diff --git a/soup_cli/commands/train.py b/soup_cli/commands/train.py index d961dac..a2c24da 100644 --- a/soup_cli/commands/train.py +++ b/soup_cli/commands/train.py @@ -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( diff --git a/tests/test_v0560.py b/tests/test_v0560.py index 48ebfe9..fe1460e 100644 --- a/tests/test_v0560.py +++ b/tests/test_v0560.py @@ -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: