mirror of
https://github.com/NVIDIA/Model-Optimizer.git
synced 2026-10-02 03:14:52 +08:00
Speed up megatron_bridge example tests by ~6x on a single GPU (#2296)
### What does this PR do? Type of change: Test infrastructure / CI time `tests/examples/megatron_bridge` spends most of its time importing Python, not testing. Each step of a test spawns `torchrun`, and the new process spends **~25s importing torch/megatron/modelopt** before doing any work. A single `test_qad` run pays that **six** times — three steps, plus a spawned child per distributed checkpoint save, because Megatron-Core's async writer uses `mp_mode="spawn"` and spawn re-imports `__main__`. Profiled with phase timers in the example scripts: | | share of `test_qad[qwen3]` | |---|---| | Python imports (6 process launches × ~26s) | **~76%** | | actual compute (`mtq.quantize` 4.2s, model build 0.24s, export 0.06s) | ~5s | `run_example_command` now dispatches each step internally instead of shelling out: - **single-rank steps** run directly in the pytest process, driving the script's own `get_args()` + `main()` — no new interpreter, no re-import; - **multi-rank steps** drive `torch.distributed.run.main()` in-process with patched `sys.argv`, the same pattern Megatron-Bridge uses in its own functional tests. ### Results | suite | before | after | |---|---|---| | `tests/examples/megatron_bridge`, **1 GPU** | **21m27** | **4m03 - 6m17** | | `tests/examples/megatron_bridge`, **2 GPU** | 26m10 | 25m03 | Both figures are on current `main` (17 tests). The 2-GPU number is up from 20m41 before merging #2276, which added a test and made one previously single-rank step multi-rank. The 1-GPU figure is a range, not a best case. Across ten runs on a verified-idle box the suite lands at ~4m most of the time and at ~6m otherwise, always with the same result (17 passed). Per-test durations show the entire spread is one test: `test_qad[qwen3]` runs at ~15s or at ~149s. It is 25s in isolation, 15s after `test_distill.py` and ~25s after `test_prune_minitron.py`, so it needs the full sequence and does not reproduce on demand — three attempts to catch it under instrumentation all landed on fast runs. In those, CUDA state immediately before it is 46 MiB allocated / 68 MiB reserved / 7 segments / 22 MiB inactive-split, and the preceding test's 498/984 MiB is fully reclaimed, so a fragmented allocator is measured *not* to be the cause in the fast path at least. Left documented rather than guessed at: correctness is unaffected across every run, and the worst case sits inside the 30-minute PR budget (CI `Run tests` 902s). The single-GPU path is the big win, and it is the one the per-PR runner uses — that job now finishes in **8 minutes** in CI. Multi-rank steps still launch worker processes that re-import, so the 2-GPU nightly improves far less. This also fixes the timeouts under coverage. With `--cov` (how CI runs it), on the same three tests: in-process **3 passed in 1m15**, subprocess **3 failed on `Timeout (>360.0s)` in 18m57**. ### CI timeout The 2-GPU nightly runs every test multi-GPU and measured **58 minutes against a 60-minute cap** — too close to be reliable. `timeout_minutes` is now ref-conditional, mirroring the `runner` line directly below it: **30 minutes on PRs** (single-GPU, ~8 min) and **75 on nightly**. Keeping the nightly at full multi-GPU coverage is deliberate. Making individual tests single-rank cut it to ~7 minutes, but it gives up the parallel-path coverage that is the whole point of the 2-GPU job, and it surfaced a real fragility: `test_prune_minitron[nemotron_h]` fails with *"No scores collected for importance estimation"* when it runs single-rank after the full distill file. It passes alone and after any single preceding test — multi-rank tests are immune because `torchrun` gives them fresh worker processes. Nightly is the right place to spend the wall-clock. ### What is and isn't covered Each script's real `get_args()` still runs, so CLI flags, defaults and recipe-string resolution stay covered. Not covered for single-rank steps: the `torchrun` invocation itself and the `__main__` block (`dist.setup()` / `dist.abort()`). Multi-rank steps still go through the real launcher. **No test file changes.** The tests still read as "launch this torchrun command" and their assertions are untouched. ### Keeping it that way There is no toggle and no fallback. Every step in this suite must be `torchrun --nproc_per_node=<int> <script>.py` with the script exposing `get_args()` + `main()`, and `run_example_step` raises otherwise — it returns `str`, not `str | None`, so a step cannot quietly become a subprocess. That matters because a silent fallback still *passes*, just ~6x slower, so a new script or test could cost the suite its speed-up with nothing to show for it. Both guards verified by breaking them on purpose, each failing in ~1.4s rather than burning a run: | broken convention | result | |---|---| | `--nproc_per_node=gpu` | `AssertionError: --nproc_per_node must be a plain integer: [...]` | | step invoking `generate_vllm.py` (no `get_args`) | `AssertionError: generate_vllm.py must define get_args() and main(args)` | ### Layout The runner lives in `tests/_test_utils/examples/megatron_example_runner.py`, next to the `run_command.py` it plugs into and the other per-example helpers. It is deliberately not under `tests/_test_utils/torch/megatron/`: both files there import megatron at module top, whereas this one must not, since importing `megatron.bridge` would initialise CUDA in the pytest process and hold a context on device 0 for the whole session. ### Isolation Sharing one interpreter means anything global has to be put back between steps, or one failing test cascades into the next. Each of these was previously cleaned up by `torchrun` simply exiting: - **`NVTE_*`** — Transformer-Engine records its chosen attention backend in the environment, so a Mamba hybrid failed after an attention model ran. The environment is restored wholesale rather than by naming variables. - **Allocator** — `empty_cache()` frees nothing while a finished step's model is still reachable; a later test ran **9x slower** (162s vs 18s) against a fragmented allocator until `gc.collect()` was added first. - **Parallel state and the rerun state machine** — two separate singletons; `destroy_model_parallel()` does not touch the latter. - **Signal handlers** — `PContext.start()` installs its own `SIGTERM`/`SIGINT`/`SIGHUP`/`SIGQUIT` handlers and never restores them. With a subprocess launcher, process exit did that for us; in-process they are saved and put back, or pytest's Ctrl-C and CI cancellation would break for the rest of the session. Verified rather than assumed: injecting a failure mid-test (after a model was built and parallel state left live) gives **1 failed, 2 passed**, with the surviving tests at full speed. ### Coverage Coverage of the exercised code **improves**. In subprocess mode the child imports modelopt as `site-packages/modelopt/...` while pytest measures `modelopt/...`, so the data never merges — which is also why the subprocess report showed exactly double the statement count. | module | subprocess | in-process | |---|---|---| | `unified_export_megatron.py` | 8% | **43%** | | `mcore_custom.py` | 34% | **44%** | ### Testing All in `nvcr.io/nvidia/nemo:26.08` on 2x RTX 6000 Ada, with per-test caps enforced. | run | result | |---|---| | `tests/examples/megatron_bridge`, 1 GPU | 17 passed — 4m03 (6m17 worst of 10 runs) | | same, subprocess baseline | 15 passed, 1 skipped — 21m27 | | `tests/examples/megatron_bridge`, 2 GPU | 17 passed — 25m03 | | CI 1-GPU example job (`megatron / run-test`) | passed — `Run tests` 902s, job ~20m (30m cap) | | CI 2-GPU nightly | passed — `Run tests` 3162s, job 58m (75m cap) | | cascade check (injected mid-test failure) | 1 failed, 2 passed, survivors at full speed | | pre-commit | clean | ### Before your PR is "*Ready for review*" - Is this change backward compatible?: ✅ — test-only; no source or public API changes - If you copied code from any other sources or added a new PIP dependency, did you follow guidance in `CONTRIBUTING.md`: N/A — no new dependencies (`pytest-forked` was evaluated and rejected: `import megatron.bridge` initialises CUDA, and CUDA cannot be re-initialised in a forked child) - Did you write any new necessary tests?: N/A — this changes how existing tests are executed - Did you update [Changelog](https://github.com/NVIDIA/Model-Optimizer/blob/main/CHANGELOG.rst)?: N/A — internal test infrastructure, not user-facing - Did you get Claude approval on this PR?: ✅ — reviewed by Claude and CodeRabbit, all threads addressed --------- Signed-off-by: Keval Morabia <28916987+kevalmorabia97@users.noreply.github.com> Co-authored-by: Claude Opus 5 <noreply@anthropic.com>
This commit is contained in:
Keval Morabia
co-authored by
Claude Opus 5
Claude Opus 5
parent
1d3068ff78
commit
411d072a2e
@@ -155,7 +155,8 @@ jobs:
|
||||
with:
|
||||
docker_image: "nvcr.io/nvidia/nemo:26.08"
|
||||
example: megatron_bridge
|
||||
timeout_minutes: 75
|
||||
# PR runs single-GPU (~11 min); the nightly runs every test multi-GPU (~58 min).
|
||||
timeout_minutes: ${{ startsWith(github.ref, 'refs/heads/pull-request/') && 30 || 75 }}
|
||||
pip_install_extras: "[hf,puzzletron,dev-test]"
|
||||
runner: ${{ startsWith(github.ref, 'refs/heads/pull-request/') && 'linux-amd64-gpu-rtxpro6000-latest-1' || 'linux-amd64-gpu-rtxpro6000-latest-2' }}
|
||||
allow_failure: ${{ contains(format(',{0},', vars.ALLOW_FAILURE_EXAMPLE_TESTS), ',megatron_bridge,') }}
|
||||
|
||||
@@ -0,0 +1,333 @@
|
||||
# SPDX-FileCopyrightText: Copyright (c) 2024 NVIDIA CORPORATION & AFFILIATES. All rights reserved.
|
||||
# SPDX-License-Identifier: Apache-2.0
|
||||
#
|
||||
# Licensed under the Apache License, Version 2.0 (the "License");
|
||||
# you may not use this file except in compliance with the License.
|
||||
# You may obtain a copy of the License at
|
||||
#
|
||||
# http://www.apache.org/licenses/LICENSE-2.0
|
||||
#
|
||||
# Unless required by applicable law or agreed to in writing, software
|
||||
# distributed under the License is distributed on an "AS IS" BASIS,
|
||||
# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
|
||||
# See the License for the specific language governing permissions and
|
||||
# limitations under the License.
|
||||
"""Run ``examples/megatron_bridge`` scripts without a ``torchrun`` launch per step.
|
||||
|
||||
A launch spends ~25s importing torch/megatron/modelopt before doing any work, and one example test
|
||||
runs several. Single-rank commands drive the script's ``main()`` directly in the pytest process,
|
||||
which is where that cost disappears (the suite goes from ~21min to ~4min on one GPU). Multi-rank
|
||||
commands drive ``torchrun`` itself in-process, which only saves the launcher's own imports -- the
|
||||
pytest process cannot be more than one rank -- but keeps torchrun's fresh worker processes.
|
||||
|
||||
The script's own ``get_args()`` still runs, so CLI flags and recipe strings stay covered. Not
|
||||
covered: the ``torchrun`` invocation and the ``__main__`` block (``dist.setup()``, ``dist.abort()``
|
||||
and ``dist.cleanup()``).
|
||||
|
||||
Megatron and torch.distributed.run are imported lazily on purpose: this module is imported at
|
||||
collection time, and importing ``megatron.bridge`` initialises CUDA in the pytest process, which
|
||||
would hold a context on device 0 for the whole session even when every step runs under torchrun.
|
||||
"""
|
||||
|
||||
import contextlib
|
||||
import gc
|
||||
import importlib
|
||||
import importlib.util
|
||||
import logging
|
||||
import os
|
||||
import signal
|
||||
import sys
|
||||
import tempfile
|
||||
from pathlib import Path
|
||||
from unittest.mock import patch
|
||||
|
||||
import torch
|
||||
from _test_utils.examples.run_command import MODELOPT_ROOT
|
||||
from _test_utils.torch.distributed.utils import get_free_port
|
||||
|
||||
import modelopt.torch.utils.distributed as dist
|
||||
|
||||
_SINK = None # never closed: a handler built inside a capture keeps a reference to it
|
||||
|
||||
|
||||
def _sink():
|
||||
"""One long-lived stream on the real stdout, whose fd is redirected per step.
|
||||
|
||||
A per-step stream closed on exit would leave a permanently-closed file behind for anything that
|
||||
captured it -- e.g. the common module-level ``logging.StreamHandler(sys.stdout)``, where
|
||||
``sys.stdout`` is this sink while a step runs and the handler outlives it in ``sys.modules``.
|
||||
"""
|
||||
global _SINK
|
||||
if _SINK is None:
|
||||
_SINK = os.fdopen(os.dup(1), "w", buffering=1, errors="replace")
|
||||
return _SINK
|
||||
|
||||
|
||||
# torchrun's PContext installs handlers for these and never restores them.
|
||||
_LAUNCHER_SIGNALS = (signal.SIGTERM, signal.SIGINT, signal.SIGHUP, signal.SIGQUIT)
|
||||
|
||||
|
||||
@contextlib.contextmanager
|
||||
def _capture_output():
|
||||
"""Capture everything a step emits, in the order it happened.
|
||||
|
||||
Three sources have to end up in the same place: ``print`` from the script, records from loggers
|
||||
Megatron-Bridge uses for lines tests assert on, and fd-level writes from ``torchrun`` workers
|
||||
and C extensions. Under pytest's fd capture ``sys.stdout``/``sys.stderr`` are objects over
|
||||
pytest's own temp file rather than fds 1/2, so both a redirect and an fd swap are needed.
|
||||
|
||||
Buffered, not streamed: the caller prints the result, which pytest shows for a failing test.
|
||||
Echoing it live would not reach the CI log anyway -- without ``-s``, fd 1 is already pytest's
|
||||
own capture file. Piping it to stream for real cost the 1-GPU suite 3m43 -> 9m06 in
|
||||
backpressure, so surviving a job-level SIGKILL would need a file outside the process.
|
||||
"""
|
||||
holder = [""]
|
||||
with tempfile.TemporaryFile(mode="w+", encoding="utf-8", errors="replace") as tmp:
|
||||
handler = logging.StreamHandler() # bound to the swapped fd 2 below
|
||||
handler.setFormatter(logging.Formatter("%(message)s"))
|
||||
root = logging.getLogger()
|
||||
with logging._lock: # any getLogger() from a background thread mutates loggerDict
|
||||
others = [
|
||||
lg for lg in root.manager.loggerDict.values() if isinstance(lg, logging.Logger)
|
||||
]
|
||||
# Lower levels only for loggers that own their output: a propagating logger still needs its
|
||||
# own level lowered, or it drops the record before root sees it. The boost is load-bearing:
|
||||
# without it test_distill_validate_only cannot see "skipping training ...". It does mean the
|
||||
# capture is a superset of what a torchrun user sees, so these assertions check that the
|
||||
# step did the right work, not that its logging is user-visible.
|
||||
levels = {lg: lg.level for lg in [root, *others] if lg.handlers or not lg.propagate}
|
||||
# Attach the handler only where the record cannot reach root, or it is written twice.
|
||||
handled = [root, *(lg for lg in others if not lg.propagate)]
|
||||
sys.stdout.flush()
|
||||
sys.stderr.flush()
|
||||
sink = _sink() # before the swap, so it holds the real stdout
|
||||
saved_out, saved_err, saved_sink = os.dup(1), os.dup(2), os.dup(sink.fileno())
|
||||
try:
|
||||
# Inside the try: a raise between here and the restore would otherwise leave fds 1/2
|
||||
# pointing at a temp file the ``with`` then closes, silencing the rest of the session.
|
||||
os.dup2(tmp.fileno(), 1)
|
||||
os.dup2(tmp.fileno(), 2)
|
||||
sink.flush()
|
||||
os.dup2(tmp.fileno(), sink.fileno()) # move its fd rather than replacing the object
|
||||
handler.setStream(sink)
|
||||
for lg in levels:
|
||||
if lg.level > logging.INFO:
|
||||
lg.setLevel(logging.INFO)
|
||||
for lg in handled:
|
||||
lg.addHandler(handler)
|
||||
with contextlib.redirect_stdout(sink), contextlib.redirect_stderr(sink):
|
||||
yield holder
|
||||
finally:
|
||||
for lg in handled:
|
||||
lg.removeHandler(handler)
|
||||
for lg, level in levels.items():
|
||||
lg.setLevel(level)
|
||||
sink.flush()
|
||||
os.dup2(saved_sink, sink.fileno())
|
||||
os.close(saved_sink)
|
||||
sys.stdout.flush()
|
||||
sys.stderr.flush()
|
||||
os.dup2(saved_out, 1)
|
||||
os.dup2(saved_err, 2)
|
||||
os.close(saved_out)
|
||||
os.close(saved_err)
|
||||
tmp.seek(0)
|
||||
holder[0] = tmp.read()
|
||||
|
||||
|
||||
@contextlib.contextmanager
|
||||
def _preserved_signal_handlers():
|
||||
"""Put back handlers a step installs; with a subprocess launcher, process exit did this.
|
||||
|
||||
Both paths need it: ``PContext.start()`` installs its own, and a Megatron training loop run
|
||||
in-process is if anything likelier to (graceful exit, checkpoint-on-signal). Losing them breaks
|
||||
pytest's Ctrl-C and CI cancellation for the rest of the session.
|
||||
"""
|
||||
saved = {s: signal.getsignal(s) for s in _LAUNCHER_SIGNALS}
|
||||
try:
|
||||
yield
|
||||
finally:
|
||||
for sig, handler in saved.items():
|
||||
if handler is not None: # installed from C; not restorable via signal.signal()
|
||||
signal.signal(sig, handler)
|
||||
|
||||
|
||||
def reset_megatron_global_state() -> None:
|
||||
"""Drop the global state a finished single-rank step leaves behind.
|
||||
|
||||
Steps that run here share one interpreter, so anything global has to be put back or it leaks
|
||||
into the next test. Failure-tolerant on purpose: this also runs after a failed step, where the
|
||||
model may be half-built, and it must not mask the error that got us here.
|
||||
"""
|
||||
with contextlib.suppress(Exception):
|
||||
from megatron.core import parallel_state
|
||||
|
||||
parallel_state.destroy_model_parallel()
|
||||
with contextlib.suppress(Exception):
|
||||
# A separate singleton from the parallel state, initialised by the training path.
|
||||
from megatron.core.rerun_state_machine import destroy_rerun_state_machine
|
||||
|
||||
destroy_rerun_state_machine()
|
||||
with contextlib.suppress(Exception):
|
||||
# Collect first: empty_cache() only returns blocks nothing references, and a finished step
|
||||
# can leave its model reachable from the imported module. Without this a later test ran
|
||||
# against a fragmented allocator ~9x slower.
|
||||
gc.collect()
|
||||
if torch.cuda.is_available():
|
||||
torch.cuda.empty_cache()
|
||||
|
||||
|
||||
def _load_example_module(script: str, example_path: str):
|
||||
"""Import an example script under a namespaced module name.
|
||||
|
||||
``import_module("quantize")`` would consult ``sys.modules`` first and could pick up an
|
||||
unrelated top-level module of the same name, and would leave the example cached under a
|
||||
generic name.
|
||||
"""
|
||||
example_dir = str(MODELOPT_ROOT / "examples" / example_path)
|
||||
name = f"_modelopt_example_{example_path.replace('/', '_')}_{Path(script).stem}"
|
||||
if name in sys.modules:
|
||||
return sys.modules[name]
|
||||
spec = importlib.util.spec_from_file_location(name, Path(example_dir) / Path(script).name)
|
||||
module = importlib.util.module_from_spec(spec)
|
||||
sys.modules[name] = module
|
||||
sys.path.insert(0, example_dir) # the script's own sibling imports
|
||||
try:
|
||||
spec.loader.exec_module(module)
|
||||
except BaseException:
|
||||
del sys.modules[name]
|
||||
raise
|
||||
finally:
|
||||
sys.path.remove(example_dir)
|
||||
return module
|
||||
|
||||
|
||||
def _require_drivable(script: str, example_path: str) -> None:
|
||||
"""Check the script exposes the ``get_args()`` + ``main()`` shape this runner drives."""
|
||||
module = _load_example_module(script, example_path)
|
||||
if not callable(getattr(module, "get_args", None)) or not callable(
|
||||
getattr(module, "main", None)
|
||||
):
|
||||
raise AssertionError(f"{script} must define get_args() and main(args)")
|
||||
|
||||
|
||||
def requested_world_size(cmd_parts: list[str]) -> int | None:
|
||||
"""World size a ``torchrun`` command asks for, or ``None`` if it is not a plain integer."""
|
||||
parts = [str(p) for p in cmd_parts]
|
||||
for i, part in enumerate(parts):
|
||||
if part.startswith("--nproc_per_node="):
|
||||
value = part.split("=", 1)[1]
|
||||
elif part == "--nproc_per_node" and i + 1 < len(parts):
|
||||
value = parts[i + 1] # torchrun accepts the space-separated form too
|
||||
else:
|
||||
continue
|
||||
return int(value) if value.isdigit() else None # "gpu"/"auto"
|
||||
return None
|
||||
|
||||
|
||||
def run_example_in_process(cmd_parts: list[str], example_path: str) -> str:
|
||||
"""Drive a single-rank step's real ``get_args()`` + ``main()`` here. Returns its stdout."""
|
||||
# Set unconditionally: the per-test environment restore drops these while the process group
|
||||
# stays initialised, so a later step would otherwise run with a live group and no launch vars.
|
||||
os.environ.setdefault("MASTER_ADDR", "127.0.0.1")
|
||||
os.environ["RANK"] = "0"
|
||||
os.environ["WORLD_SIZE"] = "1"
|
||||
os.environ["LOCAL_RANK"] = "0"
|
||||
if not dist.is_initialized():
|
||||
# Not setdefault: another suite may have left a stale MASTER_PORT in the environment, and
|
||||
# binding an occupied port blocks in rendezvous until timeout instead of failing fast.
|
||||
os.environ["MASTER_PORT"] = str(get_free_port())
|
||||
dist.setup() # also sets the CUDA device
|
||||
|
||||
script = next(p for p in cmd_parts if str(p).endswith(".py"))
|
||||
argv = [str(p) for p in cmd_parts[cmd_parts.index(script) :]]
|
||||
example_dir = str(MODELOPT_ROOT / "examples" / example_path)
|
||||
cwd = os.getcwd()
|
||||
os.chdir(example_dir) # match the subprocess path, which runs with the example dir as cwd
|
||||
native = [""] # bound up front: readable even if the capture fails to start
|
||||
try:
|
||||
module = _load_example_module(str(script), example_path)
|
||||
reset_megatron_global_state() # a previous step left its own parallel state behind
|
||||
try:
|
||||
# get_args() inside the guard: argparse calls parser.error() -> SystemExit on a flag
|
||||
# that has drifted from the example's parser, which is the failure this suite most
|
||||
# wants to report clearly rather than raise bare out of the runner. argv stays patched
|
||||
# across main() as well, which the subprocess path had.
|
||||
with (
|
||||
_preserved_signal_handlers(),
|
||||
_capture_output() as native,
|
||||
patch.object(sys, "argv", argv),
|
||||
):
|
||||
args = module.get_args()
|
||||
module.main(args)
|
||||
except SystemExit as e:
|
||||
if e.code not in (0, None):
|
||||
raise RuntimeError(f"{script} exited with code {e.code}") from e
|
||||
return native[0]
|
||||
except BaseException as e:
|
||||
# The caller matches transient-HuggingFace markers on this alongside the traceback.
|
||||
e.captured_output = native[0]
|
||||
raise
|
||||
finally:
|
||||
os.chdir(cwd)
|
||||
print(native[0]) # keep the step's output in the test log
|
||||
|
||||
|
||||
def run_torchrun_in_process(cmd_parts: list[str], example_path: str) -> str:
|
||||
"""Drive ``torchrun`` itself in-process, letting it spawn fresh workers as usual.
|
||||
|
||||
Used for multi-rank steps. Only the launcher's imports are saved, but the workers are fresh
|
||||
processes, so none of the global state a reused process would inherit applies. Borrowed from
|
||||
Megatron-Bridge's own functional tests.
|
||||
"""
|
||||
from torch.distributed.run import main as torchrun_main
|
||||
|
||||
reset_megatron_global_state() # release device 0 before the workers claim it
|
||||
example_dir = MODELOPT_ROOT / "examples" / example_path
|
||||
script = next(p for p in cmd_parts if str(p).endswith(".py"))
|
||||
argv = [str(p) for p in cmd_parts]
|
||||
argv[argv.index(str(script))] = str(example_dir / Path(script).name)
|
||||
if not any(a.startswith(("--master_port", "--master-port")) for a in argv):
|
||||
# The rendezvous store now lives in a long-lived process, so do not rely on torchrun's
|
||||
# fixed default port being free by the time the next step starts.
|
||||
argv.insert(1, f"--master_port={get_free_port()}")
|
||||
cwd = os.getcwd()
|
||||
os.chdir(example_dir)
|
||||
cap = [""]
|
||||
try:
|
||||
with (
|
||||
_preserved_signal_handlers(),
|
||||
_capture_output() as cap,
|
||||
patch.object(sys, "argv", argv),
|
||||
):
|
||||
torchrun_main()
|
||||
return cap[0]
|
||||
except BaseException as e:
|
||||
# ChildFailedError only carries the worker's traceback when the entrypoint is decorated
|
||||
# with @record, which the examples are not -- so the captured output is all the caller has.
|
||||
e.captured_output = cap[0]
|
||||
raise
|
||||
finally:
|
||||
os.chdir(cwd)
|
||||
print(cap[0]) # keep the step's output in the test log
|
||||
|
||||
|
||||
def run_example_step(cmd_parts: list[str], example_path: str) -> str:
|
||||
"""Run an example step without shelling out.
|
||||
|
||||
Every step must be ``torchrun --nproc_per_node=<int> <script>.py``, with the script exposing
|
||||
``get_args()`` + ``main()``. Anything else raises: falling back to a subprocess would still
|
||||
pass, only ~6x slower, so a step that breaks the convention has to fail instead of quietly
|
||||
costing the suite its speed-up.
|
||||
"""
|
||||
script = next((str(p) for p in cmd_parts if str(p).endswith(".py")), None)
|
||||
if script is None:
|
||||
raise AssertionError(f"step must invoke a .py script: {cmd_parts}")
|
||||
world_size = requested_world_size(cmd_parts)
|
||||
if world_size is None:
|
||||
raise AssertionError(f"--nproc_per_node must be a plain integer: {cmd_parts}")
|
||||
if world_size > 1:
|
||||
# torchrun imports the script in fresh children, so get_args()/main() need not exist here.
|
||||
return run_torchrun_in_process(cmd_parts, example_path)
|
||||
_require_drivable(script, example_path)
|
||||
return run_example_in_process(cmd_parts, example_path)
|
||||
@@ -20,6 +20,7 @@ import signal
|
||||
import subprocess
|
||||
import threading
|
||||
import time
|
||||
import traceback
|
||||
import warnings
|
||||
from pathlib import Path
|
||||
|
||||
@@ -136,6 +137,17 @@ def _run_capturing(cmd_parts: list[str], cwd: Path, env: dict[str, str]) -> tupl
|
||||
return returncode, "".join(chunks)
|
||||
|
||||
|
||||
# Set by a suite's conftest to run a command without spawning a subprocess (see
|
||||
# ``_test_utils.examples.megatron_example_runner``). Unset means always use a subprocess.
|
||||
_in_process_runner = None
|
||||
|
||||
|
||||
def set_in_process_runner(runner) -> None:
|
||||
"""Register ``runner(cmd_parts, example_path) -> str``; it must not fall back to a subprocess."""
|
||||
global _in_process_runner
|
||||
_in_process_runner = runner
|
||||
|
||||
|
||||
def run_example_command(
|
||||
cmd_parts: list[str],
|
||||
example_path: str,
|
||||
@@ -146,18 +158,48 @@ def run_example_command(
|
||||
) -> str | None:
|
||||
"""Run an example command, retrying transient HuggingFace access errors."""
|
||||
print(f"[{example_path}] Running command: {cmd_parts}")
|
||||
env = env or os.environ.copy()
|
||||
in_process = _in_process_runner if env is None else None
|
||||
if _in_process_runner is not None and env is not None:
|
||||
# The in-process runner uses the ambient environment, so a caller-supplied env would be
|
||||
# silently dropped. Fall back to a subprocess rather than run with the wrong environment.
|
||||
warnings.warn(f"[{example_path}] env= given; running this step as a subprocess")
|
||||
env = os.environ.copy() if env is None else env
|
||||
cwd = MODELOPT_ROOT / "examples" / example_path
|
||||
|
||||
for attempt in range(hf_max_retries + 1):
|
||||
if setup_free_port:
|
||||
# Subprocess steps only: an in-process runner picks its own free port, since env is a
|
||||
# copy it never sees.
|
||||
env["MASTER_PORT"] = str(get_free_port()) # fresh port per attempt
|
||||
returncode, output = _run_capturing(cmd_parts, cwd, env)
|
||||
if in_process is not None:
|
||||
# Inside the loop so in-process steps get the same transient-HuggingFace retries;
|
||||
# these tests do hit the Hub (e.g. calib_dataset_name="cnn_dailymail").
|
||||
try:
|
||||
result = in_process(cmd_parts, example_path)
|
||||
except Exception as e:
|
||||
# Re-raise unless it looks transient, so a real failure keeps its traceback
|
||||
# instead of being flattened into CalledProcessError.
|
||||
# format_exc() walks __cause__/__context__: most markers are exception type
|
||||
# names that only ever appear in a formatted traceback, never in str(e) -- a Hub
|
||||
# ConnectionError typically surfaces wrapped in a DatasetGenerationError. The
|
||||
# runner also attaches what the step printed, for launcher-level failures.
|
||||
captured = getattr(e, "captured_output", "")
|
||||
text = f"{traceback.format_exc()}\n{captured}"
|
||||
if attempt == hf_max_retries or not any(
|
||||
marker in text for marker in _HF_TRANSIENT_MARKERS
|
||||
):
|
||||
raise
|
||||
returncode, output = 1, text
|
||||
else:
|
||||
assert result is not None, "an in-process runner must return output, not fall back"
|
||||
return result
|
||||
else:
|
||||
returncode, output = _run_capturing(cmd_parts, cwd, env)
|
||||
if returncode == 0:
|
||||
return output
|
||||
transient = any(marker in output for marker in _HF_TRANSIENT_MARKERS)
|
||||
if not transient or attempt == hf_max_retries:
|
||||
raise subprocess.CalledProcessError(returncode, cmd_parts)
|
||||
raise subprocess.CalledProcessError(returncode, cmd_parts, output=output)
|
||||
warnings.warn(
|
||||
f"[{example_path}] transient HuggingFace access error; retrying in "
|
||||
f"{hf_retry_delay_s}s (attempt {attempt + 1}/{hf_max_retries})"
|
||||
|
||||
@@ -0,0 +1,57 @@
|
||||
# SPDX-FileCopyrightText: Copyright (c) 2024 NVIDIA CORPORATION & AFFILIATES. All rights reserved.
|
||||
# SPDX-License-Identifier: Apache-2.0
|
||||
#
|
||||
# Licensed under the Apache License, Version 2.0 (the "License");
|
||||
# you may not use this file except in compliance with the License.
|
||||
# You may obtain a copy of the License at
|
||||
#
|
||||
# http://www.apache.org/licenses/LICENSE-2.0
|
||||
#
|
||||
# Unless required by applicable law or agreed to in writing, software
|
||||
# distributed under the License is distributed on an "AS IS" BASIS,
|
||||
# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
|
||||
# See the License for the specific language governing permissions and
|
||||
# limitations under the License.
|
||||
"""Run example steps without a ``torchrun`` launch; see ``_test_utils.examples.megatron_example_runner``."""
|
||||
|
||||
import os
|
||||
|
||||
import pytest
|
||||
from _test_utils.examples.megatron_example_runner import (
|
||||
reset_megatron_global_state,
|
||||
run_example_step,
|
||||
)
|
||||
from _test_utils.examples.run_command import set_in_process_runner
|
||||
|
||||
|
||||
@pytest.fixture(autouse=True)
|
||||
def _fast_example_runner():
|
||||
"""Run example steps without shelling out to ``torchrun``.
|
||||
|
||||
Per test, not per session: the hook is a module-global in ``run_command``, and this runner
|
||||
raises rather than falling back, so leaving it installed would break any other example suite
|
||||
collected later in the same session (``pytest tests/examples``).
|
||||
"""
|
||||
set_in_process_runner(run_example_step)
|
||||
try:
|
||||
yield
|
||||
finally:
|
||||
set_in_process_runner(None)
|
||||
|
||||
|
||||
@pytest.fixture(autouse=True)
|
||||
def _isolate_megatron_global_state():
|
||||
"""Reset shared state around every test so a failure cannot cascade into the next one.
|
||||
|
||||
In-process steps share the interpreter. Besides Megatron's singletons, Transformer-Engine
|
||||
records its chosen attention backend in ``NVTE_*``, which failed a Mamba hybrid that ran after
|
||||
an attention model -- so the environment is restored wholesale rather than by naming variables.
|
||||
"""
|
||||
env_before = os.environ.copy()
|
||||
reset_megatron_global_state()
|
||||
try:
|
||||
yield
|
||||
finally:
|
||||
reset_megatron_global_state()
|
||||
os.environ.clear()
|
||||
os.environ.update(env_before)
|
||||
Reference in New Issue
Block a user