stormlog.context_profiler

Context profiler for easy function and code block profiling.

Functions

`clear_results`()	Clear global profiler results.
`get_global_profiler`([device])	Get or create the global profiler instance.
`get_profile_results`([limit])	Return recent profile results captured by the global profiler.
`get_summary`()	Get global profiler summary.
`profile_context`([name, device, profiler])	Context manager for profiling a block of code.
`profile_function`([func, name, device, profiler])	Decorator to profile a function's GPU memory usage.
`profile_model_training`(model, train_loader)	Profile an entire training loop.
`set_global_profiler`(profiler)	Set the global profiler instance.
`start_monitoring`([interval, device])	Start global memory monitoring.
`stop_monitoring`()	Stop global memory monitoring.

Classes

`MemoryProfiler`([device])	High-level memory profiler with convenient methods.
`ProfiledModule`(args, *kwargs)	Wrapper for PyTorch modules that automatically profiles forward passes.

stormlog.context_profiler.get_global_profiler(device=None)[source]

Get or create the global profiler instance.

Parameters:: device (str | int | torch.device | None)
Return type:: GPUMemoryProfiler

stormlog.context_profiler.set_global_profiler(profiler)[source]

Set the global profiler instance.

Parameters:: profiler (GPUMemoryProfiler)
Return type:: None

stormlog.context_profiler.profile_function(func=None, *, name=None, device=None, profiler=None)[source]

Decorator to profile a function’s GPU memory usage.

Can be used as @profile_function or @profile_function(name=”custom_name”)

Parameters:

func (F | None) – Function to profile (when used as @profile_function)
name (str | None) – Custom name for the profiled function
device (str | int | torch.device | None) – GPU device to use for profiling
profiler (GPUMemoryProfiler | None) – Custom profiler instance to use

Returns:

Decorated function or ProfileResult if called directly

Return type:

Callable[[F], F] | F

stormlog.context_profiler.profile_context(name='context', device=None, profiler=None)[source]

Context manager for profiling a block of code.

Parameters:

name (str) – Name for the profiled context
device (str | int | torch.device | None) – GPU device to use for profiling
profiler (GPUMemoryProfiler | None) – Custom profiler instance to use

Yields:

ProfileResult after the context exits

Return type:

Iterator[GPUMemoryProfiler]

Example

with profile_context(“model_forward”) as prof:: output = model(input)

class stormlog.context_profiler.ProfiledModule(*args, **kwargs)[source]

Bases: Module

Wrapper for PyTorch modules that automatically profiles forward passes.

Example

model = ProfiledModule(original_model, name=”my_model”) output = model(input) # Automatically profiled

Parameters:

module (torch.nn.Module)
name (str | None)
device (str | int | torch.device | None)
profiler (GPUMemoryProfiler | None)

forward(*args, **kwargs)[source]

Forward pass with automatic profiling.

Parameters:

args (Any)
kwargs (Any)

Return type:

Any

class stormlog.context_profiler.MemoryProfiler(device=None)[source]

Bases: object

High-level memory profiler with convenient methods.

This class provides a simplified interface for common profiling tasks.

Parameters:: device (str | int | torch.device | None)

start_monitoring(interval=0.1)[source]

Start continuous memory monitoring.

Parameters:: interval (float)
Return type:: None

stop_monitoring()[source]

Stop continuous memory monitoring.

Return type:: None

profile(func, *args, **kwargs)[source]

Profile a function call.

Parameters:

func (Callable[[...], Any])
args (Any)
kwargs (Any)

Return type:

ProfileResult

context(name='context')[source]

Context manager for profiling code blocks.

Parameters:: name (str)
Return type:: Iterator[None]

wrap_module(module, name=None)[source]

Wrap a PyTorch module for automatic profiling.

Parameters:

module (torch.nn.Module)
name (str | None)

Return type:

ProfiledModule

get_summary()[source]

Get profiling summary.

Return type:: Any

clear()[source]

Clear profiling results.

Return type:: None

save_results(filename)[source]

Save profiling results to file.

Parameters:: filename (str)
Return type:: None

load_results(filename)[source]

Load profiling results from file.

Parameters:: filename (str)
Return type:: Any

stormlog.context_profiler.start_monitoring(interval=0.1, device=None)[source]

Start global memory monitoring.

Parameters:

interval (float)
device (str | int | torch.device | None)

Return type:

None

stormlog.context_profiler.stop_monitoring()[source]

Stop global memory monitoring.

Return type:: None

stormlog.context_profiler.get_summary()[source]

Get global profiler summary.

Return type:: Any

stormlog.context_profiler.clear_results()[source]

Clear global profiler results.

Return type:: None

stormlog.context_profiler.get_profile_results(limit=None)[source]

Return recent profile results captured by the global profiler.

Parameters:: limit (int | None)
Return type:: List[ProfileResult]

stormlog.context_profiler.profile_model_training(model, train_loader, epochs=1, device=None)[source]

Profile an entire training loop.

Parameters:

model (torch.nn.Module) – PyTorch model to train
train_loader (Any) – DataLoader for training data
epochs (int) – Number of epochs to profile
device (str | int | torch.device | None) – GPU device to use

Returns:

Dictionary with profiling results

Return type:

Dict[str, Any]