train.py

"""Hugging Face Space training core: SFT / DPO / KTO, one step at a time,
resumable.

Design goals
------------
* **One bucket + one model per deployment (env, not per request).** The
  ``BUCKET_ID`` and ``MODEL_ID`` environment variables configure the whole
  Space — they are *not* re-entered on every request. Loading the base model
  once and keeping it resident avoids re-downloading / re-loading it on every
  call.
* **Model kept resident in RAM/VRAM.** The base model is loaded once and kept
  in a module-level cache across calls (the expensive part — download + load —
  happens a single time). The LoRA adapter weights are tiny and are reset from
  the resumed checkpoint on each call, so reuse is safe: the frozen base
  weights never change, only the adapter + optimizer state do.
* **GPU-when-available, CPU otherwise.** The device is chosen from
  ``torch.cuda.is_available()``. On a ZeroGPU Space ``@spaces.GPU`` attaches an
  A10G for the call (so we train on GPU); off ZeroGPU the decorator is a no-op
  and we train on CPU. Because ZeroGPU's GPU is *transient* (it detaches
  between calls), the resident cache is only used when the device is
  persistent (a dedicated-GPU or CPU Space); on ZeroGPU the model is loaded per
  call from the local HF cache instead.
* **No data persistence.** The uploaded "trace" is parsed into an in-memory
  ``datasets`` object for the duration of the call and discarded afterwards.
  Only weights + optimizer/scheduler state are ever written to the bucket.
* **Dynamic gradient accumulation.** The number of optimizer steps a call runs
  is derived from the number of uploaded examples and the deployment's
  gradient-accumulation size ``G`` (``GRAD_ACCUM`` env):
  ``steps = ceil(num_examples / G)`` (SFT/DPO). A single trace -> 1 step; 10
  traces with ``G=4`` -> 3 steps (4 + 4 + 2). The caller no longer chooses a
  fixed "steps per call" — it just uploads the trace.
* **Incremental, resumable training.** Each call advances training by the
  dynamic number of steps above. The HF ``Trainer`` checkpoint
  (model/adapter weights + ``optimizer.pt`` + ``scheduler.pt`` + ``rng`` +
  ``trainer_state.json``) is persisted to an HF Storage Bucket, so the next
  call resumes exactly where it left off — the optimizer does **not** restart
  from scratch. Shut the Space down, come back later, keep training.

Method data schemas
-------------------
SFT  (``trl.SFTTrainer``) — prompt/completion or plain text:
    {"messages": [{"role": "user", "content": "..."},
                  {"role": "assistant", "content": "..."}]}
    or  {"text": "a single string"}

DPO  (``trl.DPOTrainer``) — chosen/rejected pairs:
    {"prompt": "...", "chosen": "...", "rejected": "..."}
    (messages-format ``prompt``/``chosen``/``rejected`` also accepted)

KTO  (``trl.KTOTrainer``) — binary desirable/undesirable feedback:
    {"prompt": "...", "completion": "...", "desirable": true|false}
"""

from __future__ import annotations

import json
import os
import tempfile
from dataclasses import dataclass, field
from pathlib import Path
from typing import Any

import torch
from datasets import Dataset
from huggingface_hub import snapshot_download
from peft import LoraConfig, get_peft_model
from transformers import AutoModelForCausalLM, AutoTokenizer

# `spaces` lets us detect whether we are running on a ZeroGPU Space. On
# ZeroGPU the GPU is *transient* (attached only for the duration of a
# ``@spaces.GPU`` call), so a model cannot be kept resident across calls — it
# would be freed and any cached reference would dangle. We therefore only keep
# the model resident when the device is *persistent* (a dedicated-GPU or CPU
# Space). The import is optional; the platform ships its own `spaces` and, on
# any non-ZeroGPU Space, `spaces.Config.zero_gpu` is simply False.
try:
    from spaces.config import Config as _SpacesConfig

    _ZERO_GPU = bool(getattr(_SpacesConfig, "zero_gpu", False))
except Exception:  # pragma: no cover
    _ZERO_GPU = False

from store import (
    checkpoint_paths,
    pull_checkpoint,
    push_checkpoint,
    save_state_index,
)

HERE = Path(__file__).resolve().parent
LOCAL_CACHE = HERE / ".cache"

# Small, CPU-fine base model used for defaults / smoke tests.
DEFAULT_MODEL_ID = "HuggingFaceTB/SmolLM2-135M"

# Injected when a base model's tokenizer has no chat template, so SFT with
# chat-format ("messages") traces works on *any* base model, not just -Instruct
# ones. A simple, valid template transformers can apply.
DEFAULT_CHAT_TEMPLATE = (
    "{%- for message in messages -%}\n"
    "{% if message['role'] == 'user' %}User: {{ message['content'] }}\n"
    "{% elif message['role'] == 'assistant' %}Assistant: {{ message['content'] }}\n"
    "{% elif message['role'] == 'system' %}System: {{ message['content'] }}\n"
    "{% endif %}{%- endfor -%}"
)


# --------------------------------------------------------------------------- #
#  Configuration
# --------------------------------------------------------------------------- #
@dataclass
class TrainConfig:
    # One model + one bucket per deployment. These come from the Space's
    # environment variables (MODEL_ID, BUCKET_ID) and are NOT re-entered per
    # request. `resolve_*` fills them from the env when left empty, so the
    # whole deployment is configured once.
    model_id: str = ""  # e.g. "HuggingFaceTB/SmolLM2-135M" — from MODEL_ID env
    method: str = "sft"  # sft | dpo | kto
    # LoRA — keeps CPU training tractable on a ZeroGPU container.
    lora_r: int = 8
    lora_alpha: int = 16
    lora_dropout: float = 0.0
    target_modules: tuple = ("q_proj", "k_proj", "v_proj", "o_proj")
    # Optimizer / training.
    learning_rate: float = 1e-4
    # Constant LR keeps "one step at a time" stable across resumable calls;
    # a decaying schedule would require a fixed horizon and is not the goal here.
    lr_scheduler_type: str = "constant"
    optim: str = "adamw_torch"
    weight_decay: float = 0.0
    max_length: int = 256
    # Bucket + run.
    bucket_id: str = ""  # e.g. "username/my-bucket" — from BUCKET_ID env
    resume: bool = True  # resume from latest checkpoint if one exists
    # Gradient accumulation: how many micro-batches to accumulate per optimizer
    # step. The number of optimizer steps this call runs is derived from the
    # number of uploaded examples:  steps_this_call = ceil(num_examples /
    # grad_accum).  10 examples with GRAD_ACCUM=4  ->  3 steps (4+4+2).
    # 1 example  ->  1 step (single trace, single step).
    grad_accum: int = 4
    save_total_limit: int = 3  # keep only the N most-recent checkpoints


# --------------------------------------------------------------------------- #
#  Result
# --------------------------------------------------------------------------- #
@dataclass
class TrainResult:
    method: str
    model_id: str
    global_step: int
    loss: float | None = None
    checkpoint_id: str = ""
    resumed_from: str | None = None
    num_examples: int = 0
    grad_accum: int = 1
    steps_run: int = 0  # optimizer steps this call advanced
    notes: list = field(default_factory=list)
    elapsed: float = 0.0

    def as_dict(self) -> dict:
        return {
            "method": self.method,
            "model_id": self.model_id,
            "global_step": self.global_step,
            "loss": round(self.loss, 6) if self.loss is not None else None,
            "checkpoint_id": self.checkpoint_id,
            "resumed_from": self.resumed_from,
            "num_examples": self.num_examples,
            "grad_accum": self.grad_accum,
            "steps_run": self.steps_run,
            "notes": self.notes,
            "elapsed_sec": round(self.elapsed, 2),
        }


# --------------------------------------------------------------------------- #
#  Trace parsing  (the ONLY place the uploaded data is touched)
# --------------------------------------------------------------------------- #
def parse_trace(rows: list[dict], method: str) -> Dataset:
    """Turn raw uploaded rows into a ``datasets.Dataset`` for the method.

    Rows come in as a list of dicts (one JSON object per line, or a small
    JSON array). Anything that doesn't parse is dropped (logged, not fatal).
    """
    m = method.lower()
    if m == "sft":
        return _parse_sft(rows)
    if m == "dpo":
        return _parse_dpo(rows)
    if m == "kto":
        return _parse_kto(rows)
    raise ValueError(f"Unknown method: {method!r}")


def _parse_sft(rows: list[dict]) -> Dataset:
    # Prefer chat format; fall back to a single "text" field.
    conv = []
    for r in rows:
        if "messages" in r and isinstance(r["messages"], list):
            conv.append({"messages": r["messages"]})
    if conv:
        return Dataset.from_list(conv)
    texts = []
    for r in rows:
        if "text" in r:
            texts.append({"text": r["text"]})
        elif "completion" in r and "prompt" in r:
            # prompt/completion -> single text for SFT
            texts.append({"text": f"{r['prompt']}{r['completion']}"})
    if texts:
        return Dataset.from_list(texts)
    raise ValueError("SFT trace needs rows with 'messages' (chat) or 'text'.")


def _parse_dpo(rows: list[dict]) -> Dataset:
    out = []
    for r in rows:
        if "chosen" in r and "rejected" in r:
            out.append(
                {
                    "prompt": r.get("prompt", ""),
                    "chosen": r["chosen"],
                    "rejected": r["rejected"],
                }
            )
    if not out:
        raise ValueError("DPO trace needs rows with 'chosen' and 'rejected'.")
    return Dataset.from_list(out)


def _parse_kto(rows: list[dict]) -> Dataset:
    # trl's KTOTrainer expects a boolean ``label`` column (True = desirable);
    # the trace uses the friendlier ``desirable`` name, so map it through.
    out = []
    for r in rows:
        completion = r.get("completion")
        desirable = r.get("desirable", r.get("label"))
        if completion is not None and desirable is not None:
            out.append(
                {
                    "prompt": r.get("prompt", ""),
                    "completion": completion,
                    "label": bool(desirable),
                }
            )
    if not out:
        raise ValueError("KTO trace needs rows with 'completion' and 'desirable'.")
    return Dataset.from_list(out)


# --------------------------------------------------------------------------- #
#  Env + resident model cache
# --------------------------------------------------------------------------- #
def _resolve_token() -> str | None:
    # On a Space the token is set as a secret (HF_TOKEN / HUGGING_FACE_HUB_TOKEN).
    return (
        os.environ.get("HF_TOKEN")
        or os.environ.get("HUGGING_FACE_HUB_TOKEN")
        or os.environ.get("HF_HUB_TOKEN")
    )


def resolve_model_id(cfg: TrainConfig) -> str:
    """The model for this deployment, from MODEL_ID env (or the passed
    config). One model per deployment — not re-chosen per request."""
    return (cfg.model_id or os.environ.get("MODEL_ID") or DEFAULT_MODEL_ID).strip()


def resolve_bucket_id(cfg: TrainConfig) -> str:
    """The bucket for this deployment, from BUCKET_ID env (or the passed
    config). One bucket per deployment — not re-chosen per request."""
    return (cfg.bucket_id or os.environ.get("BUCKET_ID") or "").strip()


def resolve_lora_r(cfg: TrainConfig) -> int:
    """LoRA rank for this deployment, from LORA_R env (or the config default)."""
    raw = os.environ.get("LORA_R")
    return int(raw) if raw else cfg.lora_r


def resolve_lora_alpha(cfg: TrainConfig) -> int:
    """LoRA alpha for this deployment, from LORA_ALPHA env (or the config
    default)."""
    raw = os.environ.get("LORA_ALPHA")
    return int(raw) if raw else cfg.lora_alpha


