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
# ValueError

# 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

-

Summary

• with statement = Guaranteed resource cleanup
• __enter__ / __exit__ = Protocol for context managers
• @contextmanager = Decorator to write context managers
• Chaining = Multiple contexts in one with statement
• Exception handling = Cleanup even when errors occur
• Practical uses: GPU memory, model modes, checkpoints, profiling

-

Related Topics

• [01 Reference Counting & Garbage Collection](/05-py3/05-memory-&-performance/(01-reference-counting-garbage-collection/) - Resource management internals
• [04 Profiling & Performance Analysis](/05-py3/09-bytecode-and-execution/(04-profiling-performance-analysis/) - Timing and profiling contexts
• 00 Readme - Training loops with context management