Context Managers

Why Context Managers Are Essential for Scientific Code

Scientific computing involves managing resources that must be properly acquired and released:

  • Files: HDF5 datasets, CSV logs, checkpoint files
  • GPU memory: CUDA contexts, tensor allocations
  • Database connections: experiment tracking databases
  • Timing blocks: profiling sections of code
  • Random state: reproducible RNG seeds for experiments

The with statement guarantees cleanup even if exceptions occur. Without it, resource leaks accumulate silently during long-running simulations, eventually crashing your 12-hour training run.

Definition:

Context Manager Protocol

A context manager is any object that implements the context manager protocol — two dunder methods:

  • __enter__(self) — called when entering the with block; its return value is bound to the as variable
  • __exit__(self, exc_type, exc_val, exc_tb) — called when leaving the with block (whether normally or via exception)
class Timer:
    def __enter__(self):
        self.start = time.perf_counter()
        return self   # bound to the `as` variable

    def __exit__(self, exc_type, exc_val, exc_tb):
        self.elapsed = time.perf_counter() - self.start
        print(f"Elapsed: {self.elapsed:.4f}s")
        return False  # do not suppress exceptions

with Timer() as t:
    result = expensive_computation()
print(f"Stored: {t.elapsed:.4f}s")

If __exit__ returns True, the exception is suppressed. Returning False (or None) lets the exception propagate normally.

Definition:

contextlib.contextmanager — Generator-Based Context Managers

The contextlib.contextmanager decorator lets you write context managers using a generator function with a single yield:

from contextlib import contextmanager

@contextmanager
def timer(label: str = "Block"):
    start = time.perf_counter()
    try:
        yield start   # value bound to `as` variable
    finally:
        elapsed = time.perf_counter() - start
        print(f"{label}: {elapsed:.4f}s")

with timer("SVD") as t0:
    U, S, Vh = np.linalg.svd(large_matrix)

Everything before yield is the __enter__ logic. Everything after yield (in the finally block) is the __exit__ logic. The try/finally ensures cleanup even if an exception occurs.

Definition:

contextlib.suppress — Ignoring Specific Exceptions

contextlib.suppress is a context manager that silently catches specified exception types:

from contextlib import suppress
import os

# Instead of:
try:
    os.remove("temp_checkpoint.pt")
except FileNotFoundError:
    pass

# Use:
with suppress(FileNotFoundError):
    os.remove("temp_checkpoint.pt")

Use sparingly — silencing exceptions can hide bugs. Only suppress exceptions you genuinely expect and want to ignore.

Example: Timing Context Manager for Profiling

Build a reusable timing context manager that collects timing data for different code sections, supports nested blocks, and can generate a profiling report.

Solution
Implementation with statistics
import time
from collections import defaultdict
from contextlib import contextmanager

class Profiler:
    """Lightweight profiler using context managers."""

    def __init__(self):
        self.timings = defaultdict(list)
        self._stack = []

    @contextmanager
    def section(self, name: str):
        """Time a code section."""
        self._stack.append(name)
        full_name = " > ".join(self._stack)
        start = time.perf_counter()
        try:
            yield
        finally:
            elapsed = time.perf_counter() - start
            self.timings[full_name].append(elapsed)
            self._stack.pop()

    def report(self) -> str:
        """Generate a profiling report."""
        lines = ["Profiling Report", "=" * 50]
        for name, times in sorted(self.timings.items()):
            import numpy as np
            arr = np.array(times)
            lines.append(
                f"{name:40s}  "
                f"calls={len(times):4d}  "
                f"total={arr.sum():.4f}s  "
                f"mean={arr.mean():.4f}s  "
                f"std={arr.std():.4f}s"
            )
        return "\n".join(lines)
Usage in a simulation loop
prof = Profiler()

for epoch in range(100):
    with prof.section("epoch"):
        with prof.section("forward"):
            y_pred = model(X)
        with prof.section("loss"):
            loss = criterion(y_pred, y)
        with prof.section("backward"):
            loss.backward()
        with prof.section("update"):
            optimizer.step()

print(prof.report())
# Profiling Report
# ==================================================
# epoch                                     calls= 100  total=5.2341s ...
# epoch > forward                           calls= 100  total=2.1034s ...
# epoch > loss                              calls= 100  total=0.3121s ...
# epoch > backward                          calls= 100  total=1.8922s ...
# epoch > update                            calls= 100  total=0.9264s ...

Example: GPU Memory Tracking Context Manager

Write a context manager that tracks GPU memory usage before and after a code block, useful for debugging memory leaks in PyTorch training loops.

Solution
Implementation
from contextlib import contextmanager

@contextmanager
def track_gpu_memory(device: str = "cuda:0", label: str = ""):
    """Track GPU memory allocation during a code block.

    Parameters
    ----------
    device : str
        CUDA device string.
    label : str
        Label for the log message.
    """
    try:
        import torch
        torch.cuda.synchronize(device)
        mem_before = torch.cuda.memory_allocated(device)
        peak_before = torch.cuda.max_memory_allocated(device)
        torch.cuda.reset_peak_memory_stats(device)

        yield

        torch.cuda.synchronize(device)
        mem_after = torch.cuda.memory_allocated(device)
        peak = torch.cuda.max_memory_allocated(device)

        delta = mem_after - mem_before
        prefix = f"[{label}] " if label else ""
        print(f"{prefix}GPU Memory: "
              f"delta={delta / 1e6:+.1f}MB, "
              f"current={mem_after / 1e6:.1f}MB, "
              f"peak={peak / 1e6:.1f}MB")
    except ImportError:
        print("PyTorch not available, skipping GPU tracking")
        yield