def resolve_lora_dropout(cfg: TrainConfig) -> float:
    """LoRA dropout for this deployment, from LORA_DROPOUT env (or the config
    default)."""
    raw = os.environ.get("LORA_DROPOUT")
    return float(raw) if raw else cfg.lora_dropout


def resolve_target_modules(cfg: TrainConfig):
    """Target modules for this deployment, from LORA_TARGET_MODULES env
    (comma-separated, e.g. ``q_proj,k_proj,v_proj,o_proj``) or the config
    default tuple."""
    raw = os.environ.get("LORA_TARGET_MODULES")
    if raw:
        return tuple(p.strip() for p in raw.split(",") if p.strip())
    return tuple(cfg.target_modules)


def resolve_learning_rate(cfg: TrainConfig) -> float:
    """Learning rate for this deployment, from LEARNING_RATE env (or the config
    default)."""
    raw = os.environ.get("LEARNING_RATE")
    return float(raw) if raw else cfg.learning_rate


def resolve_max_length(cfg: TrainConfig) -> int:
    """Max sequence length for this deployment, from MAX_LENGTH env (or the
    config default)."""
    raw = os.environ.get("MAX_LENGTH")
    return int(raw) if raw else cfg.max_length


def resolve_grad_accum(cfg: TrainConfig) -> int:
    """Gradient-accumulation size for this deployment, from GRAD_ACCUM env
    (or the config default). The number of optimizer steps a call runs is
    derived from the number of uploaded examples:  ceil(num_examples /
    grad_accum). 10 examples with GRAD_ACCUM=4 -> 3 steps (4+4+2); 1 example
    -> 1 step."""
    raw = os.environ.get("GRAD_ACCUM")
    val = int(raw) if raw else cfg.grad_accum
    return max(1, val)


def _device_name() -> str:
    return "cuda" if torch.cuda.is_available() else "cpu"


def _device_map():
    """Where to load the model: auto-shard across every available GPU, else a
    single GPU, else CPU.

    ``device_map="auto"`` lets accelerate spread a large model across all
    visible GPUs (e.g. a 27B base model across 4 GPUs) so the *same code path*
    trains a 27B model on 4 GPUs as a 1B model on 1 GPU — no manual setup.
    With 0/1 GPUs accelerate returns a single device (no sharding, so the
    Trainer's single-device path is used). The model only spans multiple
    devices when >1 GPU is present.
    """
    if torch.cuda.is_available():
        return "auto"  # accelerate distributes across all visible GPUs
    return "cpu"


def _dtype_for_shard(model_id: str, is_sharded: bool):
    """Load a multi-GPU-sharded model in float16 so the per-GPU footprint
    roughly halves (a 27B model = ~54GB in fp32 vs ~27GB in fp16, splitting
    across 4 GPUs becomes trivially affordable). Single-device / CPU models
    stay float32 for accuracy. LoRA training is robust to a fp16 base."""
    if is_sharded and torch.cuda.is_available():
        return torch.float16
    return torch.float32


# A module-level cache of the loaded base model. Populated once on the first
# call and reused on every subsequent call, so the expensive download + load
# happens a single time per deployment. The base (frozen) weights never
# change between calls; only the LoRA adapter + optimizer state do, and those
# are reset from the resumed checkpoint on each call — so reuse is safe.
_MODEL_CACHE: dict = {}


def _model_cache_key(model_id: str) -> str:
    return model_id


