tessl/pypi-nvtx

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

—

Pending

Overview

Eval results

Files

Code Range Annotation

Name: tessl/pypi-nvtx
Author: tessl

Annotate code sections with decorators, context managers, or manual push/pop functions to create hierarchical performance traces in NVIDIA Nsight Systems. Supports custom messages, colors, domains, categories, and payloads for detailed profiling visualization.

Capabilities

Decorator Annotation

Use @nvtx.annotate to automatically annotate entire functions with NVTX ranges. The decorator handles range lifecycle automatically and provides exception safety.

@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():
    """
    Decorator for automatic function annotation.
    
    Parameters:
    - message: Function name used if None
    - color: Color for visualization (default varies by usage)
    - domain: Domain name (default "NVTX")
    - category: Category name or ID for grouping
    - payload: Numeric value associated with annotation
    """

Usage Example:

import nvtx
import time

@nvtx.annotate(color="blue", domain="my_app")
def process_data():
    # Function name "process_data" used as message
    time.sleep(0.1)

@nvtx.annotate("custom_name", color="red", category="processing")
def compute():
    time.sleep(0.2)

Context Manager Annotation

Use nvtx.annotate as a context manager to annotate specific code blocks with automatic range management and exception safety.

with nvtx.annotate(message: str = None, color: Union[str, int] = None,
                   domain: str = None, category: Union[str, int] = None,
                   payload: Union[int, float] = None):
    """
    Context manager for code block annotation.
    
    Parameters:
    - message: Description of annotated code block (default empty string)
    - color: Color for visualization (default varies by usage)
    - domain: Domain name (default "NVTX")
    - category: Category name or ID for grouping
    - payload: Numeric value associated with annotation
    """

Usage Example:

import nvtx
import time

# Annotate initialization phase
with nvtx.annotate("initialization", color="green", domain="startup"):
    config = load_config()
    setup_logging()

# Nested annotations for detailed profiling
with nvtx.annotate("data_processing", color="blue"):
    for i in range(100):
        with nvtx.annotate("process_item", color="yellow", payload=i):
            process_item(data[i])

Manual Range Control

Use push_range and pop_range for explicit control over annotation ranges when decorators or context managers are not suitable.

def push_range(message: str = None, color: Union[str, int] = "blue",
               domain: str = None, category: Union[str, int] = None,
               payload: Union[int, float] = None):
    """
    Mark the beginning of a code range.
    
    Parameters:
    - message: Description of the code range
    - color: Color for visualization (default "blue")
    - domain: Domain name (default "NVTX")
    - category: Category name or ID for grouping
    - payload: Numeric value associated with range
    """

def pop_range(domain: str = None):
    """
    Mark the end of a code range started with push_range.
    
    Parameters:
    - domain: Domain name (default "NVTX", must match push_range call)
    """

Usage Example:

import nvtx
import time

# Manual range control for complex logic
def complex_algorithm():
    nvtx.push_range("setup_phase", color="green")
    setup_data()
    nvtx.pop_range()
    
    for iteration in range(10):
        nvtx.push_range("iteration", color="blue", payload=iteration)
        if should_break():
            nvtx.pop_range()  # Important: always pop before breaking
            break
        process_iteration()
        nvtx.pop_range()
    
    nvtx.push_range("cleanup", color="red")
    cleanup_resources()
    nvtx.pop_range()

Parameter Details

Message Parameter

Type: Optional[str]
Decorator: Defaults to decorated function name
Context Manager: Defaults to empty string
Manual: Defaults to None
Caching: Messages are cached as Registered Strings for performance
Memory: Caching very large numbers of unique messages may increase memory usage

Color Parameter

Type: Optional[Union[str, int]]
Built-in Colors: "green", "blue", "yellow", "purple", "rapids", "cyan", "red", "white", "darkgreen", "orange"
Matplotlib Colors: Any matplotlib color name (requires matplotlib installation)
Hex Colors: Integer values in ARGB format (e.g., 0xFF0000 for red)
Default: Varies by function (typically "blue" for manual functions)

Domain Parameter

Type: Optional[str]
Default: "NVTX"
Purpose: Groups annotations for organization and performance
Performance: Domain-specific API provides better performance for high-frequency annotations

Category Parameter

Type: Optional[Union[str, int]]
String: Automatically assigned numeric ID on first use
Integer: Direct category ID
Purpose: Further sub-grouping within domains
Default: No category association

Payload Parameter

Type: Optional[Union[int, float]]
Purpose: Associate numeric data with annotations for analysis
Types: Supports both integer and floating-point values
Usage: Iteration counts, sizes, timing values, etc.

Error Handling

The annotation functions provide robust error handling:

Exception Safety: Decorator and context manager automatically clean up ranges even when exceptions occur
NVTX Disabled: All functions become no-ops when NVTX_DISABLE environment variable is set
Invalid Colors: Raises TypeError for unsupported colors when matplotlib is not available
Domain Mismatch: pop_range domain parameter must match corresponding push_range call

Performance Considerations

Message Caching: Repeated messages are cached as Registered Strings for efficiency
Domain API: Use domain-specific API (Domain.push_range, Domain.pop_range) for better performance in high-frequency scenarios
Overhead: Zero overhead when NVTX is disabled via environment variable
Nesting: Deep annotation nesting is supported but may impact profiler visualization