Extract the OpenTelemetry span-lifecycle helpers out of megatron/training/training.py

[deepakn94]

Summary

megatron/training/training.py is 5,604 lines, and a meaningful share of it is OpenTelemetry span bookkeeping rather than training logic. That telemetry code should live in its own module, leaving training.py with imports plus the call-site usage.

This was flagged as a planned follow-up in #6438 and deliberately deferred there; see the "Known follow-ups" section of that PR.

What is in training.py today

Measured on main at the time of writing — 196 _otel references in total:

11 span-lifecycle functions, in a mostly contiguous block at lines 383-680:

Function	Line
_otel_telemetry_active	383
_start_otel_job_spans	399
_otel_mark_goodput	504
_backdated_otel_span	519
_end_otel_startup_span	529
_start_otel_train_span	554
_reroot_otel_interval	572
_maybe_reroot_otel_interval	625
_end_otel_interval_span	637
_end_otel_train_span	655
_end_otel_job_spans	662

9 module-level _otel_* globals holding span handles, context tokens, the trace-interval counter, and idempotence flags, mutated through 9 global statements.

Further otel-only functions scattered below: the SIGTERM handler, force-flush, graceful-drain, and the exit-hook installer.

Proposed change

Move the above into a dedicated module (e.g. megatron/training/telemetry_spans.py). The globals travel with their mutators, which is cleaner than the status quo since state and the code that owns it end up together.

What stays in training.py is the call-site usage that should have been its only telemetry footprint:

• 21 with _otel_managed_span(...) blocks
• 6 _otel_sg_enabled(...) gates
• a handful of _otel_mark_goodput / _otel_set_attrs attribute stamps

Roughly 400-500 lines should move. Coupling to training state is light — mostly get_telemetry() and get_args() — so this is close to a pure move.

Why it was not done in #6438

1. It would have rewritten the exact training.py diff reviewers had already read.
2. The span lifecycle has no unit-test coverage. The four telemetry tests on main cover megatron/core/telemetry (span groups, metric instruments, no-op fallbacks), not this. A move touching global mutation ordering and atexit/SIGTERM registration idempotence would be unguarded.

Suggested approach

Land the move and its tests together, or tests first:

• Tests for the span lifecycle: startup span opens and closes once, interval re-rooting at save_interval boundaries, exit hooks installed exactly once across repeated pretrain() calls in one process, and spans still closed on the exception and SIGTERM paths.
• Then the move itself, kept as close to a pure relocation as possible so it reviews with git diff -M.

Related

• Add megatron.core.telemetry base layer and observability guide #6415 — megatron.core.telemetry base layer
• Instrument Megatron-LM with OpenTelemetry traces and metrics #6438 — the instrumentation that introduced this code