File size: 3,128 Bytes
994182c | 1 2 3 4 5 6 7 8 9 10 11 12 13 14 15 16 17 18 19 20 21 22 23 24 25 26 27 28 29 30 31 32 33 34 35 36 37 38 39 40 41 42 43 44 45 46 47 48 49 50 51 52 53 54 55 56 57 58 59 60 61 62 63 64 65 66 67 68 69 70 71 72 73 74 75 76 77 78 79 80 81 82 83 84 85 86 87 88 89 90 91 92 93 94 95 96 97 98 99 100 101 | #!/usr/bin/env python3
"""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())
|