"""Modal batch eval for Figment's merged Nemotron 4B checkpoint. Run the full 150-case v5 eval as a terminating Modal job: .venv/bin/modal run modal/eval_figment_nemotron.py """ from __future__ import annotations from datetime import UTC from datetime import datetime import json from pathlib import Path from typing import Any import modal APP_NAME = "figment-nemotron-4b-batch-eval" DEFAULT_DATASET_VERSION = "figment_sft_v5" DEFAULT_MERGED_MODEL_NAME = "figment-sft-v5-lora-merged-bf16" DEFAULT_MODEL_ID = "nvidia/NVIDIA-Nemotron-3-Nano-4B-BF16" DEFAULT_CASE_PATH = "data/eval/field_workflow_holdout_v1.jsonl" DEFAULT_CASE_COUNT = 150 DEFAULT_EVAL_GPU = "H100" DEFAULT_CUDA_ARCHITECTURES = "90" CHECKPOINT_VOLUME_NAME = "figment-checkpoints" RESULT_VOLUME_NAME = "figment-eval-results" CHECKPOINT_DIR = "/checkpoints" RESULT_DIR = "/eval_results" LLAMA_CPP_DIR = "/opt/llama.cpp" LLAMA_SERVER_BIN = f"{LLAMA_CPP_DIR}/build/bin/llama-server" LLAMA_CONVERT_SCRIPT = f"{LLAMA_CPP_DIR}/convert_hf_to_gguf.py" MODEL_SERVER_BASE_URL = "http://127.0.0.1:8001/v1" MODEL_SERVER_HOST = "127.0.0.1" MODEL_SERVER_PORT = 8001 app = modal.App(APP_NAME) checkpoint_volume = modal.Volume.from_name(CHECKPOINT_VOLUME_NAME, create_if_missing=True) result_volume = modal.Volume.from_name(RESULT_VOLUME_NAME, create_if_missing=True) eval_image = ( modal.Image.from_registry("pytorch/pytorch:2.7.1-cuda12.6-cudnn9-devel") .apt_install("git", "build-essential", "cmake", "curl", "ninja-build") .uv_pip_install( "accelerate>=1.8,<2", "einops>=0.8,<1", "ninja>=1.11,<2", "protobuf>=5,<7", "safetensors>=0.5,<1", "sentencepiece>=0.2,<1", "setuptools>=70", "torch==2.7.1", "transformers>=4.52,<5", "wheel>=0.45", ) .add_local_dir("tools/llama.cpp", LLAMA_CPP_DIR, copy=True) .run_commands( f"python -m pip install -r {LLAMA_CPP_DIR}/requirements.txt", f"cmake -S {LLAMA_CPP_DIR} -B {LLAMA_CPP_DIR}/build " f"-DGGML_CUDA=ON -DLLAMA_CURL=OFF -DCMAKE_BUILD_TYPE=Release " f"-DCMAKE_CUDA_ARCHITECTURES={DEFAULT_CUDA_ARCHITECTURES}", f"cmake --build {LLAMA_CPP_DIR}/build --config Release --target llama-server -j $(nproc)", ) .env( { "TOKENIZERS_PARALLELISM": "false", "PYTHON_DOTENV_DISABLED": "true", } ) .add_local_python_source("figment", "scripts") ) def build_eval_config( *, dataset_version: str = DEFAULT_DATASET_VERSION, merged_model_name: str = DEFAULT_MERGED_MODEL_NAME, output_name: str = "", model_id: str = DEFAULT_MODEL_ID, timeout_seconds: float = 300.0, expected_case_count: int = DEFAULT_CASE_COUNT, max_context_tokens: int = 16384, max_generation_tokens: int = 1536, ) -> dict[str, Any]: if not output_name: stamp = datetime.now(UTC).strftime("%Y%m%dT%H%M%SZ") output_name = f"{dataset_version}_field_workflow_holdout_modal_gpu_{stamp}" checkpoint_model_dir = f"{CHECKPOINT_DIR}/{dataset_version}/{merged_model_name}" return { "dataset_version": dataset_version, "merged_model_name": merged_model_name, "checkpoint_model_dir": checkpoint_model_dir, "checkpoint_artifact": ( f"{CHECKPOINT_VOLUME_NAME}:/{dataset_version}/{merged_model_name}" ), "result_volume_name": RESULT_VOLUME_NAME, "output_name": output_name, "output_dir": f"{RESULT_DIR}/{output_name}", "gguf_model_path": f"{RESULT_DIR}/model_cache/{dataset_version}/{merged_model_name}.bf16.gguf", "runtime": "llama_cpp_cuda", "cuda_architectures": DEFAULT_CUDA_ARCHITECTURES, "model_id": model_id, "base_url": MODEL_SERVER_BASE_URL, "case_paths": [f"/tmp/figment_eval_cases/{Path(DEFAULT_CASE_PATH).name}"], "timeout_seconds": timeout_seconds, "expected_case_count": expected_case_count, "max_context_tokens": max_context_tokens, "max_generation_tokens": max_generation_tokens, } @app.function( image=eval_image, gpu=DEFAULT_EVAL_GPU, cpu=8, memory=131072, ephemeral_disk=524288, volumes={ CHECKPOINT_DIR: checkpoint_volume, RESULT_DIR: result_volume, }, timeout=6 * 60 * 60, ) def run_batch_eval( config: dict[str, Any], case_files: list[dict[str, str]], protocol_cards: list[dict[str, str]], ) -> dict[str, Any]: import os import time import torch from scripts.run_local_4b_evidence import run_evidence started = datetime.now(UTC) model_dir = Path(str(config["checkpoint_model_dir"])) output_dir = Path(str(config["output_dir"])) output_dir.mkdir(parents=True, exist_ok=True) if not (model_dir / "config.json").exists(): raise FileNotFoundError(f"missing merged model config at {model_dir}") if not (model_dir / "model.safetensors.index.json").exists(): raise FileNotFoundError(f"missing merged model index at {model_dir}") staged_case_paths = _stage_case_files(case_files) _stage_protocol_cards(protocol_cards) runtime_gpu = { "cuda_available": torch.cuda.is_available(), "cuda_device_count": torch.cuda.device_count(), "cuda_device_name": torch.cuda.get_device_name(0) if torch.cuda.is_available() else "", } print(json.dumps({"runtime_gpu": runtime_gpu}, sort_keys=True), flush=True) gguf_model_path = _ensure_gguf_model(model_dir=model_dir, config=config) result_volume.commit() server = _LlamaCppServer( gguf_model_path=gguf_model_path, model_id=str(config["model_id"]), max_context_tokens=int(config["max_context_tokens"]), ) server.start() try: summary = run_evidence( base_url=str(config["base_url"]), model_id=str(config["model_id"]), output_dir=output_dir, case_paths=staged_case_paths, limit=None, timeout_seconds=float(config["timeout_seconds"]), force_eval=True, ) finally: server.stop() records_path = output_dir / "local_4b_eval.jsonl" record_count = _count_jsonl(records_path) finished = datetime.now(UTC) manifest = { "status": "completed" if record_count == int(config["expected_case_count"]) else "count_mismatch", "config": config, "runtime_gpu": runtime_gpu, "checkpoint_artifact": config["checkpoint_artifact"], "output_volume": RESULT_VOLUME_NAME, "output_dir": str(output_dir), "records_path": str(records_path), "record_count": record_count, "expected_case_count": int(config["expected_case_count"]), "started_at": started.isoformat(), "finished_at": finished.isoformat(), "runtime_seconds": round(time.monotonic() - server.started_monotonic, 3), "summary": summary, "files": sorted(path.name for path in output_dir.iterdir() if path.is_file()), } (output_dir / "modal_eval_manifest.json").write_text( json.dumps(manifest, indent=2, sort_keys=True) + "\n", encoding="utf-8", ) result_volume.commit() if record_count != int(config["expected_case_count"]): raise ValueError(f"expected {config['expected_case_count']} records, found {record_count}") return manifest def _ensure_gguf_model(*, model_dir: Path, config: dict[str, Any]) -> Path: import subprocess import sys gguf_path = Path(str(config["gguf_model_path"])) if gguf_path.exists() and gguf_path.stat().st_size > 0: print( json.dumps( { "modal_eval_gguf": "reuse_existing", "path": str(gguf_path), "bytes": gguf_path.stat().st_size, }, sort_keys=True, ), flush=True, ) return gguf_path gguf_path.parent.mkdir(parents=True, exist_ok=True) command = [ sys.executable, LLAMA_CONVERT_SCRIPT, str(model_dir), "--outfile", str(gguf_path), "--outtype", "bf16", ] print(json.dumps({"modal_eval_gguf": "convert_start", "command": command}, sort_keys=True), flush=True) completed = subprocess.run( command, check=False, text=True, stdout=subprocess.PIPE, stderr=subprocess.STDOUT, ) print(completed.stdout, flush=True) if completed.returncode != 0: raise RuntimeError(f"GGUF conversion failed with exit code {completed.returncode}") print( json.dumps( { "modal_eval_gguf": "convert_finished", "path": str(gguf_path), "bytes": gguf_path.stat().st_size, }, sort_keys=True, ), flush=True, ) return gguf_path class _LlamaCppServer: def __init__(self, *, gguf_model_path: Path, model_id: str, max_context_tokens: int) -> None: self.gguf_model_path = gguf_model_path self.model_id = model_id self.max_context_tokens = max_context_tokens self.process: Any = None self.started_monotonic = 0.0 def start(self) -> None: import subprocess import threading import time import urllib.request command = [ LLAMA_SERVER_BIN, "--model", str(self.gguf_model_path), "--ctx-size", str(self.max_context_tokens), "--host", MODEL_SERVER_HOST, "--port", str(MODEL_SERVER_PORT), "--alias", self.model_id, "--parallel", "1", "--temp", "0", "--top-p", "1", "--reasoning", "off", "--n-gpu-layers", "999", ] print(json.dumps({"modal_eval_llama_server": "start", "command": command}, sort_keys=True), flush=True) self.started_monotonic = time.monotonic() self.process = subprocess.Popen( command, stdout=subprocess.PIPE, stderr=subprocess.STDOUT, text=True, bufsize=1, ) threading.Thread(target=self._stream_logs, daemon=True).start() deadline = time.monotonic() + 180 last_error = "" while time.monotonic() < deadline: if self.process.poll() is not None: raise RuntimeError(f"llama-server exited early with code {self.process.returncode}") try: with urllib.request.urlopen(f"{MODEL_SERVER_BASE_URL}/models", timeout=2) as response: if response.status == 200: print(json.dumps({"modal_eval_llama_server": "ready"}, sort_keys=True), flush=True) return except OSError as exc: last_error = str(exc) time.sleep(1) raise RuntimeError(f"llama-server did not become ready: {last_error}") def stop(self) -> None: if self.process is None or self.process.poll() is not None: return self.process.terminate() try: self.process.wait(timeout=30) except Exception: self.process.kill() self.process.wait(timeout=10) def _stream_logs(self) -> None: if self.process is None or self.process.stdout is None: return for line in self.process.stdout: print(f"[llama-server] {line.rstrip()}", flush=True) class _OpenAICompatibleTransformersServer: def __init__( self, *, model_dir: Path, model_id: str, max_context_tokens: int, max_generation_tokens: int, ) -> None: self.model_dir = model_dir self.model_id = model_id self.max_context_tokens = max_context_tokens self.max_generation_tokens = max_generation_tokens self.httpd: Any = None self.thread: Any = None self.started_monotonic = 0.0 def start(self) -> None: import threading import time import urllib.request from http.server import ThreadingHTTPServer model_runtime = _TransformersRuntime( model_dir=self.model_dir, model_id=self.model_id, max_context_tokens=self.max_context_tokens, max_generation_tokens=self.max_generation_tokens, ) handler = _handler_factory(model_runtime) self.httpd = ThreadingHTTPServer((MODEL_SERVER_HOST, MODEL_SERVER_PORT), handler) self.thread = threading.Thread(target=self.httpd.serve_forever, daemon=True) self.started_monotonic = time.monotonic() self.thread.start() deadline = time.monotonic() + 30 last_error = "" while time.monotonic() < deadline: try: with urllib.request.urlopen(f"{MODEL_SERVER_BASE_URL}/models", timeout=2) as response: if response.status == 200: return except OSError as exc: last_error = str(exc) time.sleep(0.25) raise RuntimeError(f"model server did not become ready: {last_error}") def stop(self) -> None: if self.httpd is not None: self.httpd.shutdown() self.httpd.server_close() if self.thread is not None: self.thread.join(timeout=10) class _TransformersRuntime: def __init__( self, *, model_dir: Path, model_id: str, max_context_tokens: int, max_generation_tokens: int, ) -> None: import threading import torch from transformers import AutoModelForCausalLM from transformers import AutoTokenizer self.model_id = model_id self.max_context_tokens = max_context_tokens self.max_generation_tokens = max_generation_tokens self.lock = threading.Lock() self.request_count = 0 self.torch = torch self.tokenizer = AutoTokenizer.from_pretrained( model_dir, trust_remote_code=True, local_files_only=True, ) if self.tokenizer.pad_token is None: self.tokenizer.pad_token = self.tokenizer.eos_token self.model = AutoModelForCausalLM.from_pretrained( model_dir, trust_remote_code=True, local_files_only=True, dtype=torch.bfloat16, ) if torch.cuda.is_available(): self.model.to("cuda") self.model.eval() self.model.config.use_cache = True def models_payload(self) -> dict[str, Any]: return { "object": "list", "data": [ { "id": self.model_id, "object": "model", "owned_by": "figment-modal-batch", } ], } def chat_completion(self, body: dict[str, Any]) -> dict[str, Any]: import time import uuid messages = body.get("messages") if not isinstance(messages, list): raise ValueError("messages must be a list") requested_max_tokens = int(body.get("max_tokens") or 1024) with self.lock: self.request_count += 1 request_index = self.request_count started = time.monotonic() input_ids = self.tokenizer.apply_chat_template( messages, add_generation_prompt=True, return_tensors="pt", ) device = next(self.model.parameters()).device input_ids = input_ids.to(device) input_len = int(input_ids.shape[-1]) available_tokens = max(128, self.max_context_tokens - input_len - 8) max_new_tokens = max(1, min(requested_max_tokens, available_tokens, self.max_generation_tokens)) print( json.dumps( { "modal_eval_request": request_index, "status": "started", "prompt_tokens": input_len, "max_new_tokens": max_new_tokens, }, sort_keys=True, ), flush=True, ) attention_mask = self.torch.ones_like(input_ids) with self.torch.inference_mode(): output_ids = self.model.generate( input_ids=input_ids, attention_mask=attention_mask, max_new_tokens=max_new_tokens, do_sample=False, pad_token_id=self.tokenizer.pad_token_id, eos_token_id=self.tokenizer.eos_token_id, ) generated_ids = output_ids[0, input_len:] content = self.tokenizer.decode(generated_ids, skip_special_tokens=True).strip() print( json.dumps( { "modal_eval_request": request_index, "status": "finished", "completion_tokens": int(generated_ids.numel()), "elapsed_seconds": round(time.monotonic() - started, 3), }, sort_keys=True, ), flush=True, ) return { "id": f"chatcmpl-{uuid.uuid4().hex}", "object": "chat.completion", "created": int(time.time()), "model": self.model_id, "choices": [ { "index": 0, "message": {"role": "assistant", "content": content}, "finish_reason": "stop", } ], "usage": { "prompt_tokens": input_len, "completion_tokens": int(generated_ids.numel()), "total_tokens": input_len + int(generated_ids.numel()), }, } def _handler_factory(runtime: _TransformersRuntime) -> Any: import json from http.server import BaseHTTPRequestHandler class Handler(BaseHTTPRequestHandler): def do_GET(self) -> None: # noqa: N802 if self.path.rstrip("/") == "/v1/models": self._send_json(200, runtime.models_payload()) else: self._send_json(404, {"error": f"unknown path: {self.path}"}) def do_POST(self) -> None: # noqa: N802 if self.path.rstrip("/") != "/v1/chat/completions": self._send_json(404, {"error": f"unknown path: {self.path}"}) return try: length = int(self.headers.get("Content-Length", "0")) body = json.loads(self.rfile.read(length).decode("utf-8")) payload = runtime.chat_completion(body) except Exception as exc: # pragma: no cover - exercised inside Modal runtime. self._send_json(500, {"error": str(exc)}) return self._send_json(200, payload) def log_message(self, format: str, *args: Any) -> None: return def _send_json(self, status: int, payload: dict[str, Any]) -> None: encoded = json.dumps(payload).encode("utf-8") self.send_response(status) self.send_header("Content-Type", "application/json") self.send_header("Content-Length", str(len(encoded))) self.end_headers() self.wfile.write(encoded) return Handler def _stage_case_files(case_files: list[dict[str, str]]) -> list[Path]: case_dir = Path("/tmp/figment_eval_cases") case_dir.mkdir(parents=True, exist_ok=True) paths: list[Path] = [] for item in case_files: path = case_dir / Path(item["name"]).name path.write_text(item["text"], encoding="utf-8") paths.append(path) return paths def _stage_protocol_cards(protocol_cards: list[dict[str, str]]) -> None: from figment import retrieval card_dir = Path(retrieval.DEFAULT_CARD_DIR) card_dir.mkdir(parents=True, exist_ok=True) for path in card_dir.glob("*.json"): path.unlink() for item in protocol_cards: (card_dir / Path(item["name"]).name).write_text(item["text"], encoding="utf-8") def _count_jsonl(path: Path) -> int: if not path.exists(): return 0 return sum(1 for line in path.read_text(encoding="utf-8").splitlines() if line.strip()) def _read_case_files(paths: list[str]) -> list[dict[str, str]]: return [ {"name": Path(path).name, "text": Path(path).read_text(encoding="utf-8")} for path in paths ] def _read_protocol_cards() -> list[dict[str, str]]: return [ {"name": path.name, "text": path.read_text(encoding="utf-8")} for path in sorted(Path("data/protocol_cards").glob("*.json")) ] @app.local_entrypoint() def main( output_name: str = "", dataset_version: str = DEFAULT_DATASET_VERSION, merged_model_name: str = DEFAULT_MERGED_MODEL_NAME, cases: str = DEFAULT_CASE_PATH, timeout_seconds: float = 300.0, gpu: str = DEFAULT_EVAL_GPU, ) -> None: case_paths = [item.strip() for item in cases.split(",") if item.strip()] config = build_eval_config( dataset_version=dataset_version, merged_model_name=merged_model_name, output_name=output_name, timeout_seconds=timeout_seconds, ) if len(case_paths) != 1 or Path(case_paths[0]).name != Path(DEFAULT_CASE_PATH).name: config["case_paths"] = [f"/tmp/figment_eval_cases/{Path(path).name}" for path in case_paths] config["expected_case_count"] = sum( 1 for path in case_paths for line in Path(path).read_text(encoding="utf-8").splitlines() if line.strip() ) result = run_batch_eval.with_options(gpu=gpu).remote( config, _read_case_files(case_paths), _read_protocol_cards(), ) print(json.dumps({"modal_batch_eval": result}, indent=2, sort_keys=True))