Repository navigation
[1/8] launcher: make the regression gate machine-readable and opt-in strict #2839
There are no files selected for viewing
| Original file line number | Diff line number | Diff line change |
|---|---|---|
|
|
@@ -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 [--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, 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 <output_dir>/checkpoint-<step>/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.""" | ||
| checkpoint_states = sorted( | ||
| glob(os.path.join(output_dir, "checkpoint-*", "trainer_state.json")), | ||
| key=_checkpoint_step, | ||
| ) | ||
| if checkpoint_states: | ||
| return checkpoint_states[-1] | ||
|
Comment on lines
+70
to
75
Contributor
There was a problem hiding this comment. [SUGGESTION] The step-number sort fixes ordering among checkpoints, but the final state in Any numeric Since the point of the change is "grade the newest state", consider making the choice uniform by comparing
Contributor
Author
There was a problem hiding this comment. The HF Trainer saves a checkpoint at the final step whatever save_steps is, so the newest numbered checkpoint never lags a completed run's final state. I have tested this example (save_steps=5000, max_steps=2000). It produced checkpoint-2000, whose trainer_state.json has the same global_step (2000) as <output_dir>/trainer_state.json, and the gate graded checkpoint-2000. If training stops early, save_state() never runs, so the output_dir copy isn't newer in that case either. |
||
| # Fall back to output_dir itself | ||
|
|
@@ -81,8 +105,8 @@ def get_final_metrics(trainer_state_path): | |
| return metrics | ||
|
|
||
|
|
||
| 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: | ||
| if env_var not in thresholds: | ||
| continue | ||
| 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 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) | ||
| 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: | ||
|
Comment on lines
+199
to
207
Contributor
There was a problem hiding this comment. [SUGGESTION] Two unexpected-exception paths produce no verdict file, which the non-strict branch then swallows.
In both cases the process exits non-zero with no
Contributor
Author
There was a problem hiding this comment. Fixed the metrics read in 9a5904c. TypeError/AttributeError from a well-formed but wrong-shaped trainer_state.json now give METRICS_MISSING, with tests for [] and "loss": null. |
||
| # 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 | ||
|
|
||
|
|
||
| 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']}") | ||
|
Comment on lines
+238
to
+239
Contributor
There was a problem hiding this comment. [SUGGESTION] The human-readable log now prints the JSON-sanitized metrics, so a diverged run emits two contradictory lines:
The sanitization is only needed for # in evaluate(): keep the raw values
verdict.update({"pass": passed, "metrics": metrics, "messages": messages})
# in main(), at the write:
if args.json_path:
# Strict JSON: NaN/Infinity would break non-Python readers of the verdict.
serializable = dict(verdict)
serializable["metrics"] = {
k: v if math.isfinite(v) else None for k, v in verdict["metrics"].items()
}
with open(args.json_path, "w") as f:
json.dump(serializable, f, indent=2, allow_nan=False)This also keeps
Contributor
Author
There was a problem hiding this comment. Applied in 4101423. |
||
| 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__": | ||
| sys.exit(main()) | ||
There was a problem hiding this comment.
Choose a reason for hiding this comment
The reason will be displayed to describe this comment to others. Learn more.
Check checkpoint subdirs first, by numeric step: a lexicographic sort would rank checkpoint-8000 after checkpoint-12000.