Guard diagnose gate on distributed worker ranks (#169)

Co-authored-by: mzl2233 <mzl2233@users.noreply.github.com>
This commit is contained in:
Yixuan Xu 2026-05-15 20:27:54 +08:00 committed by GitHub
parent 7d81496c69
commit 4c2a578ac0
No known key found for this signature in database
GPG Key ID: B5690EEEBB952194
2 changed files with 21 additions and 2 deletions

View File

@ -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

View File

@ -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: