| """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) |
| |
| 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) |
| |
| 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 |
|
|
