code/self_healing_finetuning/training/hf_fingpt_train.py

from __future__ import annotations

import argparse
import hashlib
import importlib.util
import json
import os
import shutil
import sys
import tempfile
import time
from datetime import UTC, datetime
from pathlib import Path
from typing import Any


REQUIRED_LIVE_PACKAGES = ["torch", "transformers", "peft", "trl", "datasets", "accelerate", "bitsandbytes"]


def utc_now() -> str:
    return datetime.now(UTC).replace(microsecond=0).isoformat().replace("+00:00", "Z")


def load_jsonl(path: Path) -> list[dict[str, Any]]:
    if not path.exists():
        raise FileNotFoundError(f"JSONL file does not exist: {path}")
    rows: list[dict[str, Any]] = []
    with path.open("r", encoding="utf-8", newline="") as handle:
        for line_no, line in enumerate(handle, start=1):
            if not line.strip():
                continue
            try:
                item = json.loads(line)
            except json.JSONDecodeError as exc:
                raise ValueError(
                    f"invalid JSONL at {path}:{line_no}: {exc.msg} "
                    f"(column {exc.colno}, char {exc.pos})"
                ) from exc
            if not isinstance(item, dict):
                raise ValueError(f"record at {path}:{line_no} must be a JSON object")
            messages = item.get("messages")
            if not isinstance(messages, list) or not messages:
                raise ValueError(f"record at {path}:{line_no} must contain non-empty messages")
            rows.append(item)
    if not rows:
        raise ValueError(f"JSONL file has no records: {path}")
    return rows


def sha256_file(path: Path) -> str:
    digest = hashlib.sha256()
    with path.open("rb") as handle:
        for chunk in iter(lambda: handle.read(1024 * 1024), b""):
            digest.update(chunk)
    return digest.hexdigest()


def load_json(path: Path) -> dict[str, Any]:
    return json.loads(path.read_text(encoding="utf-8-sig"))


def format_messages(record: dict[str, Any], tokenizer: Any | None = None) -> str:
    messages = record["messages"]
    if tokenizer is not None and getattr(tokenizer, "chat_template", None):
        return tokenizer.apply_chat_template(messages, tokenize=False, add_generation_prompt=False)
    parts: list[str] = []
    for message in messages:
        if not isinstance(message, dict):
            raise ValueError("message must be an object")
        role = str(message.get("role", "")).strip()
        content = str(message.get("content", "")).strip()
        if not role or not content:
            raise ValueError("message role and content are required")
        parts.append(f"<|{role}|>\n{content}")
    return "\n".join(parts) + "\n<|end|>"


def package_status() -> dict[str, bool]:
    return {name: importlib.util.find_spec(name) is not None for name in REQUIRED_LIVE_PACKAGES}


def write_json(path: Path, payload: dict[str, Any]) -> None:
    path.parent.mkdir(parents=True, exist_ok=True)
    path.write_text(json.dumps(payload, indent=2, sort_keys=True) + "\n", encoding="utf-8")


class ZeroGradientStallError(RuntimeError):
    def __init__(self, *, step: int | None, zero_count: int, last_logs: dict[str, Any]) -> None:
        self.step = step
        self.zero_count = zero_count
        self.last_logs = last_logs
        super().__init__(
            "trainer reported zero gradient norm for "
            f"{zero_count} consecutive logged training events at step {step}"
        )


def trainable_parameter_summary(model: Any) -> dict[str, Any]:
    total_params = 0
    trainable_params = 0
    trainable_tensors = 0
    trainable_prefix_counts: dict[str, int] = {}
    for name, param in model.named_parameters():
        count = int(param.numel())
        total_params += count
        if not getattr(param, "requires_grad", False):
            continue
        trainable_params += count
        trainable_tensors += 1
        prefix = str(name).split(".")[0]
        trainable_prefix_counts[prefix] = trainable_prefix_counts.get(prefix, 0) + 1
    pct = (trainable_params / total_params * 100.0) if total_params else 0.0
    return {
        "schema_version": "shft_trainable_parameter_summary_v1",
        "total_params": total_params,
        "trainable_params": trainable_params,
        "trainable_tensors": trainable_tensors,
        "trainable_pct": round(pct, 6),
        "trainable_prefix_counts": trainable_prefix_counts,
        "ok": trainable_params > 0 and trainable_tensors > 0,
    }


def gradient_parameter_summary(model: Any) -> dict[str, Any]:
    trainable_tensors = 0
    grad_tensors = 0
    nonzero_grad_tensors = 0
    lora_trainable_tensors = 0
    lora_nonzero_grad_tensors = 0
    total_grad_norm = 0.0
    sample_zero_grad_trainable: list[str] = []
    sample_nonzero_grad_trainable: list[str] = []
    for name, param in model.named_parameters():
        if not getattr(param, "requires_grad", False):
            continue
        trainable_tensors += 1
        is_lora = "lora_" in str(name).lower()
        if is_lora:
            lora_trainable_tensors += 1
        grad = getattr(param, "grad", None)
        if grad is None:
            if len(sample_zero_grad_trainable) < 10:
                sample_zero_grad_trainable.append(str(name))
            continue
        grad_tensors += 1
        try:
            norm = float(grad.detach().float().norm().item())
        except Exception:
            norm = 0.0
        total_grad_norm += norm
        if norm > 0.0:
            nonzero_grad_tensors += 1
            if is_lora:
                lora_nonzero_grad_tensors += 1
            if len(sample_nonzero_grad_trainable) < 10:
                sample_nonzero_grad_trainable.append(str(name))
        elif len(sample_zero_grad_trainable) < 10:
            sample_zero_grad_trainable.append(str(name))
    return {
        "schema_version": "shft_gradient_parameter_summary_v1",
        "trainable_tensors": trainable_tensors,
        "grad_tensors": grad_tensors,
        "nonzero_grad_tensors": nonzero_grad_tensors,
        "lora_trainable_tensors": lora_trainable_tensors,
        "lora_nonzero_grad_tensors": lora_nonzero_grad_tensors,
        "total_grad_norm": total_grad_norm,
        "sample_zero_grad_trainable": sample_zero_grad_trainable,
        "sample_nonzero_grad_trainable": sample_nonzero_grad_trainable,
        "ok": nonzero_grad_tensors > 0 and (lora_trainable_tensors == 0 or lora_nonzero_grad_tensors > 0),
    }


def copy_file_with_retry(source: Path, destination: Path, *, attempts: int = 5, delay_seconds: float = 2.0) -> None:
    destination.parent.mkdir(parents=True, exist_ok=True)
    last_error: OSError | None = None
    for attempt in range(1, attempts + 1):
        try:
            shutil.copyfile(source, destination)
            return
        except OSError as exc:
            last_error = exc
            if attempt == attempts:
                break
            time.sleep(delay_seconds * attempt)
    if last_error is not None:
        raise last_error


def append_jsonl(path: Path, payload: dict[str, Any]) -> None:
    path.parent.mkdir(parents=True, exist_ok=True)
    with path.open("a", encoding="utf-8", newline="\n") as handle:
        handle.write(json.dumps(payload, sort_keys=True) + "\n")


def _jsonable(value: Any) -> Any:
    if value is None or isinstance(value, str | int | float | bool):
        return value
    if isinstance(value, dict):
        return {str(key): _jsonable(item) for key, item in value.items()}
    if isinstance(value, list | tuple):
        return [_jsonable(item) for item in value]
    try:
        return float(value)
    except (TypeError, ValueError):
        return str(value)


def _number(value: Any, default: float | None = None) -> float | None:
    try:
        return float(value)
    except (TypeError, ValueError):
        return default


def _metric_rows(log_history: list[dict[str, Any]]) -> list[dict[str, Any]]:
    rows: list[dict[str, Any]] = []
    for item in log_history:
        row = {str(key): _jsonable(value) for key, value in item.items()}
        if "eval_loss" in row:
            row["event_type"] = "eval"
        elif "loss" in row:
            row["event_type"] = "train_log"
        else:
            row["event_type"] = "trainer_event"
        rows.append(row)
    return rows


def _best_eval_row(rows: list[dict[str, Any]]) -> dict[str, Any] | None:
    eval_rows = [row for row in rows if _number(row.get("eval_loss")) is not None]
    if not eval_rows:
        return None
    return min(eval_rows, key=lambda row: _number(row.get("eval_loss"), float("inf")) or float("inf"))


def _last_row_with(rows: list[dict[str, Any]], key: str) -> dict[str, Any] | None:
    for row in reversed(rows):
        if key in row and _number(row.get(key)) is not None:
            return row
    return None