def _load_model_fresh(model_id: str, cfg: TrainConfig):
    """Download + load the base model and wrap it in a fresh LoRA adapter.

    The device is chosen automatically: ``device_map="auto"`` shards the model
    across every visible GPU (so a 27B model on 4 GPUs loads the same way a 1B
    model loads on 1 GPU), otherwise it uses a single GPU or CPU. A
    multi-GPU-sharded model is loaded in fp16 to halve the per-GPU footprint;
    single-device / CPU models stay fp32 for accuracy.
    """
    num_gpus = torch.cuda.device_count() if torch.cuda.is_available() else 0
    is_sharded = num_gpus > 1
    device_map = _device_map()
    torch_dtype = _dtype_for_shard(model_id, is_sharded)
    tok = AutoTokenizer.from_pretrained(model_id, token=_resolve_token())
    model = AutoModelForCausalLM.from_pretrained(
        model_id,
        token=_resolve_token(),
        torch_dtype=torch_dtype,
        device_map=device_map,
    )
    if tok.pad_token is None:
        tok.pad_token = tok.eos_token
    # A base (non-instruct) model may have no chat template; SFT with chat
    # traces needs one. Inject a simple default so any base model works.
    if getattr(tok, "chat_template", None) is None:
        tok.chat_template = DEFAULT_CHAT_TEMPLATE

    lora = LoraConfig(
        r=cfg.lora_r,
        lora_alpha=cfg.lora_alpha,
        lora_dropout=cfg.lora_dropout,
        target_modules=list(cfg.target_modules),
        bias="none",
        task_type="CAUSAL_LM",
    )
    model = get_peft_model(model, lora)
    model.print_trainable_parameters()
    if is_sharded:
        # Log the sharding so the user can see how the model was spread.
        devices = sorted({str(d) for d in set(model.hf_device_map.values())})
        print(f"[train] model {model_id} sharded across {len(devices)} device(s): "
              f"{', '.join(devices)} ({num_gpus} GPU, {torch_dtype})")
    return model, tok


def _model_is_cached(model_id: str) -> bool:
    return model_id in _MODEL_CACHE and _MODEL_CACHE[model_id] is not None


def get_model_and_tokenizer(model_id: str, cfg: TrainConfig):
    """Return a model (PEFT-wrapped) + tokenizer for the deployment.

    When the device is *persistent* (not ZeroGPU) the loaded base model is
    cached in a module-level dict and reused across calls — the whole point:
    the base model stays resident in RAM/VRAM instead of being re-downloaded /
    re-loaded on every request. The frozen base weights never change between
    calls; only the LoRA adapter + optimizer state do, and those are reset
    from the resumed checkpoint on each call, so reuse is safe.

    On ZeroGPU the GPU is *transient* (detached between calls), so we skip the
    cache and load per-call from the local HF cache (which still avoids a
    network re-download — only the load is repeated).

    Note: LoRA hyper-params (rank / alpha) are fixed at first-load time. To
    change them mid-deployment, restart the Space (which re-runs the cache).
    """
    # ZeroGPU's GPU is transient: never keep a cached device ref across calls.
    if _ZERO_GPU:
        return _load_model_fresh(model_id, cfg)

    cached = _MODEL_CACHE.get(model_id)
    if cached is not None:
        return cached

    model, tok = _load_model_fresh(model_id, cfg)
    _MODEL_CACHE[model_id] = (model, tok)
    return model, tok


def _build_config(cfg: TrainConfig, output_dir: Path, num_examples: int,
                  current_step: int):
    """Build a TRL training config for a single call.

    **Dynamic gradient accumulation.** The number of optimizer steps this call
    runs is derived from how many examples the uploaded trace contains and the
    deployment's gradient-accumulation size ``G`` (``GRAD_ACCUM`` env), *not*
    from a fixed "steps per call":

        steps_this_call = ceil(num_examples / (per_device_batch * accum))

    with ``accum = ceil(G / per_device_batch)``. With ``G=4`` this is exactly
    the requested behaviour: a single trace -> 1 step; 10 traces -> 3 steps
    (4 + 4 + 2). SFT/DPO use a per-device batch of 1 (so one optimizer step ==
    ``G`` examples, matching the example precisely); KTO uses a per-device
    batch of 2 (its KL term is undefined for a single example), so its per-step
    example count is ``2 * ceil(G/2)``.

    ``max_steps`` is the *absolute* target (``current_step + steps_this_call``);
    with ``resume_from_checkpoint`` the Trainer runs the remaining steps, so the
    call continues exactly where the last call left off.
    """
    import math

    use_gpu = torch.cuda.is_available()
    use_cpu = not use_gpu
    m = cfg.method.lower()
    # KTO needs an actual batch > 1 (its KL term is undefined for a single
    # example); SFT/DPO use a per-device batch of 1.
    per_device_batch = 2 if m == "kto" else 1
    accum = max(1, math.ceil(cfg.grad_accum / per_device_batch))
    examples_per_step = per_device_batch * accum
    steps_this_call = max(1, math.ceil(num_examples / examples_per_step)) if num_examples else 1
    target = current_step + steps_this_call

    common = dict(
        output_dir=str(output_dir),
        per_device_train_batch_size=per_device_batch,
        gradient_accumulation_steps=accum,
        max_steps=target,  # absolute target; resume continues from here
        learning_rate=cfg.learning_rate,
        lr_scheduler_type=cfg.lr_scheduler_type,
        optim=cfg.optim,
        weight_decay=cfg.weight_decay,
        max_grad_norm=1.0,
        logging_strategy="steps",
        logging_steps=1,
        save_strategy="steps",
        save_steps=1,
        save_total_limit=cfg.save_total_limit,
        save_only_model=False,  # MUST keep optimizer.pt / scheduler.pt / rng
        use_cpu=use_cpu,        # GPU when available, else CPU
        report_to="none",
        dataloader_num_workers=0,
        remove_unused_columns=False,
        ignore_data_skip=True,  # robust resume for tiny/arbitrary datasets
    )
    if m == "sft":
        from trl import SFTConfig

        args = SFTConfig(**common, max_length=cfg.max_length)
    elif m == "dpo":
        from trl import DPOConfig

        args = DPOConfig(**common, max_length=cfg.max_length)
    elif m == "kto":
        from trl import KTOConfig

        args = KTOConfig(**common, max_length=cfg.max_length)
    else:
        raise ValueError(f"Unknown method {cfg.method!r}")
    return args, steps_this_call


def _build_trainer(method: str, model, tok, dataset, args):
    if method == "sft":
        from trl import SFTTrainer

        return SFTTrainer(model=model, args=args, train_dataset=dataset, processing_class=tok)
    if method == "dpo":
        from trl import DPOTrainer

        # ref_model=None + a PEFT model => trl uses the adapter-disabled model
        # as the reference (standard LoRA-DPO pattern).
        return DPOTrainer(model=model, args=args, train_dataset=dataset, processing_class=tok)
    if method == "kto":
        from trl import KTOTrainer

        return KTOTrainer(model=model, args=args, train_dataset=dataset, processing_class=tok)
    raise ValueError(f"Unknown method {method!r}")


# --------------------------------------------------------------------------- #
#  The main entry point
# --------------------------------------------------------------------------- #
def _steps_for_call(cfg: TrainConfig, num_examples: int) -> int:
    """Dynamic optimizer-step count for a single call.

    ``steps = ceil(num_examples / (per_device_batch * accum))`` with
    ``accum = ceil(cfg.grad_accum / per_device_batch)``. SFT/DPO use a
    per-device batch of 1, so ``steps = ceil(num_examples / grad_accum)`` —
    exactly the requested behaviour (10 examples, grad_accum=4 -> 3 steps;
    1 example -> 1 step). KTO uses a per-device batch of 2. A call with no
    examples still runs 1 step (keeps the optimizer warm / makes the call
    idempotent)."""
    import math

    m = cfg.method.lower()
    per_device_batch = 2 if m == "kto" else 1
    accum = max(1, math.ceil(cfg.grad_accum / per_device_batch))
    examples_per_step = per_device_batch * accum
    if not num_examples:
        return 1
    return max(1, math.ceil(num_examples / examples_per_step))


def run_one_call(
    rows: list[dict],
    cfg: TrainConfig,
    *,
    log=print,
) -> TrainResult:
    """Parse the uploaded trace, run ``steps_per_call`` optimizer steps for the
    chosen method (resuming from the latest checkpoint if one exists), then
    persist the resulting weights + optimizer state to the bucket.

    The uploaded ``rows`` (the training data) are never written anywhere —
    only the trained weights and optimizer/scheduler state are persisted.
    """
    import time

    start = time.time()
    # One bucket + one model per deployment: resolve from the env when the
    # caller did not pin them (the UI no longer asks for them per request).
    cfg.model_id = resolve_model_id(cfg)
    cfg.bucket_id = resolve_bucket_id(cfg)
    # LoRA hyper-params are also fixed per deployment via env (one model + one
    # bucket + one LoRA config per deployment). They are resolved here so the
    # resident model cache (keyed by model_id) wraps the *deployment's* LoRA
    # config at first-load, and every later call reuses the same cached model.
    cfg.lora_r = resolve_lora_r(cfg)
    cfg.lora_alpha = resolve_lora_alpha(cfg)
    cfg.lora_dropout = resolve_lora_dropout(cfg)
    cfg.target_modules = resolve_target_modules(cfg)
    cfg.learning_rate = resolve_learning_rate(cfg)
    cfg.max_length = resolve_max_length(cfg)
    cfg.grad_accum = resolve_grad_accum(cfg)
    m = cfg.method.lower()
    dataset = parse_trace(rows, m)
    num_examples = len(dataset)
    log(f"[train] model={cfg.model_id}  bucket={cfg.bucket_id or '(none)'}  "
        f"parsed {num_examples} {m} examples from trace  "
        f"(grad_accum={cfg.grad_accum})")

    notes: list[str] = []

    # --- resolve resume checkpoint -------------------------------------- #
    local_cache = LOCAL_CACHE
    local_cache.mkdir(parents=True, exist_ok=True)
    resume_dir = None
    resumed_from = None
    if cfg.resume:
        resume_dir, resumed_from = _resolve_resume(
            cfg, local_cache, notes
        )
        if resume_dir:
            log(f"[train] resuming from {resumed_from}")
    current_step = _read_global_step(resume_dir) if resume_dir else 0

    # --- dynamic gradient accumulation: how many steps this call runs -- #
    steps_this_call = _steps_for_call(cfg, num_examples)
    target = current_step + steps_this_call
    output_dir = local_cache / f"run-{target}"
    log(
        f"[train] {num_examples} example(s) / grad_accum={cfg.grad_accum} "
        f"-> {steps_this_call} step(s): {current_step} -> {target}"
    )

    args, steps_this_call = _build_config(cfg, output_dir, num_examples, current_step)

    # --- build model + trainer (resident cache keeps the base model in
    #     RAM/VRAM across calls; the adapter+optimizer state is reset from the
    #     resumed checkpoint below) ------------------------- #
    model, tok = get_model_and_tokenizer(cfg.model_id, cfg)
    trainer = _build_trainer(m, model, tok, dataset, args)

    # --- train (resume continues from current_step to target) ---------- #
    train_args = {}
    if resume_dir:
        train_args["resume_from_checkpoint"] = resume_dir
    log(f"[train] training steps {current_step} -> {target}")
    train_res = trainer.train(**train_args)

    # --- read back the new global step + loss -------------------------- #
    new_step = current_step
    loss = None
    try:
        st = trainer.state
        new_step = int(getattr(st, "global_step", current_step) or current_step)
    except Exception:
        pass
    logs = getattr(train_res, "log", None) or []
    if logs:
        for entry in logs:
            if "loss" in entry:
                loss = float(entry["loss"])
    if loss is None:
        try:
            metrics = getattr(train_res, "metrics", None) or {}
            train_metrics = metrics.get("train", {})
            if "train_loss" in train_metrics:
                loss = float(train_metrics["train_loss"])
        except Exception:
            pass

    # --- persist the latest checkpoint to the bucket ------------------- #
    checkpoint_id = f"checkpoint-{new_step}"
    saved_dir = _find_checkpoint_dir(output_dir, new_step) or output_dir
    if cfg.bucket_id:
        try:
            push_checkpoint(cfg.bucket_id, saved_dir, checkpoint_id)
            save_state_index(
                cfg.bucket_id,
                {
                    "global_step": new_step,
                    "method": m,
                    "model_id": cfg.model_id,
                    "loss": loss,
                },
            )
            log(f"[train] pushed {checkpoint_id} to bucket {cfg.bucket_id}")
        except Exception as exc:
            notes.append(f"push failed: {exc}")
            log(f"[train] PUSH FAILED: {exc}")
    else:
        notes.append("no bucket_id set — weights NOT persisted")

    result = TrainResult(
        method=m,
        model_id=cfg.model_id,
        global_step=new_step,
        loss=loss,
        checkpoint_id=checkpoint_id,
        resumed_from=resumed_from,
        num_examples=num_examples,
        grad_accum=cfg.grad_accum,
        steps_run=new_step - current_step,
        notes=notes,
        elapsed=time.time() - start,
    )
    log(f"[train] done: {result.as_dict()}")
    return result


