scripts/preflight_training_v1.py

#!/usr/bin/env python3
"""Fail-fast checks for ControlAI v1 MLX training configurations and artifacts."""

from __future__ import annotations

import argparse
import hashlib
import importlib.metadata
import json
import shutil
import sys
from pathlib import Path

import yaml
from transformers import AutoConfig, AutoTokenizer


PROJECT_ROOT = Path(__file__).resolve().parents[1]
STAGES = {
    "sft": PROJECT_ROOT / "configs" / "lora_controlai_sft_v1.yaml",
    "dapt": PROJECT_ROOT / "configs" / "lora_controlai_dapt_v1.yaml",
    "sft-after-dapt": PROJECT_ROOT / "configs" / "lora_controlai_sft_after_dapt_v1.yaml",
}


def qwen_lora_parameter_count(architecture, config: dict) -> int:
    """Return exact LoRA A/B parameter count for the configured Qwen projections."""
    hidden = architecture.hidden_size
    head_dim = architecture.head_dim
    attention = architecture.num_attention_heads * head_dim
    key_value = architecture.num_key_value_heads * head_dim
    intermediate = architecture.intermediate_size
    dimensions = {
        "self_attn.q_proj": (hidden, attention),
        "self_attn.k_proj": (hidden, key_value),
        "self_attn.v_proj": (hidden, key_value),
        "self_attn.o_proj": (attention, hidden),
        "mlp.gate_proj": (hidden, intermediate),
        "mlp.up_proj": (hidden, intermediate),
        "mlp.down_proj": (intermediate, hidden),
    }
    lora = config["lora_parameters"]
    keys = lora.get("keys")
    if not keys:
        raise ValueError("lora_parameters.keys must explicitly pin target projections")
    unknown = sorted(set(keys) - dimensions.keys())
    if unknown:
        raise ValueError(f"unsupported LoRA projection keys: {unknown}")
    per_block = int(lora["rank"]) * sum(
        dimensions[key][0] + dimensions[key][1] for key in keys
    )
    return int(config["num_layers"]) * per_block


def rows(path: Path) -> list[dict]:
    with path.open(encoding="utf-8") as stream:
        return [json.loads(line) for line in stream if line.strip()]


def main() -> int:
    parser = argparse.ArgumentParser(description=__doc__)
    parser.add_argument("stage", choices=STAGES)
    parser.add_argument("--allow-existing-output", action="store_true")
    args = parser.parse_args()
    errors: list[str] = []
    warnings: list[str] = []

    checksum_manifest = PROJECT_ROOT / "configs" / "training_artifacts_v1.sha256"
    for line in checksum_manifest.read_text(encoding="utf-8").splitlines():
        if not line.strip():
            continue
        expected, relative = line.split(maxsplit=1)
        artifact = PROJECT_ROOT / relative.strip()
        if not artifact.exists():
            errors.append(f"checksummed artifact is missing: {artifact}")
            continue
        digest = hashlib.sha256()
        with artifact.open("rb") as stream:
            for block in iter(lambda: stream.read(1024 * 1024), b""):
                digest.update(block)
        if digest.hexdigest() != expected:
            errors.append(f"artifact checksum changed: {relative.strip()}")

    config_path = STAGES[args.stage]
    config = yaml.safe_load(config_path.read_text(encoding="utf-8"))
    data_dir = PROJECT_ROOT / config["data"]
    train_path, valid_path = data_dir / "train.jsonl", data_dir / "valid.jsonl"
    for path in (train_path, valid_path):
        if not path.exists() or path.stat().st_size == 0:
            errors.append(f"missing or empty data split: {path}")
    train = rows(train_path) if train_path.exists() else []
    valid = rows(valid_path) if valid_path.exists() else []
    optimizer_updates = config["iters"] // config["grad_accumulation_steps"]
    schedule = config.get("lr_schedule")
    if schedule:
        warmup = int(schedule.get("warmup", 0))
        decay_steps = int(schedule["arguments"][1])
        if abs((warmup + decay_steps) - optimizer_updates) > 2:
            errors.append(
                "learning-rate schedule length is not expressed in optimizer updates: "
                f"warmup+decay={warmup + decay_steps}, updates={optimizer_updates}"
            )

    trainable_parameters = None
    try:
        architecture = AutoConfig.from_pretrained(config["model"], local_files_only=True)
        if config["num_layers"] > architecture.num_hidden_layers:
            errors.append(
                f"num_layers={config['num_layers']} exceeds model depth {architecture.num_hidden_layers}"
            )
        trainable_parameters = qwen_lora_parameter_count(architecture, config)
    except Exception as exc:
        errors.append(f"model/config is not available locally: {exc}")

    try:
        tokenizer = AutoTokenizer.from_pretrained(config["model"], local_files_only=True)
        sample = train[: min(128, len(train))]
        lengths = []
        for row in sample:
            if "messages" in row:
                lengths.append(len(tokenizer.apply_chat_template(row["messages"], return_dict=False)))
            else:
                lengths.append(len(tokenizer.encode(row["text"], add_special_tokens=False)))
        if lengths and max(lengths) > config["max_seq_length"]:
            errors.append(f"sample contains {max(lengths)} tokens, above max_seq_length")
    except Exception as exc:
        errors.append(f"tokenizer preflight failed: {exc}")

    adapter_path = PROJECT_ROOT / config["adapter_path"]
    if adapter_path.exists() and any(adapter_path.iterdir()) and not args.allow_existing_output:
        errors.append(
            f"output already contains files: {adapter_path}; choose a new path or pass --allow-existing-output"
        )
    resume = config.get("resume_adapter_file")
    if resume and not (PROJECT_ROOT / resume).exists():
        errors.append(f"required resume adapter is missing: {PROJECT_ROOT / resume}")

    free_gb = shutil.disk_usage(PROJECT_ROOT).free / 1e9
    if free_gb < 10:
        warnings.append(f"only {free_gb:.1f} GB disk space is free")
    if sys.prefix == sys.base_prefix:
        errors.append("virtual environment is not active")
    try:
        mlx_lm_version = importlib.metadata.version("mlx-lm")
    except importlib.metadata.PackageNotFoundError:
        errors.append("mlx-lm is not installed in this environment")
        mlx_lm_version = "missing"

    print(f"stage: {args.stage}")
    print(f"config: {config_path.relative_to(PROJECT_ROOT)}")
    print(f"mlx-lm: {mlx_lm_version}")
    print(f"model: {config['model']}")
    print(f"train/valid rows: {len(train):,}/{len(valid):,}")
    print(f"iterations: {config['iters']:,}; gradient accumulation: {config['grad_accumulation_steps']}")
    print(f"optimizer updates: {optimizer_updates:,}")
    print(f"adapter: rank {config['lora_parameters']['rank']}, layers {config['num_layers']}")
    if trainable_parameters is not None:
        print(
            f"trainable LoRA parameters: {trainable_parameters:,} "
            f"(~{trainable_parameters * 4 / 1_000_000:.1f} MB as float32 weights)"
        )
    print(f"free disk: {free_gb:.1f} GB")
    for warning in warnings:
        print(f"WARNING: {warning}")
    if errors:
        print("preflight failed:", file=sys.stderr)
        for error in errors:
            print(f"- {error}", file=sys.stderr)
        return 1
    print("preflight passed")
    print(f"train command: python -m mlx_lm.lora --config {config_path.relative_to(PROJECT_ROOT)}")
    return 0


if __name__ == "__main__":
    raise SystemExit(main())