Usage
with track_gpu_memory(label="Model init"):
    model = LargeModel().cuda()

with track_gpu_memory(label="Forward pass"):
    output = model(input_tensor)

# [Model init] GPU Memory: delta=+450.2MB, current=450.2MB, peak=450.2MB
# [Forward pass] GPU Memory: delta=+128.5MB, current=578.7MB, peak=612.3MB

Example: Temporary Random Seed Context Manager

Write a context manager that temporarily sets a specific random seed for reproducibility, then restores the original RNG state when the block exits.

Solution
Implementation for NumPy
from contextlib import contextmanager
import numpy as np

@contextmanager
def temporary_seed(seed: int):
    """Temporarily set NumPy random seed, restore state on exit.

    Parameters
    ----------
    seed : int
        Random seed for the block.
    """
    state = np.random.get_state()
    np.random.seed(seed)
    try:
        yield
    finally:
        np.random.set_state(state)

# Usage:
np.random.seed(0)
print(np.random.rand())  # 0.5488...

with temporary_seed(42):
    print(np.random.rand())  # 0.3745... (seed=42)

print(np.random.rand())  # 0.7152... (continues from seed=0)
Why this matters

In scientific computing, you often need reproducible random numbers for a specific section (e.g., generating test data) without affecting the global RNG state used by the rest of the simulation. This context manager provides that isolation.

Context Manager Execution Flow

Context Manager Execution Flow
The lifecycle of a context manager: enter is called when entering the with block, the body executes, and exit is called on both normal exit and exception paths. If exit returns True, the exception is suppressed.

Context Manager Timing Comparison

Compare the overhead of different resource management patterns: manual try/finally, class-based context managers, and contextlib.contextmanager. Vary the workload to see when context manager overhead becomes negligible.

Parameters

Common Mistake: Accidentally Suppressing Exceptions in exit

Mistake:

Returning True from __exit__ suppresses any exception that occurred in the with block:

class BadContext:
    def __enter__(self):
        return self

    def __exit__(self, exc_type, exc_val, exc_tb):
        print("Cleaning up")
        return True   # SUPPRESSES ALL EXCEPTIONS!

with BadContext():
    raise ValueError("Critical error!")

# The ValueError is silently swallowed — program continues
print("This runs!")

Correction:

Return False (or None) unless you intentionally want to suppress specific exceptions:

class GoodContext:
    def __exit__(self, exc_type, exc_val, exc_tb):
        print("Cleaning up")
        return False  # let exceptions propagate

# Or suppress only specific exceptions:
class SelectiveContext:
    def __exit__(self, exc_type, exc_val, exc_tb):
        if exc_type is FileNotFoundError:
            return True   # suppress only this type
        return False      # propagate everything else

Context Managers You Already Use

Many Python standard library and scientific library objects are context managers:

# File I/O
with open("results.csv", "w") as f:
    f.write("snr,ber\n")

# HDF5 files
with h5py.File("data.h5", "r") as f:
    dataset = f["signals"][:]

# PyTorch inference mode
with torch.no_grad():
    predictions = model(test_data)

# Matplotlib figures
with plt.ioff():
    fig, ax = plt.subplots()
    ax.plot(data)
    fig.savefig("plot.png")

# Temporary directories
with tempfile.TemporaryDirectory() as tmpdir:
    save_checkpoint(model, f"{tmpdir}/model.pt")
    validate(f"{tmpdir}/model.pt")
# Directory and contents automatically deleted

Why This Matters: Context Managers in Production Research

Context managers are indispensable in production research pipelines:

  • Experiment tracking: with mlflow.start_run() as run: ensures experiment metadata is always properly closed
  • Distributed training: with torch.distributed.init_process_group(): manages multi-GPU coordination
  • Database connections: with sqlite3.connect("results.db") as conn: ensures transactions are committed or rolled back
  • Temporary computation precision: with torch.cuda.amp.autocast(): for mixed-precision training

In Chapter 6, we will use context managers extensively for managing simulation resources and experiment lifecycle.

See full treatment in Chapter 6

Context Manager Patterns for Scientific Computing

python
Production-ready context managers: Timer, Profiler, TemporarySeed, GPUMemoryTracker, and WorkingDirectory. Both class-based and generator-based implementations. 
# Code from: ch02/python/context_managers.py
# Load from backend supplements endpoint

Resource Management in Scientific Pipelines

python
Examples of using context managers for file I/O (HDF5, CSV), database connections, and temporary directories in scientific workflows. 
# Code from: ch02/python/resource_management.py
# Load from backend supplements endpoint

Quick Check

Which dunder methods must an object implement to be used as a context manager with the with statement?

__init__ and __del__

__enter__ and __exit__

__open__ and __close__

__start__ and __stop__

Correction:
__enter__ and __exit__

These two methods define the context manager protocol.

context manager

An object that implements __enter__ and __exit__ methods, enabling use with the with statement for automatic resource management. Guarantees cleanup code runs even if exceptions occur.

Related: decorator

contextlib

Python standard library module providing utilities for working with context managers: @contextmanager (generator-based), suppress (exception ignoring), redirect_stdout, and ExitStack (dynamic context manager composition).

Key Takeaway

Use context managers for any acquire/release pattern. The with statement guarantees cleanup even on exceptions. Use contextlib.contextmanager for simple cases (generator-based) and class-based context managers when you need state or reusability. In scientific code, timing blocks, GPU memory tracking, and temporary RNG seeds are prime use cases.

DecoratorsChapter Summary