# --------------------------------------------------------------------------- #
#  helpers
# --------------------------------------------------------------------------- #
def _resolve_resume(cfg: TrainConfig, local_cache: Path, notes: list) -> tuple[str | None, str | None]:
    """Decide which checkpoint to resume from.

    1. Prefer a fresh local checkpoint from a previous call (same session) —
       so back-to-back calls resume even without a bucket.
    2. Otherwise pull the latest checkpoint from the bucket.
    """
    # 1) fresh local checkpoint from a prior call
    local = _find_any_checkpoint(local_cache)
    if local is not None:
        return str(local), local.name
    # 2) latest from the bucket
    if cfg.bucket_id:
        try:
            latest = _latest_checkpoint_prefix(cfg.bucket_id)
        except Exception as exc:  # pragma: no cover
            latest = None
            notes.append(f"could not list bucket: {exc}")
        if latest:
            local_dir = local_cache / f"pull-{latest}"
            if pull_checkpoint(cfg.bucket_id, latest, local_dir):
                return str(local_dir), latest
            notes.append(f"could not download {latest}")
    return None, None


def _find_any_checkpoint(root: Path):
    """Return the local ``.../checkpoint-N`` dir with the highest N under
    ``root`` (searches run-*/checkpoint-* and top-level checkpoint-*)."""
    best = None
    for cand in list(root.glob("checkpoint-*")) + list(root.glob("run-*/checkpoint-*")):
        s = _dir_step(cand)
        if best is None or s > _dir_step(best):
            best = cand
    return best


def _latest_checkpoint_prefix(bucket_id: str) -> str | None:
    try:
        from store import list_latest_checkpoint

        return list_latest_checkpoint(bucket_id)
    except Exception:
        return None


def _read_global_step(checkpoint_dir: str | None) -> int:
    if not checkpoint_dir:
        return 0
    p = Path(checkpoint_dir) / "trainer_state.json"
    try:
        data = json.loads(p.read_text())
        return int(data.get("global_step", 0))
    except Exception:
        return 0


def _find_checkpoint_dir(output_dir: Path, step: int) -> Path | None:
    cand = output_dir / f"checkpoint-{step}"
    if cand.exists():
        return cand
    # the trainer may have saved under a different name; take the max one
    best = None
    for d in output_dir.glob("checkpoint-*"):
        try:
            s = int(d.name.split("-")[-1])
        except Exception:
            continue
        if best is None or s > int(_dir_step(best)):
            best = d
    return best


def _dir_step(d: Path) -> int:
    try:
        return int(d.name.split("-")[-1])
    except Exception:
        return -1


def current_progress(bucket_id: str) -> dict:
    """Lightweight status: latest checkpoint + global step from the index."""
    try:
        from store import load_state_index

        idx = load_state_index(bucket_id)
    except Exception:
        idx = {}
    latest = None
    try:
        latest = _latest_checkpoint_prefix(bucket_id)
    except Exception:
        pass
    return {
        "global_step": idx.get("global_step", 0),
        "latest_checkpoint": latest,
        "checkpoints": checkpoint_paths(bucket_id),
    }