From a3810823d14f656afb8ee0070c269c9be8f0464f Mon Sep 17 00:00:00 2001
From: Alpamys <vpn.alpamys@gmail.com>
Date: Fri, 15 May 2026 17:29:26 +0500
Subject: [PATCH] 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.
---
 soup_cli/commands/train.py | 16 ++++++++++++++--
 tests/test_v0560.py        |  9 +++++++++
 2 files changed, 23 insertions(+), 2 deletions(-)

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: