Add always-on critical-path timing for trainers

[micahtyong]

Summary

A client step time today is one number that mixes "my work took this long on the GPU" with "the trainer was busy with somebody else's LoRA". Nothing separates those unless OTLP is configured against a tracing backend (off by default) — which is why "is this Lilo overhead or fair sharing?" keeps being unanswerable after the fact.

This adds an exporter-free breakdown the trainer always keeps, readable per model, so a training loop can log it to W&B per step (or stdout when there is no W&B).

One instrumentation path, not three

CriticalPath (lilo/telemetry/critical_path.py) is an implementation of the engine's existing Observer protocol — the same callbacks that feed TrainerTelemetry's OTLP spans. The engine carries no timing code of its own: no submitted_at fields, no timings.record(...) call sites.

class CriticalPath:                       # Observer
    def register_model(self, model_id, spec):   pending[(model_id, "accept")] = now
    def begin(self, operation):                 pending[(model_id, request_id)] = now
    def span(self, model, name, lane, t0, t1=None, **attrs):
        # "forward_backward"          -> forward_backward.execute
        # "accept"                    -> accept.execute (+ first_model_ready_s gauge)
        # "capture:<kind>"            -> <kind>.capture
        # "persist:<kind>"            -> <kind>.persist
        # first span for a pending command also records <kind>.queue_wait = t0 - submitted
    def forget_model(self, model_id):           drop the model's series + pending entries

The only engine change is passing request_ids= on the execute span (needed because a coalesced batch can hold several commands for one model, so seq_ids doesn't zip with models), plus Observers(*observers), a fan-out so the Modal serve path runs both. Engine(observer=x) composes x with the timing observer rather than replacing it, so timing can't be turned off by accident:

engine.observer = Observers(trainer_telemetry, engine.timings)

CriticalPath is a bounded in-memory accumulator: (model_id, phase) -> {count, total_s, mean_s, max_s, mean_batch}, plus single-valued startup gauges. Every record credits both the model's series and an ALL_MODELS aggregate. execute is recorded on the failure path too. mean_batch is the number of commands coalesced into one backend execution — the multi-LoRA batching factor.

Startup is three gauges measured from process start, set once: trainer.backend_ready_s (Megatron answering /healthz), trainer.serving_ready_s, trainer.first_model_ready_s.

Reading it

• GET /api/v1/timing?model_id=...&reset=true on the engine, proxied by the control plane through engine_for(model_id), so a client hits its own trainer.
• lilo.timing.log_critical_path(model_id, step=step) — flattens the snapshot to lilo/<phase>.<stat> scalars, logs to the active W&B run if one exists (discovered via sys.modules, so W&B stays a non-dependency), prints a JSON line otherwise, and swallows every exception (fetch and run.log) so instrumentation can't fail a training step. reset=True by default, making each call the interval since the last.
• A trainer nobody polls for the aggregate prints one lilo_critical_path JSON line every 5 min, so container logs alone explain a step.

Off the training thread

log_critical_path returns immediately; the HTTP fetch and the log run on one daemon worker. Correctness guards, since a late-arriving log is the classic way to corrupt a W&B run:

• One in flight. A call made while the previous one is still fetching is dropped, not queued (_Reporter.submit), so a slow control plane costs at most one outstanding request and can't pile up threads. The next successful call covers both intervals because counters weren't reset in between.
• Never touches the global step. W&B rejects step= values behind the run's current step, which a background thread can't avoid. Values go out with commit=False and a lilo/step column, with define_metric("lilo/*", step_metric="lilo/step") (once per run) so charts plot against the step they describe:

  run.log({**metrics, "lilo/step": step}, commit=False)

• Run captured in the caller. The active run is resolved on the training thread at call time, not on the worker, so a run finishing mid-fetch can't make the worker log to a different run.
• flush_critical_path(timeout) joins the in-flight report; call it before wandb.finish(). background=False gives the old blocking behavior and returns the metrics.

Unchanged and worth knowing: reset=true is destructive, so two readers of the same model (e.g. a background fetch plus a manual curl) split the interval between them.

Semantics worth knowing

• queue_wait is submit → execution start. A pipelined client's own backlog counts toward it, not just neighbors.
• execute for a coalesced batch is the batch's wall time, credited to every model in it.
• A model's series are evicted on unload (forget_model), so a long-lived multi-tenant trainer never exhausts MAX_SERIES; the aggregate keeps them. The * aggregate is exempt from the cap (it's bounded by the phase count), so per-model churn can never drop trainer-wide totals.

Nothing here is configurable or optional: no env var, no exporter, no payload capture, bounded at MAX_SERIES per-model entries.

Link to Devin session: https://modal.devinenterprise.com/sessions/2ea9fb5c53e5482c86a85dffa85b2c7c
Open in Devin Desktop: https://modal.devinenterprise.com/desktop/session/2ea9fb5c53e5482c86a85dffa85b2c7c?variant=devin
Requested by: @micahtyong

[devin-ai-integration]

I'll fix CI failures and address comments from users with write access that start with 'DevinAI' or '@devin'.

• Disable automatic comment, CI, and merge conflict monitoring

[micahtyong]

/devin review

[devin-ai-integration]

Starting Devin Review.

[devin-ai-integration]

Devin Review found 5 potential issues.