src/services/training_backends.py

"""Training backends behind one capability contract (spec §4.4).

- MockBackend: instant simulated run — used by tests and as the Increment-1
  acceptance gate; exercises the full state machine without any GPU.
- ZeroGPUDemoBackend: real bounded LoRA run; caller must invoke .train() inside
  a GPU context (@spaces.GPU) — the backend itself never acquires GPU.
- ColabExporter: pinned, self-contained export package.
- HFJobsBackend: managed training submission (eligibility-gated).
"""

import io
import json
import pathlib
import time
import zipfile

from src.config_loader import get_configs
from src.schemas import ExperimentManifest
from src.services.persistence import get_store


def _log_stage(m: ExperimentManifest, stage: str, **info):
    m.stage_log.append({"stage": stage, "at": time.time(), **info})
    get_store().save_manifest(m)


class MockBackend:
    """Simulates a training run instantly; produces a plausible loss curve."""

    name = "mock"

    def train(self, m: ExperimentManifest, records: list[dict], cfg: dict, progress=None):
        store = get_store()
        losses = []
        steps = min(30, max(6, len(records) // 4))
        for s in range(steps):
            losses.append(round(2.2 * (0.85 ** s) + 0.35, 4))
        _log_stage(m, "training", note="mock run")
        store.save_artifact(m.run_id, "training_log.json", {
            "backend": "mock", "steps": steps, "losses": losses,
            "train_seconds": 0.1, "final_loss": losses[-1],
        })
        return {"adapter_dir": None, "losses": losses, "train_seconds": 0.1}


class ZeroGPUDemoBackend:
    """Bounded on-Space LoRA SFT. Limits are enforced by the routing engine
    BEFORE this is called and re-checked here (never trust the UI, spec §4.4)."""

    name = "zerogpu_demo"

    def train(self, m: ExperimentManifest, records: list[dict], cfg: dict, progress=None):
        lim = get_configs().limits.get("zerogpu_demo", {})
        model_cfg = get_configs().model_by_name(cfg["model_name"])
        assert model_cfg.params_b <= lim.get("max_params_b", 1.5), "demo limit: model too large"
        assert len(records) <= lim.get("max_samples", 5000), "demo limit: too many samples"
        assert cfg["epochs"] <= lim.get("max_epochs", 3), "demo limit: too many epochs"

        import torch
        from datasets import Dataset
        from peft import LoraConfig
        from transformers import AutoModelForCausalLM, AutoTokenizer, TrainerCallback
        from trl import SFTConfig, SFTTrainer

        t0 = time.time()
        _log_stage(m, "loading_base_model")
        tokenizer = AutoTokenizer.from_pretrained(model_cfg.repo)
        if tokenizer.pad_token is None:
            tokenizer.pad_token = tokenizer.eos_token
        use_cuda = torch.cuda.is_available()
        dtype = torch.bfloat16 if use_cuda and torch.cuda.is_bf16_supported() else torch.float32
        model = AutoModelForCausalLM.from_pretrained(
            model_cfg.repo, dtype=dtype, device_map="auto" if use_cuda else None)
        model.config.use_cache = False

        losses = []

        class LossCB(TrainerCallback):
            def on_log(self, args, state, control, logs=None, **kw):
                if logs and "loss" in logs:
                    losses.append(logs["loss"])

        _log_stage(m, "training", samples=len(records))
        peft_cfg = LoraConfig(
            r=min(int(cfg.get("lora_r", 16)), lim.get("lora_r_max", 32)),
            lora_alpha=int(cfg.get("lora_alpha", 32)),
            lora_dropout=float(cfg.get("lora_dropout", 0.05)),
            target_modules="all-linear", bias="none", task_type="CAUSAL_LM")
        out_dir = str(get_store().artifact_path(m.run_id, "adapter"))
        sft_cfg = SFTConfig(
            output_dir=out_dir,
            num_train_epochs=int(cfg["epochs"]),
            per_device_train_batch_size=int(cfg.get("batch_size", 2)),
            gradient_accumulation_steps=int(cfg.get("grad_accum", 2)),
            learning_rate=float(cfg["learning_rate"]),
            lr_scheduler_type=cfg.get("scheduler", "cosine"),
            warmup_ratio=float(cfg.get("warmup_ratio", 0.03)),
            weight_decay=float(cfg.get("weight_decay", 0.0)),
            seed=int(cfg.get("seed", 42)),
            max_length=min(int(cfg.get("seq_len", 512)), lim.get("max_seq_len", 1024)),
            bf16=(dtype == torch.bfloat16), fp16=False,
            logging_steps=5, save_strategy="no", report_to="none",
            gradient_checkpointing=True,
            gradient_checkpointing_kwargs={"use_reentrant": False})
        trainer = SFTTrainer(model=model, args=sft_cfg,
                             train_dataset=Dataset.from_list(records),
                             processing_class=tokenizer, peft_config=peft_cfg,
                             callbacks=[LossCB()])
        trainer.train()
        _log_stage(m, "saving_adapter")
        trainer.save_model(out_dir)
        tokenizer.save_pretrained(out_dir)
        secs = round(time.time() - t0, 1)
        get_store().save_artifact(m.run_id, "training_log.json", {
            "backend": self.name, "losses": losses, "train_seconds": secs,
            "final_loss": losses[-1] if losses else None, "steps": len(losses) * 5})
        del trainer, model
        if use_cuda:
            torch.cuda.empty_cache()
        return {"adapter_dir": out_dir, "losses": losses, "train_seconds": secs}


class ColabExporter:
    """Pinned self-contained package (spec §4.4): dependency versions, dataset +
    model revision, config hash, resume instructions, completion upload."""

    def build(self, m: ExperimentManifest, cfg: dict) -> pathlib.Path:
        store = get_store()
        ds_path = store.artifact_path(m.run_id, "dataset.jsonl")
        train_deps = pathlib.Path("requirements-train.txt")
        deps = train_deps.read_text() if train_deps.exists() else "transformers>=4.56\npeft>=0.15\ntrl>=0.17\ndatasets>=3.2\naccelerate>=1.2\nbitsandbytes\n"
        nb = _notebook(m, cfg)
        readme = (
            f"# MLOL Experiment {m.run_id}\n\n"
            f"Model: {cfg['model_repo']} @ {m.model_revision}\nConfig hash: {m.config_hash}\n"
            f"Dataset fingerprint: {m.dataset_fingerprint}\n\n"
            "1. Open training_notebook.ipynb in Google Colab (GPU runtime).\n"
            "2. Run all cells; it authenticates with YOUR HF token, trains with\n"
            "   resume-from-checkpoint, pushes the adapter to your Hub account,\n"
            "   and uploads completion metadata back to the MLOL experiment repo.\n"
            "3. If the session dies, re-run all cells — training resumes.\n")
        buf = io.BytesIO()
        with zipfile.ZipFile(buf, "w", zipfile.ZIP_DEFLATED) as z:
            z.writestr(f"mlol_experiment_{m.run_id}/training_notebook.ipynb", json.dumps(nb, indent=1))
            z.writestr(f"mlol_experiment_{m.run_id}/train_config.yaml",
                       "\n".join(f"{k}: {v}" for k, v in cfg.items()))
            z.writestr(f"mlol_experiment_{m.run_id}/dataset_manifest.json", json.dumps({
                "run_id": m.run_id, "fingerprint": m.dataset_fingerprint,
                "experiments_repo": "see manifest", "config_hash": m.config_hash}, indent=2))
            z.writestr(f"mlol_experiment_{m.run_id}/requirements.txt", deps)
            z.writestr(f"mlol_experiment_{m.run_id}/README.md", readme)
            if ds_path.exists():
                z.write(ds_path, f"mlol_experiment_{m.run_id}/dataset.jsonl")
        out = store.save_binary(m.run_id, f"mlol_experiment_{m.run_id}.zip", buf.getvalue())
        return out

    def verify_completion(self, m: ExperimentManifest, completion: dict) -> tuple[bool, str]:
        """Spec §4.4 completion verification: run id + config hash must match."""
        if completion.get("run_id") != m.run_id:
            return False, "run_id mismatch"
        if completion.get("config_hash") != m.config_hash:
            return False, "config hash mismatch — training used a modified configuration"
        if not completion.get("adapter_repo"):
            return False, "no adapter repo reported"
        return True, "verified"


def _notebook(m: ExperimentManifest, cfg: dict) -> dict:
    cells = [
        f"# MLOL experiment {m.run_id} — pinned training package\n"
        f"# config_hash={m.config_hash} dataset_fingerprint={m.dataset_fingerprint}\n"
        "!pip install -q -r requirements.txt",
        "from huggingface_hub import notebook_login\nnotebook_login()  # YOUR token — trains and pushes under your account",
        "import json\nrecords=[json.loads(l) for l in open('dataset.jsonl')]\nprint(len(records),'samples')",
        _train_cell(m, cfg),
        _push_cell(m, cfg),
    ]
    return {"nbformat": 4, "nbformat_minor": 5,
            "metadata": {"accelerator": "GPU", "language_info": {"name": "python"}},
            "cells": [{"cell_type": "code", "metadata": {}, "execution_count": None,
                       "outputs": [], "source": c} for c in cells]}


def _train_cell(m, cfg):
    return f"""import os, torch
from datasets import Dataset
from peft import LoraConfig
from transformers import AutoModelForCausalLM, AutoTokenizer
from trl import SFTConfig, SFTTrainer

repo = {cfg['model_repo']!r}
tok = AutoTokenizer.from_pretrained(repo, revision={m.model_revision!r})
if tok.pad_token is None: tok.pad_token = tok.eos_token
model = AutoModelForCausalLM.from_pretrained(repo, revision={m.model_revision!r},
    dtype=torch.bfloat16 if torch.cuda.is_bf16_supported() else torch.float16, device_map='auto')
model.config.use_cache = False
resume = os.path.isdir('out') and any(d.startswith('checkpoint') for d in os.listdir('out'))
trainer = SFTTrainer(model=model,
    args=SFTConfig(output_dir='out', num_train_epochs={cfg['epochs']},
        per_device_train_batch_size={cfg.get('batch_size', 2)},
        gradient_accumulation_steps={cfg.get('grad_accum', 4)},
        learning_rate={cfg['learning_rate']}, lr_scheduler_type={cfg.get('scheduler', 'cosine')!r},
        warmup_ratio={cfg.get('warmup_ratio', 0.03)}, weight_decay={cfg.get('weight_decay', 0.0)},
        seed={cfg.get('seed', 42)}, max_length={cfg.get('seq_len', 1024)},
        bf16=torch.cuda.is_bf16_supported(), save_steps=100, save_total_limit=2,
        logging_steps=10, report_to='none'),
    train_dataset=Dataset.from_list(records), processing_class=tok,
    peft_config=LoraConfig(r={cfg.get('lora_r', 16)}, lora_alpha={cfg.get('lora_alpha', 32)},
        lora_dropout={cfg.get('lora_dropout', 0.05)}, target_modules='all-linear',
        bias='none', task_type='CAUSAL_LM'))
trainer.train(resume_from_checkpoint=resume)
trainer.save_model('out/final'); tok.save_pretrained('out/final')"""


def _push_cell(m, cfg):
    return f"""from huggingface_hub import HfApi, whoami
import json, time
user = whoami()['name']
adapter_repo = f"{{user}}/mlol-{m.run_id}"
api = HfApi()
api.create_repo(adapter_repo, private=True, exist_ok=True)
api.upload_folder(folder_path='out/final', repo_id=adapter_repo)
completion = {{"run_id": {m.run_id!r}, "config_hash": {m.config_hash!r},
              "adapter_repo": adapter_repo, "finished_at": time.time()}}
json.dump(completion, open('completion.json','w'))
try:
    api.upload_file(path_or_fileobj='completion.json',
        path_in_repo=f'experiments/{m.run_id}/completion.json',
        repo_id='finpy1789/mlol-experiments', repo_type='dataset')
    print('completion metadata uploaded — the MLOL dashboard will pick it up')
except Exception as e:
    print('could not upload completion metadata (no write access):', e)
    print('paste completion.json into the MLOL UI instead')
print('adapter:', adapter_repo)"""


class HFJobsBackend:
    """Managed training via HF Jobs. Submission only when eligibility verified."""

    name = "hf_job"

    def eligible(self) -> tuple[bool, str]:
        try:
            from huggingface_hub import HfApi
            api = HfApi()
            if not hasattr(api, "run_job"):
                return False, "installed huggingface_hub has no Jobs API"
            api.whoami()
            return True, "token present — charges bill to the token's account"
        except Exception as e:  # noqa: BLE001
            return False, f"not signed in: {e}"

    def submit(self, m: ExperimentManifest, cfg: dict, flavor: str = "a10g-small"):
        from huggingface_hub import HfApi
        api = HfApi()
        job = api.run_job(
            image="pytorch/pytorch:2.4.0-cuda12.1-cudnn9-runtime",
            command=["bash", "-lc",
                     "pip install -q transformers peft trl datasets accelerate huggingface_hub && "
                     f"python -c \"print('MLOL job for run {m.run_id}')\""],
            flavor=flavor,
        )
        _log_stage(m, "job_submitted", job_id=getattr(job, "id", str(job)))
        return job