Context Managers

TL;DR

A context manager is an object that defines what happens when a with block starts and ends. The with statement guarantees the cleanup runs — even if an exception is raised inside the block, even if the program is being killed by KeyboardInterrupt. That’s the whole point: it’s the language’s answer to “always close the file, always release the lock, always restore the GPU state.”

Two forms exist. The protocol form: a class with __enter__ and __exit__. The decorator form: a generator function with @contextlib.contextmanager and one yield inside. The decorator form is the one you’ll write 95% of the time.

ML codebases lean on context managers for: file handles, DB sessions, PyTorch’s torch.no_grad(), torch.cuda.amp.autocast, MLflow runs, W&B runs, profilers, temporary directories, tempfile.NamedTemporaryFile, locks, and any “set state, do work, restore state” idiom.

The picture in your head

with X as y: is a try / finally in nicer clothes:

with open("data.csv") as f:
    process(f)

# is exactly equivalent to:
f = open("data.csv").__enter__()
try:
    process(f)
finally:
    open("data.csv").__exit__(*exc_info_or_Nones)   # actually called on the same obj

The cleanup always runs. That guarantee is the value. Without with, you’d write the try/finally by hand every time, and forget on the twentieth file you opened.

The protocol — __enter__ and __exit__

Two methods. __enter__ runs at block start; its return value is bound to the as variable. __exit__ runs at block end with three arguments describing any exception (or three Nones if the block exited normally). Returning True from __exit__ suppresses the exception; usually you return None (let it propagate).

class TimingContext:
    def __init__(self, label: str):
        self.label = label

    def __enter__(self):
        self.t0 = time.perf_counter()
        return self                     # bound to `as` variable

    def __exit__(self, exc_type, exc_val, exc_tb):
        elapsed = time.perf_counter() - self.t0
        print(f"{self.label}: {elapsed:.3f}s")
        return None                     # don't suppress exceptions

with TimingContext("training epoch") as t:
    train_one_epoch(model, loader)
# prints: training epoch: 12.345s

If train_one_epoch raises, __exit__ still runs (and the timing still prints), then the exception propagates. That’s the whole guarantee.

The easier form — @contextlib.contextmanager

Most context managers fit this template:

from contextlib import contextmanager

@contextmanager
def timing(label):
    t0 = time.perf_counter()
    try:
        yield                              # block body runs here
    finally:
        elapsed = time.perf_counter() - t0
        print(f"{label}: {elapsed:.3f}s")

with timing("training epoch"):
    train_one_epoch(model, loader)

Same behaviour, three lines fewer. The try / finally is doing the real work; the decorator turns the generator into a context manager.

If you yield value, that value is what as x binds to:

@contextmanager
def temp_file(suffix=".txt"):
    f = tempfile.NamedTemporaryFile(suffix=suffix, delete=False)
    try:
        yield f.name
    finally:
        os.unlink(f.name)

with temp_file(".pt") as path:
    torch.save(model.state_dict(), path)
    # path exists here
# path is deleted here

This is the canonical “set up resource, hand it to the user, clean up” pattern.

Setting and restoring state

The single most useful application of context managers in ML is “temporarily change a global, restore it on exit, no matter what.”

@contextmanager
def temp_seed(seed):
    """Make the next block deterministic without polluting the global RNG."""
    state = random.getstate()
    np_state = np.random.get_state()
    torch_state = torch.get_rng_state()
    random.seed(seed)
    np.random.seed(seed)
    torch.manual_seed(seed)
    try:
        yield
    finally:
        random.setstate(state)
        np.random.set_state(np_state)
        torch.set_rng_state(torch_state)

with temp_seed(42):
    sample = sample_from_model()
# RNG state is fully restored here, regardless of whether sampling raised

This is exactly how torch.no_grad() is implemented. It saves the “grad enabled” flag, sets it to False, runs the block, restores the flag.

Multiple context managers in one with

# Three resources, all cleaned up even if any raise
with open("input.csv") as fin, open("output.csv", "w") as fout, lock:
    process(fin, fout)

Equivalent to nested with blocks, but flatter. Use it when the managers don’t depend on each other’s return values.

If they do, nest:

with mlflow.start_run() as run:
    with timing(f"run-{run.info.run_id}"):
        train(...)

The standard-library context managers worth knowing

ExitStack is the under-known one and worth a closer look.

ExitStack — variable-number resources

When you don’t know how many resources you need until runtime — say, opening N files where N comes from a config — you can’t write with open(...), open(...), ... because you don’t know N at write time. ExitStack solves this.

from contextlib import ExitStack

