#!/usr/bin/env python3 """End-to-end pipeline: repo URL -> served, repo-personalized model. Stages (each written to ``/status.json`` as it runs so the Next.js app can poll progress): clone -> git clone the repo (shallow) embed -> 6-view frozen-encoder embedding (12288-d) model -> ensure a local snapshot of the base model (first run only) generate -> hypernetwork emits a LoRA adapter for this repo merge -> merge the adapter into the frozen base model serve -> launch an OpenAI-compatible server (vLLM, else fallback) ready -> endpoint is live on the recorded port Each heavy stage runs as its own subprocess so peak memory stays bounded (the ~1GB encoder and the ~10GB base model are never resident together). Usage: python pipeline.py --job-id --repo-url https://github.com/owner/repo """ from __future__ import annotations import argparse import json import os import shutil import subprocess import sys import time from pathlib import Path import config STAGES = ["clone", "embed", "model", "generate", "merge", "serve", "ready"] def write_status(ws: Path, **fields) -> None: status_path = ws / "status.json" cur = {} if status_path.exists(): try: cur = json.loads(status_path.read_text()) except json.JSONDecodeError: cur = {} cur.update(fields) cur["updated_at"] = time.time() status_path.write_text(json.dumps(cur, indent=2)) def run_step(cmd: list[str], log: Path) -> None: """Run a step, streaming combined output to a per-job log file. Raises on non-zero exit.""" with open(log, "a") as lf: lf.write(f"\n$ {' '.join(cmd)}\n") lf.flush() proc = subprocess.run(cmd, stdout=lf, stderr=subprocess.STDOUT, cwd=str(config.ENGINE_DIR)) if proc.returncode != 0: raise RuntimeError(f"step failed ({proc.returncode}): {' '.join(cmd)} — see {log}") def clone_repo(repo_url: str, dest: Path, log: Path) -> None: if dest.exists(): shutil.rmtree(dest, ignore_errors=True) run_step(["git", "clone", "--depth", "80", "--quiet", repo_url, str(dest)], log) def serve_available() -> str: """Pick the serving backend: vLLM if importable, else the transformers fallback (Apple Silicon / no-GPU machines).""" try: import vllm # noqa: F401 return "vllm" except Exception: # noqa: BLE001 return "fallback" def launch_server(merged_dir: Path, adapter_dir: Path, port: int, ws: Path, log: Path, merged: bool = True) -> dict: """Start the inference server. ``merged`` selects between serving the self-contained merged model and serving the frozen base with the generated adapter loaded on top (the ``--no-merge`` path).""" backend = serve_available() env = dict(os.environ) logf = open(log, "a") env["PORT"] = str(port) env["SERVED_NAME"] = "memory-lora" if backend == "vllm": if merged: env["MODE"] = "merged" env["MODEL"] = str(merged_dir) else: env["MODE"] = "lora" env["BASE"] = config.BASE_MODEL env["ADAPTER"] = str(adapter_dir) proc = subprocess.Popen( ["bash", str(config.ENGINE_DIR / "serve_vllm.sh")], env=env, stdout=logf, stderr=subprocess.STDOUT, start_new_session=True, ) else: cmd = [sys.executable, str(config.ENGINE_DIR / "serve_fallback.py"), "--port", str(port), "--served-name", "memory-lora"] if merged: cmd += ["--model", str(merged_dir)] else: cmd += ["--base", config.BASE_MODEL, "--adapter", str(adapter_dir)] proc = subprocess.Popen( cmd, env=env, stdout=logf, stderr=subprocess.STDOUT, start_new_session=True, ) return {"backend": backend, "pid": proc.pid, "port": port, "mode": "merged" if merged else "lora"} def wait_healthy(port: int, timeout: float = 900.0) -> bool: import urllib.request deadline = time.time() + timeout url = f"http://127.0.0.1:{port}/v1/models" while time.time() < deadline: try: with urllib.request.urlopen(url, timeout=3) as r: if r.status == 200: return True except Exception: # noqa: BLE001 time.sleep(3) return False def main() -> None: ap = argparse.ArgumentParser() ap.add_argument("--job-id", required=True) ap.add_argument("--repo-url", required=True) ap.add_argument("--checkpoint", default=str(config.DEFAULT_CHECKPOINT)) ap.add_argument("--port", type=int, default=config.SERVE_PORT) ap.add_argument("--device", default=config.DEVICE) ap.add_argument("--no-merge", action="store_true", default=True, help="serve base+adapter (default): one resident model can " "answer as both base and adapted, which the side-by-side " "comparison needs, and it skips a ~10GB merged copy") ap.add_argument("--merge", dest="no_merge", action="store_false", help="write a self-contained merged model instead") args = ap.parse_args() ws = config.workspace(args.job_id) log = ws / "pipeline.log" clone_dir = ws / "repo" emb_path = ws / "embedding.npy" views_path = ws / "views.json" adapter_dir = ws / "adapter" merged_dir = ws / "merged" # Record our PID so the app can tell "still working" from "died without # writing a terminal status" (a killed or crashed pipeline would otherwise # leave the job showing as running forever). write_status(ws, job_id=args.job_id, repo_url=args.repo_url, stages=STAGES, stage="clone", state="running", error=None, pipeline_pid=os.getpid()) try: # 1. clone clone_repo(args.repo_url, clone_dir, log) # 2. embed (subprocess: encoder only) write_status(ws, stage="embed") run_step([sys.executable, str(config.ENGINE_DIR / "build_embedding.py"), "--repo", str(clone_dir), "--out", str(emb_path), "--out-views", str(views_path), "--device", args.device], log) # 3. base model: ensure a local snapshot exists. The base model is not # part of this repo -- it is pulled separately from the official HF # repository. Doing it as its own stage (instead of letting # transformers download lazily mid-merge) keeps the UI honest about # what is happening during a multi-GB first-run download. if config.resolve_base_model() == config.BASE_MODEL_ID: write_status(ws, stage="model") run_step([sys.executable, str(config.ENGINE_DIR / "fetch_base_model.py"), "--repo", config.BASE_MODEL_ID], log) # 4-5. generate adapter + merge (subprocess: base model only) write_status(ws, stage="generate") gen_cmd = [sys.executable, str(config.ENGINE_DIR / "generate_and_merge.py"), "--embedding", str(emb_path), "--checkpoint", args.checkpoint, "--adapter-out", str(adapter_dir), "--device", args.device] if args.no_merge: gen_cmd.append("--no-merge") else: gen_cmd += ["--merged-out", str(merged_dir)] write_status(ws, stage="merge") run_step(gen_cmd, log) # 5. serve write_status(ws, stage="serve") srv = launch_server(merged_dir, adapter_dir, args.port, ws, log, merged=not args.no_merge) write_status(ws, server=srv) # 6. ready if wait_healthy(args.port): write_status(ws, stage="ready", state="ready", endpoint=f"http://127.0.0.1:{args.port}", views=json.loads(views_path.read_text()) if views_path.exists() else None) else: write_status(ws, state="error", error="server did not become healthy in time") except Exception as e: # noqa: BLE001 write_status(ws, state="error", error=str(e)) print(f"[pipeline] ERROR: {e}", file=sys.stderr) sys.exit(1) if __name__ == "__main__": main()