NVIDIA / NVIDIA/Megatron-LM

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

Open
#6,631 0 comments 0 reactions 0 assignees View on GitHub
enhancement
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

Open the contributing 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

Get new issues in your inbox

A short digest of beginner-friendly GitHub issues.