def build_trainer_metrics_summary(
    *,
    rows: list[dict[str, Any]],
    selected_checkpoint: dict[str, Any],
    overfit_tolerance: float,
) -> dict[str, Any]:
    best_eval = _best_eval_row(rows)
    final_eval = _last_row_with(rows, "eval_loss")
    final_train = _last_row_with(rows, "loss")
    best_eval_loss = _number(best_eval.get("eval_loss")) if best_eval else None
    final_eval_loss = _number(final_eval.get("eval_loss")) if final_eval else None
    final_train_loss = _number(final_train.get("loss")) if final_train else None
    selected_metric = _number(selected_checkpoint.get("selection_metric_value"))
    train_eval_gap = None
    if final_train_loss is not None and final_eval_loss is not None:
        train_eval_gap = round(final_eval_loss - final_train_loss, 6)
    late_eval_regression = None
    if best_eval_loss is not None and final_eval_loss is not None:
        late_eval_regression = round(final_eval_loss - best_eval_loss, 6)
    token_accuracy_rows = [
        row
        for row in rows
        if _number(row.get("mean_token_accuracy")) is not None or _number(row.get("eval_mean_token_accuracy")) is not None
    ]
    first_token_accuracy = None
    final_token_accuracy = None
    if token_accuracy_rows:
        first_token_accuracy = _number(
            token_accuracy_rows[0].get("mean_token_accuracy")
            or token_accuracy_rows[0].get("eval_mean_token_accuracy")
        )
        final_token_accuracy = _number(
            token_accuracy_rows[-1].get("mean_token_accuracy")
            or token_accuracy_rows[-1].get("eval_mean_token_accuracy")
        )
    overfit_flags: list[str] = []
    if late_eval_regression is not None and late_eval_regression > overfit_tolerance:
        overfit_flags.append("late_eval_loss_regression")
    return {
        "schema_version": "shft_trainer_metrics_summary_v1",
        "created_at": utc_now(),
        "metric_row_count": len(rows),
        "eval_row_count": len([row for row in rows if _number(row.get("eval_loss")) is not None]),
        "train_log_row_count": len([row for row in rows if _number(row.get("loss")) is not None]),
        "best_eval_step": best_eval.get("step") if best_eval else None,
        "best_eval_loss": best_eval_loss,
        "final_eval_step": final_eval.get("step") if final_eval else None,
        "final_eval_loss": final_eval_loss,
        "final_train_step": final_train.get("step") if final_train else None,
        "final_train_loss": final_train_loss,
        "train_eval_loss_gap": train_eval_gap,
        "late_eval_loss_regression": late_eval_regression,
        "first_token_accuracy": first_token_accuracy,
        "final_token_accuracy": final_token_accuracy,
        "selected_checkpoint": selected_checkpoint,
        "selected_metric_value": selected_metric,
        "overfit_tolerance": overfit_tolerance,
        "overfit_detected": bool(overfit_flags),
        "overfit_flags": overfit_flags,
    }


def validate_dataset_provenance(
    args: argparse.Namespace,
    train_rows: list[dict[str, Any]],
    valid_rows: list[dict[str, Any]],
) -> dict[str, Any]:
    errors: list[str] = []
    checks: dict[str, Any] = {}
    dataset_arg = getattr(args, "dataset_dir", None)
    dataset_dir = Path(dataset_arg) if dataset_arg else None
    manifest: dict[str, Any] | None = None
    train_sha = sha256_file(args.train_jsonl)
    valid_sha = sha256_file(args.valid_jsonl)

    if dataset_dir is not None:
        manifest_path = dataset_dir / "dataset_manifest.json"
        if not manifest_path.exists():
            errors.append(f"missing dataset manifest: {manifest_path}")
        else:
            manifest = load_json(manifest_path)
            split_counts = manifest.get("split_counts", {})
            split_hashes = manifest.get("split_sha256", {})
            expected_train_count = int(split_counts.get("train", -1))
            expected_valid_count = int(split_counts.get("valid", -1))
            checks["manifest_train_records"] = {
                "ok": expected_train_count == len(train_rows),
                "detail": f"{len(train_rows)} == {expected_train_count}",
            }
            checks["manifest_valid_records"] = {
                "ok": expected_valid_count == len(valid_rows),
                "detail": f"{len(valid_rows)} == {expected_valid_count}",
            }
            if expected_train_count != len(train_rows):
                errors.append(f"manifest_train_records: {len(train_rows)} != {expected_train_count}")
            if expected_valid_count != len(valid_rows):
                errors.append(f"manifest_valid_records: {len(valid_rows)} != {expected_valid_count}")

            manifest_train_sha = split_hashes.get("train")
            manifest_valid_sha = split_hashes.get("valid")
            if manifest_train_sha:
                checks["manifest_train_sha256"] = {
                    "ok": train_sha == manifest_train_sha,
                    "detail": f"{train_sha} == {manifest_train_sha}",
                }
                if train_sha != manifest_train_sha:
                    errors.append("manifest_train_sha256 mismatch")
            if manifest_valid_sha:
                checks["manifest_valid_sha256"] = {
                    "ok": valid_sha == manifest_valid_sha,
                    "detail": f"{valid_sha} == {manifest_valid_sha}",
                }
                if valid_sha != manifest_valid_sha:
                    errors.append("manifest_valid_sha256 mismatch")

            expected_manifest_sha = (getattr(args, "expected_dataset_manifest_sha256", "") or "").strip()
            if expected_manifest_sha:
                actual_manifest_sha = sha256_file(manifest_path)
                checks["expected_dataset_manifest_sha256"] = {
                    "ok": actual_manifest_sha == expected_manifest_sha,
                    "detail": f"{actual_manifest_sha} == {expected_manifest_sha}",
                }
                if actual_manifest_sha != expected_manifest_sha:
                    errors.append("expected_dataset_manifest_sha256 mismatch")

    expected_hashes = {
        "train": (getattr(args, "expected_train_sha256", "") or "").strip(),
        "valid": (getattr(args, "expected_valid_sha256", "") or "").strip(),
    }
    actual_hashes = {"train": train_sha, "valid": valid_sha}
    for split, expected in expected_hashes.items():
        if not expected:
            continue
        ok = actual_hashes[split] == expected
        checks[f"expected_{split}_sha256"] = {"ok": ok, "detail": f"{actual_hashes[split]} == {expected}"}
        if not ok:
            errors.append(f"expected_{split}_sha256 mismatch")

    test_hash = ""
    expected_test_sha = (getattr(args, "expected_test_sha256", "") or "").strip()
    if dataset_dir is not None and (dataset_dir / "test.jsonl").exists():
        test_hash = sha256_file(dataset_dir / "test.jsonl")
        if expected_test_sha:
            ok = test_hash == expected_test_sha
            checks["expected_test_sha256"] = {"ok": ok, "detail": f"{test_hash} == {expected_test_sha}"}
            if not ok:
                errors.append("expected_test_sha256 mismatch")

    return {
        "ok": not errors,
        "errors": errors,
        "checks": checks,
        "dataset_dir": str(dataset_dir) if dataset_dir else None,
        "train_jsonl": str(args.train_jsonl),
        "valid_jsonl": str(args.valid_jsonl),
        "actual_split_counts": {"train": len(train_rows), "valid": len(valid_rows)},
        "actual_split_sha256": {"train": train_sha, "valid": valid_sha, "test": test_hash or None},
        "manifest_split_counts": (manifest or {}).get("split_counts") if manifest else None,
        "manifest_split_sha256": (manifest or {}).get("split_sha256") if manifest else None,
    }


def build_plan(args: argparse.Namespace, train_rows: list[dict[str, Any]], valid_rows: list[dict[str, Any]]) -> dict[str, Any]:
    missing = [name for name, present in package_status().items() if not present]
    min_steps = int(os.environ.get("SHFT_MIN_PRODUCTION_STEPS", "100"))
    min_train_records = int(os.environ.get("SHFT_MIN_PRODUCTION_TRAIN_RECORDS", "100"))
    readiness_warnings: list[str] = []
    if args.max_steps < min_steps:
        readiness_warnings.append(f"max_steps={args.max_steps} is below production minimum {min_steps}; this is a smoke run")
    if len(train_rows) < min_train_records:
        readiness_warnings.append(
            f"train_records={len(train_rows)} is below production minimum {min_train_records}; corpus is too small for a durable role adapter"
        )
    return {
        "run_id": args.run_id,
        "model_candidate": args.model_candidate,
        "finetune_start_policy": getattr(args, "finetune_start_policy", "bootstrap"),
        "start_adapter": getattr(args, "start_adapter", args.model_candidate),
        "base_model_id": args.base_model_id,
        "train_jsonl": str(args.train_jsonl),
        "valid_jsonl": str(args.valid_jsonl),
        "output_dir": str(args.output_dir),
        "train_records": len(train_rows),
        "valid_records": len(valid_rows),
        "dataset_provenance": validate_dataset_provenance(args, train_rows, valid_rows),
        "hyperparameters": {
            "max_steps": args.max_steps,
            "per_device_train_batch_size": args.per_device_train_batch_size,
            "gradient_accumulation_steps": args.gradient_accumulation_steps,
            "learning_rate": args.learning_rate,
            "lora_r": args.lora_r,
            "lora_alpha": args.lora_alpha,
            "lora_dropout": args.lora_dropout,
            "max_seq_length": args.max_seq_length,
            "logging_steps": getattr(args, "logging_steps", 5),
            "checkpoint_steps": getattr(args, "checkpoint_steps", 50),
            "eval_steps": getattr(args, "eval_steps", getattr(args, "checkpoint_steps", 50)),
            "save_total_limit": getattr(args, "save_total_limit", 4),
            "metric_for_best_model": getattr(args, "metric_for_best_model", "eval_loss"),
            "greater_is_better": getattr(args, "greater_is_better", False),
            "overfit_tolerance": getattr(args, "overfit_tolerance", 0.10),
        },
        "quantization": {
            "load_in_4bit": True,
            "bnb_4bit_quant_type": "nf4",
            "bnb_4bit_compute_dtype": "bfloat16",
        },
        "required_packages": package_status(),
        "missing_packages": missing,
        "readiness": {
            "production_candidate": not readiness_warnings,
            "warnings": readiness_warnings,
            "minimums": {"max_steps": min_steps, "train_records": min_train_records},
        },
        "dry_run": args.dry_run,
        "created_at": utc_now(),
    }


def run_live_training(args: argparse.Namespace, plan: dict[str, Any]) -> dict[str, Any]:
    missing = plan["missing_packages"]
    if missing:
        return {
            "status": "blocked_missing_dependencies",
            "missing_packages": missing,
            "note": "Install the live ML stack or use an HF Jobs image that includes these packages.",
        }

    import torch
    from datasets import Dataset
    from peft import LoraConfig, PeftModel, get_peft_model, prepare_model_for_kbit_training
    from transformers import AutoModelForCausalLM, AutoTokenizer, BitsAndBytesConfig, TrainerCallback
    from trl import SFTConfig, SFTTrainer

    class JsonlMetricsCallback(TrainerCallback):
        def __init__(self, path: Path) -> None:
            self.path = path
            self.started = time.time()

        def _write(self, event_type: str, state: Any, logs: dict[str, Any] | None = None) -> None:
            row = {
                "schema_version": "shft_trainer_metrics_v1",
                "event_type": event_type,
                "timestamp": utc_now(),
                "elapsed_seconds": round(time.time() - self.started, 3),
                "step": getattr(state, "global_step", None),
                "epoch": getattr(state, "epoch", None),
            }
            if logs:
                row.update({str(key): _jsonable(value) for key, value in logs.items()})
            append_jsonl(self.path, row)

        def on_log(self, args: Any, state: Any, control: Any, logs: dict[str, Any] | None = None, **kwargs: Any) -> None:
            self._write("trainer_log", state, logs)

        def on_evaluate(self, args: Any, state: Any, control: Any, metrics: dict[str, Any] | None = None, **kwargs: Any) -> None:
            self._write("eval", state, metrics)

        def on_save(self, args: Any, state: Any, control: Any, **kwargs: Any) -> None:
            self._write("save", state, None)

        def on_train_end(self, args: Any, state: Any, control: Any, **kwargs: Any) -> None:
            self._write("train_end", state, None)

    class ZeroGradientStallCallback(TrainerCallback):
        def __init__(self, path: Path, *, max_zero_logs: int = 3, hard_halt_enabled: bool = True) -> None:
            self.path = path
            self.max_zero_logs = max(1, max_zero_logs)
            self.hard_halt_enabled = hard_halt_enabled
            self.zero_count = 0

        def on_log(self, args: Any, state: Any, control: Any, logs: dict[str, Any] | None = None, **kwargs: Any) -> None:
            if not logs or "grad_norm" not in logs:
                return
            try:
                grad_norm = float(logs["grad_norm"])
            except (TypeError, ValueError):
                return
            if grad_norm > 0.0:
                self.zero_count = 0
                return
            self.zero_count += 1
            step = getattr(state, "global_step", None)
            row = {
                "schema_version": "shft_zero_gradient_stall_v1",
                "event_type": "zero_gradient_observed",
                "timestamp": utc_now(),
                "step": step,
                "zero_count": self.zero_count,
                "max_zero_logs": self.max_zero_logs,
                "hard_halt_enabled": self.hard_halt_enabled,
                "grad_norm": grad_norm,
                "logs": {str(key): _jsonable(value) for key, value in logs.items()},
            }
            append_jsonl(self.path, row)
            if self.hard_halt_enabled and self.zero_count >= self.max_zero_logs:
                halt = (
                    "[SHFT TRAIN STALL] ACTION=HALT_ZERO_GRADIENT_TRAINING "
                    f"step={step} consecutive_zero_grad_logs={self.zero_count} "
                    "reason=trainer_reported_zero_grad_norm_repeatedly"
                )
                print(halt)
                print(halt, file=sys.stderr)
                raise ZeroGradientStallError(step=step, zero_count=self.zero_count, last_logs=logs)

    train_rows = load_jsonl(args.train_jsonl)
    valid_rows = load_jsonl(args.valid_jsonl)
    tokenizer = AutoTokenizer.from_pretrained(args.base_model_id, use_fast=True)
    if tokenizer.pad_token is None:
        tokenizer.pad_token = tokenizer.eos_token

    train_dataset = Dataset.from_list([{"text": format_messages(row, tokenizer)} for row in train_rows])
    valid_dataset = Dataset.from_list([{"text": format_messages(row, tokenizer)} for row in valid_rows])

    quantization = BitsAndBytesConfig(
        load_in_4bit=True,
        bnb_4bit_quant_type="nf4",
        bnb_4bit_compute_dtype=torch.bfloat16,
    )
    model = AutoModelForCausalLM.from_pretrained(
        args.base_model_id,
        quantization_config=quantization,
        device_map="auto",
        trust_remote_code=False,
    )
    if getattr(model, "config", None) is not None:
        model.config.use_cache = False
    model = prepare_model_for_kbit_training(model)
    if hasattr(model, "enable_input_require_grads"):
        model.enable_input_require_grads()
    start_adapter = args.start_adapter.strip() if args.start_adapter else ""
    if start_adapter:
        model = PeftModel.from_pretrained(model, start_adapter, is_trainable=True)
    else:
        lora_config = LoraConfig(
            r=args.lora_r,
            lora_alpha=args.lora_alpha,
            lora_dropout=args.lora_dropout,
            bias="none",
            task_type="CAUSAL_LM",
            target_modules=["q_proj", "k_proj", "v_proj", "o_proj", "gate_proj", "up_proj", "down_proj"],
        )
        model = get_peft_model(model, lora_config)

    trainable_summary = trainable_parameter_summary(model)
    write_json(args.output_dir / "trainable_parameters.json", trainable_summary)
    print("[SHFT TRAINABLE PARAMETERS] " + json.dumps(trainable_summary, sort_keys=True))
    if not trainable_summary["ok"]:
        print("[SHFT TRAIN STALL] ACTION=HALT_NO_TRAINABLE_PARAMETERS", file=sys.stderr)
        return {
            "status": "blocked_no_trainable_parameters",
            "run_id": args.run_id,
            "completed_at": utc_now(),
            "trainable_parameters": trainable_summary,
            "trainable_parameters_path": str(args.output_dir / "trainable_parameters.json"),
        }

    preflight_path = args.output_dir / "gradient_preflight.json"
    try:
        model.train()
        model.zero_grad(set_to_none=True)
        smoke_text = train_dataset[0]["text"]
        encoded = tokenizer(
            smoke_text,
            return_tensors="pt",
            truncation=True,
            max_length=args.max_seq_length,
        )
        if int(encoded["input_ids"].numel()) < 2:
            raise ValueError("gradient preflight sample tokenized to fewer than 2 tokens")
        labels = encoded["input_ids"].clone()
        encoded["labels"] = labels
        first_param = next(param for param in model.parameters() if getattr(param, "requires_grad", False))
        encoded = {key: value.to(first_param.device) for key, value in encoded.items()}
        with torch.enable_grad():
            outputs = model(**encoded)
            loss = outputs.loss
            loss.backward()
        gradient_summary = gradient_parameter_summary(model)
        preflight = {
            "schema_version": "shft_gradient_preflight_v1",
            "run_id": args.run_id,
            "created_at": utc_now(),
            "base_model_id": args.base_model_id,
            "max_seq_length": args.max_seq_length,
            "sample_token_count": int(encoded["input_ids"].numel()),
            "loss": float(loss.detach().float().item()),
            "trainable_parameters": trainable_summary,
            "gradient_parameters": gradient_summary,
            "ok": bool(gradient_summary["ok"]),
            "quit_policy": "halt before paid full training unless a one-batch backward pass proves nonzero trainable LoRA gradients",
        }
        write_json(preflight_path, preflight)
        print("[SHFT GRADIENT PREFLIGHT] " + json.dumps(preflight, sort_keys=True))
        if not preflight["ok"]:
            halt = (
                "[SHFT TRAIN STALL] ACTION=HALT_GRADIENT_PREFLIGHT_FAILED "
                f"nonzero_grad_tensors={gradient_summary['nonzero_grad_tensors']} "
                f"lora_nonzero_grad_tensors={gradient_summary['lora_nonzero_grad_tensors']} "
                "reason=one_batch_backward_did_not_prove_trainable_lora_gradients"
            )
            print(halt)
            print(halt, file=sys.stderr)
            return {
                "status": "blocked_gradient_preflight_failed",
                "run_id": args.run_id,
                "completed_at": utc_now(),
                "trainable_parameters": trainable_summary,
                "trainable_parameters_path": str(args.output_dir / "trainable_parameters.json"),
                "gradient_preflight": preflight,
                "gradient_preflight_path": str(preflight_path),
                "quit_policy": preflight["quit_policy"],
            }
    except Exception as exc:
        preflight = {
            "schema_version": "shft_gradient_preflight_v1",
            "run_id": args.run_id,
            "created_at": utc_now(),
            "base_model_id": args.base_model_id,
            "trainable_parameters": trainable_summary,
            "ok": False,
            "error": str(exc),
            "quit_policy": "halt before paid full training when gradient preflight cannot complete",
        }
        write_json(preflight_path, preflight)
        print("[SHFT TRAIN STALL] ACTION=HALT_GRADIENT_PREFLIGHT_ERROR reason=" + str(exc), file=sys.stderr)
        return {
            "status": "blocked_gradient_preflight_error",
            "run_id": args.run_id,
            "completed_at": utc_now(),
            "trainable_parameters": trainable_summary,
            "trainable_parameters_path": str(args.output_dir / "trainable_parameters.json"),
            "gradient_preflight": preflight,
            "gradient_preflight_path": str(preflight_path),
            "quit_policy": preflight["quit_policy"],
        }
    finally:
        model.zero_grad(set_to_none=True)

    checkpoint_steps = max(1, min(int(args.checkpoint_steps), int(args.max_steps)))
    eval_steps = max(1, min(int(args.eval_steps or args.checkpoint_steps), int(args.max_steps)))
    if checkpoint_steps % eval_steps != 0:
        checkpoint_steps = eval_steps
    logging_steps = max(1, int(args.logging_steps))
    training_args = SFTConfig(
        output_dir=str(args.output_dir),
        max_steps=args.max_steps,
        per_device_train_batch_size=args.per_device_train_batch_size,
        gradient_accumulation_steps=args.gradient_accumulation_steps,
        learning_rate=args.learning_rate,
        logging_steps=logging_steps,
        save_strategy="steps",
        save_steps=checkpoint_steps,
        eval_strategy="steps",
        eval_steps=eval_steps,
        save_total_limit=args.save_total_limit,
        load_best_model_at_end=True,
        metric_for_best_model=args.metric_for_best_model,
        greater_is_better=args.greater_is_better,
        report_to=[],
        bf16=True,
        gradient_checkpointing=True,
        gradient_checkpointing_kwargs={"use_reentrant": False},
        dataset_text_field="text",
        max_length=args.max_seq_length,
    )
    metrics_path = args.output_dir / "trainer_metrics.jsonl"
    local_metrics_dir = Path(tempfile.gettempdir()) / "linvest21_shft_metrics" / args.run_id
    local_metrics_path = local_metrics_dir / "trainer_metrics.jsonl"
    if local_metrics_path.exists():
        local_metrics_path.unlink()
    trainer = SFTTrainer(
        model=model,
        args=training_args,
        train_dataset=train_dataset,
        eval_dataset=valid_dataset,
        processing_class=tokenizer,
        callbacks=[
            JsonlMetricsCallback(local_metrics_path),
            ZeroGradientStallCallback(local_metrics_path, hard_halt_enabled=True),
        ],
    )
    try:
        result = trainer.train()
    except ZeroGradientStallError as exc:
        if local_metrics_path.exists():
            copy_file_with_retry(local_metrics_path, metrics_path)
        return {
            "status": "blocked_zero_gradient_stall",
            "run_id": args.run_id,
            "completed_at": utc_now(),
            "stall_reason": str(exc),
            "stall_step": exc.step,
            "consecutive_zero_grad_logs": exc.zero_count,
            "last_logs": {str(key): _jsonable(value) for key, value in exc.last_logs.items()},
            "trainable_parameters": trainable_summary,
            "trainable_parameters_path": str(args.output_dir / "trainable_parameters.json"),
            "trainer_metrics": str(metrics_path),
            "quit_policy": "halt paid training when trainable parameters exist but gradient norm is repeatedly zero",
        }
    metric_rows = _metric_rows(getattr(trainer.state, "log_history", []))
    if not local_metrics_path.exists() and metric_rows:
        for row in metric_rows:
            append_jsonl(local_metrics_path, row)
    if local_metrics_path.exists():
        copy_file_with_retry(local_metrics_path, metrics_path)
    best_checkpoint = getattr(trainer.state, "best_model_checkpoint", None)
    best_metric = getattr(trainer.state, "best_metric", None)
    best_eval = _best_eval_row(metric_rows)
    if best_metric is None and best_eval is not None:
        best_metric = best_eval.get(args.metric_for_best_model)
    if not best_checkpoint and best_eval is not None and best_eval.get("step") is not None:
        candidate = args.output_dir / f"checkpoint-{best_eval['step']}"
        if candidate.exists():
            best_checkpoint = str(candidate)
    final_eval = _last_row_with(metric_rows, args.metric_for_best_model)
    selected_checkpoint = {
        "schema_version": "shft_selected_checkpoint_v1",
        "run_id": args.run_id,
        "created_at": utc_now(),
        "selection_metric": args.metric_for_best_model,
        "selection_metric_value": _number(best_metric),
        "selected_checkpoint": best_checkpoint,
        "selected_step": best_eval.get("step") if best_eval else None,
        "final_step": getattr(trainer.state, "global_step", None),
        "final_metric_value": _number(final_eval.get(args.metric_for_best_model)) if final_eval else None,
        "greater_is_better": args.greater_is_better,
        "candidate_adapter_dir": str(args.output_dir / "adapter"),
        "selection_policy": "trainer_best_model_checkpoint_by_validation_metric",
    }
    selected_metric = _number(selected_checkpoint.get("selection_metric_value"))
    final_metric = _number(selected_checkpoint.get("final_metric_value"))
    selected_checkpoint["selected_vs_final_delta"] = (
        round(final_metric - selected_metric, 6)
        if selected_metric is not None and final_metric is not None
        else None
    )
    selected_checkpoint["overfit_detected"] = (
        selected_checkpoint["selected_vs_final_delta"] is not None
        and selected_checkpoint["selected_vs_final_delta"] > args.overfit_tolerance
    )
    write_json(args.output_dir / "selected_checkpoint.json", selected_checkpoint)
    summary = build_trainer_metrics_summary(
        rows=metric_rows,
        selected_checkpoint=selected_checkpoint,
        overfit_tolerance=args.overfit_tolerance,
    )
    write_json(args.output_dir / "trainer_metrics_summary.json", summary)
    trainer.save_model(str(args.output_dir / "adapter"))
    tokenizer.save_pretrained(str(args.output_dir / "adapter"))
    return {
        "status": "completed",
        "run_id": args.run_id,
        "completed_at": utc_now(),
        "train_loss": getattr(result, "training_loss", None),
        "adapter_dir": str(args.output_dir / "adapter"),
        "trainer_metrics": str(metrics_path),
        "trainer_metrics_summary": str(args.output_dir / "trainer_metrics_summary.json"),
        "selected_checkpoint": selected_checkpoint,
        "selected_checkpoint_path": str(args.output_dir / "selected_checkpoint.json"),
        "overfit_detected": summary["overfit_detected"],
        "overfit_flags": summary["overfit_flags"],
    }


