training/modal_app.py

"""Modal backend adapter: runs the training pipeline on a cloud GPU.

The pipeline itself (SFT warmup -> two-phase GRPO -> save/eval/push) lives in
``training/pipeline.py`` and never imports ``modal`` — this module is only the
Modal provisioning shell (image, Volume, Secrets, function registration, CLI),
per the ports/adapters seam in ``training/ports.py``. A future on-prem k8s or
GCP Vertex backend reuses the pipeline and replaces only this file's role.

The repo's flat package is pip-installed as ``sql_env`` inside the image (the
same import paths the notebook uses), so the container runs the real training
modules: ``sql_env.training.pipeline``, ``.config``, ``.data_loading``,
``.notebook_pipeline``, ``.trl_adapter``, ``sql_env.scripts.validate_questions``.

--------------------------------------------------------------------------- ---
Run commands
--------------------------------------------------------------------------- ---
    # 0. One-time: create the secrets the app references (see "Secrets" below)
    modal secret create huggingface-secret HF_TOKEN=hf_xxx
    modal secret create wandb-secret WANDB_API_KEY=...          # optional

    # 1. Cheap smoke test FIRST (tiny model, few steps, small GPU) — validates
    #    plumbing/imports/volume/secrets before paying for big-GPU hours.
    modal run training/modal_app.py --smoke

    # 1b. Resume drill (validates crash recovery for cents): Ctrl-C the smoke
    #     after its first "checkpoint" log line, then relaunch with:
    modal run training/modal_app.py --smoke --resume auto

    # 2. Full Qwen3-4B + LoRA + vLLM run on an A100. Crash/timeout-safe: the
    #    run checkpoints every save_steps and a relaunch of the SAME command
    #    resumes from the last checkpoint instead of starting over.
    modal run --detach training/modal_app.py --config configs/modal_a100.json

    # 3. Override the GPU (e.g. cheaper/bigger) without editing the config.
    modal run --detach training/modal_app.py --config configs/modal_a100.json --gpu L40S

    # 4. Evaluate a checkpoint (local dir on the Volume, or a HF Hub repo id
    #    in the config's output_dir) — the success gate.
    modal run training/modal_app.py --eval-only --config configs/modal_a100.json
    # Render the SFT data as real tokenized input + loss mask (CPU, near-free):
    modal run training/modal_app.py --inspect-sft --config configs/modal_a100.json
    modal run training/modal_app.py --eval-only \
        --config configs/eval_0p6b_hub.json --gpu T4

    # 5. Push the trained adapter / merged model to the HF Hub afterwards.
    modal run training/modal_app.py --config configs/modal_a100.json --push

    # 6. Start over deliberately (archive run_state.json, ignore checkpoints).
    modal run --detach training/modal_app.py \
        --config configs/modal_a100.json --resume fresh

--------------------------------------------------------------------------- ---
Required secrets (Modal Secret names referenced below)
--------------------------------------------------------------------------- ---
    huggingface-secret  -> HF_TOKEN        (required: model download + Hub push)
    wandb-secret        -> WANDB_API_KEY   (optional: experiment tracking)

--------------------------------------------------------------------------- ---
Notes
--------------------------------------------------------------------------- ---
* ``modal`` and all training deps (torch/trl/peft/vllm/...) are NOT installed
  locally. Every heavy import is guarded or deferred inside a function so this
  file imports cleanly and lints under ``uv run ruff check``.
* A single ``modal.Volume`` holds both the HF cache (``HF_HOME``) and the
  checkpoints/output dir, so weights download once and outputs survive
  container teardown. Mid-run GRPO checkpoints are made durable on every save
  via ``ModalVolumeStore.persist()`` (= ``volume.commit()``), so a crashed or
  timed-out run resumes from its last checkpoint (``resume: "auto"``) instead
  of repeating paid GPU hours. See training/pipeline.py for the resume design.
"""

from __future__ import annotations

import json
import logging
from typing import TYPE_CHECKING, Any

if TYPE_CHECKING:  # pragma: no cover - import only for type checkers
    import modal  # noqa: F401 — referenced only in type-checker context

logger = logging.getLogger(__name__)

# --------------------------------------------------------------------------- #
# Constants — image/volume/secret wiring
# --------------------------------------------------------------------------- #
APP_NAME = "sqlenv-grpo"

# Where the persistent Volume is mounted inside the container.
VOLUME_MOUNT = "/vol"
# HF_HOME lives on the Volume so model weights download exactly once.
HF_HOME_DIR = f"{VOLUME_MOUNT}/hf-cache"
# Default output/checkpoint dir on the Volume (configs/modal_a100.json mirrors this).
DEFAULT_OUTPUT_DIR = f"{VOLUME_MOUNT}/checkpoints/grpo_run"

# Default full-run GPU; overridable via --gpu. Smoke runs use SMOKE_GPU below.
# "A100" alone resolves to the 40GB card; we default to 80GB because vLLM
# colocate + LoRA + KV cache is tight on 40GB and OOMs if utilization is high
# (the guide marks 80GB as "Recommended (headroom)"). Cheaper alternatives and
# their trade-offs: docs/guides/modal-rl-training.md + the F001 playbook's GPU
# economics section.
DEFAULT_GPU = "A100-80GB"
# L4, not T4: the smoke trains in bf16, which Turing (sm75) lacks natively —
# observed as CUBLAS_STATUS_ALLOC_FAILED at the first SFT matmul on the
# torch-2.11/CUDA-13 image. L4 (sm89) has native bf16 AND is the same Ada
# architecture as the L40S pilot target, so the smoke validates the real arch.
SMOKE_GPU = "L4"

# Generous timeout for a full SFT + two-phase GRPO run (seconds). With resume
# enabled, hitting it is recoverable: relaunch and the run continues. Raised to 12h
# after v3 (expanded data + 2 epochs ~= 7h) hit the old 6h cap mid-Phase-2.
TRAIN_TIMEOUT = 12 * 60 * 60  # 12h
SMOKE_TIMEOUT = 30 * 60  # 30m

# --------------------------------------------------------------------------- #
# !!! VERSION RECONCILIATION — the #1 cause of a failed FIRST run. !!!
# The ML pins are SINGLE-SOURCED in pyproject.toml ([project.optional-dependencies]
# training + training-gpu); this image reads them via _image_requirements() so
# there is no second copy to drift. VALIDATE the (transformers, trl, vllm) trio
# with `modal run training/modal_app.py --smoke` (T4, cents) BEFORE any paid A100
# run: the smoke uses the SAME image and the SAME GRPOTrainer(environment_factory=)
# agent path with vLLM OFF, so a too-old/too-new transformers fails CHEAPLY at
# trainer construction instead of after model load on a paid GPU.
#
# Constraints (encode any EXACT pin you validate back into pyproject):
#  - The agentic environment_factory path needs transformers >= 5.2.0
#    (notebooks/train_grpo.ipynb cell 2); a 4.x backtrack rejects it.
#  - transformers and vLLM MUST resolve TOGETHER — an open transformers floor
#    co-installed with vLLM can backtrack to an incompatible pair. If smoke
#    fails at GRPOTrainer construction, pin an EXACT (transformers, vllm) pair
#    in pyproject — or drop vLLM (training-gpu) and set use_vllm=false — re-smoke.
# --------------------------------------------------------------------------- #


# Mirror of pyproject's base + training + training-gpu deps. Used ONLY as a
# fallback if the pyproject read fails (e.g. an odd container import path). The
# client-side image BUILD reads pyproject (the real single source); the
# container's image is already built so the exact list is moot there — this
# just stops a read failure from silently disabling every @app.function via the
# broad except below. test_image_requirements_match_fallback pins the two so
# they can never drift.
_FALLBACK_REQUIREMENTS = [
    "pydantic>=2.0.0",
    "requests>=2.31.0",
    "torch>=2.6.0",
    "transformers>=5.2.0",
    "trl>=0.29.0",
    "peft>=0.13.0",
    "accelerate>=0.34.0",
    "datasets>=3.0.0",
    "huggingface_hub>=0.37",
    "wandb>=0.16",
    "vllm>=0.8.0",
]


