refactor(train): harden diagnose-gate rank guard from PR #169

PR #169 wired LOCAL_RANK==0 guard on _run_diagnose_gate so distributed
launches only run the gate on one worker per machine. Two minor polish
items on top of the merged version:

- Wrap the int() parse in try/except ValueError. A malformed LOCAL_RANK
  (garbage value from a misconfigured launcher) would previously crash
  the post-training gate. Falling back to True is safer than silently
  skipping the gate -- over-running is recoverable, under-running hides
  failures.
- Expand the docstring to explain why we use LOCAL_RANK (per-machine)
  rather than RANK (global): the gate reads the local output_dir, so
  one gate per machine is the right granularity for typical single-
  machine multi-GPU runs. Documents the choice for future readers.
- Add a focused test (test_diagnose_gate_handles_malformed_local_rank)
  asserting the safe fallback path.
This commit is contained in:
Alpamys 2026-05-15 17:29:26 +05:00
parent 4c2a578ac0
commit a3810823d1
2 changed files with 23 additions and 2 deletions

View File

@ -952,8 +952,20 @@ def train(
def _should_run_diagnose_gate_on_rank() -> bool:
"""Return true only for rank 0 in distributed launches."""
return int(os.environ.get("LOCAL_RANK", "0")) == 0
"""Return True only for LOCAL_RANK=0 in distributed launches.
Uses LOCAL_RANK (per-machine rank) -- not RANK (global rank across all
nodes) -- because the diagnose gate reads the local training output
directory. We want one gate per machine, not one across the whole
cluster. For typical single-machine multi-GPU runs both are equivalent.
Defaults to True (run gate) on any parse error: a malformed env var is
safer to over-run than to silently skip.
"""
try:
return int(os.environ.get("LOCAL_RANK", "0")) == 0
except ValueError:
return True
def _run_diagnose_gate(

View File

@ -721,6 +721,15 @@ class TestTrainDiagnoseGate:
monkeypatch.setenv("LOCAL_RANK", "0")
assert _should_run_diagnose_gate_on_rank() is True
def test_diagnose_gate_handles_malformed_local_rank(
self, monkeypatch: pytest.MonkeyPatch,
) -> None:
"""Garbage LOCAL_RANK falls back to True -- safer to over-run than skip."""
from soup_cli.commands.train import _should_run_diagnose_gate_on_rank
monkeypatch.setenv("LOCAL_RANK", "not-an-int")
assert _should_run_diagnose_gate_on_rank() is True
def test_run_diagnose_gate_rejects_non_dict_payload(
self, tmp_path: Path
) -> None: