tessl/pypi-nvtx

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

—

Pending

Overview

Eval results

Files

Domain Management

Name: tessl/pypi-nvtx
Author: tessl

Create and manage custom domains for organizing annotations and achieving better performance through domain-specific API usage. Domains provide namespace separation and optimized APIs for high-frequency annotation scenarios.

Capabilities

Domain Creation and Access

Use nvtx.get_domain to create or retrieve domain objects for organized and high-performance annotation.

def get_domain(name: str = None) -> Union[Domain, DummyDomain]:
    """
    Get or create a Domain object for a domain name.
    
    Parameters:
    - name: Domain name (default creates default "NVTX" domain)
    
    Returns:
    - Domain object if NVTX is enabled
    - DummyDomain object if NVTX is disabled (no-op operations)
    """

Usage Example:

import nvtx

# Create custom domains for different components
ui_domain = nvtx.get_domain("UI")
network_domain = nvtx.get_domain("Network") 
compute_domain = nvtx.get_domain("Compute")

# Default domain
default_domain = nvtx.get_domain()  # or nvtx.get_domain(None)

Domain-Specific API

Domain objects provide optimized methods for high-frequency annotation scenarios with reduced overhead compared to global functions.

class Domain:
    def get_registered_string(self, string: str) -> RegisteredString:
        """
        Register a string for efficient reuse within this domain.
        
        Parameters:
        - string: String to register
        
        Returns:
        - RegisteredString object for efficient reuse
        """
    
    def get_category_id(self, name: str) -> int:
        """
        Get or create a category ID for the given name.
        
        Parameters:
        - name: Category name
        
        Returns:
        - Numeric category ID for use in annotations
        """
    
    def get_event_attributes(self, message: str = None, 
                           color: Union[str, int] = None,
                           category: Union[str, int] = None, 
                           payload: Union[int, float] = None) -> EventAttributes:
        """
        Create an EventAttributes object for this domain.
        
        Parameters:
        - message: Event message (auto-registered if string)
        - color: Color specification
        - category: Category name (auto-converted to ID) or ID
        - payload: Numeric payload value
        
        Returns:
        - EventAttributes object for use with domain methods
        """
    
    def mark(self, attributes: EventAttributes):
        """
        Mark an instantaneous event using pre-created attributes.
        
        Parameters:
        - attributes: EventAttributes from get_event_attributes()
        """
    
    def push_range(self, attributes: EventAttributes):
        """
        Start a code range using pre-created attributes.
        
        Parameters:
        - attributes: EventAttributes from get_event_attributes()
        """
    
    def pop_range(self):
        """
        End the most recent code range started with push_range().
        """
    
    def start_range(self, attributes: EventAttributes) -> int:
        """
        Start a process range using pre-created attributes.
        
        Parameters:
        - attributes: EventAttributes from get_event_attributes()
        
        Returns:
        - Range ID for use with end_range()
        """
    
    def end_range(self, range_id: int):
        """
        End a process range started with start_range().
        
        Parameters:
        - range_id: Range ID returned by start_range()
        """

Usage Example:

import nvtx

# Create domain for high-frequency annotations
compute_domain = nvtx.get_domain("Compute")

# Pre-create attributes for reuse
loop_attrs = compute_domain.get_event_attributes("inner_loop", color="blue")
checkpoint_attrs = compute_domain.get_event_attributes("checkpoint", color="green")

def high_frequency_function():
    # High-performance annotation using domain API
    compute_domain.push_range(loop_attrs)
    
    for i in range(1000000):
        if i % 100000 == 0:
            compute_domain.mark(checkpoint_attrs)
        # ... computation ...
    
    compute_domain.pop_range()

String and Category Management

Domains provide efficient management of strings and categories through registration and caching.

class RegisteredString:
    """
    Registered string handle for efficient reuse within a domain.
    
    Attributes:
    - string: The original string value
    - domain: Domain object this string is registered with
    - handle: Low-level string handle for C API
    """
    string: str
    domain: Domain
    handle: StringHandle

