File size: 7,486 Bytes
69928d6 | 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 102 103 104 105 106 107 108 109 110 111 112 113 114 115 116 117 118 119 120 121 122 123 124 125 126 127 128 129 130 131 132 133 134 135 136 137 138 139 140 141 142 143 144 145 146 147 148 149 150 151 152 153 154 155 156 157 158 159 160 161 162 163 164 165 166 167 168 169 170 171 172 173 174 175 176 177 178 179 180 | """Compose existing data, training, and raw evaluation stages without supplied weights."""
from __future__ import annotations
import json
import os
import shutil
import subprocess
from copy import deepcopy
from pathlib import Path
from typing import Any
from omegaconf import DictConfig, OmegaConf
from softasp_core.config.validation import validate_config
from softasp_core.training.schedule import CORRECTED_PROTOCOL
from softasp.cli.train import run_training
from softasp.instantiations.runtime import evaluate as evaluate_release
from softasp.sudoku_backbone import extract_backbone
from softasp.sudoku_reporting import file_sha256
from softasp.sudoku_scratch_data import prepare_scratch_data
def validate_scratch_config(cfg: DictConfig) -> None:
"""Fail before downloads/training when the requested workflow is incompatible."""
if (
cfg.dataset.name != "sudoku"
or list(cfg.dataset.splits) != ["train", "val", "test"]
or cfg.model.name != "softasp_sudoku"
or cfg.model.params.backbone != "default"
or cfg.model.params.oracle_perception
or cfg.backbone.model.name != "perception_only_sudoku"
or cfg.backbone.model.params.d_model != cfg.model.params.d_model
or cfg.release_profile != "sudoku_release1"
or int(os.environ.get("WORLD_SIZE", "1")) != 1
):
raise ValueError(
"From-scratch workflow requires single-process Sudoku AS2 and matching default CNN"
)
if cfg.pipeline.data_mode not in ("generate", "existing"):
raise ValueError("pipeline.data_mode must be generate or existing")
subset = cfg.preprocessing.subset_size
if subset is not None and (
isinstance(subset, bool) or not isinstance(subset, int) or subset < 1
):
raise ValueError("preprocessing.subset_size must be a positive integer or null")
name = cfg.pipeline.name
if not isinstance(name, str) or not name or name in (".", "..") or Path(name).name != name:
raise ValueError("pipeline.name must be a single directory name")
if cfg.pipeline.data_mode == "generate" and (
cfg.data_artifact.path or cfg.data_artifact.sha256 or cfg.data_artifact.protocol_id
):
raise ValueError("Generate mode cannot use an existing data_artifact; select existing mode")
for model, trainer in ((cfg.model, cfg.trainer), (cfg.backbone.model, cfg.backbone.trainer)):
if model.pretrained_from or trainer.resume_from:
raise ValueError("From-scratch training cannot accept pretrained_from or resume_from")
if (
trainer.protocol != CORRECTED_PROTOCOL
or trainer.seed != 42
or trainer.distributed != "false"
or trainer.compile
or trainer.gradient_accumulation_steps != 1
or trainer.epochs < 1
or trainer.eval_on_start
or not 1 <= trainer.eval_every <= trainer.epochs
or trainer.batch_size_per_gpu < 1
or trainer.num_workers < 0
):
raise ValueError(
"From-scratch training requires seed 42, corrected single-device training "
"and post-training validation"
)
validate_config(
OmegaConf.create({"trainer": OmegaConf.to_container(trainer, resolve=True)})
)
def _stage_config(cfg: DictConfig, root: Path, *, backbone: bool) -> DictConfig:
stage = deepcopy(cfg)
if backbone:
stage.model = deepcopy(cfg.backbone.model)
stage.trainer = deepcopy(cfg.backbone.trainer)
# This stage intentionally trains a fresh perception model, not AS2.
stage.release_profile = None
del stage.backbone
del stage.pipeline
stage.paths.output_dir = str(root / ("backbone" if backbone else "as2"))
return stage
def train_from_scratch(config: DictConfig) -> Path:
"""Run a fresh two-stage training workflow and record exact artifact lineage."""
cfg = deepcopy(config)
OmegaConf.resolve(cfg)
validate_scratch_config(cfg)
root = Path(cfg.pipeline.run_dir).expanduser().resolve()
models = Path(cfg.pipeline.model_dir).expanduser().resolve()
if root == models or root in models.parents or models in root.parents:
raise ValueError("Workflow run_dir and model_dir must be separate directories")
for directory in (root, models):
if directory.exists():
raise FileExistsError(f"Use a fresh workflow directory: {directory}")
root.mkdir(parents=True, exist_ok=False)
models.mkdir(parents=True, exist_ok=False)
manifest_path = root / "manifest.json"
source_root = Path(__file__).resolve().parents[2]
try:
revision = subprocess.check_output(
["git", "rev-parse", "HEAD"], cwd=source_root, text=True, stderr=subprocess.DEVNULL
).strip()
dirty = bool(
subprocess.check_output(
["git", "status", "--porcelain"], cwd=source_root, text=True
).strip()
)
except (OSError, subprocess.CalledProcessError):
revision, dirty = None, None
record: dict[str, Any] = {
"status": "running",
"supplied_weights": False,
"git_commit": revision,
"git_dirty": dirty,
"uv_lock_sha256": (
file_sha256(source_root / "uv.lock")
if (source_root / "uv.lock").is_file() else None
),
"configuration": OmegaConf.to_container(cfg, resolve=True),
"stages": {},
}
def save() -> None:
manifest_path.write_text(json.dumps(record, indent=2, allow_nan=False) + "\n")
save()
try:
prepare_scratch_data(cfg)
record["dataset"] = OmegaConf.to_container(cfg.resolved_dataset, resolve=True)
OmegaConf.save(cfg, root / "resolved.yaml", resolve=True)
save()
backbone_cfg = _stage_config(cfg, root, backbone=True)
backbone_run = run_training(backbone_cfg)
backbone_path = models / "backbone.pt"
record["stages"]["backbone"] = extract_backbone(backbone_run / "best.pt", backbone_path)
save()
as2_cfg = _stage_config(cfg, root, backbone=False)
as2_cfg.model.pretrained_from = str(backbone_path)
# The release guard checks the complete tensor set, dtype, shape, and finiteness.
as2_run = run_training(as2_cfg)
selected = as2_run / "best.pt"
checkpoint = models / "as2.pt"
with selected.open("rb") as source, checkpoint.open("xb") as target:
shutil.copyfileobj(source, target)
record["stages"]["as2"] = {
"source_checkpoint": str(selected),
"path": str(checkpoint),
"sha256": file_sha256(checkpoint),
"backbone_sha256": file_sha256(backbone_path),
}
save()
evaluation_cfg = deepcopy(as2_cfg)
evaluation_cfg.dataset.splits = ["test"]
evaluation_cfg.evaluator.checkpoint = str(checkpoint)
evaluation_cfg.paths.output_dir = str(root / "evaluation")
reports = evaluate_release(evaluation_cfg)
record["stages"]["evaluation"] = {
"report_dir": str(root / "evaluation" / "eval"),
"metrics": reports["test"]["stages"],
}
record["status"] = "complete"
save()
except BaseException as error:
record["status"] = "failed"
record["error"] = f"{type(error).__name__}: {error}"
save()
raise
return manifest_path
|