Context Managers & Resource Management¶
Overview¶
Context managers are essential for ML infrastructure: - GPU memory allocation/deallocation - File handling (datasets, checkpoints) - Distributed training setup/teardown - Automatic differentiation context - Profiling sessions
The with statement ensures resources are properly released even if errors occur.
The Context Manager Protocol¶
Core Concept¶
A context manager implements two methods:
class ContextManager:
def __enter__(self):
"""Called when entering 'with' block."""
# Setup: acquire resources
return self # or any resource object
def __exit__(self, exc_type, exc_val, exc_tb):
"""Called when exiting 'with' block."""
# Teardown: release resources
# Called even if exception occurs!
return False # Propagate exception
Simple Example¶
# File handling uses context managers
with open("data.txt") as f: # Calls __enter__
data = f.read()
# Calls __exit__ (file closed automatically)
# Equivalent to:
f = open("data.txt")
try:
data = f.read()
finally:
f.close() # Guaranteed to run
PyTorch Context Managers¶
GPU/Device Context¶
import torch
# Allocate tensor on CPU
x = torch.randn(3, 224, 224)
# Use GPU temporarily without permanently moving tensors
if torch.cuda.is_available():
with torch.cuda.device(0): # Use GPU 0
x_gpu = x.to("cuda") # Move to GPU 0
y = x_gpu * 2 # Compute on GPU 0
# GPU 0 memory is used here
# Back to CPU context
# Or specify device in tensor creation
with torch.cuda.device(0):
a = torch.randn(3, 3, device="cuda") # Created on GPU 0
No-Grad Context (Inference)¶
model = torch.nn.Linear(10, 5)
# During inference, don't compute gradients
with torch.no_grad(): # Disable autograd
predictions = model(input_data)
# Faster, uses less memory (no gradient history)
# Gradient is still computed here (if model.train())
loss = model(input_data).mean()
loss.backward()
Why this matters: Disabling gradients saves ~2x memory and 2-3x speed during inference.
Gradient Context Variations¶
# Disable gradient computation (inference)
with torch.no_grad():
output = model(x)
# Enable gradient computation (even if model.eval())
with torch.enable_grad():
output = model(x)
# Set gradient computation explicitly
with torch.set_grad_enabled(False):
output = model(x)
Autograd Context¶
Backward Context¶
x = torch.randn(3, requires_grad=True)
# By default, autograd tracks operations
y = x * 2
z = y.sum()
z.backward() # Computes gradients
print(x.grad) # [2, 2, 2]
Custom Autograd Functions¶
class CustomFunction(torch.autograd.Function):
@staticmethod
def forward(ctx, input):
ctx.save_for_backward(input)
return input * 2
@staticmethod
def backward(ctx, grad_output):
input, = ctx.saved_tensors
return grad_output * 2
# Use in forward pass (context managers enable this)
def forward(x):
return CustomFunction.apply(x)
Writing Custom Context Managers¶
Using Class Syntax¶
class DeviceContext:
"""Move tensors to device temporarily."""
def __init__(self, device):
self.device = device
self.original_devices = {}
def __enter__(self):
# Save original devices and move
return self
def __exit__(self, exc_type, exc_val, exc_tb):
# Restore original devices
pass
# Usage
with DeviceContext("cuda"):
model.to("cuda")
output = model(input_data)
Using Decorator Syntax (contextlib)¶
from contextlib import contextmanager
@contextmanager
def timer(name):
"""Context manager for timing code blocks."""
import time
start = time.time()
print(f"[{name}] Starting...")
try:
yield # Code in 'with' block runs here
finally:
elapsed = time.time() - start
print(f"[{name}] Done in {elapsed:.3f}s")
# Usage
with timer("Model inference"):
output = model(input_data)
# Output:
# [Model inference] Starting...
# [Model inference] Done in 0.123s
Practical Example: GPU Memory Context¶
@contextmanager
def gpu_memory_context(device_id):
"""Context that clears GPU memory before and after."""
import torch
# Setup: Clear cache
torch.cuda.empty_cache()
initial_memory = torch.cuda.memory_allocated(device_id)
try:
yield
finally:
# Teardown: Report memory usage and clear
final_memory = torch.cuda.memory_allocated(device_id)
peak_memory = torch.cuda.max_memory_allocated(device_id)
print(f"Memory used: {(final_memory - initial_memory) / 1e9:.2f}GB")
print(f"Peak memory: {peak_memory / 1e9:.2f}GB")
torch.cuda.reset_peak_memory_stats(device_id)
torch.cuda.empty_cache()
# Usage
with gpu_memory_context(0):
model = torch.nn.Linear(10000, 10000).to("cuda")
output = model(torch.randn(100, 10000, device="cuda"))
Practical ML Patterns¶
Pattern 1: Training/Evaluation Mode¶
# Custom context for model mode switching
@contextmanager
def eval_mode(model):
"""Temporarily set model to eval mode."""
was_training = model.training
model.eval()
try:
with torch.no_grad(): # Chain contexts
yield
finally:
model.train(was_training) # Restore original mode
# Usage
model = torch.nn.Linear(10, 5)
# Training
model.train()
output = model(input_data)
loss = output.mean()
loss.backward() # Gradients computed
# Validation (no grad, eval mode)
with eval_mode(model):
predictions = model(val_data) # No gradients, dropout disabled
Pattern 2: Distributed Training Context¶
import torch.distributed as dist
@contextmanager
def ddp_context(rank, world_size):
"""Setup and teardown distributed training."""
# Setup: Initialize process group
os.environ['MASTER_ADDR'] = 'localhost'
os.environ['MASTER_PORT'] = '12355'
dist.init_process_group(
backend='nccl',
rank=rank,
world_size=world_size
)
try:
yield
finally:
# Teardown: Cleanup distributed training
dist.destroy_process_group()
# Usage
def run_ddp(rank, world_size):
with ddp_context(rank, world_size):
model = torch.nn.Linear(10, 5)
model = torch.nn.parallel.DistributedDataParallel(model)
# Train model...
Pattern 3: Checkpoint Context¶
@contextmanager
def checkpoint_context(model, save_path, freq=1):
"""Save checkpoints periodically."""
import json
from datetime import datetime
checkpoint_dir = Path(save_path)
checkpoint_dir.mkdir(exist_ok=True)
metrics = []
epoch = 0
try:
yield metrics
finally:
# Save final checkpoint
checkpoint = {
"state_dict": model.state_dict(),
"metrics": metrics,
"timestamp": datetime.now().isoformat(),
}
torch.save(checkpoint, checkpoint_dir / "final.pt")
print(f"Checkpoint saved to {checkpoint_dir}")
# Usage
model = torch.nn.Linear(10, 5)
with checkpoint_context(model, "./checkpoints") as metrics:
for epoch in range(10):
loss = train_epoch(model)
metrics.append({"epoch": epoch, "loss": float(loss)})
print(f"Epoch {epoch}: loss={loss:.4f}")
Nested Context Managers¶
Multiple Contexts¶
# Nest contexts explicitly
with open("log.txt", "w") as f:
with torch.no_grad():
with torch.cuda.device(0):
output = model(input_data)
f.write(f"Output shape: {output.shape}\n")
# Or use comma syntax (Python 3.1+)
with open("log.txt", "w") as f, torch.no_grad(), torch.cuda.device(0):
output = model(input_data)
f.write(f"Output shape: {output.shape}\n")
Context Manager Stack¶
from contextlib import ExitStack
# Dynamically manage multiple resources
with ExitStack() as stack:
# Open multiple files conditionally
if config.save_output:
output_file = stack.enter_context(open("output.txt", "w"))
if config.save_log:
log_file = stack.enter_context(open("log.txt", "w"))
# Use contexts
with torch.no_grad():
output = model(input_data)
if config.save_output:
output_file.write(str(output))
Error Handling in Context Managers¶
Exception Propagation¶
class ErrorHandlingContext:
def __enter__(self):
print("Setting up...")
return self
def __exit__(self, exc_type, exc_val, exc_tb):
print(f"Cleaning up... (Exception: {exc_type})")
# Return False to propagate exception (default)
# Return True to suppress exception
return False
# Exception propagates
with ErrorHandlingContext():
raise ValueError("Something went wrong!")
# Output:
# Setting up...
# Cleaning up... (Exception: <class 'ValueError'>)
# ValueError: Something went wrong!
# Suppress specific exceptions
class SuppressionContext:
def __exit__(self, exc_type, exc_val, exc_tb):
# Suppress only ValueError
if exc_type is ValueError:
print("Suppressed ValueError")
return True
return False
with SuppressionContext():
raise ValueError("This won't propagate")
print("Execution continues!")
Practical: Robust Inference¶
@contextmanager
def robust_inference(model, timeout=30):
"""Inference with timeout and error handling."""
import signal
def timeout_handler(signum, frame):
raise TimeoutError("Inference timeout")
# Set timeout
signal.signal(signal.SIGALRM, timeout_handler)
signal.alarm(timeout)
try:
yield
except TimeoutError as e:
print(f"Inference timed out: {e}")
raise
finally:
signal.alarm(0) # Cancel alarm
# Usage
try:
with robust_inference(model, timeout=30):
output = model(input_data)
except TimeoutError:
print("Fallback: Use cached prediction")
Performance: Context Manager Overhead¶
import timeit
# No context manager
def simple_op():
x = torch.randn(1000, 1000)
return x @ x
# With context manager
def with_context_op():
with torch.no_grad():
x = torch.randn(1000, 1000)
return x @ x
# Benchmark
simple_time = timeit.timeit(simple_op, number=100)
context_time = timeit.timeit(with_context_op, number=100)
print(f"Simple: {simple_time:.3f}s")
print(f"With context: {context_time:.3f}s")
# Result: Context overhead is negligible (<1%)
Summary¶
withstatement = Guaranteed resource cleanup__enter__/__exit__= Protocol for context managers@contextmanager= Decorator to write context managers- Chaining = Multiple contexts in one
withstatement - Exception handling = Cleanup even when errors occur
- Practical uses: GPU memory, model modes, checkpoints, profiling
Related Topics¶
- 01 Reference Counting & Garbage Collection - Resource management internals
- 04 Profiling & Performance Analysis - Timing and profiling contexts
- 00 Readme - Training loops with context management