feedmill-auto-label/backend/training.py

"""Fine-tune the project's base model on its master dataset (REQ-060…065).

The default is old + new together: the master dataset already accumulates every
merged batch, so a run sees the whole history. Training on the newest batch
alone is what makes a model quietly forget what it used to know, so it is not
what happens here.
"""

import json
import os
import shutil
import time
from typing import Optional

from backend import (augment, base_dataset, config, dataset, datasets, db, evaluate,
                     hardware, jobs, projects)

PRETRAINED = {"bbox": "yolo11n.pt", "polygon": "yolo11n-seg.pt"}


class TrainingError(Exception):
    pass


def models_dir(project_slug: str) -> str:
    return os.path.join(config.project_dir(project_slug), "models")


#: Rough bytes one decoded 640px training image occupies in the RAM cache,
#: measured against the run that OOMed: 10.2 GB across 15,774 images.
_BYTES_PER_CACHED_IMAGE = 700_000


def _cache_mode(train_images: int, job) -> object:
    """Pick Ultralytics' `cache` argument for the memory this host actually has.

    RAM caching is a large speedup and worth taking when it fits. It is only
    taken with three times the headroom the raw estimate asks for: the run that
    died had ~10 GB of cache on a 30 GB host and still lost, because the
    dataloader workers fork after the cache is built and their copy-on-write
    pages are what turn "just fits" into a kill. Two-times headroom would have
    green-lit exactly the run that failed.
    """
    needed = train_images * _BYTES_PER_CACHED_IMAGE
    try:
        import psutil
        available = psutil.virtual_memory().available
    except Exception:
        available = 0

    if available == 0:
        job.log(f"Image cache: disk (cannot read free memory; {train_images} images)")
        return "disk"
    if needed * 3 <= available:
        job.log(f"Image cache: RAM (~{needed / 1e9:.1f} GB of "
                f"{available / 1e9:.1f} GB free)")
        return "ram"
    job.log(f"Image cache: disk (RAM cache would need ~{needed / 1e9:.1f} GB, "
            f"only {available / 1e9:.1f} GB free)")
    return "disk"


def start(project_id: int, epochs: int = 50, overrides: Optional[dict] = None,
          batch_ids: Optional[list] = None, class_ids: Optional[list] = None,
          dataset_ids: Optional[list] = None,
          base_dataset_ids: Optional[list] = None) -> dict:
    project = projects.get(project_id)
    if project is None:
        raise TrainingError("No such project")

    bases = list(base_dataset_ids or [])

    # No dataset picked means "everything this project has", which is what the
    # single-dataset app always did. Base datasets are opt-in, so an empty pick
    # never silently drags them in.
    chosen = list(dataset_ids or [])
    if not chosen and not bases:
        chosen = [item["id"] for item in datasets.listing(project_id)]
    if not chosen and not bases:
        raise TrainingError(
            "This project has no dataset yet — approve and merge a batch before training"
        )

    items = datasets.combined_items(project_id, chosen) if chosen else []
    base_train = sum(base_dataset.get(bid)["image_count"] for bid in bases
                     if base_dataset.get(bid) is not None)
    counts = {"train": sum(1 for i in items if i["split"] == "train") + base_train,
              "val": sum(1 for i in items if i["split"] == "val")}
    if counts["train"] == 0:
        raise TrainingError(
            "The chosen dataset(s) hold no training images — merge a batch before training"
        )
    # A base dataset is train-only, so it can never supply the val split that
    # REQ-063's base-vs-new comparison is measured on.
    if counts["val"] == 0:
        raise TrainingError(
            "Nothing to validate on — a base dataset only contributes training images, "
            "so pick at least one of this project's own datasets too"
        )

    settings = hardware.resolve(overrides, epochs)
    job = jobs.create(
        "train",
        params={"project_id": project_id, "settings": settings, "batch_ids": batch_ids,
                "class_ids": class_ids, "dataset_ids": chosen,
                "base_dataset_ids": bases},
        project_id=project_id,
        message=f"{counts['train']} train / {counts['val']} val",
    )
    return job.to_dict()


def listing(project_id: int) -> list:
    with db.cursor() as cur:
        cur.execute(
            "SELECT * FROM model_versions WHERE project_id = ? ORDER BY version DESC",
            (project_id,),
        )
        rows = []
        for row in cur.fetchall():
            item = dict(row)
            item["metrics"] = json.loads(item["metrics"] or "null")
            item["base_metrics"] = json.loads(item["base_metrics"] or "null")
            rows.append(item)
        return rows


def get_version(model_id: int) -> Optional[dict]:
    with db.cursor() as cur:
        cur.execute("SELECT * FROM model_versions WHERE id = ?", (model_id,))
        row = cur.fetchone()
        return dict(row) if row else None


def promote(model_id: int) -> dict:
    """Make a trained version the project's base model for the next round (REQ-064)."""
    version = get_version(model_id)
    if version is None:
        raise TrainingError("No such model version")
    project = projects.get(version["project_id"])
    base_path = os.path.join(config.project_dir(project["slug"]), "base", "model.pt")
    os.makedirs(os.path.dirname(base_path), exist_ok=True)
    shutil.copyfile(version["weights_path"], base_path)
    with db.cursor() as cur:
        cur.execute(
            "UPDATE projects SET base_model_path = ?, base_model_kind = 'trained' WHERE id = ?",
            (base_path, project["id"]),
        )
    return projects.get(project["id"])


def _next_version(cur, project_id: int) -> int:
    cur.execute(
        "SELECT COALESCE(MAX(version), 0) + 1 FROM model_versions WHERE project_id = ?",
        (project_id,),
    )
    return cur.fetchone()[0]


@jobs.handler("train")
def _run_train(job) -> None:
    from ultralytics import YOLO

    os.environ["ULTRALYTICS_OFFLINE"] = "true"
    os.environ["YOLO_OFFLINE"] = "true"

    project = projects.get(job.params["project_id"])
    settings = job.params["settings"]
    batch_ids = job.params.get("batch_ids")
    class_ids = job.params.get("class_ids")
    dataset_ids = job.params.get("dataset_ids") or []
    data_yaml = dataset.write_data_yaml(project, dataset_ids, batch_ids=batch_ids,
                                        selected_class_ids=class_ids, require_val=True,
                                        base_dataset_ids=job.params.get("base_dataset_ids") or [])

    # SAM3 and a training run must not hold VRAM at the same time (REQ-065).
    from backend.sam3_engine import release_engine

    if release_engine():
        job.log("Released SAM3 from VRAM before training")

    start_point = project["base_model_path"] or PRETRAINED[project["label_type"]]
    if not project["base_model_path"]:
        job.log(f"No base model on this project — starting from {start_point}")
    job.log(f"Fine-tuning {os.path.basename(start_point)} for {settings['epochs']} epoch(s) "
            f"(batch={settings['batch']}, imgsz={settings['imgsz']}, "
            f"device={settings['device']})")

    with db.cursor() as cur:
        version = _next_version(cur, project["id"])
    out_dir = os.path.join(models_dir(project["slug"]), str(version))
    os.makedirs(out_dir, exist_ok=True)

    model = YOLO(start_point)

    def on_epoch(trainer):
        # trainer.epoch is 0-based; report a human-facing 1-based count.
        epoch = getattr(trainer, 'epoch', 0) + 1
        total = getattr(trainer, 'epochs', settings["epochs"])
        job.progress(epoch, total, f"epoch {epoch}/{total}")
        if job.cancelled:
            trainer.stop_training = True

    model.add_callback("on_fit_epoch_end", on_epoch)
    job.progress(0, settings["epochs"])

    import torch
    if torch.cuda.is_available():
        torch.backends.cudnn.benchmark = True

    # REQ-110: explicit rather than inherited. An untouched project gets MEDIUM,
    # which is Ultralytics' own default set, so this changes nothing by itself.
    augmentation = augment.get(project["id"])
    job.log(f"Augmentation: {augmentation['preset']} — "
            + ", ".join(f"{k}={v:g}" for k, v in sorted(augmentation["settings"].items())))

    # `cache="ram"` used to be hardcoded. It holds the whole training set in
    # memory, which was invisible at a few hundred images and fatal at fifteen
    # thousand: the run below died mid-epoch with no traceback, killed by the
    # host OOM killer, because 10 GB of cache plus per-worker copies did not fit
    # in 30 GB. RAM caching is now earned, not assumed (CLAUDE.md §9).
    train_list = os.path.join(os.path.dirname(data_yaml), "selected_train.txt")
    with open(train_list, encoding="utf-8") as handle:
        train_images = sum(1 for line in handle if line.strip())
    cache_mode = _cache_mode(train_images, job)

    keep_run_dir = False
    try:
        model.train(
            data=data_yaml,
            epochs=settings["epochs"],
            imgsz=settings["imgsz"],
            batch=settings["batch"],
            **augmentation["settings"],
            device=settings["device"],
            workers=settings.get("workers", 8),
            cache=cache_mode,
            project=os.path.join(out_dir, "runs"),
            name="train",
            exist_ok=True,
            amp=True,
            plots=False,
            verbose=False,
        )

        produced = os.path.join(out_dir, "runs", "train", "weights", "best.pt")
        if not os.path.isfile(produced):
            raise TrainingError("Training finished without producing best.pt")
        weights = os.path.join(out_dir, "best.pt")
        shutil.copyfile(produced, weights)

        job.log("Validating the base model and the new one on the same val set…")
        comparison = evaluate.compare(
            project["base_model_path"] or None, weights, data_yaml,
            [item["name"] for item in project["classes"]],
            imgsz=settings["imgsz"], device=settings["device"], batch=settings["batch"],
        )
        with open(os.path.join(out_dir, "metrics.json"), "w", encoding="utf-8") as handle:
            json.dump(comparison, handle, indent=2)

        with db.cursor() as cur:
            cur.execute(
                """INSERT INTO model_versions (project_id, version, weights_path,
                                               parent_model_path, metrics, base_metrics,
                                               created_at, augment)
                   VALUES (?, ?, ?, ?, ?, ?, ?, ?)""",
                (project["id"], version, weights, project["base_model_path"],
                 json.dumps(comparison["new"]), json.dumps(comparison["base"]), time.time(),
                 json.dumps(augmentation["settings"])),
            )

        new = comparison["new"]
        if comparison["delta"]:
            delta = comparison["delta"]
            job.log(f"v{version}: mAP50 {new['map50']:.4f} ({delta['map50']:+.4f} vs base), "
                    f"mAP50-95 {new['map50_95']:.4f} ({delta['map50_95']:+.4f})")
        else:
            job.log(f"v{version}: mAP50 {new['map50']:.4f}, mAP50-95 {new['map50_95']:.4f} "
                    f"— {comparison['skipped']}")
    except Exception:
        # Keep the runs directory on failure: results.csv and logs are the only record
        # of why training failed (REQ-006, REQ-064).
        keep_run_dir = True
        raise
    finally:
        # Delete run directory on success or cancellation. job.cancelled finishes training
        # without raising an exception, so it takes this delete path.
        if not keep_run_dir:
            shutil.rmtree(os.path.join(out_dir, "runs"), ignore_errors=True)