nsight-python 0.9.4__py3-none-any.whl

nsight/collection/core.py

@@ -0,0 +1,399 @@
+# SPDX-FileCopyrightText: Copyright (c) 2025 NVIDIA CORPORATION & AFFILIATES. All rights reserved.
+# SPDX-License-Identifier: Apache-2.0
+
+import abc
+import dataclasses
+import functools
+import importlib.util
+import inspect
+import os
+import time
+from collections.abc import Callable, Sequence
+from typing import Any
+
+import pandas as pd
+
+from nsight import exceptions, thermovision, transformation, utils
+
+
+def _sanitize_configs(
+    func: Callable[..., Any],
+    *args: Any,
+    configs: Sequence[Sequence[Any]] | None = None,
+    decorator_configs: Sequence[Sequence[Any]] | None = None,
+    **kwargs: Any,
+) -> list[tuple[Any, ...]]:
+    """
+    Sanitizes and validates configuration inputs for a profile-decorated function.
+
+    This function ensures that the provided configurations are consistent and
+    handles different cases of passing configurations.
+
+    1. As regular args+kw - A single configuration -> the function arguments
+    2. As configs=, a list of configurations at functionn call time
+    3. As decorator_configs=, a list of configurations at decoration time
+
+    Args:
+        *args: Positional arguments that may contain configuration data.
+        configs: A list of configurations provided at runtime.
+        decorator_configs: A list of configurations provided
+            at decoration time.
+        **kwargs: Keyword arguments that may contain additional configuration data.
+
+    Returns:
+        A sanitized list of configurations.
+
+    Raises:
+        exceptions.ProfilerException: If no configurations are provided or if
+            configurations are provided both at decoration time and runtime.
+        AssertionError: If `configs` is not a list when provided.
+
+    Notes:
+        - If `args` are provided, `configs` and `decorator_configs` must be `None`.
+        - If `configs` is provided, `decorator_configs` must be `None`, and vice versa.
+        - The function combines `args` and `kwargs` into a single list if `args` are provided.
+        - The function assumes that `kwargs` keys are in the expected order when combining.
+    """
+    if len(args) > 0:
+        # We do not expect any configs in this case
+        assert configs is None and decorator_configs is None
+        # kwargs not supported here yet
+        if len(kwargs) != 0:
+            raise exceptions.ProfilerException(
+                f"Keyword arguments are not supported yet: {list(kwargs.keys())}"
+            )
+
+        configs = [list(args)]
+
+    if configs is None:
+        if decorator_configs is None:
+            raise exceptions.ProfilerException(
+                "You have provided no configs. Provide configs at decoration time or at runtime."
+            )
+        configs = decorator_configs
+
+    else:
+        if decorator_configs is not None:
+            raise exceptions.ProfilerException(
+                "You have provided configs at decoration time and at runtime. Provide configs at decoration time or at runtime."
+            )
+        assert isinstance(configs, list), f"configs must be a list, got {type(configs)}"
+
+    # Validate that all configs have the same number of arguments
+    if len(configs) == 0:
+        raise exceptions.ProfilerException("configs list cannot be empty")
+    config_lengths = [len(config) for config in configs]
+    if not all(length == config_lengths[0] for length in config_lengths):
+        raise exceptions.ProfilerException(
+            f"All configs must have the same number of arguments. Found lengths: {config_lengths}"
+        )
+    first_config_arg_count = config_lengths[0]
+
+    # Validate that the number of args matches the number of function parameters
+    sig = inspect.signature(func)
+    expected_arg_count = len(sig.parameters)
+    if first_config_arg_count != expected_arg_count:
+        raise exceptions.ProfilerException(
+            f"Configs have {first_config_arg_count} arguments, but function expects {expected_arg_count}"
+        )
+
+    return configs  # type: ignore[return-value]
+
+
+def run_profile_session(
+    func: Callable[..., Any],
+    configs: Sequence[Sequence[Any]],
+    runs: int,
+    output_progress: bool,
+    output_detailed: bool,
+    thermal_control: bool,
+) -> None:
+
+    if output_progress:
+        print("")
+        print("")
+
+    if thermal_control:
+        thermovision_initialized = thermovision.init()
+
+    total_configs = len(configs)
+    total_runs = total_configs * runs  # Total runs executed
+    curr_config = 0
+    curr_run = 0
+    total_time: float = 0
+    bar_length = 100
+    progress_time: float = 0
+
+    # overwrite flag: we do not overwrite when output mode is detailed
+    overwrite_output = not output_detailed
+
+    for c in configs:
+        curr_config += 1
+
+        if output_progress:
+            utils.print_config(total_configs, curr_config, c, overwrite_output)
+
+        for i in range(runs):
+            start_time = time.time()
+            curr_run += 1
+            if thermal_control:
+                if thermovision_initialized:
+                    thermovision.throttle_guard()
+
+            # Check if func supports the input configs
+            if len(inspect.signature(func).parameters) != len(c):
+                raise exceptions.ProfilerException(
+                    f"Function '{func.__name__}' does not support the input configuration"
+                )
+
+            # Run the function with the config
+            func(*c)
+
+            elapsed_time = time.time() - start_time
+            if curr_run > 1:
+                total_time += elapsed_time
+                avg_time_per_run = total_time / curr_run
+            else:
+                avg_time_per_run = elapsed_time  # Use first run's time only
+
+            # Update time estimates every half second
+            if time.time() - progress_time > 0.5:
+                if output_progress:
+                    utils.print_progress_bar(
+                        total_runs,
+                        curr_run,
+                        bar_length,
+                        avg_time_per_run,
+                        overwrite_output,
+                    )
+                progress_time = time.time()
+
+    # Update progress bar at end so it shows 100%
+    if output_progress:
+        utils.print_progress_bar(
+            total_runs,
+            curr_run,
+            bar_length,
+            avg_time_per_run,
+            overwrite_output,
+        )
+
+
+@dataclasses.dataclass
+class ProfileSettings:
+    """
+    Class to hold profile settings for Nsight Python.
+    """
+
+    configs: Sequence[Sequence[Any]] | None
+    """
+    A list of configurations to run the
+    function with. Each configuration is a tuple of arguments for the
+    decorated function. Nsight Python invokes the decorated function
+    ``len(configs) * runs`` times. If the configs are not provided at
+    decoration time, they must be provided when calling the decorated function.
+    """
+
+    runs: int
+    """Number of times each configuration should be executed."""
+
+    output_progress: bool
+    """
+    Will display a progress bar on stdout during profiling
+    """
+
+    output_detailed: bool
+    """
+    Will display a progress bar, detailed output for each config along with the profiler logs
+    """
+
+    derive_metric: Callable[..., float] | None
+    """
+    A function to transform the collected metric.
+    This can be used to compute derived metrics like TFLOPs that cannot
+    be captured by ncu directly. The function takes the metric value and
+    the arguments of the profile-decorated function and returns the new
+    metric. See the examples for concrete use cases.
+    """
+
+    normalize_against: str | None
+    """
+    Annotation name to normalize metrics against.
+    This is useful to compute relative metrics like speedup.
+    """
+
+    thermal_control: bool
+    """
+    Toggles whether to enable thermal control.
+    """
+
+    output_prefix: str | None
+    """
+    All intermediate profiler files are created with this prefix
+    """
+
+    output_csv: bool
+    """
+    Controls whether to output raw and processed profiling data to CSV files
+    """
+
+
+class ProfileResults:
+    """
+    Class to hold profile results for Nsight Python
+    """
+
+    def __init__(self, results: pd.DataFrame):
+        """
+        Initialize a ProfileResults object.
+
+        Args:
+            results: Processed profiling results.
+        """
+        self._results = results
+
+    def to_dataframe(self) -> pd.DataFrame:
+        """
+        Returns the processed profiling data as a pandas DataFrame.
+
+        This DataFrame contains aggregated statistics across multiple runs for each
+        configuration and annotation combination. The data is equivalent to what is
+        written to the ``processed_data-<function_name>-<run_id>.csv`` file when
+        ``output_csv=True``.
+
+        Returns:
+            Processed profiling data with the following columns:
+
+                - ``Annotation``: Name of the annotated region being profiled
+                - ``<param_name>``: One column for each parameter of the decorated function
+                - ``AvgValue``: Average metric value across all runs
+                - ``StdDev``: Standard deviation of the metric across runs
+                - ``MinValue``: Minimum metric value observed
+                - ``MaxValue``: Maximum metric value observed
+                - ``NumRuns``: Number of runs used for aggregation
+                - ``CI95_Lower``: Lower bound of the 95% confidence interval
+                - ``CI95_Upper``: Upper bound of the 95% confidence interval
+                - ``RelativeStdDevPct``: Standard deviation as a percentage of the mean
+                - ``StableMeasurement``: Boolean indicating if the measurement is stable (low variance). The measurement is stable if ``RelativeStdDevPct`` < 2 % .
+                - ``Metric``: The metric being collected
+                - ``Transformed``: Name of the function used to transform the metric (specified via ``derive_metric``), or ``False`` if no transformation was applied. For lambda functions, this shows ``"<lambda>"``
+                - ``Kernel``: Name of the GPU kernel(s) launched
+                - ``GPU``: GPU device name
+                - ``Host``: Host machine name
+                - ``ComputeClock``: GPU compute clock frequency
+                - ``MemoryClock``: GPU memory clock frequency
+        """
+        return self._results
+
+
+class NsightCollector(abc.ABC):
+    @abc.abstractmethod
+    def collect(
+        self,
+        func: Callable[..., Any],
+        configs: Sequence[Sequence[Any]],
+        settings: ProfileSettings,
+    ) -> pd.DataFrame | None:
+        """
+        Collects profiling data for the given function and configurations.
+
+        Args:
+            func: The function to be profiled.
+            configs: List of configurations for profiling.
+            settings: Settings for profiling.
+
+        Returns:
+            Collected profiling data.
+        """
+        pass
+
+
+class NsightProfiler:
+    """
+    A decorator class for profiling functions using Nsight Python's profiling tools.
+
+    This class allows you to wrap a function with profiling capabilities,
+    collecting performance data and saving the results to CSV files. It uses
+    a collector to gather raw profiling data and processes it according to
+    the provided settings.
+
+    Attributes:
+        settings: Configuration settings for profiling,
+            including normalization, and other options.
+        collector: The collector responsible for gathering
+            raw profiling data.
+
+    Methods:
+        __call__(func):
+            Wraps the given function with profiling logic. Collects raw
+            profiling data, processes it, and saves the results to CSV files.
+            Returns the processed data.
+    """
+
+    def __init__(self, settings: ProfileSettings, collector: NsightCollector):
+        self.settings = settings
+        self.collector = collector
+
+    def __call__(
+        self, func: Callable[..., Any]
+    ) -> Callable[..., ProfileResults | None]:
+        func._nspy_ncu_run_id = 0  # type: ignore[attr-defined]
+
+        @functools.wraps(func)
+        def wrapper(
+            *args: Any, configs: Sequence[Sequence[Any]] | None = None, **kwargs: Any
+        ) -> ProfileResults | None:
+
+            tag = f"{func.__name__}-{func._nspy_ncu_run_id}"  # type: ignore[attr-defined]
+
+            configs = _sanitize_configs(
+                func,
+                *args,
+                configs=configs,
+                decorator_configs=self.settings.configs,
+                **kwargs,
+            )
+
+            raw_df = self.collector.collect(func, configs, self.settings)
+
+            if raw_df is not None:
+                processed = transformation.aggregate_data(
+                    raw_df,
+                    func,
+                    self.settings.normalize_against,
+                    self.settings.output_progress,
+                )
+
+                # Save to CSV if enabled
+                if self.settings.output_csv:
+                    raw_csv_path = (
+                        f"{self.settings.output_prefix}profiled_data-{tag}.csv"
+                    )
+                    processed_csv_path = (
+                        f"{self.settings.output_prefix}processed_data-{tag}.csv"
+                    )
+
+                    raw_df.to_csv(
+                        raw_csv_path,
+                        index=False,
+                    )
+                    processed.to_csv(
+                        processed_csv_path,
+                        index=False,
+                    )
+
+                    if self.settings.output_progress:
+                        print(
+                            f"[NSIGHT-PYTHON] Refer to {raw_csv_path} for the raw profiling data"
+                        )
+                        print(
+                            f"[NSIGHT-PYTHON] Refer to {processed_csv_path} for the processed profiling data"
+                        )
+
+                func._nspy_ncu_run_id += 1  # type: ignore[attr-defined]
+
+                return ProfileResults(results=processed)
+
+            return None
+
+        return wrapper

