Extract the OpenTelemetry span-lifecycle helpers out of megatron/training/training.py
- Dominant language
- Python
- Stars
- 17.9k
- Forks
- 4.5k
- Avg merge
- 4d 6h
- Merged PRs (30d)
- 271
Description
## 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](https://github.com/NVIDIA/Megatron-LM/pull/6438) 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
- #6415 — `megatron.core.telemetry` base layer
- #6438 — the instrumentation that introduced this code
Contributor guide
Research direction
Start by reading the listed span-lifecycle functions and module-level globals in megatron/training/training.py, then inspect the existing telemetry tests under megatron/core/telemetry. Add coverage for startup and interval spans, repeated exit-hook installation, exception, and SIGTERM paths before relocating the helpers to megatron/training/telemetry_spans.py. Done means training.py retains only the listed call-site usage and the lifecycle tests pass.
Written by the indexing model from the issue text.
Assessment
- Tech stack
- python
- Domain
- observability, testing
- Issue type
- Refactor
- Difficulty
- 4/5
- Estimated time
- 3-5 days
- Activity status
- Active
- Clarity
- Clearly specified
- Newbie friendliness
- 52/100