8. Context managers and exceptions

Exceptions are how Python reports failure, and context managers guarantee cleanup when failure happens. Together they decide whether a crashed job leaves behind open files, held locks or half-written checkpoints.

The exception model

• Raising creates an exception object and unwinds the call stack until a matching except clause catches it.
• If nothing catches it, Python prints a traceback and exits. Read tracebacks bottom-up: the last line is the error, and the frames above show the path to it. Find the lowest frame in your own code.
• The hierarchy: BaseException contains Exception, which contains ValueError, TypeError, KeyError, OSError and the rest. KeyboardInterrupt and SystemExit sit outside Exception, so except Exception does not swallow Ctrl-C.

The full try statement

try:
    resp = call_model(prompt)
except (TimeoutError, ConnectionError) as e:
    log.warning("transient failure: %s", e)
    raise                         # re-raise with the original traceback
except ValueError:
    resp = fallback_answer()
else:
    cache.put(prompt, resp)       # runs only if nothing was raised
finally:
    metrics.flush()               # always runs

• Catch the narrowest exception you can actually handle. Never write a bare except:.
• except Exception: pass hides bugs. If you must catch broadly at a boundary, log with log.exception(...), which records the traceback.
• raise ConfigError("bad model name") from e chains exceptions so the root cause stays visible.

Custom exceptions

class RateLimited(Exception):
    def __init__(self, retry_after: float):
        super().__init__(f"rate limited; retry after {retry_after}s")
        self.retry_after = retry_after

Create one per failure mode your caller handles differently: retryable versus permanent, user error versus system error.

Retry with exponential backoff

This is the staple of AI systems code, because model APIs fail transiently all the time.

import random, time

def with_backoff(fn, *, retries=5, base=0.5, cap=30.0,
                 retry_on=(TimeoutError, ConnectionError)):
    for attempt in range(retries):
        try:
            return fn()
        except retry_on:
            if attempt == retries - 1:
                raise
            delay = min(cap, base * 2 ** attempt)
            time.sleep(delay * random.uniform(0.5, 1.5))   # jitter spreads out retries

Libraries such as tenacity package this pattern. Most official LLM SDKs already retry some errors internally, so check yours before stacking a second retry layer on top.

Exception groups (3.11+)

An ExceptionGroup bundles several failures, for example from concurrent tasks, and except* handles them by type. You will meet them with asyncio.TaskGroup (Part 12).

Context managers

with guarantees that setup is paired with teardown:

with open("out.jsonl", "w", encoding="utf-8") as f:
    f.write(line)
# the file is closed here, even if write() raised

The protocol: __enter__ runs at the start, and its return value is bound by as. __exit__(exc_type, exc, tb) runs at the end, with or without an exception. Returning True from __exit__ suppresses the exception, which you will rarely want.

The easiest way to write your own is a generator with @contextmanager:

from contextlib import contextmanager
import time

@contextmanager
def timer(label):
    t0 = time.perf_counter()
    try:
        yield                      # the with-block runs here
    finally:
        print(f"{label}: {time.perf_counter() - t0:.2f}s")

with timer("eval"):
    evaluate(model)

Where context managers appear in AI code:

• torch.no_grad() and torch.inference_mode() switch off gradient tracking inside the block.
• torch.autocast(...) enables mixed precision inside the block.
• model.eval() is not a context manager. It flips a flag that stays flipped until you call model.train().
• HTTP clients (with httpx.Client() as client:), database sessions, locks, temporary directories and tracing spans.
• contextlib.ExitStack handles a variable number of contexts; contextlib.suppress(FileNotFoundError) ignores one specific error.

Questions to ask:

• If this line raises, what is left open or half-written?
• Is this exception retryable, and is something else already retrying it?
• Is this except clause hiding a bug?