nsight-python 1.0.0__py3-none-any.whl → 1.1.1__py3-none-any.whl

nsight/__init__.py

@@ -1,16 +1,40 @@
 # SPDX-FileCopyrightText: Copyright (c) 2025 NVIDIA CORPORATION & AFFILIATES. All rights reserved.
 # SPDX-License-Identifier: Apache-2.0
 
+import os
 from importlib.metadata import version
 
-from nsight import analyze
+from nsight import analyze, exceptions, experimental
 from nsight.annotation import annotate
-from nsight.collection.ncu import try_init_injection
+from nsight.collection import ncu
+from nsight.info_collector import CollectionScope, InfoCollector
+from nsight.tools_manager import Tool, activate, deactivate, get_active_tool
 from nsight.utils import VerbosityLevel, col_panel, row_panel
 
-# Load injection library and symbols for NCU attach (if available)
-try_init_injection()
+# NVIDIA Nsight Compute must attach before CUDA is initialized, so by default it
+# is activated here -- importing nsight early is enough to satisfy that. Set
+# NSPY_NCU_INIT_AT_IMPORT=0 to disable this and activate a tool yourself.
+if os.environ.get("NSPY_NCU_INIT_AT_IMPORT", "1") != "0":
+    try:
+        activate(Tool.NCU)
+    except (exceptions.ProfilerException, exceptions.NCUNotAvailableError) as exc:
+        # Importing nsight must keep working when the NVIDIA Nsight Compute
+        # injection fails to load, so record the error rather than raising it
+        # here. NCUCollector.collect re-raises it late, when a decorated
+        # function is first called.
+        ncu.injection_load_error = exc
 
 __version__ = version("nsight-python")
 
-__all__ = ["analyze", "annotate", "VerbosityLevel"]
+__all__ = [
+    "CollectionScope",
+    "InfoCollector",
+    "Tool",
+    "VerbosityLevel",
+    "activate",
+    "analyze",
+    "annotate",
+    "deactivate",
+    "experimental",
+    "get_active_tool",
+]

nsight/analyze.py

@@ -16,6 +16,8 @@ from numpy.typing import NDArray
 
 import nsight.collection as collection
 import nsight.visualization as visualization
+from nsight.info_collector import CollectionScope, InfoCollector
+from nsight.tools_manager import Tool, activate, get_active_tool
 from nsight.utils import VerbosityLevel
 
 
@@ -54,6 +56,14 @@ def kernel(
     combine_kernel_metrics: Callable[[float, float], float] | None = None,
     output_prefix: str | None = None,
     output_csv: bool = False,
+    info_collectors: (
+        Sequence[
+            InfoCollector
+            | tuple[str, Callable[..., Any], str | CollectionScope]
+            | Callable[..., Any]
+        ]
+        | None
+    ) = None,
 ) -> Callable[[Callable[..., None]], Callable[..., collection.core.ProfileResults]]: ...
 
 