paths = ["a.csv", "b.csv", "c.csv"]   # could be any length

with ExitStack() as stack:
    files = [stack.enter_context(open(p)) for p in paths]
    # all files are open here; all get closed at end of block
    merged = list(itertools.chain.from_iterable(files))

Every enter_context(...) registers a cleanup that runs when the ExitStack exits, in reverse order. It’s the “with” equivalent of a dynamic try/finally.

A worked ML example — a training-run context

A common pattern: every training run wraps a few cross-cutting things. Telemetry, profiling, deterministic seeds, MLflow logging. A single context manager hides the boilerplate.

import logging
from contextlib import ExitStack, contextmanager

import mlflow
import torch

logger = logging.getLogger(__name__)

@contextmanager
def training_run(name: str, *, seed: int, profile: bool = False):
    """Set up everything; tear it all down even if training crashes."""
    with ExitStack() as stack:
        # MLflow run — logged on success or failure
        run = stack.enter_context(mlflow.start_run(run_name=name))
        mlflow.log_param("seed", seed)

        # Deterministic seeds
        stack.enter_context(temp_seed(seed))

        # Optional profiling
        if profile:
            prof = stack.enter_context(
                torch.profiler.profile(
                    activities=[torch.profiler.ProfilerActivity.CUDA]
                )
            )
        else:
            prof = None

        logger.info("training run %s started (mlflow_id=%s)", name, run.info.run_id)
        try:
            yield {"run": run, "prof": prof}
        except Exception:
            mlflow.set_tag("status", "failed")
            raise
        else:
            mlflow.set_tag("status", "ok")

with training_run("baseline-v3", seed=42, profile=True) as ctx:
    train(model, loader)
    # MLflow run, profiling, seed restoration all cleaned up here regardless

If train raises, mlflow.set_tag("status", "failed") runs, the profiler stops, the seed gets restored, the MLflow run gets closed. None of that is your call site’s responsibility — the context manager owns it.

Async context managers

Use async with on objects that define __aenter__ and __aexit__, or generators decorated with @contextlib.asynccontextmanager.

from contextlib import asynccontextmanager

@asynccontextmanager
async def http_client(base_url: str):
    client = httpx.AsyncClient(base_url=base_url, timeout=30)
    try:
        yield client
    finally:
        await client.aclose()

async with http_client("https://api.openai.com") as client:
    resp = await client.post("/v1/chat/completions", json=...)

The httpx.AsyncClient itself is already an async context manager — you’d more typically write async with httpx.AsyncClient(...) as client: directly. The wrapper here is illustrative.

Common gotchas

• __exit__ returning truthy swallows the exception. Return None (or no return) unless you really mean to suppress.
• Context manager objects are typically single-use. Calling with on the same instance twice is undefined behaviour. Most managers (file handles, locks) are single-use; create a new one each time.
• @contextmanager generators must yield exactly once. Two yields raise RuntimeError. Zero yields raise RuntimeError. The yield must be inside try if you want guaranteed cleanup.
• Reraising from __exit__. If you do work in __exit__ that itself might fail, wrap it in try/except so you don’t mask the original exception. Or use contextlib.suppress carefully.
• GPU contexts and with. torch.cuda.device(0) is a context manager that switches the default device for the block — useful, but silently scoped. Errors from forgetting to enter it manifest as RuntimeError from kernels running on the wrong device.

Where this shows up in real ML codebases

• with torch.no_grad(): — disable autograd for an inference block; restored on exit. Used in every eval loop.
• with torch.cuda.amp.autocast(): — mixed-precision context.
• with model.train(): / with model.eval(): — actually these aren’t context managers in PyTorch (just methods); but plenty of wrapper libraries make them so for symmetry.
• with mlflow.start_run(): — log everything between the lines to the same MLflow run.
• with wandb.init(...): — same idea.
• with tempfile.TemporaryDirectory() as d: — a scratch dir for intermediate artifacts; auto-cleaned.
• with httpx.AsyncClient() as client: — HTTP client lifecycle.
• with sqlalchemy.Session() as session: — DB session that commits on success, rolls back on error.

The defensive habit: if a function “must always be paired with a cleanup call,” that’s a context manager waiting to be written.

Resources

• PEP 343 — The “with” Statement — peps.python.org — the original spec.
• PEP 492 — Coroutines with async and await — peps.python.org — introduces async with / __aenter__ / __aexit__.
• Python docs — contextlib — docs.python.org — contextmanager, ExitStack, suppress, redirect_stdout, closing.
• Python docs — with statement — docs.python.org — the language reference.