or run

npx @tessl/cli init
Log in

Version

Tile

Overview

Evals

Files

docs

annotation.mddomains.mdevents-ranges.mdindex.mdprofiling.md
tile.json

tessl/pypi-nvtx

Python code annotation library for NVIDIA Tools Extension enabling performance profiling with NVIDIA Nsight Systems

Workspace
tessl
Visibility
Public
Created
Last updated
Describes
pypipkg:pypi/nvtx@0.2.x

To install, run

npx @tessl/cli install tessl/pypi-nvtx@0.2.0

index.mddocs/

NVTX

Python code annotation library for NVIDIA Tools Extension that enables developers to add custom annotations, ranges, and markers to their Python code for profiling and performance analysis with NVIDIA Nsight Systems. The library provides decorators, context managers, and manual functions for easy integration with NVIDIA's profiling ecosystem.

Package Information

  • Package Name: nvtx
  • Language: Python
  • Installation: pip install nvtx or conda install -c conda-forge nvtx

Core Imports

import nvtx

For direct access to specific functions:

from nvtx import annotate, enabled, mark, push_range, pop_range, start_range, end_range, get_domain, Domain, Profile

For type annotations:

from typing import Union, Optional, Tuple, Any

Basic Usage

import nvtx
import time

# Using as a decorator
@nvtx.annotate(color="blue")
def my_function():
    for i in range(5):
        with nvtx.annotate("my_loop", color="red"):
            time.sleep(0.1)

# Using context manager
with nvtx.annotate("initialization", color="green"):
    data = initialize_data()

# Manual range control
nvtx.push_range("processing", color="yellow")
process_data(data)
nvtx.pop_range()

# Instantaneous markers
nvtx.mark("checkpoint", color="purple")

my_function()

Architecture

NVTX provides multiple annotation approaches optimized for different use cases:

  • High-level API: Global functions (nvtx.annotate, nvtx.mark) for simple usage with automatic domain management
  • Domain-specific API: Domain objects for better performance and organization when creating many annotations
  • Automatic profiling: Profile-based annotation of all function calls with configurable detail levels
  • Cython bindings: Low-level C extension interface for optimal performance in performance-critical code

The library automatically handles NVTX_DISABLE environment variable and gracefully falls back to no-op implementations for zero overhead when profiling is disabled.

Capabilities

Code Range Annotation

Annotate code sections with decorators, context managers, or manual push/pop functions. Supports custom messages, colors, domains, categories, and payloads for detailed profiling visualization.

@nvtx.annotate(message: str = None, color: Union[str, int] = None, 
               domain: str = None, category: Union[str, int] = None, 
               payload: Union[int, float] = None)
def function(): ...

with nvtx.annotate(message: str = None, color: Union[str, int] = None,
                   domain: str = None, category: Union[str, int] = None,
                   payload: Union[int, float] = None): ...

def push_range(message: str = None, color: Union[str, int] = "blue",
               domain: str = None, category: Union[str, int] = None,
               payload: Union[int, float] = None): ...

def pop_range(domain: str = None): ...

Code Range Annotation

Instantaneous Events and Process Ranges

Mark specific points in execution and create cross-process or cross-thread ranges that can span multiple function calls.

def mark(message: str = None, color: Union[str, int] = "blue",
         domain: str = None, category: Union[str, int] = None,
         payload: Union[int, float] = None): ...

def start_range(message: str = None, color: Union[str, int] = None,
                domain: str = None, category: Union[str, int] = None,
                payload: Union[int, float] = None) -> Optional[Tuple[int, int]]: ...

def end_range(range_id: Optional[Tuple[int, int]]): ...

Events and Process Ranges

Domain Management

Create and manage custom domains for organizing annotations and achieving better performance through domain-specific API usage.

def get_domain(name: str = None) -> Union[Domain, DummyDomain]: ...

class Domain:
    def get_registered_string(self, string: str) -> RegisteredString: ...
    def get_category_id(self, name: str) -> int: ...
    def get_event_attributes(self, message: str = None, color: Union[str, int] = None,
                           category: Union[str, int] = None, 
                           payload: Union[int, float] = None) -> EventAttributes: ...
    def mark(self, attributes: EventAttributes): ...
    def push_range(self, attributes: EventAttributes): ...
    def pop_range(self): ...
    def start_range(self, attributes: EventAttributes) -> int: ...
    def end_range(self, range_id: int): ...

Domain Management

Automatic Profiling

Automatically annotate all function calls in your program with configurable detail levels, including line numbers and C function support.

class Profile:
    def __init__(self, linenos: bool = True, annotate_cfuncs: bool = True): ...
    def enable(self): ...
    def disable(self): ...

Automatic Profiling

Utility Functions

Status and Color Support

def enabled() -> bool: ...

The library includes built-in color support for common colors and matplotlib integration:

  • Built-in colors: "green", "blue", "yellow", "purple", "rapids", "cyan", "red", "white", "darkgreen", "orange"
  • Matplotlib colors: Any matplotlib color name when matplotlib is installed
  • Hex colors: Direct integer color values in ARGB format (e.g., 0xFF0000FF for blue, 0xFFFF0000 for red)

Types

class EventAttributes:
    """Event attributes for NVTX annotations."""
    message: RegisteredString
    color: Union[str, int]
    category: int
    payload: Union[int, float]

class RegisteredString:
    """Registered string handle for efficient reuse."""
    string: str
    domain: Domain
    handle: StringHandle

class StringHandle:
    """Low-level string handle for C API interaction."""
    c_obj: Any  # C object handle

class DummyDomain:
    """No-op domain replacement when NVTX is disabled."""
    def get_registered_string(self, string: str): ...
    def get_category_id(self, name: str): ...
    def get_event_attributes(self, message: str = None, color: Union[str, int] = None,
                           category: Union[str, int] = None, 
                           payload: Union[int, float] = None): ...
    def mark(self, attributes: EventAttributes): ...
    def push_range(self, attributes: EventAttributes): ...
    def pop_range(self): ...
    def start_range(self, attributes: EventAttributes) -> int: ...
    def end_range(self, range_id: int): ...