def main() -> int:
    parser = argparse.ArgumentParser(description="Linvest21 SHFT Hugging Face QLoRA/PEFT continuation trainer.")
    parser.add_argument("--run-id", required=True)
    parser.add_argument("--model-candidate", required=True)
    parser.add_argument("--start-adapter")
    parser.add_argument("--finetune-start-policy", default="bootstrap", choices=["bootstrap", "continue-best"])
    parser.add_argument("--base-model-id", default="meta-llama/Meta-Llama-3-8B")
    parser.add_argument("--dataset-dir")
    parser.add_argument("--train-jsonl")
    parser.add_argument("--valid-jsonl")
    parser.add_argument("--output-dir", required=True)
    parser.add_argument("--max-steps", type=int, default=20)
    parser.add_argument("--per-device-train-batch-size", type=int, default=1)
    parser.add_argument("--gradient-accumulation-steps", type=int, default=8)
    parser.add_argument("--learning-rate", type=float, default=0.00008)
    parser.add_argument("--lora-r", type=int, default=16)
    parser.add_argument("--lora-alpha", type=int, default=32)
    parser.add_argument("--lora-dropout", type=float, default=0.05)
    parser.add_argument("--max-seq-length", type=int, default=2048)
    parser.add_argument("--logging-steps", type=int, default=5)
    parser.add_argument("--checkpoint-steps", type=int, default=int(os.environ.get("SHFT_CHECKPOINT_STEPS", "50")))
    parser.add_argument("--eval-steps", type=int, default=int(os.environ.get("SHFT_EVAL_STEPS", "50")))
    parser.add_argument("--save-total-limit", type=int, default=int(os.environ.get("SHFT_SAVE_TOTAL_LIMIT", "4")))
    parser.add_argument("--metric-for-best-model", default=os.environ.get("SHFT_METRIC_FOR_BEST_MODEL", "eval_loss"))
    parser.add_argument("--greater-is-better", action="store_true")
    parser.add_argument("--overfit-tolerance", type=float, default=float(os.environ.get("SHFT_OVERFIT_TOLERANCE", "0.10")))
    parser.add_argument("--dry-run", action="store_true")
    parser.add_argument("--expected-dataset-manifest-sha256", default="")
    parser.add_argument("--expected-train-sha256", default="")
    parser.add_argument("--expected-valid-sha256", default="")
    parser.add_argument("--expected-test-sha256", default="")
    args = parser.parse_args()
    dataset_dir = Path(args.dataset_dir) if args.dataset_dir else None
    args.train_jsonl = Path(args.train_jsonl) if args.train_jsonl else (dataset_dir / "train.jsonl" if dataset_dir else None)
    args.valid_jsonl = Path(args.valid_jsonl) if args.valid_jsonl else (dataset_dir / "valid.jsonl" if dataset_dir else None)
    args.output_dir = Path(args.output_dir)
    if args.train_jsonl is None or args.valid_jsonl is None:
        raise SystemExit("--dataset-dir or both --train-jsonl and --valid-jsonl are required")

    train_rows = load_jsonl(args.train_jsonl)
    valid_rows = load_jsonl(args.valid_jsonl)
    args.output_dir.mkdir(parents=True, exist_ok=True)
    plan = build_plan(args, train_rows, valid_rows)
    write_json(args.output_dir / "training_plan.json", plan)

    if not plan["dataset_provenance"]["ok"]:
        result = {
            "status": "blocked_dataset_provenance_mismatch",
            "run_id": args.run_id,
            "completed_at": utc_now(),
            "errors": plan["dataset_provenance"]["errors"],
            "plan_path": str(args.output_dir / "training_plan.json"),
        }
        write_json(args.output_dir / "training_result.json", result)
        print(json.dumps(result, indent=2))
        return 2

    if args.dry_run:
        result = {
            "status": "dry_run_validated",
            "run_id": args.run_id,
            "completed_at": utc_now(),
            "formatted_preview": format_messages(train_rows[0])[:500],
            "plan_path": str(args.output_dir / "training_plan.json"),
        }
        write_json(args.output_dir / "training_result.json", result)
        print(json.dumps(result, indent=2))
        return 0

    result = run_live_training(args, plan)
    write_json(args.output_dir / "training_result.json", result)
    print(json.dumps(result, indent=2))
    return 0 if result["status"] == "completed" else 2


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