Usage Example:

import nvtx

# Create domain and register frequently used strings
domain = nvtx.get_domain("MyApp")

# Register strings for efficient reuse
func_name = domain.get_registered_string("process_data")
loop_name = domain.get_registered_string("inner_loop")

# Register categories
data_category = domain.get_category_id("DataProcessing")
ui_category = domain.get_category_id("UserInterface")

# Use registered strings and categories efficiently
attrs = domain.get_event_attributes(
    message=func_name.string,
    color="blue", 
    category=data_category
)
domain.push_range(attrs)
# ... work ...
domain.pop_range()

Disabled Domain Handling

When NVTX is disabled, get_domain returns a DummyDomain that provides no-op implementations of all methods.

class DummyDomain:
    """
    No-op domain replacement when NVTX is disabled.
    All methods are safe to call but perform no operations.
    """
    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): ...

Performance Optimization Patterns

Attribute Pre-creation

For maximum performance in high-frequency scenarios, pre-create EventAttributes objects:

import nvtx

# Setup phase
domain = nvtx.get_domain("HighFreq")
attrs = domain.get_event_attributes("hot_path", color="red")

# Hot path - minimal overhead
def hot_function():
    domain.push_range(attrs)
    # ... critical code ...
    domain.pop_range()

Batch Registration

import nvtx

class PerformanceCritical:
    def __init__(self):
        self.domain = nvtx.get_domain("Critical")
        
        # Pre-register all strings and categories
        self.strings = {
            'init': self.domain.get_registered_string("initialization"),
            'compute': self.domain.get_registered_string("computation"),
            'cleanup': self.domain.get_registered_string("cleanup")
        }
        
        self.categories = {
            'setup': self.domain.get_category_id("Setup"),
            'work': self.domain.get_category_id("Work"),
            'teardown': self.domain.get_category_id("Teardown")
        }
        
        # Pre-create attributes
        self.attrs = {
            'init': self.domain.get_event_attributes(
                message=self.strings['init'].string,
                category=self.categories['setup']
            ),
            'compute': self.domain.get_event_attributes(
                message=self.strings['compute'].string,
                category=self.categories['work']
            )
        }
    
    def run(self):
        # Maximum performance annotation
        self.domain.push_range(self.attrs['init'])
        self.initialize()
        self.domain.pop_range()
        
        self.domain.push_range(self.attrs['compute'])
        self.compute()
        self.domain.pop_range()

Domain Organization Strategies

Component-Based Domains

Organize domains by major application components:

import nvtx

# Create domains for different components
ui_domain = nvtx.get_domain("UI")
database_domain = nvtx.get_domain("Database")
network_domain = nvtx.get_domain("Network")
compute_domain = nvtx.get_domain("Compute")

def handle_request():
    # UI layer
    ui_attrs = ui_domain.get_event_attributes("request_handling")
    ui_domain.push_range(ui_attrs)
    
    # Database layer
    db_attrs = database_domain.get_event_attributes("query_execution")
    database_domain.push_range(db_attrs)
    execute_query()
    database_domain.pop_range()
    
    ui_domain.pop_range()

Layer-Based Domains

Organize domains by architectural layers:

import nvtx

# Architectural layer domains
presentation_domain = nvtx.get_domain("Presentation")
business_domain = nvtx.get_domain("Business")
data_domain = nvtx.get_domain("Data")

Thread-Based Domains

Use domains to track work across different threads:

import nvtx
import threading

# Thread-specific domains
main_domain = nvtx.get_domain("MainThread")
worker_domain = nvtx.get_domain("WorkerThread")
io_domain = nvtx.get_domain("IOThread")

Error Handling and Edge Cases

Disabled NVTX: All domain methods become no-ops when NVTX_DISABLE is set
Invalid Attributes: Passing None attributes to domain methods is safe (no-op)
Thread Safety: Domain objects and their methods are thread-safe
Memory Management: Registered strings and categories are cached efficiently
Domain Reuse: Multiple calls to get_domain with same name return the same object