nsight/collection/ncu.py

@@ -0,0 +1,268 @@
+# SPDX-FileCopyrightText: Copyright (c) 2025 NVIDIA CORPORATION & AFFILIATES. All rights reserved.
+# SPDX-License-Identifier: Apache-2.0
+
+"""
+Collection utilities for profiling Nsight Python runs using NVIDIA Nsight Compute (ncu).
+
+This module contains logic for launching NVIDIA Nsight Compute with appropriate settings.
+NCU is instructed to profile specific code sections marked by NVTX ranges - the
+Nsight Python annotations.
+"""
+
+import os
+import subprocess
+import sys
+from collections.abc import Callable, Sequence
+from typing import Any, Literal
+
+import pandas as pd
+
+from nsight import exceptions, extraction, utils
+from nsight.collection import core
+from nsight.exceptions import NCUErrorContext
+
+
+def launch_ncu(
+    report_path: str,
+    name: str,
+    metric: str,
+    cache_control: Literal["none", "all"],
+    clock_control: Literal["none", "base"],
+    replay_mode: Literal["kernel", "range"],
+    verbose: bool,
+) -> str | None:
+    """
+    Launch NVIDIA Nsight Compute to profile the current script with specified options.
+
+    Args:
+        report_path: Path to write report file to.
+        metric: Specific metric to collect.
+        cache_control: Select cache control option
+        clock_control: Select clock control option
+        replay_mode: Select replay mode option
+        verbose: If False, log is written to a file (ncu_log.txt)
+
+    Raises:
+        NCUNotAvailableError: If NCU is not available on the system.
+        SystemExit: If profiling fails due to an error from NVIDIA Nsight Compute.
+
+    Returns:
+        path to the NVIDIA Nsight Compute log file
+        Produces NVIDIA Nsight Compute report file with profiling data.
+    """
+    assert report_path.endswith(".ncu-rep")
+
+    # Determine the script being executed
+    script_path = os.path.abspath(sys.argv[0])
+    script_args = " ".join(sys.argv[1:])
+
+    # Set an environment variable to detect recursive calls
+    env = os.environ.copy()
+    env["NSPY_NCU_PROFILE"] = name
+
+    if cache_control not in ("none", "all"):
+        raise ValueError("cache_control must be 'none', or 'all'")
+    if clock_control not in ("none", "base"):
+        raise ValueError("clock_control must be 'none', or 'base'")
+    if replay_mode not in ("kernel", "range"):
+        raise ValueError("replay_mode must be 'kernel', or 'range'")
+
+    cache = f"--cache-control {cache_control}"
+    clocks = f"--clock-control {clock_control}"
+    replay = f"--replay-mode {replay_mode}"
+    log_path = os.path.splitext(report_path)[0] + ".log"
+    log = f"--log-file {log_path}"
+    nvtx = f'--nvtx --nvtx-include "regex:{utils.NVTX_DOMAIN}@.+/"'
+
+    # Construct the ncu command
+    ncu_command = f"""ncu {log} {cache} {clocks} {replay} {nvtx} --metrics {metric} -f -o {report_path} {sys.executable} {script_path} {script_args}"""
+
+    # Check if ncu is available on the system
+    ncu_available = False
+    try:
+        subprocess.run(
+            ["ncu", "--version"],
+            check=True,
+            stdout=subprocess.DEVNULL,
+            stderr=subprocess.DEVNULL,
+        )
+        ncu_available = True
+    except:
+        ncu_available = False
+
+    if ncu_available:
+        try:
+            subprocess.run(
+                ncu_command,
+                shell=True,
+                check=True,
+                env=env,
+            )
+
+            return log_path
+        except subprocess.CalledProcessError as e:
+            log_parser = utils.NCULogParser()
+            error_logs = log_parser.get_logs(log_path, "ERROR")
+
+            # Create error context
+            error_context = NCUErrorContext(
+                errors=error_logs,
+                log_file_path=log_path,
+                metric=metric,
+            )
+
+            error_message = utils.format_ncu_error_message(error_context)
+            print(error_message)
+            sys.exit(1)
+    else:
+        subprocess.run([sys.executable, script_path], env=env)
+        raise exceptions.NCUNotAvailableError(
+            "Nsight Compute CLI (ncu) is not available on this system. Profiling will not be performed.\n"
+            "Please install Nsight Compute CLI."
+        )
+
+
+class NCUCollector(core.NsightCollector):
+    """
+    NCU collector for Nsight Python.
+
+    Args:
+        metric: Metric to collect from
+            NVIDIA Nsight Compute. By default we collect kernel runtimes in nanoseconds.
+            A list of supported metrics can be found with ``ncu --list-metrics``.
+        ignore_kernel_list: List of kernel names to ignore.
+            If you call a library within a ``annotation`` context, you might not have
+            precise control over which and how many kernels are being launched.
+            If some of these kernels should be ignored in the Nsight Python profile, their
+            their names can be blacklisted. Default: ``None``
+        combine_kernel_metrics: By default, Nsight Python
+            expects one kernel launch per annotation. In case an annotated region launches
+            multiple kernels, instead of failing the profiling run, you can specify
+            how to summarize the collected metrics into a single number. For example,
+            if we profile runtime and want to sum the times of all kernels we can specify
+            ``combine_kernel_metrics = lambda x, y: x + y``. The function should take
+            two arguments and return a single value. Default: ``None``.
+        clock_control: Select clock_control option
+            control in NVIDIA Nsight Compute. If ``None``, we launch ``ncu --clock-control none ...``.
+            For more details, see the NVIDIA Nsight Compute Profiling Guide:
+            https://docs.nvidia.com/nsight-compute/ProfilingGuide/index.html#clock-control
+            Default: ``None``
+        cache_control: Select cache_control option
+            control in NVIDIA Nsight Compute. If ``None``, we launch ``ncu --cache-control none ...``.
+            For more details, see the NVIDIA Nsight Compute Profiling Guide:
+            https://docs.nvidia.com/nsight-compute/ProfilingGuide/index.html#cache-control
+            Default: ``all``
+        replay_mode: Select replay mode option
+            control in NVIDIA Nsight Compute. If ``None``, we launch ``ncu --replay-mode kernel ...``.
+            For more details, see the NVIDIA Nsight Compute Profiling Guide:
+            https://docs.nvidia.com/nsight-compute/ProfilingGuide/index.html#replay
+            Default: ``kernel``
+    """
+
+    def __init__(
+        self,
+        metric: str = "gpu__time_duration.sum",
+        ignore_kernel_list: Sequence[str] | None = None,
+        combine_kernel_metrics: Callable[[float, float], float] | None = None,
+        clock_control: Literal["base", "none"] = "none",
+        cache_control: Literal["all", "none"] = "all",
+        replay_mode: Literal["kernel", "range"] = "kernel",
+    ):
+        if clock_control not in ("none", "base"):
+            raise ValueError("clock_control must be 'none', or 'base'")
+        if cache_control not in ("none", "all"):
+            raise ValueError("cache_control must be 'none', or 'all'")
+        if replay_mode not in ("kernel", "range"):
+            raise ValueError("replay_mode must be 'kernel', or 'range'")
+
+        self.metric = metric
+        self.ignore_kernel_list = ignore_kernel_list or []
+        self.combine_kernel_metrics = combine_kernel_metrics
+        self.clock_control = clock_control
+        self.cache_control = cache_control
+        self.replay_mode = replay_mode
+
+    def collect(
+        self,
+        func: Callable[..., Any],
+        configs: Sequence[Sequence[Any]],
+        settings: core.ProfileSettings,
+    ) -> pd.DataFrame | None:
+        """
+        Collects profiling data using NVIDIA Nsight Compute.
+
+        Args:
+            func: The function to profile.
+            configs: List of configurations to run the function with.
+            settings: Profiling settings.
+
+        Returns:
+            Collected profiling data.
+        """
+
+        # Check if the script is already running under ncu
+        if "NSPY_NCU_PROFILE" not in os.environ:
+
+            tag = f"{func.__name__}-{func._nspy_ncu_run_id}"  # type: ignore[attr-defined]
+            report_path = f"{settings.output_prefix}ncu-output-{tag}.ncu-rep"
+
+            # Launch NVIDIA Nsight Compute
+            log_path = launch_ncu(
+                report_path,
+                func.__name__,
+                self.metric,
+                self.cache_control,
+                self.clock_control,
+                self.replay_mode,
+                settings.output_detailed,
+            )
+
+            if settings.output_progress:
+                print("[NSIGHT-PYTHON] Profiling completed successfully !")
+                print(
+                    f"[NSIGHT-PYTHON] Refer to {report_path} for the NVIDIA Nsight Compute CLI report"
+                )
+                print(
+                    f"[NSIGHT-PYTHON] Refer to {log_path} for the NVIDIA Nsight Compute CLI logs"
+                )
+
+            # Extract raw data
+            df = extraction.extract_df_from_report(
+                report_path,
+                self.metric,
+                configs,  # type: ignore[arg-type]
+                settings.runs,
+                func,
+                settings.derive_metric,
+                self.ignore_kernel_list,  # type: ignore[arg-type]
+                settings.output_progress,
+                self.combine_kernel_metrics,
+            )
+            return df
+
+        else:
+            # If NSPY_NCU_PROFILE is set, just run the function normally
+            name = os.environ["NSPY_NCU_PROFILE"]
+
+            # If this is not the function we are profiling, stop
+            if func.__name__ != name:
+                return None
+
+            if settings.output_progress:
+                utils.print_header(
+                    f"Profiling {name}",
+                    f"{len(configs)} configurations, {settings.runs} runs each",
+                )
+
+            core.run_profile_session(
+                func,
+                configs,
+                settings.runs,
+                settings.output_progress,
+                settings.output_detailed,
+                settings.thermal_control,
+            )
+
+            # Exit after profiling to prevent the rest of the script from running
+            # Use os._exit() instead of sys.exit() to avoid pytest catching SystemExit
+            os._exit(0)