| |
| """Check required artifacts before a training run starts.""" |
|
|
| from __future__ import annotations |
|
|
| import argparse |
| from pathlib import Path |
| from typing import Any |
|
|
| import yaml |
|
|
|
|
| DEFAULT_BASELINE_REPORT = Path("reports/cybergym/qwen36_27b_base/cybergym_base_level1_sample.md") |
| DEFAULT_FROZEN_TASKS = Path("reports/cybergym/frozen_level1_baseline_tasks.txt") |
|
|
|
|
| def parse_args() -> argparse.Namespace: |
| parser = argparse.ArgumentParser() |
| parser.add_argument("--config", required=True, help="Training config YAML.") |
| parser.add_argument( |
| "--allow-missing-cybergym-baseline", |
| action="store_true", |
| help="Allow local dry-runs before the remote CyberGym baseline exists.", |
| ) |
| return parser.parse_args() |
|
|
|
|
| def read_yaml(path: Path) -> dict[str, Any]: |
| with path.open("r", encoding="utf-8") as fh: |
| payload = yaml.safe_load(fh) or {} |
| if not isinstance(payload, dict): |
| raise TypeError(f"Expected a YAML mapping in {path}") |
| return payload |
|
|
|
|
| def existing_path(path: Path) -> Path | None: |
| candidates = [path] |
| if not path.is_absolute(): |
| candidates.append(Path("/workspace") / path) |
| candidates.append(Path("/workspace/infosec") / path) |
| for candidate in candidates: |
| if candidate.is_file(): |
| return candidate |
| return None |
|
|
|
|
| def main() -> int: |
| args = parse_args() |
| config_path = Path(args.config) |
| config = read_yaml(config_path) |
|
|
| missing: list[str] = [] |
| warnings: list[str] = [] |
|
|
| run_stage = str(config.get("run", {}).get("stage", "unknown")) |
| baseline_required = run_stage in {"pilot_qlora", "stage1_lora_sft", "stage1_dpo"} |
|
|
| gates = config.get("gates", {}) or {} |
| baseline_report = Path(gates.get("cybergym_baseline_report", DEFAULT_BASELINE_REPORT)) |
| frozen_tasks = Path(gates.get("frozen_cybergym_tasks", DEFAULT_FROZEN_TASKS)) |
|
|
| if baseline_required and not args.allow_missing_cybergym_baseline: |
| if existing_path(baseline_report) is None: |
| missing.append(str(baseline_report)) |
| if existing_path(frozen_tasks) is None: |
| missing.append(str(frozen_tasks)) |
|
|
| data = config.get("data", {}) or {} |
| for key in ("train_jsonl", "validation_jsonl"): |
| value = data.get(key) |
| if value and existing_path(Path(value)) is None: |
| missing.append(str(value)) |
|
|
| if data.get("decontamination_required"): |
| report = data.get("decontamination_report") |
| if not report: |
| missing.append("data.decontamination_report (decontamination_required=true but unset)") |
| elif existing_path(Path(report)) is None: |
| missing.append(str(report)) |
|
|
| print(f"Config: {config_path}") |
| print(f"Stage: {run_stage}") |
|
|
| if warnings: |
| print("Warnings:") |
| for warning in warnings: |
| print(f" - {warning}") |
|
|
| if missing: |
| print("Missing required files:") |
| for path in missing: |
| print(f" - {path}") |
| return 2 |
|
|
| print("All configured training gates passed.") |
| return 0 |
|
|
|
|
| if __name__ == "__main__": |
| raise SystemExit(main()) |
|
|