def _image_requirements() -> list[str]:
    """The image's pip requirements, read from pyproject (single source).

    base deps + the CPU-installable `training` extra + the CUDA-only
    `training-gpu` extra. The editable install in the image is ``--no-deps``,
    so every runtime dep must be listed in this explicit (cache-friendly) layer.

    NEVER raises: it runs at import time inside ``_build_modal_objects``, and a
    failure there is swallowed by the broad except that disables ALL Modal
    functions (surfacing as a baffling "module has no attribute X" in the
    container). If pyproject can't be located/parsed, fall back to the mirror.
    """
    import os  # noqa: PLC0415
    import tomllib  # noqa: PLC0415 — py3.11 stdlib
    from pathlib import Path  # noqa: PLC0415

    # Try hard to find pyproject across client + container import paths.
    roots = [os.environ.get("PROJECT_ROOT"), _project_root_str(), os.getcwd(), "/"]
    for root in roots:
        if not root:
            continue
        pp = Path(root) / "pyproject.toml"
        try:
            if pp.is_file():
                proj = tomllib.loads(pp.read_text())["project"]
                extras = proj["optional-dependencies"]
                return [
                    *proj["dependencies"],
                    *extras["training"],
                    *extras["training-gpu"],
                ]
        except Exception:  # noqa: BLE001 - never let a parse error escape
            continue
    return list(_FALLBACK_REQUIREMENTS)


# --------------------------------------------------------------------------- #
# Modal object construction — guarded so the module imports without `modal`.
# --------------------------------------------------------------------------- #
def _build_modal_objects() -> tuple[Any, Any, Any, list[Any]]:
    """Construct (app, image, volume, secrets).

    Imported lazily so this file parses/lints without ``modal`` installed.
    """
    import modal  # noqa: PLC0415 — intentional deferred heavy import

    # Install the local repo as the `sql_env` package (matches notebook imports),
    # plus the GRPO acceleration extras (peft/vllm) on top of the training extra.
    image = (
        modal.Image.debian_slim(python_version="3.11")
        .apt_install("git", "build-essential")
        # ML pins single-sourced from pyproject (training + training-gpu extras).
        .pip_install(*_image_requirements())
        .env(
            {
                # Pin HF cache onto the Volume so weights download once.
                "HF_HOME": HF_HOME_DIR,
                "HF_HUB_ENABLE_HF_TRANSFER": "0",
                # Reduce CUDA fragmentation across the SFT -> GRPO handoff
                # (same flag the notebook sets in its Configuration cell).
                "PYTORCH_CUDA_ALLOC_CONF": "expandable_segments:True",
                # SQLite/duckdb ship with Python stdlib; no extra install needed.
            }
        )
        # Ship the repo into the image and INSTALL it as the `sql_env` package.
        # `sql_env` is a package-dir alias for the repo root (see pyproject
        # [tool.setuptools] package-dir = {"sql_env" = "."}), so `import sql_env`
        # only resolves once the package is installed — PYTHONPATH alone is NOT
        # enough. We bake the repo in with copy=True (so the editable install can
        # see it at build time) and `pip install -e` it exactly as `uv sync` does
        # locally. The heavy .pip_install layer above stays cached; only this
        # small copy + metadata-only editable install re-runs when source changes.
        .add_local_dir(
            _project_root_str(),
            remote_path="/root/analyst-buddy",
            copy=True,
            ignore=[
                "**/.git",
                "**/.venv",
                "**/__pycache__",
                "**/.pytest_cache",
                "**/.ruff_cache",
                "**/*.egg-info",
                "**/outputs",
            ],
        )
        .run_commands("pip install --no-deps -e /root/analyst-buddy")
        .env(
            {
                "PYTHONPATH": "/root/analyst-buddy",
                # The pipeline resolves the repo root from this (ports rule:
                # container paths are adapter config, not pipeline logic).
                "PROJECT_ROOT": "/root/analyst-buddy",
                # Group all runs under one W&B project; SFT + the two GRPO
                # phases appear as separate runs inside it.
                "WANDB_PROJECT": "analyst-buddy-grpo",
            }
        )
    )

    volume = modal.Volume.from_name("sqlenv-grpo-vol", create_if_missing=True)

    secrets = [
        modal.Secret.from_name("huggingface-secret"),
        # Live training curves (reward/loss per logging_steps) in the W&B web
        # UI — the "inspect the graph continuously" requirement. Create once:
        #   modal secret create wandb-secret WANDB_API_KEY=...
        # Enabled per-run via report_to="wandb" in the config (smoke forces it
        # off). To drop W&B entirely, remove this line + the config key.
        modal.Secret.from_name("wandb-secret"),
    ]

    app = modal.App(APP_NAME)
    return app, image, volume, secrets


