nsight-python 0.9.6__py3-none-any.whl → 1.1.0__py3-none-any.whl

nsight/__init__.py

@@ -1,12 +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.utils import col_panel, row_panel
+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
+
+# 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"]
+__all__ = [
+    "CollectionScope",
+    "InfoCollector",
+    "Tool",
+    "VerbosityLevel",
+    "activate",
+    "analyze",
+    "annotate",
+    "deactivate",
+    "experimental",
+    "get_active_tool",
+]

nsight/analyze.py

@@ -7,7 +7,7 @@ import os
 import tempfile
 import warnings
 from collections.abc import Callable, Iterable, Sequence
-from typing import Any, Literal, overload
+from typing import Any, Literal, Mapping, overload
 
 import matplotlib
 import matplotlib.figure
@@ -16,6 +16,9 @@ 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
 
 
 # Overload 1: When used without parentheses: @kernel
@@ -32,9 +35,14 @@ def kernel(
     *,
     configs: Iterable[Any] | None = None,
     runs: int = 1,
-    derive_metric: Callable[..., float | dict[str, float]] | None = None,
+    derive_metric: (
+        Callable[
+            ..., float | tuple[float, str] | Mapping[str, float | tuple[float, str]]
+        ]
+        | None
+    ) = None,
     normalize_against: str | None = None,
-    output: Literal["quiet", "progress", "verbose"] = "progress",
+    verbosity: VerbosityLevel = VerbosityLevel.INFO,
     metrics: Sequence[str] = ["gpu__time_duration.sum"],
     ignore_kernel_list: Sequence[str] | None = None,
     clock_control: Literal["base", "none"] = "none",
@@ -44,9 +52,18 @@ def kernel(
     thermal_wait: int | None = None,
     thermal_cont: int | None = None,
     thermal_timeout: int | None = None,
+    thermal_device: int | None = None,
     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]]: ...
 
 
@@ -56,9 +73,14 @@ def kernel(
     *,
     configs: Iterable[Any] | None = None,
     runs: int = 1,
-    derive_metric: Callable[..., float | dict[str, float]] | None = None,
+    derive_metric: (
+        Callable[
+            ..., float | tuple[float, str] | Mapping[str, float | tuple[float, str]]
+        ]
+        | None
+    ) = None,
     normalize_against: str | None = None,
-    output: Literal["quiet", "progress", "verbose"] = "progress",
+    verbosity: VerbosityLevel = VerbosityLevel.INFO,
     metrics: Sequence[str] = ["gpu__time_duration.sum"],
     ignore_kernel_list: Sequence[str] | None = None,
     clock_control: Literal["base", "none"] = "none",
@@ -68,15 +90,24 @@ def kernel(
     thermal_wait: int | None = None,
     thermal_cont: int | None = None,
     thermal_timeout: int | None = None,
+    thermal_device: int | None = None,
     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)
@@ -93,6 +124,29 @@ 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
+        signature. If the function has ``*args`` or ``**kwargs``, they are tolerated
+        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:
@@ -105,30 +159,39 @@ def kernel(
 
                    configs = [1, 2, 3, 4]
 
+            If configs provide fewer values than the function has parameters,
+            the remaining parameters must have default values, which will be used
+            automatically.
+
             If the configs are not provided at decoration time, they must be provided when calling the decorated function.
         runs:  Number of times each configuration should be executed.
         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:
 
             - A single value (float/int): Will be added as a new metric with the function name as the metric name. For lambda functions, the metric name will be ``"<lambda>"``.
-            - A dictionary: Keys will be used as metric names, values as metric values.
+            - A tuple (value, unit): A single metric value with its unit specified. The value will be added with the function name as the metric name, and the unit will be stored in the ``Unit`` column of the dataframe. Example: ``(tflops_value, "TFLOPS")``.
+            - A dictionary: Keys will be used as metric names, values can be either:
+                - Scalar values (float/int): Metric values without units. ``np.nan`` will be added to the ``Unit`` column and a warning will be issued.
+                - Tuples (value, unit): Metric values with units specified. The unit will be stored in the ``Unit`` column of the dataframe. Example: ``{"TFLOPS": (tflops_value, "TFLOPS"), "ArithIntensity": (intensity_value, "FLOP/byte")}``.
 
             The parameter order requirements for the custom function:
 
             - First several arguments: Must exactly match the order of metrics declared in the @kernel decorator. These arguments will receive the actual measured values of those metrics.
             - Remaining arguments: Must exactly match the signature of the decorated function. In other words, the original function's parameters are passed in order.
 
+            **Note on units**: When units are not provided (returning scalar values instead of tuples), ``np.nan`` will be added to the ``Unit`` column in the dataframe and a warning will be issued. To avoid warnings and ensure proper unit tracking, return tuples ``(value, unit)`` for all derived metrics.
+
             See the examples for concrete use cases.
         normalize_against:
             Annotation name to normalize metrics against (computes current / baseline).
             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``
@@ -138,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.
@@ -152,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.
@@ -159,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
@@ -193,11 +266,22 @@ def kernel(
         thermal_timeout: Maximum wait time in seconds for GPU to cool down.
             Default: 180 seconds
 
-        output: Controls the verbosity level of the output.
+        thermal_device: CUDA device ordinal to monitor for thermal throttling.
+
+            If set, monitor that CUDA device. If ``None``, monitor the current
+            CUDA device context. Device numbering follows CUDA runtime ordinals and
+            honors ``CUDA_VISIBLE_DEVICES``.
+
+            Default: ``None``
 
-            - ``"quiet"``: Suppresses all output.
-            - ``"progress"``: Shows a progress bar along with details about profiling and data extraction progress.
-            - ``"verbose"``: Displays the progress bar, configuration-specific logs, and profiler logs.
+        verbosity: Controls the verbosity level of the output. Selecting level N
+            enables output at all levels with value <= N. Default: ``VerbosityLevel.INFO``.
+
+            - ``VerbosityLevel.SILENT``: Suppresses all output.
+            - ``VerbosityLevel.INFO``: Shows a progress bar, profiling completion messages,
+              and report file paths.
+            - ``VerbosityLevel.DEBUG``: Adds configuration-specific logs, the full NCU CLI
+              command, and enables NCU verbose output.
 
         output_prefix: When specified, all intermediate profiler files are created with this prefix.
             For example, if `output_prefix="/home/user/run1_"`, the profiler will generate:
@@ -221,6 +305,7 @@ def kernel(
                 - ``Annotation``: Name of the annotated region being profiled
                 - ``Value``: Raw metric values collected by the profiler
                 - ``Metric``: The metrics being collected (e.g., ``gpu__time_duration.sum``) 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
                 - ``GPU``: GPU device name
                 - ``Host``: Host machine name
@@ -239,18 +324,79 @@ 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
                 - ``GPU``: GPU device name
                 - ``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:
@@ -263,46 +409,111 @@ def kernel(
 
     def _create_profiler() -> collection.core.NsightProfiler:
         """Helper to create the profiler with the given settings."""
-        if output not in ("quiet", "progress", "verbose"):
-            raise ValueError("output must be 'quiet', 'progress' or 'verbose'")
-
-        output_progress = output == "progress" or output == "verbose"
-        output_detailed = output == "verbose"
-
         # Create the output paths needed for the ncu report, ncu logs and the CSVs
         prefix = output_prefix
-        if "NSPY_NCU_PROFILE" not in os.environ:
-            if prefix is None:
-                prefix = tempfile.mkdtemp(prefix="nspy_")
-                prefix = os.path.join(
-                    prefix, ""
-                )  # Adds a trailing forward/backward slash
-            else:
-                os.makedirs(os.path.dirname(prefix) or ".", exist_ok=True)
+        if prefix is None:
+            prefix = tempfile.mkdtemp(prefix="nspy_")
+            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,
-            output_progress=output_progress,
-            output_detailed=output_detailed,
+            verbosity=verbosity,
             derive_metric=derive_metric,
             normalize_against=normalize_against,
             thermal_mode=thermal_mode,
             thermal_wait=thermal_wait,
             thermal_cont=thermal_cont,
             thermal_timeout=thermal_timeout,
+            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:
@@ -377,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,
@@ -443,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"
 
@@ -451,30 +667,29 @@ def plot(
         def wrapper(*args: Any, **kwargs: Any) -> Any:
             result = func(*args, **kwargs)
 
-            if "NSPY_NCU_PROFILE" not in os.environ:
-                # Legalize the user-provided metric for plotting
-                legalized_metric = _legalize_metric(result, metric)
-
-                visualization.visualize(
-                    result.to_dataframe(),
-                    metric=legalized_metric,
-                    row_panels=row_panels,
-                    col_panels=col_panels,
-                    x_keys=x_keys,
-                    print_data=print_data,
-                    title=title,
-                    filename=filename,
-                    ylabel=ylabel or "",
-                    annotate_points=annotate_points,
-                    show_avg=show_avg,
-                    show_geomean=show_geomean,
-                    plot_type=plot_type,
-                    plot_width=plot_width,
-                    plot_height=plot_height,
-                    variant_fields=variant_fields,
-                    variant_annotations=variant_annotations,
-                    plot_callback=plot_callback,
-                )
+            # Legalize the user-provided metric for plotting
+            legalized_metric = _legalize_metric(result, metric)
+
+            visualization.visualize(
+                result.to_dataframe(),
+                metric=legalized_metric,
+                row_panels=row_panels,
+                col_panels=col_panels,
+                x_keys=x_keys,
+                print_data=print_data,
+                title=title,
+                filename=filename,
+                ylabel=ylabel or "",
+                annotate_points=annotate_points,
+                show_avg=show_avg,
+                show_geomean=show_geomean,
+                plot_type=plot_type,
+                plot_width=plot_width,
+                plot_height=plot_height,
+                variant_fields=variant_fields,
+                variant_annotations=variant_annotations,
+                plot_callback=plot_callback,
+            )
             return result
 
         return wrapper

nsight/annotation.py

@@ -9,6 +9,7 @@ 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 +89,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
@@ -129,6 +130,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 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.
 
     """
 
@@ -140,9 +145,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():
@@ -160,14 +163,17 @@ 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:

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"]