training_code/ml/training/modal_lora_train.py

"""
Modal GPU LoRA trainer for LandfillSentry LFM2-VL domain adaptation.

Usage:
    modal run ml/training/modal_lora_train.py --config-json '{"epochs": 1, "max_steps": 24}'
"""

from __future__ import annotations

import base64
import io
import json
import os
import random
import zipfile
from datetime import datetime, timezone
from pathlib import Path
from typing import Any, Dict, Iterable, List

import modal


def _gpu_from_env():
    gpu_name = os.getenv("MODAL_GPU", "L4").strip().upper()
    allowed = {"T4", "L4", "A10G", "A100"}
    return gpu_name if gpu_name in allowed else "L4"


APP_NAME = os.getenv("MODAL_APP_NAME", "landfillsentry-lora-train")
VOLUME_NAME = os.getenv("MODAL_ARTIFACT_VOLUME", "landfillsentry-model-artifacts")
ARTIFACT_ROOT = Path("/artifacts")
REMOTE_PROJECT_ROOT = Path("/workspace")

image = (
    modal.Image.debian_slim(python_version="3.11")
    .pip_install(
        "accelerate>=1.0.0",
        "huggingface_hub>=1.0.0",
        "peft>=0.13.0",
        "pillow>=10.0.0",
        "safetensors>=0.4.0",
        "sentencepiece>=0.2.0",
        "torch>=2.0.0",
        "torchvision>=0.20.0",
        "transformers>=5.0.0",
    )
    .add_local_dir("apps", "/workspace/apps", copy=True)
    .add_local_dir("ml", "/workspace/ml", copy=True)
    .add_local_dir("data/cache/assets", "/workspace/data/cache/assets", copy=True)
    .add_local_dir("data/labels", "/workspace/data/labels", copy=True)
    .add_local_dir("data/manifests", "/workspace/data/manifests", copy=True)
)
volume = modal.Volume.from_name(VOLUME_NAME, create_if_missing=True)
app = modal.App(APP_NAME)
_secret_values = {
    key: value
    for key in ("HF_TOKEN", "HUGGINGFACE_TOKEN")
    if (value := os.getenv(key, "").strip())
}
hf_secrets = [modal.Secret.from_dict(_secret_values)] if _secret_values else []


def _utc_now() -> str:
    return datetime.now(timezone.utc).isoformat()


def _normalize_config(config: Dict[str, Any]) -> Dict[str, Any]:
    return {
        "model_id": str(config.get("model_id", "LiquidAI/LFM2.5-VL-450M")),
        "revision": str(config.get("revision", "main")),
        "epochs": int(config.get("epochs", 1)),
        "learning_rate": float(config.get("learning_rate", 2e-4)),
        "lora_r": int(config.get("lora_r", 8)),
        "lora_alpha": int(config.get("lora_alpha", 16)),
        "lora_dropout": float(config.get("lora_dropout", 0.05)),
        "max_steps": int(config.get("max_steps", 24)),
        "max_train_samples": int(config.get("max_train_samples", 49)),
        "max_eval_samples": int(config.get("max_eval_samples", 8)),
        "seed": int(config.get("seed", 2103)),
        "dataset_manifest_path": str(config.get("dataset_manifest_path", "data/manifests/dataset_manifest_v1.json")),
        "dataset_split_path": str(config.get("dataset_split_path", "data/manifests/dataset_splits_v1.json")),
        "hf_token": str(config.get("hf_token") or os.getenv("HF_TOKEN") or os.getenv("HUGGINGFACE_TOKEN") or ""),
    }


def _project_path(path: str | Path) -> Path:
    candidate = Path(path)
    if candidate.is_absolute():
        return candidate
    return REMOTE_PROJECT_ROOT / candidate


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


def _load_samples(config: Dict[str, Any]) -> Dict[str, List[Dict[str, Any]]]:
    manifest = _load_json(_project_path(config["dataset_manifest_path"]))
    by_id = {sample["sample_id"]: sample for sample in manifest.get("samples", [])}
    split_doc = _load_json(_project_path(config["dataset_split_path"]))
    split_ids = split_doc.get("splits", {})
    return {
        split: [by_id[sample_id] for sample_id in sample_ids if sample_id in by_id]
        for split, sample_ids in split_ids.items()
    }


def _panel_image_path(sample: Dict[str, Any]) -> Path:
    panel_path = _project_path(sample["panel_artifact_path"])
    panel = _load_json(panel_path)
    slots = panel.get("slots", {})
    image_ref = slots.get("current_rgb") or slots.get("spectral_composite") or slots.get("mapbox_context")
    if not image_ref:
        raise ValueError(f"panel has no usable image slot: {panel_path}")
    return _project_path(image_ref)


def _prompt_text(sample: Dict[str, Any]) -> tuple[str, str]:
    panel = _load_json(_project_path(sample["panel_artifact_path"]))
    metadata_block = panel.get("metadata_text") or (
        f"site_id={sample['site_id']}; sample_id={sample['sample_id']}; "
        f"source_ref={sample.get('provenance', {}).get('source_ref', 'unknown')}"
    )
    system_prompt = (
        "You are LandfillSentry Incident Assistant. Return JSON only. "
        "Do not include markdown or narrative. Use enum values exactly as specified."
    )
    user_prompt = (
        "Interpret the evidence panel and produce an incident object using this schema: "
        "{incident_id, site_id, job_id, analysis_time, plume_likely, confidence, bbox_norm, "
        "likely_source_zone, persistence_score, priority_tier, severity_tier, review_status, "
        "feedback_status, evidence_summary, recommended_followup, model_version}. "
        f"Metadata: {metadata_block}."
    )
    return system_prompt, user_prompt


def _target_json(sample: Dict[str, Any], model_id: str) -> str:
    annotation = sample["annotation"]
    priority = str(annotation.get("priority_tier", "medium"))
    severity = "critical" if priority == "urgent" else "high" if priority == "high" else "medium" if priority == "medium" else "low"
    payload = {
        "incident_id": f"train_{sample['sample_id']}",
        "site_id": sample["site_id"],
        "job_id": sample["sample_id"],
        "analysis_time": sample.get("provenance", {}).get("created_at", _utc_now()),
        "plume_likely": bool(annotation.get("plume_likely", True)),
        "confidence": 0.85 if annotation.get("plume_likely", True) else 0.35,
        "bbox_norm": annotation.get("bbox_norm", [0.2, 0.2, 0.5, 0.5]),
        "likely_source_zone": annotation.get("likely_source_zone", "perimeter_or_unknown"),
        "persistence_score": 0.74 if annotation.get("plume_likely", True) else 0.18,
        "priority_tier": priority,
        "severity_tier": severity,
        "review_status": "needs_review",
        "feedback_status": "unreviewed",
        "evidence_summary": f"Domain-labeled landfill sample {sample['sample_id']} indicates {priority} review priority.",
        "recommended_followup": "Review the evidence panel, confirm source zone, and compare with recent temporal context.",
        "model_version": model_id,
    }
    return json.dumps(payload, sort_keys=True, separators=(",", ":"))


def _conversation(sample: Dict[str, Any], image: Any, model_id: str, include_answer: bool) -> List[Dict[str, Any]]:
    system_prompt, user_prompt = _prompt_text(sample)
    messages = [
        {"role": "system", "content": [{"type": "text", "text": system_prompt}]},
        {
            "role": "user",
            "content": [
                {"type": "image", "image": image},
                {"type": "text", "text": user_prompt},
            ],
        },
    ]
    if include_answer:
        messages.append({"role": "assistant", "content": [{"type": "text", "text": _target_json(sample, model_id)}]})
    return messages


def _move_to_device(batch: Dict[str, Any], device: str) -> Dict[str, Any]:
    moved: Dict[str, Any] = {}
    for key, value in batch.items():
        moved[key] = value.to(device) if hasattr(value, "to") else value
    return moved


def _encode_sample(processor: Any, sample: Dict[str, Any], model_id: str) -> Dict[str, Any]:
    from PIL import Image

    image_path = _panel_image_path(sample)
    image = Image.open(image_path).convert("RGB")
    prompt_messages = _conversation(sample, image, model_id, include_answer=False)
    full_messages = _conversation(sample, image, model_id, include_answer=True)
    full = processor.apply_chat_template(
        full_messages,
        add_generation_prompt=False,
        return_tensors="pt",
        return_dict=True,
        tokenize=True,
    )
    prompt = processor.apply_chat_template(
        prompt_messages,
        add_generation_prompt=True,
        return_tensors="pt",
        return_dict=True,
        tokenize=True,
    )
    labels = full["input_ids"].clone()
    prompt_len = min(prompt["input_ids"].shape[1], labels.shape[1] - 1)
    labels[:, :prompt_len] = -100
    full["labels"] = labels
    return full


def _linear_target_modules(model: Any) -> List[str]:
    import torch

    preferred = {"q_proj", "k_proj", "v_proj", "o_proj"}
    found = set()
    fallback = set()
    for name, module in model.named_modules():
        if not isinstance(module, torch.nn.Linear):
            continue
        short = name.rsplit(".", 1)[-1]
        if short in preferred:
            found.add(short)
        elif len(fallback) < 8:
            fallback.add(short)
    return sorted(found or fallback)


def _average_loss(
    model: Any,
    processor: Any,
    samples: Iterable[Dict[str, Any]],
    model_id: str,
    device: str,
) -> float | None:
    import torch

    losses: List[float] = []
    model.eval()
    with torch.no_grad():
        for sample in samples:
            batch = _move_to_device(_encode_sample(processor, sample, model_id), device)
            loss = model(**batch).loss
            losses.append(float(loss.detach().cpu()))
    model.train()
    if not losses:
        return None
    return sum(losses) / len(losses)


def _archive_dir_b64(path: Path) -> str:
    buffer = io.BytesIO()
    with zipfile.ZipFile(buffer, mode="w", compression=zipfile.ZIP_DEFLATED) as zf:
        for item in sorted(path.rglob("*")):
            if item.is_file():
                zf.write(item, item.relative_to(path).as_posix())
    return base64.b64encode(buffer.getvalue()).decode("ascii")


@app.function(image=image, gpu=_gpu_from_env(), timeout=60 * 10)
def gpu_smoke() -> Dict[str, Any]:
    import torch

    cuda_available = bool(torch.cuda.is_available())
    return {
        "cuda_available": cuda_available,
        "device_name": torch.cuda.get_device_name(0) if cuda_available else "cpu",
        "torch_version": torch.__version__,
    }


@app.function(
    image=image,
    gpu=_gpu_from_env(),
    timeout=60 * 120,
    volumes={str(ARTIFACT_ROOT): volume},
    secrets=hf_secrets,
)
def run_lora_training(config: Dict[str, Any]) -> Dict[str, Any]:
    import torch
    from peft import LoraConfig, get_peft_model
    from transformers import AutoModelForImageTextToText, AutoProcessor

    cfg = _normalize_config(config)
    random.seed(cfg["seed"])
    torch.manual_seed(cfg["seed"])
    if torch.cuda.is_available():
        torch.cuda.manual_seed_all(cfg["seed"])

    samples_by_split = _load_samples(cfg)
    train_samples = list(samples_by_split.get("train", []))[: cfg["max_train_samples"]]
    eval_samples = (list(samples_by_split.get("validation", [])) or list(samples_by_split.get("test", [])))[: cfg["max_eval_samples"]]
    if not train_samples:
        raise RuntimeError("no train samples found in frozen split manifest")

    device = "cuda" if torch.cuda.is_available() else "cpu"
    dtype = torch.float16 if device == "cuda" else torch.float32
    token = cfg["hf_token"] or None
    load_kwargs: Dict[str, Any] = {
        "revision": cfg["revision"],
        "trust_remote_code": True,
        "dtype": dtype,
    }
    processor_kwargs: Dict[str, Any] = {"revision": cfg["revision"], "trust_remote_code": True}
    if token:
        load_kwargs["token"] = token
        processor_kwargs["token"] = token

    processor = AutoProcessor.from_pretrained(cfg["model_id"], **processor_kwargs)
    model = AutoModelForImageTextToText.from_pretrained(cfg["model_id"], **load_kwargs)
    model.to(device)
    model.config.use_cache = False

    target_modules = _linear_target_modules(model)
    peft_config = LoraConfig(
        r=cfg["lora_r"],
        lora_alpha=cfg["lora_alpha"],
        target_modules=target_modules,
        lora_dropout=cfg["lora_dropout"],
        bias="none",
        task_type="CAUSAL_LM",
    )
    model = get_peft_model(model, peft_config)
    model.train()

    eval_loss_before = _average_loss(model, processor, eval_samples, cfg["model_id"], device)
    optimizer = torch.optim.AdamW(model.parameters(), lr=cfg["learning_rate"])
    train_losses: List[float] = []
    completed_steps = 0

    for _epoch in range(cfg["epochs"]):
        for sample in train_samples:
            if cfg["max_steps"] > 0 and completed_steps >= cfg["max_steps"]:
                break
            optimizer.zero_grad(set_to_none=True)
            batch = _move_to_device(_encode_sample(processor, sample, cfg["model_id"]), device)
            loss = model(**batch).loss
            loss.backward()
            optimizer.step()
            train_losses.append(float(loss.detach().cpu()))
            completed_steps += 1
        if cfg["max_steps"] > 0 and completed_steps >= cfg["max_steps"]:
            break

    eval_loss_after = _average_loss(model, processor, eval_samples, cfg["model_id"], device)

    timestamp = datetime.now(timezone.utc).strftime("%Y%m%dT%H%M%SZ")
    run_id = f"lora_run_{timestamp}"
    run_dir = ARTIFACT_ROOT / run_id
    checkpoint_dir = run_dir / "checkpoint-lora-v1"
    checkpoint_dir.mkdir(parents=True, exist_ok=True)
    model.save_pretrained(checkpoint_dir)
    processor.save_pretrained(checkpoint_dir / "processor")

    manifest = _load_json(_project_path(cfg["dataset_manifest_path"]))
    metadata = {
        "adapter_status": "fine_tuned",
        "base_model": cfg["model_id"],
        "base_revision": cfg["revision"],
        "completed_steps": completed_steps,
        "created_at": _utc_now(),
        "dataset_manifest_checksum": manifest.get("manifest_checksum"),
        "dataset_sample_count": manifest.get("sample_count"),
        "eval_loss_after": eval_loss_after,
        "eval_loss_before": eval_loss_before,
        "eval_loss_delta": None
        if eval_loss_before is None or eval_loss_after is None
        else eval_loss_before - eval_loss_after,
        "lora_alpha": cfg["lora_alpha"],
        "lora_dropout": cfg["lora_dropout"],
        "lora_r": cfg["lora_r"],
        "max_steps": cfg["max_steps"],
        "target_modules": target_modules,
        "train_loss_mean": None if not train_losses else sum(train_losses) / len(train_losses),
        "train_sample_count_used": len(train_samples),
        "validation_sample_count_used": len(eval_samples),
    }
    (checkpoint_dir / "landfillsentry_adapter_metadata.json").write_text(
        json.dumps(metadata, indent=2, sort_keys=True),
        encoding="utf-8",
    )
    (checkpoint_dir / "README.md").write_text(
        "# LandfillSentry LFM2-VL LoRA Adapter\n\n"
        "PEFT LoRA adapter trained on the frozen LandfillSentry satellite evidence-panel dataset.\n\n"
        f"- Base model: `{cfg['model_id']}@{cfg['revision']}`\n"
        f"- Completed optimizer steps: `{completed_steps}`\n"
        f"- Validation loss before: `{eval_loss_before}`\n"
        f"- Validation loss after: `{eval_loss_after}`\n"
        f"- Dataset checksum: `{manifest.get('manifest_checksum')}`\n",
        encoding="utf-8",
    )

    run_manifest = {
        "adapter_artifact_ref": f"modal-volume://{VOLUME_NAME}/{run_id}/checkpoint-lora-v1",
        "artifact_volume": VOLUME_NAME,
        "checkpoint_dir": str(checkpoint_dir),
        "config": {k: v for k, v in cfg.items() if k != "hf_token"},
        "created_at": _utc_now(),
        "manifest_path": str(run_dir / "run_manifest.json"),
        "metrics": metadata,
        "model_id": cfg["model_id"],
        "revision": cfg["revision"],
        "run_id": run_id,
        "status": "ok",
        "training_mode": "peft_lora_supervised",
    }
    (run_dir / "run_manifest.json").write_text(json.dumps(run_manifest, indent=2, sort_keys=True), encoding="utf-8")
    volume.commit()

    return {
        **run_manifest,
        "adapter_archive_b64": _archive_dir_b64(checkpoint_dir),
    }


def _extract_adapter_archive(result: Dict[str, Any], output_root: Path) -> Path | None:
    archive_b64 = result.pop("adapter_archive_b64", None)
    if not archive_b64:
        return None
    run_id = str(result.get("run_id", "lora_run_latest"))
    out_dir = output_root / run_id / "checkpoint-lora-v1"
    out_dir.mkdir(parents=True, exist_ok=True)
    data = base64.b64decode(archive_b64.encode("ascii"))
    with zipfile.ZipFile(io.BytesIO(data), mode="r") as zf:
        zf.extractall(out_dir)
    return out_dir


@app.local_entrypoint()
def main(config_json: str = "", output_dir: str = "data/processed/hf_adapter_trained") -> None:
    config = json.loads(config_json) if config_json else {}
    smoke = gpu_smoke.remote()
    print("GPU smoke:", smoke)
    result = run_lora_training.remote(config)
    local_adapter_dir = _extract_adapter_archive(result, Path(output_dir))
    if local_adapter_dir is not None:
        result["local_adapter_dir"] = str(local_adapter_dir)
        print("LOCAL_ADAPTER_DIR:", local_adapter_dir)
    print("Training result:", result)
    print("TRAINING_RESULT_JSON:", json.dumps(result, sort_keys=True))