@@ -84,12 +94,20 @@ def kernel(
     combine_kernel_metrics: Callable[[float, float], float] | None = None,
     output_prefix: str | None = None,
     output_csv: bool = False,
+    info_collectors: (
+        Sequence[
+            InfoCollector
+            | tuple[str, Callable[..., Any], str | CollectionScope]
+            | Callable[..., Any]
+        ]
+        | None
+    ) = None,
 ) -> (
     Callable[..., collection.core.ProfileResults]
     | Callable[[Callable[..., None]], Callable[..., collection.core.ProfileResults]]
 ):
     """
-    A decorator that collects profiling data using NVIDIA Nsight Compute.
+    A decorator that collects profiling data from GPU kernels.
 
     Can be used with or without parentheses:
         - ``@nsight.analyze.kernel`` (no parentheses)
@@ -106,6 +124,14 @@ def kernel(
         - ``**kwargs``: Original function keyword arguments
         - Returns ``ProfileResults`` object containing profiling data
 
+    .. warning::
+        Decorating a function changes its return value. Any value returned by the
+        original function is ignored, and calling the decorated function returns a
+        ``ProfileResults`` object instead. If the original function returns a
+        non-``None`` value, Nsight Python emits a ``RuntimeWarning``. Keep any
+        computation whose result you need in an undecorated helper and use the
+        decorated function only as a profiling wrapper.
+
     .. note::
         Configs are passed to the decorated function as positional arguments.
         Both positional and keyword-only parameters are supported in the function
@@ -113,6 +139,14 @@ def kernel(
         but will always be empty during profiling. Parameters with default values
         may be omitted from configs — missing values will be filled from the defaults.
 
+    Profiling uses NVIDIA Nsight Compute by default. Activating it is automatic
+    -- it happens during ``import nsight`` -- and needs no attention from you.
+    For more control over when NVIDIA Nsight Compute is loaded, see the
+    :doc:`/tools` page.
+
+    ``@nsight.analyze.kernel`` supports collection of
+    ``"gpu__time_duration.sum"`` via the low-overhead, experimental CUPTI
+    backend. See the :doc:`/tools` page for how to enable it.
 
     Parameters:
         configs: An iterable of configurations to run the function with. Each configuration can be either:
@@ -134,7 +168,7 @@ def kernel(
         derive_metric:
             A function to transform the collected metrics.
             This can be used to compute derived metrics like TFLOPs that cannot
-            be captured by ncu directly. The function takes the metric values and
+            be captured directly. The function takes the metric values and
             the arguments of the profile-decorated function and returns the new
             metrics. Return value can be either:
 
@@ -157,7 +191,7 @@ def kernel(
             Combine with ``derive_metric`` to compute speedup (reciprocal of normalized value).
         metrics: The metrics to collect. By default, kernel runtimes in nanoseconds
             are collected. Default: ``["gpu__time_duration.sum"]``. To see the available
-            metrics on your system, use the command: ``ncu --query-metrics``.
+            metrics on your system when using ncu, use the command: ``ncu --query-metrics``.
         ignore_kernel_list:
             List of kernel names to ignore. If you call a library within an annotated range 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 profile, their names can be provided in this parameter. Default: ``None``
@@ -167,13 +201,19 @@ def kernel(
             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``.
+            two arguments and return a single value. Each annotation must still launch
+            the same number of kernels across configurations and repeated runs;
+            different counts may cause an error or incorrect results. Default: ``None``.
         clock_control: Control the behavior of the GPU clocks during profiling. Allowed values:
 
             - ``"base"``: GPC and memory clocks are locked to their respective base frequency during profiling. This has no impact on thermal throttling. Note that actual clocks might still vary, depending on the level of driver support for this feature. As an alternative, use nvidia-smi to lock the clocks externally and set this option to ``"none"``.
             - ``"none"``: No GPC or memory frequencies are changed during profiling.
 
             Default: ``"none"``
+
+            Refer the `NVIDIA Nsight Compute documentation on Clock Control <https://docs.nvidia.com/nsight-compute/ProfilingGuide/index.html#clock-control>`_ for more details.
+
+            Ignored when the CUPTI backend is active.
         cache_control: Control the behavior of the GPU caches during profiling. Allowed values:
 
             - ``"all"``: All GPU caches are flushed before each kernel replay iteration during profiling. While metric values in the execution environment of the application might be slightly different without invalidating the caches, this mode offers the most reproducible metric results across the replay passes and also across multiple runs of the target application.
@@ -181,6 +221,8 @@ def kernel(
 
             Default: ``"all"``
 
+            Ignored when the CUPTI backend is active.
+
         replay_mode: Mechanism used for replaying a kernel launch multiple times to collect selected metrics. Allowed values:
 
             - ``"kernel"``:  Replay individual kernel launches  during the execution of the application.
@@ -188,6 +230,8 @@ def kernel(
 
             Default: ``"kernel"``
 
+            Ignored when the CUPTI backend is active.
+
         thermal_mode: Controls GPU thermal management mode. Default: ``"auto"``
 
             Monitors GPU thermal headroom (T.Limit) and pauses profiling when GPU gets
@@ -280,8 +324,8 @@ def kernel(
                 - ``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 % .
+                - ``RelativeStdDevPct``: Standard deviation as a percentage of the absolute mean
+                - ``StableMeasurement``: Nullable Boolean: true if ``RelativeStdDevPct`` < 2%, false otherwise, and ``pd.NA`` when relative variability is undefined (for example, with one valid sample or a zero mean).
                 - ``Metric``: The metrics being collected and the metrics being derived
                 - ``Unit``: Unit of measurement for the metric value (e.g., ``ns`` for nanoseconds)
                 - ``Kernel``: Name of the GPU kernel(s) launched
@@ -289,10 +333,70 @@ def kernel(
                 - ``Host``: Host machine name
                 - ``ComputeClock``: GPU compute clock frequency
                 - ``MemoryClock``: GPU memory clock frequency
+        info_collectors: Optional list of custom information collectors.
+            Collectors should be functions decorated with
+            ``@nsight.experimental.collect(scope=...)``.
+            The function name will be used as the column name in the DataFrame.
+
+            **Decorated functions** (recommended)::
+
+                @nsight.experimental.collect(scope=nsight.CollectionScope.ONCE)
+                def driver_version():
+                    from cuda.core import system
+                    return ".".join(
+                        str(part) for part in system.get_user_mode_driver_version()
+                    )
+
+                @nsight.experimental.collect(scope=nsight.CollectionScope.CONFIG)
+                def matrix_elements(n, dtype):
+                    return n * n
+
+                @nsight.experimental.collect(scope=nsight.CollectionScope.RUN)
+                def gpu_temp_c(n, dtype):
+                    return get_gpu_temperature()
+
+                @nsight.experimental.collect(
+                    scope=nsight.CollectionScope.ANNOTATION
+                )
+                def ann_timestamp(annotation_name, *config_args):
+                    return time.time()
+
+                @nsight.analyze.kernel(
+                    info_collectors=[
+                        driver_version, matrix_elements, gpu_temp_c, ann_timestamp
+                    ]
+                )
+                def my_benchmark(n, dtype):
+                    ...
+
+            Collectors may also be supplied as ``(name, callback, scope)``
+            tuples, where ``scope`` is ``"once"``, ``"config"``, ``"run"``,
+            or ``"annotation"`` (or a :class:`nsight.CollectionScope`).
+
+            Column naming: ``once``-scope and non-numeric collectors produce a
+            single column named after the function. Numeric ``config``/``run``/
+            ``annotation``-scope collectors are aggregated across runs into four
+            columns ``<name>_Avg``, ``<name>_Std``, ``<name>_Min``, ``<name>_Max``.
     """
     # Strip whitespace
     metrics = [m.strip() for m in metrics]
 
