code/models/perf/benchmarking_utils.py

# SPDX-FileCopyrightText: © 2023 Tenstorrent USA, Inc.

# SPDX-License-Identifier: Apache-2.0

import os
import pickle
import re
from datetime import datetime
from typing import List

import pytz
from loguru import logger

# Decouple dependency of model tests on infra folder unless running in CI
IS_CI_ENV = os.getenv("CI") == "true"
if IS_CI_ENV:
    from infra.data_collection.pydantic_models import BenchmarkMeasurement, PartialBenchmarkRun
else:
    logger.warning("Skipping import of pydantic_models for benchmarking since not running in CI environment")

# Map tt-metal device names to unified names used in benchmark database
UNIFIED_DEVICE_NAME_MAP = {
    # Blackhole
    "P100": "p100",
    "P150": "p150",
    "P300": "p300",
    "P150x4": "p150x4",
    "P150x8": "p150x8",
    "BHGLX": "galaxy_bh",
    # Wormhole
    "N150": "n150",
    "N300": "n300",
    "N150x4": "n150x4",
    "T3K": "t3k",
    "TG": "galaxy",
}


class BenchmarkProfiler:
    def __init__(self):
        self.start_times = dict()
        self.end_times = dict()

    def __call__(self, step_name: str, iteration: int = 0):
        # Return a context manager for this step
        return self.StepContext(self, step_name, iteration)

    class StepContext:
        def __init__(self, profiler, step_name: str, iteration: int):
            self.profiler = profiler
            self.step_name = step_name
            self.iteration = iteration

        def __enter__(self):
            self.profiler.start(self.step_name, self.iteration)
            return self.profiler

        def __exit__(self, exc_type, exc_val, exc_tb):
            self.profiler.end(self.step_name, self.iteration)
            return False

    def start(self, step_name: str, iteration: int = 0):
        self.start_times[(iteration, step_name)] = datetime.now(tz=pytz.UTC)

    def end(self, step_name: str, iteration: int = 0):
        self.end_times[(iteration, step_name)] = datetime.now(tz=pytz.UTC)

    def get_duration(self, step_name: str, iteration: int = 0):
        start_time = self.start_times[(iteration, step_name)]
        end_time = self.end_times[(iteration, step_name)]
        return (end_time - start_time).total_seconds()

    def get_str_start(self, step_name: str, iteration: int = 0):
        return self._get_str_ts(self.start_times[(iteration, step_name)])

    def get_str_end(self, step_name: str, iteration: int = 0):
        return self._get_str_ts(self.end_times[(iteration, step_name)])

    def contains_step(self, step_name: str, iteration: int = 0):
        return (iteration, step_name) in self.start_times and (iteration, step_name) in self.end_times

    def _get_str_ts(self, timestamp):
        return timestamp.strftime("%Y-%m-%dT%H:%M:%S%z")

    def _get_step_durations(self, step_name: str, start_iteration: int = 0):
        """Helper method to collect all durations for a given step across iterations."""
        durations = []
        iteration = start_iteration
        while self.contains_step(step_name, iteration):
            durations.append(self.get_duration(step_name, iteration))
            iteration += 1
        return durations

    def get_duration_average(self, step_name: str, start_iteration: int = 0):
        durations = self._get_step_durations(step_name, start_iteration)
        return sum(durations) / len(durations)

    def get_duration_sum(self, step_name: str, start_iteration: int = 0):
        durations = self._get_step_durations(step_name, start_iteration)
        return sum(durations)


class BenchmarkData:
    def __init__(self):
        if IS_CI_ENV:
            self.measure_data: List[BenchmarkMeasurement] = []
            self.output_folder = "generated/benchmark_data/"

    def add_measurement(
        self,
        profiler: BenchmarkProfiler,
        iteration: int,
        step_name: str,
        name: str,
        value: float,
        step_warm_up_num_iterations: int = None,
        target: float = None,
        device_power: float = None,
        device_temperature: float = None,
    ):
        if IS_CI_ENV:  # no-op if not running in CI environment
            """
            Measurement data contains records and attributes for each measurement performed at each iteration and each step of the benchmark run.
            The triad of fields (iteration, step_name, name) must be unique.
            """
            assert None not in [profiler, iteration, step_name, name, value], "Missing required fields"
            assert profiler.contains_step(
                step_name, iteration
            ), f"Completed step '{step_name}' for iteration {iteration} not found in profiler"
            self.measure_data.append(
                BenchmarkMeasurement(
                    step_start_ts=profiler.get_str_start(step_name, iteration),
                    step_end_ts=profiler.get_str_end(step_name, iteration),
                    iteration=iteration,
                    step_name=step_name,
                    name=name,
                    value=value,
                    step_warm_up_num_iterations=step_warm_up_num_iterations,
                    target=target,
                    device_power=device_power,
                    device_temperature=device_temperature,
                )
            )

    def save_partial_run_json(
        self,
        profiler: BenchmarkProfiler,  # must contain a "run" step for the entire run
        run_type: str,
        ml_model_name: str,
        ml_model_type: str = None,
        device_name: str = "",
        num_layers: int = None,
        batch_size: int = None,
        config_params: dict = None,
        precision: str = None,
        device_hostname: str = None,
        dataset_name: str = None,
        profiler_name: str = None,
        input_sequence_length: int = None,
        output_sequence_length: int = None,
        image_dimension: str = None,
        perf_analysis: bool = None,
        training: bool = None,
    ):
        if IS_CI_ENV:  # no-op if not running in CI environment
            assert None not in [profiler, run_type, ml_model_name], "Missing required fields"
            assert profiler.contains_step("run"), "Run step not found in profiler"

            run_start_ts = profiler.get_str_start("run")
            run_end_ts = profiler.get_str_end("run")

            device_info = None
            if device_name:
                device_name = UNIFIED_DEVICE_NAME_MAP.get(device_name, device_name)
                device_info = {"device_name": device_name}

            # Standardize model name - convert lowercase and replace whitespace and underscores with hyphens
            # strip org name if present (e.g. meta-llama/Llama-3.1-70B -> Llama-3.1-70B)
            ml_model_name = ml_model_name.rstrip("/").split("/", 1)[-1]
            ml_model_name = re.sub(r"[\s_]+", "-", ml_model_name.lower())

            partial_benchmark_run = PartialBenchmarkRun(
                run_start_ts=run_start_ts,
                run_end_ts=run_end_ts,
                run_type=run_type,
                device_hostname=device_hostname,
                device_info=device_info,
                ml_model_name=ml_model_name,
                ml_model_type=ml_model_type,
                num_layers=num_layers,
                batch_size=batch_size,
                config_params=config_params,
                precision=precision,
                dataset_name=dataset_name,
                profiler_name=profiler_name,
                input_sequence_length=input_sequence_length,
                output_sequence_length=output_sequence_length,
                image_dimension=image_dimension,
                perf_analysis=perf_analysis,
                training=training,
                measurements=self.measure_data,
            )

            pkl_data = pickle.dumps(partial_benchmark_run)

            filename = os.path.join(self.output_folder, f"partial_run_{run_start_ts}.pkl")
            parent_dir = os.path.dirname(filename)
            if parent_dir != "":
                # exist_ok=True avoids a race when parallel MPI ranks create the directory concurrently
                os.makedirs(parent_dir, exist_ok=True)
            with open(filename, "wb") as f:
                f.write(pkl_data)
            logger.info(f"Run and measurement data saved to {filename}")
        else:
            logger.info("Skipping saving partial benchmark data PKL since not running in CI environment")