diff --git a/plugins/modelopt/skills/speculative-decoding/references/algorithms/dflash.md b/plugins/modelopt/skills/speculative-decoding/references/algorithms/dflash.md index 94909bbe946..b0b72b7f45c 100644 --- a/plugins/modelopt/skills/speculative-decoding/references/algorithms/dflash.md +++ b/plugins/modelopt/skills/speculative-decoding/references/algorithms/dflash.md @@ -115,7 +115,7 @@ Export is automatic: after training, rank 0 exports every `checkpoint-` to | Task | Log evidence | Artifact | | --- | --- | --- | | Hidden-state dump (offline) | vLLM extraction completes over the input data | Dump directory populated | -| Training | `Training time: N seconds`, then `=== Exporting: → ===` and `=== Regression Check (...) ===` | `/checkpoint-*/trainer_state.json`, `/exported-checkpoint-*` | +| Training | `Training time: N seconds`, then `=== Exporting: → ===` and `=== Regression Check (...) ===` | `/checkpoint-*/trainer_state.json`, `/exported-checkpoint-*`, `/regression_check.json` | | Smoke test | `Auto-detected draft model: ...`, `Server ready after Ns` | Smoke-test log with responses | | AR eval | Per-category MT-Bench AR output | AR results | | Benchmark | `Average_AL` in the saved results | JSON under `--save_dir` | @@ -137,9 +137,12 @@ convergence baseline (8×B200, bs=1, seq_len=4096, 5-layer draft, block_size=16, samples, 1 epoch ≈ 12,500 steps) is in the YAML header — compare against it when judging whether a run under-trained. -Note: `check_regression.py` is invoked with `|| true`, and it only warns when no -`trainer_state.json` exists. A green Slurm exit is not proof the gate ran — confirm -the `=== Regression Check ===` block is present in the log. +Judge the run by `/regression_check.json`: it passed only if `pass` is true +and `thresholds` is non-empty; no file means the check never ran (training failed). +`failure_class` is `TRAIN_REGRESSION` (a metric missed its threshold or is NaN/inf) or +`METRICS_MISSING` (no readable metrics, or an invalid threshold). Only +`REGRESSION_STRICT=1` (DFlash2 online) fails the task on either; elsewhere a green Slurm +exit proves nothing. Any value but `0`/`1` fails before training. **2. Smoke test** — `MIN_ACCEPTANCE_LENGTH` env var (Qwen3-8B online uses 1.4) with `NUM_SPEC_TOKENS` speculative tokens. diff --git a/plugins/modelopt/skills/speculative-decoding/references/algorithms/dflash2.md b/plugins/modelopt/skills/speculative-decoding/references/algorithms/dflash2.md index 66f59bbaa3c..dfab614aecb 100644 --- a/plugins/modelopt/skills/speculative-decoding/references/algorithms/dflash2.md +++ b/plugins/modelopt/skills/speculative-decoding/references/algorithms/dflash2.md @@ -124,9 +124,11 @@ task_2 (`vllm_smoke_test.sh`, gate `MIN_ACCEPTANCE_LENGTH=1.2`); the online exam ships no such step, so benchmark its export separately on the serving stack (vLLM PR #52816) or via the offline acceptance-length harness. -The training-regression gate from `dflash.md` (`MAX_FINAL_LOSS`, `MIN_FINAL_ACC` via -`check_regression.py`) applies; the online example sets `MAX_FINAL_LOSS=5.0` and -`MIN_FINAL_ACC=0.15` for its 2000-step convergence check. +The training-regression gate from `dflash.md` applies. The online example sets +`MAX_FINAL_LOSS=7.0` (DFlash2's loss runs higher than DFlash's, so 5.0 fails healthy +runs) and `MIN_FINAL_ACC=0.15` for its 2000-step convergence check, and runs it strict +(`REGRESSION_STRICT=1`): a missed threshold, or no metrics to check, fails the +training task. ## Known failures diff --git a/tools/launcher/common/check_regression.py b/tools/launcher/common/check_regression.py index 08002336dc8..47292de44f8 100644 --- a/tools/launcher/common/check_regression.py +++ b/tools/launcher/common/check_regression.py @@ -26,27 +26,51 @@ passing unevaluated. MAX_FINAL_PERPLEXITY: Final perplexity must be below this value +An empty variable counts as unset, and one that isn't a finite number fails the check. +Once any threshold is set, a missing or unreadable trainer_state.json, or one without +training logs, also fails the check, and so does a NaN or infinite value for a metric +that has a threshold. + Usage: - python check_regression.py /path/to/output_dir + python check_regression.py /path/to/output_dir [--json verdict.json] Or from a shell script: python common/check_regression.py ${OUTPUT_DIR} + --json also writes the verdict as {"pass", "failure_class", "trainer_state", + "metrics", "thresholds", "messages"}, so a pipeline gate can read it without + parsing logs. failure_class is null, "TRAIN_REGRESSION" or "METRICS_MISSING". It is + strict JSON: a non-finite metric is written as null. + Exit codes: 0 — all checks pass (or no thresholds set) - 1 — regression detected + 1 — regression detected, or thresholds set but no metrics to check """ +import argparse import json +import math import os import sys from glob import glob +from typing import Any + +THRESHOLD_VARS = ("MAX_FINAL_LOSS", "MIN_FINAL_ACC", "MAX_FINAL_PERPLEXITY") + + +def _checkpoint_step(trainer_state_path): + """Step number of /checkpoint-/trainer_state.json, or -1.""" + name = os.path.basename(os.path.dirname(trainer_state_path)) + step = name.removeprefix("checkpoint-") + return int(step) if step.isdigit() else -1 def find_trainer_state(output_dir): """Find the latest trainer_state.json in the output directory.""" - # Check checkpoint subdirs first (sorted by step number) - checkpoint_states = sorted(glob(os.path.join(output_dir, "checkpoint-*", "trainer_state.json"))) + checkpoint_states = sorted( + glob(os.path.join(output_dir, "checkpoint-*", "trainer_state.json")), + key=_checkpoint_step, + ) if checkpoint_states: return checkpoint_states[-1] # Fall back to output_dir itself @@ -81,8 +105,8 @@ def get_final_metrics(trainer_state_path): return metrics -def check_regression(metrics): - """Check metrics against environment variable thresholds. Returns (passed, messages).""" +def check_regression(metrics, thresholds): + """Check metrics against {env_var: threshold}. Returns (passed, messages).""" checks = [ ( "MAX_FINAL_LOSS", @@ -108,10 +132,9 @@ def check_regression(metrics): messages = [] for env_var, metric_key, check_fn, fail_msg in checks: - thresh_str = os.environ.get(env_var) - if thresh_str is None: + if env_var not in thresholds: continue - thresh = float(thresh_str) + thresh = thresholds[env_var] val = metrics.get(metric_key) if val is None: # Setting a threshold asserts the run must be checked against it. If the @@ -123,6 +146,10 @@ def check_regression(metrics): ) passed = False continue + if not math.isfinite(val): + messages.append(f"REGRESSION: {metric_key} is {val}, not a finite number") + passed = False + continue if check_fn(val, thresh): messages.append(f"PASS: {metric_key}={val:.3f} (threshold: {env_var}={thresh})") else: @@ -132,40 +159,96 @@ def check_regression(metrics): return passed, messages -def main(): - """Entry point for regression check CLI.""" - if len(sys.argv) < 2: - print("Usage: python check_regression.py ") - sys.exit(0) - - output_dir = sys.argv[1] - - # Skip if no thresholds set - if not any( - os.environ.get(v) for v in ["MAX_FINAL_LOSS", "MIN_FINAL_ACC", "MAX_FINAL_PERPLEXITY"] - ): - return +def parse_thresholds(environ): + """Returns ({env_var: value}, [error]) for the set threshold variables.""" + thresholds, errors = {}, [] + for v in THRESHOLD_VARS: + raw = environ.get(v) + if not raw: + continue + try: + value = float(raw) + except ValueError: + value = math.nan + if math.isfinite(value): + thresholds[v] = value + else: + errors.append(f"REGRESSION: {v}={raw!r} is not a finite number") + return thresholds, errors + + +def evaluate(output_dir): + """Run the check on output_dir. Returns the verdict that --json writes.""" + thresholds, errors = parse_thresholds(os.environ) + verdict: dict[str, Any] = { + "pass": True, + "failure_class": None, + "trainer_state": None, + "metrics": {}, + "thresholds": thresholds, + "messages": [], + } + if errors: + # The recipe's thresholds are unusable, so the run cannot be checked against them. + verdict.update({"pass": False, "failure_class": "METRICS_MISSING", "messages": errors}) + return verdict + if not thresholds: + return verdict trainer_state = find_trainer_state(output_dir) - if not trainer_state: - print(f"WARNING: No trainer_state.json found in {output_dir}, skipping regression check") - return - - print(f"=== Regression Check ({trainer_state}) ===") - metrics = get_final_metrics(trainer_state) + verdict["trainer_state"] = trainer_state + where = trainer_state or f"{output_dir} (no trainer_state.json)" + try: + metrics = get_final_metrics(trainer_state) if trainer_state else {} + missing = f"no training logs to check in {where}" + # TypeError/AttributeError: valid JSON of the wrong shape, e.g. "loss": null or a list. + except (OSError, ValueError, TypeError, AttributeError) as e: + metrics, missing = {}, f"cannot read {where}: {e}" if not metrics: - print("No training logs found in trainer_state.json") - return + # Thresholds were set, so an unreadable run must not pass as unchecked. + verdict.update( + { + "pass": False, + "failure_class": "METRICS_MISSING", + "messages": [f"REGRESSION: {missing}"], + } + ) + return verdict + + passed, messages = check_regression(metrics, thresholds) + verdict.update({"pass": passed, "metrics": metrics, "messages": messages}) + if not passed: + verdict["failure_class"] = "TRAIN_REGRESSION" + return verdict - print(f"Final metrics: {metrics}") - passed, messages = check_regression(metrics) - for msg in messages: - print(f" {msg}") - if not passed: - sys.exit(1) - print("Regression check PASSED") +def main(argv=None): + """Entry point for regression check CLI.""" + parser = argparse.ArgumentParser(description="Regression check for training jobs.") + parser.add_argument("output_dir", nargs="?", help="Training output directory") + parser.add_argument("--json", dest="json_path", help="Create a JSON file storing the verdict") + args = parser.parse_args(argv) + if args.output_dir is None: + parser.print_usage() + return 0 + + verdict = evaluate(args.output_dir) + if verdict["thresholds"] or verdict["messages"]: + print(f"=== Regression Check ({verdict['trainer_state'] or args.output_dir}) ===") + if verdict["metrics"]: + print(f"Final metrics: {verdict['metrics']}") + for msg in verdict["messages"]: + print(f" {msg}") + print(f"Regression check {'PASSED' if verdict['pass'] else 'FAILED'}") + + if args.json_path: + # Strict JSON, serialized before the file is opened so a failure cannot truncate it. + metrics = {k: v if math.isfinite(v) else None for k, v in verdict["metrics"].items()} + text = json.dumps({**verdict, "metrics": metrics}, indent=2, allow_nan=False) + with open(args.json_path, "w") as f: + f.write(text) + return 0 if verdict["pass"] else 1 if __name__ == "__main__": - main() + sys.exit(main()) diff --git a/tools/launcher/common/specdec/dflash_online_training.sh b/tools/launcher/common/specdec/dflash_online_training.sh index efc38be3b17..e70800e069a 100644 --- a/tools/launcher/common/specdec/dflash_online_training.sh +++ b/tools/launcher/common/specdec/dflash_online_training.sh @@ -38,6 +38,25 @@ SCRIPT_DIR="$(dirname "$(readlink -f "$0")")" source ${SCRIPT_DIR}/../service_utils.sh +OUTPUT_DIR=$(python3 -c " +import sys +for arg in sys.argv[1:]: + if arg.startswith('training.output_dir='): + print(arg.split('=', 1)[1]) + break +" "$@") +# A verdict left in a reused output_dir by an earlier run must never be read as this run's, +# so remove it before anything below can stop the script. +if [ "${SLURM_PROCID:-0}" = "0" ] && [ -n "$OUTPUT_DIR" ]; then + rm -f "${OUTPUT_DIR}/regression_check.json" +fi + +# Checked before any work: a typo such as "true" must not cost a full training run. +case "${REGRESSION_STRICT:-0}" in + 0|1) ;; + *) echo "ERROR: REGRESSION_STRICT must be 0 or 1, got '${REGRESSION_STRICT}'" >&2; exit 1 ;; +esac + pip install -r modules/Model-Optimizer/examples/speculative_decoding/requirements.txt pip install huggingface-hub>=1.2.1 export PATH=$PATH:/workspace/.local/bin @@ -117,20 +136,23 @@ fi set -x start_time=$(date +%s) -accelerate launch $ACCEL_CONFIG_ARGS --mixed_precision "${MIXED_PRECISION:-bf16}" $MULTI_NODE_ARGS $MAIN_PY "$@" +# The ERR trap only records a failure and lets the script go on, so test the status here: +# a crashed run must not export or write a regression verdict from whatever checkpoints +# and trainer_state.json it, or an earlier run, left behind. +if accelerate launch $ACCEL_CONFIG_ARGS --mixed_precision "${MIXED_PRECISION:-bf16}" $MULTI_NODE_ARGS $MAIN_PY "$@"; then + TRAIN_STATUS=0 +else + # The if suppresses the ERR trap, so report as it would; $? is still the launch status. + error_handler "$0" "$LINENO" + TRAIN_STATUS=1 +fi echo "Training time: $(( $(date +%s) - start_time )) seconds" set +x +if [ "$TRAIN_STATUS" != 0 ]; then + echo "ERROR: training failed; skipping export and regression check." >&2 # Export last checkpoint to deployment format (rank 0 only, single GPU) -if [ "${SLURM_PROCID:-0}" = "0" ]; then - OUTPUT_DIR=$(python3 -c " -import sys -for arg in sys.argv[1:]: - if arg.startswith('training.output_dir='): - print(arg.split('=', 1)[1]) - break -" "$@") - +elif [ "${SLURM_PROCID:-0}" = "0" ]; then if [ -n "$OUTPUT_DIR" ]; then # Export all checkpoints that don't already have an exported version EXPORTED=0 @@ -162,8 +184,15 @@ for arg in sys.argv[1:]: fi fi - # Regression check (uses env vars MAX_FINAL_LOSS, MIN_FINAL_ACC, etc.) + # Regression check (uses env vars MAX_FINAL_LOSS, MIN_FINAL_ACC, etc.). The verdict + # lands next to the exports for downstream gates. A regression only warns unless + # REGRESSION_STRICT=1, which fails the task (via the ERR trap) so Slurm reports it. if [ -n "$OUTPUT_DIR" ]; then - python3 common/check_regression.py "${OUTPUT_DIR}" || true + REGRESSION_ARGS=("${OUTPUT_DIR}" --json "${OUTPUT_DIR}/regression_check.json") + if [ "${REGRESSION_STRICT:-0}" = 1 ]; then + python3 common/check_regression.py "${REGRESSION_ARGS[@]}" + else + python3 common/check_regression.py "${REGRESSION_ARGS[@]}" || true + fi fi fi diff --git a/tools/launcher/examples/Qwen/Qwen3-8B/hf_online_dflash2.yaml b/tools/launcher/examples/Qwen/Qwen3-8B/hf_online_dflash2.yaml index e300e0665b4..6d08abba99c 100644 --- a/tools/launcher/examples/Qwen/Qwen3-8B/hf_online_dflash2.yaml +++ b/tools/launcher/examples/Qwen/Qwen3-8B/hf_online_dflash2.yaml @@ -98,8 +98,11 @@ pipeline: # Sliding-window draft attention, matching the published DFlash2 drafters. - dflash.dflash_swa_window_size=2048 environment: - - MAX_FINAL_LOSS: "5.0" + # DFlash2's loss runs higher than DFlash's; 5.0 fails healthy runs. + - MAX_FINAL_LOSS: "7.0" - MIN_FINAL_ACC: "0.15" + # Fail the task on a regression instead of only logging it. + - REGRESSION_STRICT: "1" slurm_config: _factory_: "slurm_factory" nodes: 1 diff --git a/tools/launcher/tests/test_check_regression.py b/tools/launcher/tests/test_check_regression.py index 93c7998a837..0832fced283 100644 --- a/tools/launcher/tests/test_check_regression.py +++ b/tools/launcher/tests/test_check_regression.py @@ -16,6 +16,7 @@ """Regression-gate behaviour, especially the unevaluated-threshold case.""" import importlib.util +import json from pathlib import Path import pytest @@ -33,30 +34,40 @@ def _clear_thresholds(monkeypatch): monkeypatch.delenv(var, raising=False) -def test_threshold_with_no_metric_fails(monkeypatch): +def test_threshold_with_no_metric_fails(): """A configured threshold whose metric never reached trainer_state.json must fail. The DFlash/EAGLE trainers print accuracy every log step but did not persist it, so this path used to warn and pass -- stamping PASS on a run nothing had checked. """ - monkeypatch.setenv("MIN_FINAL_ACC", "0.15") - passed, messages = check_regression_mod.check_regression({"loss": 0.4}) + passed, messages = check_regression_mod.check_regression({"loss": 0.4}, {"MIN_FINAL_ACC": 0.15}) assert not passed assert any("never evaluated" in m for m in messages) def test_unset_threshold_is_not_a_failure(): """Only a threshold the operator actually set is enforced.""" - passed, messages = check_regression_mod.check_regression({"loss": 0.4}) + passed, messages = check_regression_mod.check_regression({"loss": 0.4}, {}) assert passed assert not messages +def test_empty_threshold_var_is_unset(tmp_path, monkeypatch): + """An empty MIN_FINAL_ACC disables that check instead of crashing on float('').""" + monkeypatch.setenv("MAX_FINAL_LOSS", "5.0") + monkeypatch.setenv("MIN_FINAL_ACC", "") + (tmp_path / "trainer_state.json").write_text('{"log_history": [{"loss": 1.0}]}') + verdict = check_regression_mod.evaluate(str(tmp_path)) + assert verdict["pass"] + assert verdict["thresholds"] == {"MAX_FINAL_LOSS": 5.0} + + @pytest.mark.parametrize(("acc", "expected"), [(0.52, True), (0.10, False)]) -def test_accuracy_threshold_is_enforced_when_present(monkeypatch, acc, expected): +def test_accuracy_threshold_is_enforced_when_present(acc, expected): """A present metric is compared against the threshold in both directions.""" - monkeypatch.setenv("MIN_FINAL_ACC", "0.15") - passed, _ = check_regression_mod.check_regression({"loss": 0.4, "accuracy": acc}) + passed, _ = check_regression_mod.check_regression( + {"loss": 0.4, "accuracy": acc}, {"MIN_FINAL_ACC": 0.15} + ) assert passed is expected @@ -69,3 +80,56 @@ def test_metrics_pick_up_any_accuracy_key(tmp_path): metrics = check_regression_mod.get_final_metrics(state) assert metrics["loss"] == pytest.approx(0.39) assert metrics["accuracy"] == pytest.approx(0.5153) + + +def test_latest_checkpoint_is_chosen_by_step_number(tmp_path): + """checkpoint-12000 is newer than checkpoint-8000, though it sorts first as a string.""" + for step in (8000, 12000): + ckpt = tmp_path / f"checkpoint-{step}" + ckpt.mkdir() + (ckpt / "trainer_state.json").write_text("{}") + latest = check_regression_mod.find_trainer_state(str(tmp_path)) + assert Path(latest).parent.name == "checkpoint-12000" + + +@pytest.mark.parametrize( + ("state_text", "exit_code", "failure_class"), + [ + (None, 1, "METRICS_MISSING"), # no trainer_state.json at all + ('{"log_hist', 1, "METRICS_MISSING"), # truncated by a job killed mid-save + ("[]", 1, "METRICS_MISSING"), # valid JSON, wrong shape + ('{"log_history": [{"loss": null}]}', 1, "METRICS_MISSING"), + ('{"log_history": [{"loss": NaN}]}', 1, "TRAIN_REGRESSION"), # a diverged run + # An unchecked non-finite metric still serializes, as null. + ('{"log_history": [{"loss": 4.5, "acc": Infinity}]}', 0, None), + (json.dumps({"log_history": [{"loss": 6.5}]}), 1, "TRAIN_REGRESSION"), + (json.dumps({"log_history": [{"loss": 4.5}]}), 0, None), + ], +) +def test_cli_exit_code_and_json_verdict( + tmp_path, monkeypatch, state_text, exit_code, failure_class +): + """With a threshold set, the CLI fails both a regression and a run it cannot read.""" + monkeypatch.setenv("MAX_FINAL_LOSS", "5.0") + if state_text is not None: + ckpt = tmp_path / "checkpoint-2000" + ckpt.mkdir() + (ckpt / "trainer_state.json").write_text(state_text) + verdict_path = tmp_path / "verdict.json" + assert check_regression_mod.main([str(tmp_path), "--json", str(verdict_path)]) == exit_code + # Strict JSON: NaN/Infinity would break non-Python readers of the verdict. + verdict = json.loads(verdict_path.read_text(), parse_constant=pytest.fail) + assert verdict["pass"] is (exit_code == 0) + assert verdict["failure_class"] == failure_class + assert verdict["thresholds"] == {"MAX_FINAL_LOSS": 5.0} + + +@pytest.mark.parametrize("raw", ["7,0", "nan", "inf"]) +def test_invalid_threshold_writes_failed_verdict(tmp_path, monkeypatch, raw): + """A malformed or non-finite threshold fails as a readable verdict, not a traceback.""" + monkeypatch.setenv("MAX_FINAL_LOSS", raw) + verdict_path = tmp_path / "verdict.json" + assert check_regression_mod.main([str(tmp_path), "--json", str(verdict_path)]) == 1 + verdict = json.loads(verdict_path.read_text(), parse_constant=pytest.fail) + assert verdict["failure_class"] == "METRICS_MISSING" + assert verdict["thresholds"] == {} diff --git a/tools/launcher/tests/test_dflash_online_training.py b/tools/launcher/tests/test_dflash_online_training.py new file mode 100644 index 00000000000..1a065f48abb --- /dev/null +++ b/tools/launcher/tests/test_dflash_online_training.py @@ -0,0 +1,124 @@ +# SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. +# SPDX-License-Identifier: Apache-2.0 +# +# Licensed under the Apache License, Version 2.0 (the "License"); +# you may not use this file except in compliance with the License. +# You may obtain a copy of the License at +# +# http://www.apache.org/licenses/LICENSE-2.0 +# +# Unless required by applicable law or agreed to in writing, software +# distributed under the License is distributed on an "AS IS" BASIS, +# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +# See the License for the specific language governing permissions and +# limitations under the License. + +"""dflash_online_training.sh: the regression verdict, its strict mode and its validation.""" + +import json +import os +import shutil +import subprocess +from pathlib import Path + +import pytest + +_COMMON = Path(__file__).resolve().parents[1] / "common" +_EXPORT = "modules/Model-Optimizer/examples/speculative_decoding/scripts/export_hf_checkpoint.py" + + +def _run(tmp_path, train_exit=0, loss=4.0, **env_overrides): + """Run the script in a sandbox: stub pip/nvidia-smi/accelerate/export/torch, real gate. + + Returns (process, verdict or None, whether training was launched). + """ + (tmp_path / "common" / "specdec").mkdir(parents=True) + for name in ("service_utils.sh", "check_regression.py", "specdec/dflash_online_training.sh"): + shutil.copy(_COMMON / name, tmp_path / "common" / name) + (tmp_path / "modules/Model-Optimizer/modelopt").mkdir( + parents=True + ) # service_utils appends here + export = tmp_path / _EXPORT + export.parent.mkdir(parents=True) + export.write_text( + "import os, sys\n" + "d = sys.argv[sys.argv.index('--export_path') + 1]\n" + "os.makedirs(d, exist_ok=True)\n" + "open(os.path.join(d, 'config.json'), 'w').write('{}')\n" + ) + # The script counts GPUs with torch; without it that line trips the ERR trap and every + # run exits 1, which would hide the exit code under test. + pystub = tmp_path / "pystub" + pystub.mkdir() + (pystub / "torch.py").write_text( + "class cuda:\n @staticmethod\n def device_count():\n return 1\n" + ) + launched = tmp_path / "launched" + bin_dir = tmp_path / "bin" + bin_dir.mkdir() + for tool, body in ( + ("pip", "exit 0"), + ("nvidia-smi", "echo 1"), + ("accelerate", f"touch {launched}; exit {train_exit}"), + ): + (bin_dir / tool).write_text(f"#!/bin/bash\n{body}\n") + (bin_dir / tool).chmod(0o755) + + # A checkpoint with the given loss, and a passing verdict from an earlier run. + out = tmp_path / "out" + (out / "checkpoint-2000").mkdir(parents=True) + (out / "checkpoint-2000" / "trainer_state.json").write_text( + json.dumps({"log_history": [{"loss": loss}]}) + ) + (out / "regression_check.json").write_text('{"pass": true}') + + env = { + **os.environ, + "PATH": f"{bin_dir}:{os.environ['PATH']}", + "PYTHONPATH": str(pystub), + "MAX_FINAL_LOSS": "5.0", + **env_overrides, + } + env.pop("SLURM_PROCID", None) + proc = subprocess.run( + ["bash", "common/specdec/dflash_online_training.sh", f"training.output_dir={out}"], + cwd=tmp_path, + env=env, + capture_output=True, + text=True, + ) + verdict = out / "regression_check.json" + return proc, json.loads(verdict.read_text()) if verdict.exists() else None, launched.exists() + + +@pytest.mark.parametrize("train_exit", [0, 1]) +def test_failed_training_writes_no_verdict(tmp_path, train_exit): + """A crash must not let a surviving checkpoint, or an old verdict, pass the gate.""" + proc, verdict, _ = _run(tmp_path, train_exit) + if train_exit: + assert verdict is None, verdict + assert proc.returncode == 1, proc.stderr + assert "training failed; skipping export and regression check" in proc.stderr + else: # control: the same sandbox grades the checkpoint when training succeeds + assert verdict is not None, proc.stderr + assert verdict["pass"] is True + assert verdict["trainer_state"].endswith("checkpoint-2000/trainer_state.json") + assert proc.returncode == 0, proc.stderr + assert "[ERROR]" not in proc.stderr + + +@pytest.mark.parametrize(("strict", "exit_code"), [("1", 1), ("0", 0)]) +def test_strict_mode_fails_the_task_on_regression(tmp_path, strict, exit_code): + """REGRESSION_STRICT=1 turns a failed check into a failed task; 0 only records it.""" + proc, verdict, _ = _run(tmp_path, loss=6.0, REGRESSION_STRICT=strict) + assert verdict["failure_class"] == "TRAIN_REGRESSION", proc.stderr + assert proc.returncode == exit_code, proc.stderr + + +def test_invalid_regression_strict_fails_before_training(tmp_path): + """A typo fails at once, without training, and leaves no earlier run's verdict behind.""" + proc, verdict, launched = _run(tmp_path, REGRESSION_STRICT="true") + assert proc.returncode == 1 + assert "REGRESSION_STRICT must be 0 or 1, got 'true'" in proc.stderr + assert not launched + assert verdict is None