+    # These are decorator parameters, so reject illegal values now rather than
+    # when the decorated function is first called. Whether the active tool
+    # supports a legal value is a separate check, made by the collector.
+    if clock_control not in ("base", "none"):
+        raise ValueError(
+            f"Invalid clock_control {clock_control!r}. Expected 'base' or 'none'."
+        )
+    if cache_control not in ("all", "none"):
+        raise ValueError(
+            f"Invalid cache_control {cache_control!r}. Expected 'all' or 'none'."
+        )
+    if replay_mode not in ("kernel", "range"):
+        raise ValueError(
+            f"Invalid replay_mode {replay_mode!r}. Expected 'kernel' or 'range'."
+        )
+
     # Validate thermal parameters if both are provided
     if thermal_wait is not None and thermal_cont is not None:
         if thermal_cont <= thermal_wait:
@@ -309,10 +413,72 @@ def kernel(
         prefix = output_prefix
         if prefix is None:
             prefix = tempfile.mkdtemp(prefix="nspy_")
-            prefix = os.path.join(prefix, "")  # Adds a trailing forward/backward slash
+            prefix = os.path.join(prefix, "")  # Adds a trailing path separator
         else:
             os.makedirs(os.path.dirname(prefix) or ".", exist_ok=True)
 
+        # Normalize info_collectors: convert decorated functions to
+        # (name, callback, scope) tuples with a canonical string scope.
+        valid_scopes = {s.value for s in CollectionScope}
+
+        def _canonical_scope(scope: Any, label: str) -> str:
+            if isinstance(scope, CollectionScope):
+                return scope.value
+            if isinstance(scope, str) and scope in valid_scopes:
+                return scope
+            raise ValueError(
+                f"Invalid scope {scope!r} for collector {label}. "
+                f"Expected one of {sorted(valid_scopes)} or a CollectionScope."
+            )
+
+        normalized_collectors = None
+        if info_collectors:
+            normalized_collectors = []
+            for collector in info_collectors:
+                if isinstance(collector, InfoCollector):
+                    # An InfoCollector instance (the exported public class).
+                    normalized_collectors.append(
+                        (
+                            collector.name,
+                            collector.callback,
+                            _canonical_scope(collector.scope, repr(collector.name)),
+                        )
+                    )
+                elif isinstance(collector, tuple):
+                    # Raw tuple format: (name, callback, scope)
+                    if len(collector) != 3:
+                        raise ValueError(
+                            "Tuple-format info collector must be "
+                            f"(name, callback, scope); got {collector!r}."
+                        )
+                    name, callback, scope = collector
+                    if not isinstance(name, str) or not callable(callback):
+                        raise TypeError(
+                            "Tuple-format info collector must be "
+                            f"(str name, callable callback, scope); got {collector!r}."
+                        )
+                    normalized_collectors.append(
+                        (name, callback, _canonical_scope(scope, repr(name)))
+                    )
+                elif callable(collector):
+                    # Decorated function: extract metadata
+                    if hasattr(collector, "_nsight_collector_scope"):
+                        name = getattr(collector, "_nsight_collector_name")
+                        scope = getattr(collector, "_nsight_collector_scope")
+                        normalized_collectors.append(
+                            (name, collector, _canonical_scope(scope, repr(name)))
+                        )
+                    else:
+                        raise ValueError(
+                            f"Collector function {collector.__name__} is missing "
+                            "@nsight.experimental.collect(scope=...) decorator"
+                        )
+                else:
+                    raise TypeError(
+                        f"Invalid collector type: {type(collector)}. "
+                        "Expected tuple (name, callback, scope) or decorated function."
+                    )
+
         settings = collection.core.ProfileSettings(
             configs=configs,
             runs=runs,
@@ -326,16 +492,28 @@ def kernel(
             thermal_device=thermal_device,
             output_prefix=prefix,
             output_csv=output_csv,
+            info_collectors=normalized_collectors,
         )
-        ncu = collection.ncu.NCUCollector(
-            metrics=metrics,
-            ignore_kernel_list=ignore_kernel_list,
-            combine_kernel_metrics=combine_kernel_metrics,
-            clock_control=clock_control,
-            cache_control=cache_control,
-            replay_mode=replay_mode,
-        )
-        return collection.core.NsightProfiler(settings, ncu)
+
+        def get_collector() -> collection.core.NsightCollector:
+            if get_active_tool() == Tool.CUPTI:
+                # clock_control, cache_control and replay_mode have no
+                # meaning for CUPTI, so the collector does not take them.
+                return collection.cupti.CUPTICollector(
+                    metrics=metrics,
+                    ignore_kernel_list=ignore_kernel_list,
+                    combine_kernel_metrics=combine_kernel_metrics,
+                )
+            return collection.ncu.NCUCollector(
+                metrics=metrics,
+                ignore_kernel_list=ignore_kernel_list,
+                combine_kernel_metrics=combine_kernel_metrics,
+                clock_control=clock_control,
+                cache_control=cache_control,
+                replay_mode=replay_mode,
+            )
+
+        return collection.core.NsightProfiler(settings, get_collector)
 
     # Support both @kernel and @kernel() syntax
     if _func is None:
@@ -410,7 +588,7 @@ def plot(
     ylabel: str | None = None,
     annotate_points: bool = False,
     show_aggregate: str | None = None,
-    plot_type: str = "line",
+    plot_type: Literal["line", "bar"] = "line",
     plot_width: int = 6,
     plot_height: int = 4,
     row_panels: Sequence[str] | None = None,
@@ -476,6 +654,11 @@ def plot(
         variant_fields: List of config fields to use as variant fields (lines).
         variant_annotations: List of annotated range names for which to apply variant splitting. The provided strings must each match one of the names defined using nsight.annotate.
     """
+    # Reject illegal values now rather than when the decorated function is
+    # first called.
+    if plot_type not in ("line", "bar"):
+        raise ValueError(f"Invalid plot_type {plot_type!r}. Expected 'line' or 'bar'.")
+
     show_avg = show_aggregate == "avg"
     show_geomean = show_aggregate == "geomean"
 

nsight/annotation.py

@@ -3,12 +3,14 @@
 
 import functools
 import threading
+import warnings
 from collections.abc import Callable
 from typing import Any
 
 import nvtx
 
 import nsight.utils as utils
+from nsight import info_collector
 from nsight.exceptions import CUDA_CORE_UNAVAILABLE_MSG
 
 # Thread-local storage for tracking annotations
@@ -88,7 +90,7 @@ def is_active_annotation(name: str) -> bool:
     return False
 
 
-class annotate(nvtx.annotate):  # type: ignore[misc]
+class annotate:
     """
     A decorator/context-manager hybrid for marking profiling regions.
     The encapsulated code will be profiled and associated with an NVTX
@@ -111,7 +113,9 @@ class annotate(nvtx.annotate):  # type: ignore[misc]
         ignore_failures: Flag indicating whether to ignore
             failures in the annotate context. If set to ``True``, any exceptions
             raised within the context will be ignored, and the profiling will
-            continue. The measured metric for this run will be set to NaN.
+            continue. The measured metric for this run will be set to NaN, and
+            a ``RuntimeWarning`` naming the annotation and the ignored exception
+            is emitted.
             Default: ``False``
 
     Raises:
@@ -129,10 +133,10 @@ class annotate(nvtx.annotate):  # type: ignore[misc]
         Nested annotations are currently not supported. However, since each
         annotation is expected to contain a single kernel launch by default,
         nested annotations should not be necessary in typical usage scenarios.
-        Entering and exiting the annotation context pushes and pops an NVTX
-        range, so calling ``nvtx.pop_range()`` inside the annotation context
-        without a corresponding push would pop the annotation context's own
-        range instead, resulting in incorrect profiling results.
+        Entering and exiting the annotation context starts and ends an NVTX
+        range that is visible across all threads in the profiled process, so
+        kernels launched from worker threads within the annotated region are
+        also captured.
 
     """
 
@@ -144,9 +148,7 @@ class annotate(nvtx.annotate):  # type: ignore[misc]
         if ignore_failures and not utils.CUDA_CORE_AVAILABLE:
             raise ImportError(CUDA_CORE_UNAVAILABLE_MSG)
 
-        super().__init__(name, domain=utils.NVTX_DOMAIN)
-
-    def __enter__(self) -> nvtx.annotate:
+    def __enter__(self) -> "annotate":
 
         # Check for nested annotations
         if _is_in_annotation():
@@ -164,19 +166,32 @@ class annotate(nvtx.annotate):  # type: ignore[misc]
 
         add_active_annotation(self.name)
         _set_in_annotation(True)
-        return super().__enter__()
+        # Collect annotation-scope information
+        info_collector.collect_for_annotation(self.name)
+        self._range_id = nvtx.start_range(message=self.name, domain=utils.NVTX_DOMAIN)
+        return self
 
     def __exit__(self, exc_type: Any, exc_value: Any, traceback: Any) -> bool:
         try:
             if exc_type and self.ignore_failures:
                 utils.launch_dummy_kernel_module()
         finally:
-            super().__exit__(exc_type, exc_value, traceback)
+            nvtx.end_range(self._range_id)
             _set_in_annotation(False)
 
         if exc_type and not self.ignore_failures:
             return False  # propagate the exception
 
+        if exc_type:
+            # The failed run only shows up as NaN in the results, so report why
+            # it failed rather than discarding the exception.
+            warnings.warn(
+                f"Ignored failure in annotation '{self.name}': "
+                f"{exc_type.__name__}: {exc_value}",
+                category=RuntimeWarning,
+                stacklevel=2,
+            )
+
         return True
 
     def __call__(self, func: Callable[..., Any]) -> Callable[..., Any]:

nsight/collection/__init__.py

@@ -4,7 +4,8 @@
 import functools
 
 import nsight.collection.core as core
+import nsight.collection.cupti as cupti
 import nsight.collection.ncu as ncu
 import nsight.utils as utils
 
-__all__ = ["ncu", "core"]
+__all__ = ["ncu", "cupti", "core"]