def _project_root_str() -> str:
    """Absolute path to the repo root (the dir containing pyproject.toml)."""
    from pathlib import Path  # noqa: PLC0415

    here = Path(__file__).resolve()
    for parent in here.parents:
        if (parent / "pyproject.toml").exists():
            return str(parent)
    # Fallback: parent of the training/ dir.
    return str(here.parent.parent)


# Build the Modal objects at import time *only when modal is available*. When it
# is not (local lint/CI), expose ``None`` placeholders so the file still imports.
try:  # pragma: no cover - exercised only in the Modal runtime
    app, image, volume, secrets = _build_modal_objects()
    _MODAL_AVAILABLE = True
except ImportError:
    # modal not installed (local lint/CI) — expected; stay quiet.
    app = image = volume = None  # type: ignore[assignment]
    secrets = []  # type: ignore[assignment]
    _MODAL_AVAILABLE = False
except Exception as _exc:  # pragma: no cover - runtime only
    # modal IS installed but object construction failed (e.g. a bad image spec
    # or a config read). Do NOT fail silently: this disables EVERY @app.function
    # and surfaces in the container as a baffling "module 'modal_app' has no
    # attribute <fn>". Make it loud so the real cause is visible in the logs.
    import sys as _sys

    print(f"FATAL: _build_modal_objects() failed: {_exc!r}", file=_sys.stderr)
    app = image = volume = None  # type: ignore[assignment]
    secrets = []  # type: ignore[assignment]
    _MODAL_AVAILABLE = False


class ModalVolumeStore:
    """ArtifactStore adapter: checkpoint durability via ``Volume.commit()``.

    The pipeline calls ``persist()`` on every Trainer checkpoint save and at
    every stage boundary, so mid-run state survives a crash/timeout and a
    relaunch can resume. Failures are logged, never raised — a flaky commit
    must not kill a healthy training step (the next save retries anyway).
    """

    def __init__(self, vol: Any) -> None:
        self._vol = vol

    def persist(self) -> None:
        if self._vol is None:
            return
        try:
            self._vol.commit()
            logger.info("Volume committed.")
        except Exception as exc:  # pragma: no cover - runtime only
            logger.warning("Volume commit failed: %s", exc)


def _write_config_tmp(config_json: str) -> str:
    """Write the serialized config to a temp file, return its path.

    We pass config *contents* (not a path) into Modal functions so they don't
    depend on the local file being present in the container image.
    """
    import tempfile  # noqa: PLC0415

    with tempfile.NamedTemporaryFile(
        "w", suffix=".json", delete=False, encoding="utf-8"
    ) as tmp:
        tmp.write(config_json)
        return tmp.name


# --------------------------------------------------------------------------- #
# Modal functions — only registered when modal is importable.
# --------------------------------------------------------------------------- #
if _MODAL_AVAILABLE:

    @app.function(
        image=image,
        gpu=DEFAULT_GPU,
        volumes={VOLUME_MOUNT: volume},
        secrets=secrets,
        timeout=TRAIN_TIMEOUT,
    )
    def train_remote(
        config_json: str,
        *,
        do_push: bool = False,
    ) -> dict[str, Any]:
        """Full GPU training run. ``config_json`` is the serialized config dict."""
        from sql_env.training.pipeline import run_training  # noqa: PLC0415

        return run_training(
            _write_config_tmp(config_json),
            smoke=False,
            do_push=do_push,
            store=ModalVolumeStore(volume),
        )

    @app.function(
        image=image,
        gpu=SMOKE_GPU,
        volumes={VOLUME_MOUNT: volume},
        secrets=secrets,
        timeout=SMOKE_TIMEOUT,
    )
    def smoke_remote(config_json: str) -> dict[str, Any]:
        """Cheap plumbing check on a tiny model/GPU before paying for A100 hours."""
        from sql_env.training.pipeline import run_training  # noqa: PLC0415

        return run_training(
            _write_config_tmp(config_json),
            smoke=True,
            do_push=False,
            store=ModalVolumeStore(volume),
        )

    @app.function(
        image=image,
        gpu=DEFAULT_GPU,
        volumes={VOLUME_MOUNT: volume},
        secrets=secrets,
        timeout=60 * 60,
    )
    def push_remote(config_json: str) -> dict[str, Any]:
        """Push an already-trained checkpoint (from the Volume) to the HF Hub.

        Loads the saved model+tokenizer from ``output_dir`` and pushes to the
        config's ``hf_repo``. Use after a training run if you skipped --push.
        """
        from transformers import (  # noqa: PLC0415
            AutoModelForCausalLM,
            AutoTokenizer,
        )

        from sql_env.training.backends.hf_hub import HfHubRegistry  # noqa: PLC0415

        cfg = json.loads(config_json)
        output_dir = cfg["output_dir"]
        repo = cfg.get("hf_repo")
        if not repo:
            raise ValueError("config has no 'hf_repo' to push to")
        model = AutoModelForCausalLM.from_pretrained(output_dir)
        tokenizer = AutoTokenizer.from_pretrained(output_dir)
        HfHubRegistry().publish(
            model,
            tokenizer,
            repo,
            merged=cfg.get("push_merged", False),
            private=cfg.get("hf_private", True),
        )
        return {"hf_repo": repo, "output_dir": output_dir}

    @app.function(
        image=image,
        gpu=DEFAULT_GPU,
        volumes={VOLUME_MOUNT: volume},
        secrets=secrets,
        # 2h: a full-set ×k eval of a model that flails (e.g. the BASE model,
        # which burns its whole step budget ~9 steps/episode) runs ~9 s/episode
        # and blew the old 1h cap at ~430/498. 2h covers the slowest eval.
        timeout=2 * 60 * 60,
    )
    def evaluate_remote(config_json: str) -> dict[str, Any]:
        """Measure success_rate of a checkpoint on the held-out eval set.

        ``output_dir`` may be a Volume checkpoint dir or a HF Hub repo id
        (e.g. the 0.6B fallback). Use --gpu T4 for small models — the default
        A100 is wasted on a 0.6B eval.
        """
        from sql_env.training.pipeline import run_eval  # noqa: PLC0415

        return run_eval(json.loads(config_json))

    @app.function(
        image=image,
        gpu=DEFAULT_GPU,
        volumes={VOLUME_MOUNT: volume},
        secrets=secrets,
        timeout=30 * 60,
    )
    def probe_remote(config_json: str) -> dict[str, Any]:
        """H1 capability probe: few-shot the BASE model on join questions.

        Loads the untrained ``model_name`` (no checkpoint) and tests whether
        multi-hop JOIN synthesis is latent — the fork the plateau diagnosis turns
        on. Cheap (one small model, ~12 single-turn generations). Use --gpu T4/L4
        for the 1.7B base; the default A100 is overkill.
        """
        from sql_env.scripts.probe_base_model import run_probe  # noqa: PLC0415

        return run_probe(json.loads(config_json))

    @app.function(
        image=image,
        volumes={VOLUME_MOUNT: volume},
        secrets=secrets,
        timeout=15 * 60,
    )
    def inspect_sft_remote(config_json: str, n: int = 3) -> str:
        """Render the SFT dataset as the REAL tokenized input + loss mask.

        CPU-only (no GPU) — tokenizer-only, so it's near-free. Lets you SEE the
        exact model input (incl. masked train_on_loss=False recovery turns)
        without transformers on the laptop.
        """
        from sql_env.training.pipeline import render_sft_dataset  # noqa: PLC0415

        return render_sft_dataset(json.loads(config_json), n=n)

    @app.local_entrypoint()
    def main(
        config: str = "configs/modal_a100.json",
        smoke: bool = False,
        gpu: str = "",
        push: bool = False,
        push_only: bool = False,
        eval_only: bool = False,
        inspect_sft: bool = False,
        probe: bool = False,
        n: int = 3,
        resume: str = "",
        max_steps: int = 0,
    ) -> None:
        """CLI entrypoint: ``modal run training/modal_app.py [--smoke] [--push] ...``.

        Args:
            config: Path to a training config JSON (default configs/modal_a100.json).
            smoke: Run the cheap tiny-model plumbing check instead of full training.
            gpu: Override the GPU (e.g. "L40S"); applies to training AND
                --eval-only (smoke always uses the T4).
            push: After training, push the result to the config's hf_repo.
            push_only: Skip training; load the already-trained checkpoint from
                the config's output_dir (Volume) and push it to hf_repo. Use
                after a run launched without --push. No GPU compute needed —
                pass --gpu T4 to avoid wasting the default A100.
            eval_only: Skip training; evaluate the checkpoint named by the
                config's output_dir (Volume dir or Hub repo id) and print
                success_rate (the success gate).
            inspect_sft: Skip training; render the SFT dataset as the real
                tokenized model input + loss mask (CPU-only, near-free). Use
                --n to set how many examples (default 3).
            probe: Skip training; few-shot the BASE (untrained) model_name on
                join questions to test whether multi-hop JOIN synthesis is
                latent (the plateau-diagnosis fork). Cheap — use --gpu L4/T4.
            n: Number of SFT examples to render with --inspect-sft.
            resume: Override the config's resume mode: "auto" (continue from
                the last checkpoint/stage in output_dir — the default for real
                runs) or "fresh" (archive run state and start over; the default
                for smoke runs).
            max_steps: Bounded-spend cap per GRPO phase (the pilot pattern):
                launch with e.g. --max-steps 30, inspect reward/parse-rate/
                step-time, then relaunch without it — resume continues from the
                pilot's checkpoint, so the pilot spend is never wasted. 0 = no
                cap (the config may still set its own).

        Note: ``--gpu`` is honored by re-binding the function's GPU via
        ``.with_options`` so you don't have to edit the decorator.
        """
        with open(config, encoding="utf-8") as fh:
            config_json = fh.read()

        if resume or max_steps or smoke:
            cfg = json.loads(config_json)
            if resume:
                if resume not in ("auto", "fresh"):
                    raise ValueError(
                        f"--resume must be 'auto' or 'fresh', got {resume!r}"
                    )
                cfg["resume"] = resume
            elif smoke:
                # A smoke must re-validate training from scratch unless the
                # operator explicitly asked for the resume drill (--resume
                # auto). Without this, the base config's resume:"auto" (meant
                # for real runs) leaks through the smoke overrides and a
                # repeated smoke silently skips completed stages.
                cfg["resume"] = "fresh"
            if max_steps:
                cfg["max_steps"] = int(max_steps)
            config_json = json.dumps(cfg)

        if smoke:
            result = smoke_remote.remote(config_json)
            print(f"Smoke run complete: {result}")
            return

        if push_only:
            # Publish an already-trained checkpoint (no training, no GPU compute).
            fn = push_remote.with_options(gpu=gpu) if gpu else push_remote
            result = fn.remote(config_json)
            print(f"Push complete: {result}")
            return

        if eval_only:
            fn = evaluate_remote.with_options(gpu=gpu) if gpu else evaluate_remote
            result = fn.remote(config_json)
            print(f"Eval complete: {result}")
            return

        if inspect_sft:
            # CPU-only, tokenizer render of the SFT data (incl. loss mask).
            print(inspect_sft_remote.remote(config_json, n))
            return

        if probe:
            # H1: few-shot the BASE model on joins (is the capability latent?).
            fn = probe_remote.with_options(gpu=gpu) if gpu else probe_remote
            result = fn.remote(config_json)
            print(f"Probe complete: {json.dumps(result, indent=2)}")
            return

        fn = train_remote
        if gpu:
            fn = train_remote.with_options(gpu=gpu)
        result = fn.remote(config_json, do_push=push)
        print(f"Training complete: {result}")