Skip to main content

Python module

profiler

Performance profiling and tracing utilities for MAX.

This module provides tools for profiling and tracing MAX operations to analyze performance characteristics. Profiling captures timing information for code execution, which helps identify bottlenecks and optimize your models.

To enable profiling, set the MODULAR_ENABLE_PROFILING=1 environment variable before running your code. Without this variable, profiling calls will be no-ops with minimal overhead.

The profiler supports three usage patterns:

  1. Context manager: Use Tracer as a context manager to profile a code block.
  2. Decorator: Use @traced to profile entire functions.
  3. Manual stack: Use Tracer methods to explicitly control profiling spans.

Tracer

class max.profiler.Tracer(message=None, color='modular_purple')

A stack-based profiling manager for creating nested profiling spans.

Manages a stack of profiling spans that allows for nested tracing without requiring deeply nested with Trace(name): statements. This is especially useful when you need to dynamically create and manage profiling spans based on runtime conditions or when profiling spans don’t align with your code’s block structure.

The Tracer can be used both as a context manager and as a manual stack manager. As a context manager, it ensures all pushed spans are properly closed when the context exits.

from max.profiler import Tracer

tracer = Tracer("parent_operation", color="modular_purple")
tracer.push("child_operation")
# ... perform work ...
tracer.pop()

# Context manager with manual stack
with Tracer("parent_operation", color="modular_purple") as tracer:
    # The parent span is named "parent_operation"
    tracer.push("child_operation")
    # ... perform work ...
    tracer.pop()
    # All spans are automatically closed on context exit

Initializes the tracer stack.

Creates an empty trace stack and optionally pushes an initial profiling span if a message is provided.

Parameters:

  • message (str | None) – The name of the initial profiling span. If None, no initial span is created.
  • color (str) – The color of the profiling span for visualization tools.

cleanup()

cleanup()

Closes all remaining profiling spans.

Pops and closes all profiling spans that were pushed onto the stack. This method is automatically called when the tracer is used as a context manager or when the object is deleted.

Return type:

None

mark()

mark()

Marks the current profiling span with a timestamp.

Records a timestamp event within the current profiling span. This is useful for marking significant events or milestones within a longer operation.

Raises:

AssertionError – If the stack is empty when mark is called.

Return type:

None

next()

next(message, color='modular_purple')

Transitions to the next profiling span.

Pops the current profiling span and immediately pushes a new one with the specified message. This is a convenience method for sequential operations at the same nesting level.

Parameters:

  • message (str) – The name of the new profiling span.
  • color (str) – The color of the profiling span for visualization tools.

Return type:

None

pop()

pop(exc_type=None, exc_value=None, traceback=None)

Pops a profiling span off the stack and closes it.

Removes the most recently pushed profiling span from the stack and closes it, recording its execution time. Exception information can be passed through for proper error handling in context managers.

Parameters:

  • exc_type (type[BaseException] | None) – The exception type if an exception occurred, or None.
  • exc_value (BaseException | None) – The exception instance if an exception occurred, or None.
  • traceback (TracebackType | None) – The traceback object if an exception occurred, or None.

Return type:

None

push()

push(message=None, color='modular_purple')

Pushes a new profiling span onto the stack.

Creates and activates a new profiling span. If profiling is disabled or no message is provided, pushes a None placeholder to maintain stack consistency.

Parameters:

  • message (str | None) – The name of the profiling span. If None, no span is created.
  • color (str) – The color of the profiling span for visualization tools.

Return type:

None

traced()

max.profiler.traced(func=None, *, message=None, color='modular_purple')

Decorator for creating a profiling span for a function.

Creates a profiling span that measures the execution time of the decorated function. This is useful for identifying performance bottlenecks without modifying the function’s internal code. The decorator supports both synchronous and asynchronous functions.

from max.profiler import traced

# Decorator with custom span name
@traced(message="inference", color="red")
def run_model() -> None:
    # The profiling span is named "inference"
    model.execute()

# Decorator with default span name (uses function name)
@traced
def preprocess_data() -> None:
    # The profiling span is named "preprocess_data"
    data.normalize()

Parameters:

  • func (_FuncType | None) – The function to profile.
  • message (str | None) – The name of the profiling span. If None, uses the function name.
  • color (str) – The color of the profiling span for visualization tools.

Returns:

The decorated function wrapped in a trace object.

Return type:

Callable

Was this page helpful?