"""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 ". 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}")