From 4c2a578ac0aa68b2b48a84096c0dcc99e51b3640 Mon Sep 17 00:00:00 2001
From: Yixuan Xu <109468061+mzl2233@users.noreply.github.com>
Date: Fri, 15 May 2026 20:27:54 +0800
Subject: [PATCH] Guard diagnose gate on distributed worker ranks (#169)

Co-authored-by: mzl2233 <mzl2233@users.noreply.github.com>
---
 soup_cli/commands/train.py | 11 +++++++++--
 tests/test_v0560.py        | 12 ++++++++++++
 2 files changed, 21 insertions(+), 2 deletions(-)

diff --git a/soup_cli/commands/train.py b/soup_cli/commands/train.py
index 26154f8..d961dac 100644
--- a/soup_cli/commands/train.py
+++ b/soup_cli/commands/train.py
@@ -937,7 +937,7 @@ def train(
     )
 
     # --- v0.56.0 --diagnose-gate: post-training failure-mode check ---
-    if diagnose_gate:
+    if diagnose_gate and _should_run_diagnose_gate_on_rank():
         try:
             _run_diagnose_gate(
                 diagnose_gate, run_id, cfg.base, result["output_dir"]
@@ -951,6 +951,11 @@ def train(
             raise typer.Exit(1) from exc
 
 
+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
+
+
 def _run_diagnose_gate(
     evidence_path: str, run_id: str, base: str, adapter: str
 ) -> None:
@@ -959,7 +964,9 @@ def _run_diagnose_gate(
     Loads a JSON ``evidence`` file with optional per-mode scores and
     refuses to mark the run successful if any mode comes back MAJOR.
     Missing modes fall back to a neutral OK score so partial evidence
-    still produces a useful report card.
+    still produces a useful report card. The train command only calls
+    this helper on LOCAL_RANK=0 so distributed runs do not execute the
+    gate once per worker.
     """
     import json
 
diff --git a/tests/test_v0560.py b/tests/test_v0560.py
index d9b387b..48ebfe9 100644
--- a/tests/test_v0560.py
+++ b/tests/test_v0560.py
@@ -709,6 +709,18 @@ class TestTrainDiagnoseGate:
         from soup_cli.commands.train import _run_diagnose_gate
         _run_diagnose_gate(str(evidence), "run1", "base", "adapter")
 
+    def test_diagnose_gate_skips_nonzero_local_rank(self, monkeypatch: pytest.MonkeyPatch) -> None:
+        from soup_cli.commands.train import _should_run_diagnose_gate_on_rank
+
+        monkeypatch.setenv("LOCAL_RANK", "1")
+        assert _should_run_diagnose_gate_on_rank() is False
+
+    def test_diagnose_gate_runs_on_rank_zero(self, monkeypatch: pytest.MonkeyPatch) -> None:
+        from soup_cli.commands.train import _should_run_diagnose_gate_on_rank
+
+        monkeypatch.setenv("LOCAL_RANK", "0")
+        assert _should_run_diagnose_gate_on_rank() is True
+
     def test_run_diagnose_gate_rejects_non_dict_payload(
         self, tmp_path: Path
     ) -> None: