# DeepSpeed

A small helper that annotates a live Ingero eBPF trace with the training
step and epoch from a DeepSpeed loop, so a recorded trace can be sliced
per step / per epoch after the run.

DeepSpeed has no callback API for a standalone loop (you drive
`engine.step()` yourself), so this is a helper you call once per step
rather than a registered callback. **If you train DeepSpeed through
`transformers.Trainer`, use the `hf-trainer` integration instead** - the
Trainer callback already fires.

## What it does

`IngeroDeepSpeed` connects once to the agent's annotation socket and,
each time you call `.step(engine)`, writes one NDJSON annotation scoped
to the training process PID with the engine's `global_steps` (and an
optional epoch). `.epoch(n, phase=...)` marks epoch boundaries. The
agent joins those labels to the eBPF event stream; `ingero query` and
`ingero explain` then slice the trace by step or epoch.

## Requirements

The agent must be running with the annotation ingest socket enabled:

    sudo ingero trace --record --annotate

If it is not, the helper is a no-op: it logs one line and every call
returns without sending. It never crashes or slows the loop.

## Use in your own loop

    from ingero_deepspeed import IngeroDeepSpeed

    ingero = IngeroDeepSpeed(run_id="my-run")
    for epoch in range(epochs):
        ingero.epoch(epoch)
        for batch in loader:
            loss = engine(batch)
            engine.backward(loss)
            engine.step()
            ingero.step(engine, epoch=epoch)
        ingero.epoch(epoch, phase="epoch_end")
    ingero.close()

`IngeroDeepSpeed` is also a context manager (`with IngeroDeepSpeed() as
ingero:`).

## Rate cap

The agent enforces a per-connection annotation rate cap
(`AnnotationConnRateLimit`, see `pkg/contract/annotate.go`). For very
fast loops, annotate at a coarser boundary so every emitted label is
recorded.

## Files

| File | Purpose |
|------|---------|
| `ingero_annotate.py` | Framework-agnostic NDJSON socket writer. Shared verbatim with the lightning / ray / hf-trainer integrations. |
| `ingero_deepspeed.py` | The `IngeroDeepSpeed` loop helper. |
| `train_example.py` | A runnable end-to-end example. |
| `test_*.py`, `conftest.py` | pytest suite (fake socket server, no agent or DeepSpeed needed). |

## Tests

    pip install pytest
    pytest
