diff --git a/.env.example b/.env.example index 6311c1e..f256e04 100644 --- a/.env.example +++ b/.env.example @@ -4,3 +4,4 @@ OPENAI_API_KEY= # Job Server (optional) # API_KEY= # JOB_STORE_PATH=jobs.db +# BREV_TOKEN= diff --git a/Containerfile b/Containerfile index 3bcd5a1..bc660bd 100644 --- a/Containerfile +++ b/Containerfile @@ -10,8 +10,7 @@ RUN apt-get update && apt-get install -y --no-install-recommends curl git && \ curl -sL https://mirror.openshift.com/pub/openshift-v4/clients/ocp/stable/openshift-client-linux.tar.gz \ | tar xzf - -C /usr/local/bin oc kubectl && \ curl -sL https://dl.min.io/client/mc/release/linux-amd64/mc -o /usr/local/bin/mc && \ - chmod +x /usr/local/bin/mc && \ - apt-get remove -y curl && apt-get autoremove -y && rm -rf /var/lib/apt/lists/* + chmod +x /usr/local/bin/mc RUN pip install --upgrade pip uv @@ -21,4 +20,7 @@ USER 1001 RUN uv sync --no-cache +# Install Brev +RUN bash -c "$(curl -fsSL https://raw.githubusercontent.com/brevdev/brev-cli/main/bin/install-latest.sh)" + CMD ["echo", "Image is live!"] diff --git a/README.md b/README.md index 224101b..1057efd 100644 --- a/README.md +++ b/README.md @@ -23,6 +23,7 @@ Reproducible benchmarks for coding agents and models using Harbor - [Queue Service](#queue-service) - [Set up the service](#set-up-the-service) - [Use the service](#use-the-service) + - [Brev Integration](#brev-integration) - [Harbor Command Examples](#harbor-command-examples) - [Claude Code vLLM](#claude-code-vllm) - [Codex vLLM](#codex-vllm) @@ -144,6 +145,7 @@ If you want to see a preview of Harbor command that would be run for a given set The queue service is a FastAPI application that can be deployed on OpenShift to queue and run benchmarks automatically. Benchmark results are stored to MinIO for later review. +The queue service can also be configured to spin up/down models on demand using GPU-enabled VMs from [NVIDIA's Brev](https://developer.nvidia.com/brev). ```mermaid sequenceDiagram @@ -181,7 +183,7 @@ sequenceDiagram oc apply -f deploy/harbor-orchestrator-sa.yml oc apply -f deploy/harbor-task-sa.yml ``` -4. Create a secret file named `job-queue-secret` with an `API_KEY` and apply it: +4. Create a secret file named `job-queue-secret` with an `API_KEY` (and optional `BREV_TOKEN`) and apply it: ```yaml apiVersion: v1 kind: Secret @@ -189,6 +191,7 @@ sequenceDiagram name: job-queue-secret stringData: API_KEY: + # BREV_TOKEN: type: Opaque ``` 5. Create the queue service: @@ -243,6 +246,41 @@ Cancel a running or queued job: curl -X DELETE $JOB_QUEUE_URL/jobs/ -H "X-API-Key: " ``` +### Brev Integration + +The queue service can automatically manage [NVIDIA Brev](https://developer.nvidia.com/brev) GPU VMs to serve models on demand. When a job is submitted with `server_url` set to `"brev"`, the service will: + +1. Create a Brev VM instance and set up a port-forward to it +2. Start the requested model as a Docker container on the VM +3. Wait for the model's `/health` endpoint to respond +4. Run the benchmark job +5. Stop the model container when the job completes + +If the next job in the queue uses the same model, the model container is kept running to avoid unnecessary restart cycles. Jobs are automatically reordered so that jobs using the same model run consecutively. When the queue is empty, the Brev instance is deleted. + +**Prerequisites:** + +- Add your `BREV_TOKEN` to the `job-queue-secret` (see [Set up the service](#set-up-the-service) step 4) +- Ensure your model is configured in `src/coding_agent_bench/brev.py` under `MODEL_CONFIGS` + +**Usage:** + +```sh +curl -X POST $JOB_QUEUE_URL/jobs \ + -H "Content-Type: application/json" \ + -H "X-API-Key: " \ + -d '{ + "job_name": "my-benchmark", + "agent": "pi", + "dataset": "swe-bench/swe-bench-verified", + "model_name": "RedHatAI/Qwen3.6-27B-FP8", + "server_url": "brev", + "n_concurrent": 16 + }' +``` + +To bypass Brev and use a specific server, set `server_url` to the server URL as usual. + ## Harbor Command Examples **Prerequisites:** diff --git a/deploy/job-queue-service.yml b/deploy/job-queue-service.yml index 19d09a2..8bccfe0 100644 --- a/deploy/job-queue-service.yml +++ b/deploy/job-queue-service.yml @@ -58,9 +58,16 @@ spec: secretKeyRef: name: job-queue-secret key: API_KEY + - name: BREV_TOKEN + valueFrom: + secretKeyRef: + name: job-queue-secret + key: BREV_TOKEN ports: - containerPort: 8000 name: http + - containerPort: 9000 + name: brev-model imagePullPolicy: Always volumeMounts: - mountPath: /app/data @@ -89,6 +96,24 @@ spec: targetPort: 8000 protocol: TCP --- +kind: Service +apiVersion: v1 +metadata: + name: brev-model-service + labels: + app: job-queue + component: brev-model +spec: + selector: + app: job-queue + component: api + type: ClusterIP + ports: + - name: http + port: 80 + targetPort: 9000 + protocol: TCP +--- apiVersion: route.openshift.io/v1 kind: Route metadata: diff --git a/src/coding_agent_bench/api.py b/src/coding_agent_bench/api.py index 52c10f5..b728e41 100644 --- a/src/coding_agent_bench/api.py +++ b/src/coding_agent_bench/api.py @@ -11,6 +11,9 @@ from enum import Enum from pathlib import Path +import atexit + +from coding_agent_bench.brev import BrevInstance from coding_agent_bench.builder import SupportedAgent, HarborCommandBuilder from coding_agent_bench.job import OpenshiftJob import json @@ -21,10 +24,11 @@ logger = logging.getLogger(__name__) -_job_queue: list[tuple[str, list[str]]] = [] +_job_queue: list[tuple[str, list[str], str]] = [] _job_event = asyncio.Event() _active_job: tuple[str, asyncio.Task, OpenshiftJob] | None = None _shutting_down = False +_brev_instance: BrevInstance | None = None db_path = Path(os.environ.get("JOB_STORE_PATH", "jobs.db")) @@ -161,13 +165,32 @@ async def _verify_api_key(key: str = Depends(_api_key_header)) -> str: return key +def _brev_emergency_cleanup(*_args) -> None: + """Last-resort synchronous cleanup registered via atexit.""" + if _brev_instance is not None: + _brev_instance.destroy_sync() + + @asynccontextmanager async def lifespan(_app: FastAPI) -> AsyncIterator[None]: global _shutting_down + + logger.info("Server starting up") + + logger.info("Cleaning up orphaned Brev instances from previous runs") + BrevInstance.cleanup_orphaned() + + logger.info("Marking orphaned jobs (queued/running) as failed") job_store.mark_orphaned() worker_task = asyncio.create_task(_worker()) cleanup_task = asyncio.create_task(_build_pod_cleanup_loop()) + + atexit.register(_brev_emergency_cleanup) + + logger.info("Server ready — worker and cleanup tasks started") yield + + logger.info("Server shutting down") _shutting_down = True worker_task.cancel() cleanup_task.cancel() @@ -176,6 +199,11 @@ async def lifespan(_app: FastAPI) -> AsyncIterator[None]: await task except asyncio.CancelledError: pass + if _brev_instance is not None: + logger.info("Destroying Brev instance during shutdown") + await _brev_instance.destroy() + logger.info("Brev instance destroyed") + logger.info("Server shutdown complete") app = FastAPI(lifespan=lifespan) @@ -238,9 +266,56 @@ async def _best_effort_cleanup(oj: OpenshiftJob, signal: bool = False) -> str | return "; ".join(errors) if errors else None -async def _run_job(job_id: str, command: list[str]): +def _substitute_server_url(command: list[str], new_url: str) -> list[str]: + """Replace the --server-url value in a CLI command list.""" + result = list(command) + for i, arg in enumerate(result): + if arg == "--server-url" and i + 1 < len(result): + result[i + 1] = new_url + break + return result + + +def _next_brev_model() -> str | None: + """Return the model_name of the next queued Brev job, or None.""" + for _, cmd, model in _job_queue: + for i in range(len(cmd) - 1): + if cmd[i] == "--server-url" and cmd[i + 1] == "brev": + return model + return None + + +async def _brev_post_job_cleanup(model_name: str) -> None: + """Stop the current model if the next job needs a different one, + and destroy the instance if the queue is empty.""" + global _brev_instance + if _brev_instance is None: + return + + next_model = _next_brev_model() + if next_model == model_name: + logger.info("Brev: next job uses same model %s — keeping it running", model_name) + return + + logger.info("Brev: stopping model %s (next model: %s)", model_name, next_model or "none") + await _brev_instance.stop_model(model_name) + if next_model is None: + logger.info("Brev: no more queued Brev jobs — destroying instance") + await _brev_instance.destroy() + _brev_instance = None + logger.info("Brev: instance destroyed") + + +async def _run_job(job_id: str, command: list[str], model_name: str): """Run and monitor an Openshift Job.""" - global _active_job + global _active_job, _brev_instance + + uses_brev = any( + command[i] == "--server-url" and i + 1 < len(command) and command[i + 1] == "brev" + for i in range(len(command)) + ) + + logger.info("Job %s: starting (model=%s, uses_brev=%s)", job_id, model_name, uses_brev) oj = OpenshiftJob(job_name=job_id) task = asyncio.current_task() @@ -248,12 +323,32 @@ async def _run_job(job_id: str, command: list[str]): _active_job = (job_id, task, oj) try: + if uses_brev: + if _brev_instance is None: + logger.info("Job %s: creating new Brev instance", job_id) + _brev_instance = BrevInstance() + logger.info("Job %s: ensuring Brev instance is running", job_id) + await _brev_instance.ensure_running() + if _brev_instance._current_model != model_name: + if _brev_instance._current_model is not None: + logger.info("Job %s: stopping previous Brev model %s", job_id, _brev_instance._current_model) + await _brev_instance.stop_model(_brev_instance._current_model) + logger.info("Job %s: starting Brev model %s", job_id, model_name) + await _brev_instance.start_model(model_name) + logger.info("Job %s: Brev model ready, substituting server URL", job_id) + command = _substitute_server_url(command, _brev_instance.server_url) + + logger.info("Job %s: creating OpenShift job", job_id) job_spec = oj._job_spec(command) await oj._run_oc_command( ["apply", "-f", "-"], stdin_data=json.dumps(job_spec).encode(), ) + + logger.info("Job %s: waiting for pod to be ready", job_id) await oj._wait_for_job_pod_ready() + + logger.info("Job %s: pod ready — status → RUNNING", job_id) job_store.update_status(job_id, JobStatus.RUNNING) while True: @@ -266,41 +361,59 @@ async def _run_job(job_id: str, command: list[str]): if pods: phase = pods[0].get("status", {}).get("phase", "") if phase == "Succeeded": + logger.info("Job %s: pod succeeded — cleaning up", job_id) cleanup_err = await _best_effort_cleanup(oj) + if uses_brev: + await _brev_post_job_cleanup(model_name) job_store.update_status( job_id, JobStatus.COMPLETED, error=f"cleanup failed: {cleanup_err}" if cleanup_err else None, ) + logger.info("Job %s: status → COMPLETED", job_id) return if phase in ("Failed", "Unknown", "Error"): reason = pods[0].get("status", {}).get("reason", "") message = pods[0].get("status", {}).get("message", "") + logger.error("Job %s: pod entered %s (reason=%s, message=%s)", job_id, phase, reason, message) cleanup_err = await _best_effort_cleanup(oj) + if uses_brev: + await _brev_post_job_cleanup(model_name) error = f"{phase}: reason={reason}, message={message}" if cleanup_err: error += f"; cleanup failed: {cleanup_err}" job_store.update_status(job_id, JobStatus.FAILED, error=error) + logger.info("Job %s: status → FAILED", job_id) return await asyncio.sleep(5) except asyncio.CancelledError: + logger.info("Job %s: cancelled (shutting_down=%s)", job_id, _shutting_down) + if uses_brev: + await _brev_post_job_cleanup(model_name) if _shutting_down: cleanup_err = await _best_effort_cleanup(oj) error = "Server shut down" if cleanup_err: error += f"; cleanup failed: {cleanup_err}" job_store.update_status(job_id, JobStatus.FAILED, error=error) + logger.info("Job %s: status → FAILED (server shutdown)", job_id) raise + logger.info("Job %s: signalling pod and cleaning up", job_id) cleanup_err = await _best_effort_cleanup(oj, signal=True) error = f"cleanup failed: {cleanup_err}" if cleanup_err else None job_store.update_status(job_id, JobStatus.CANCELLED, error=error) + logger.info("Job %s: status → CANCELLED", job_id) except Exception as e: + logger.exception("Job %s: unexpected error", job_id) + if uses_brev: + await _brev_post_job_cleanup(model_name) cleanup_err = await _best_effort_cleanup(oj) error = str(e) if cleanup_err: error += f"; cleanup failed: {cleanup_err}" job_store.update_status(job_id, JobStatus.FAILED, error=error) + logger.info("Job %s: status → FAILED", job_id) finally: _active_job = None @@ -308,15 +421,19 @@ async def _run_job(job_id: str, command: list[str]): async def _worker(): """Process jobs from the queue one at a time.""" + logger.info("Worker started — waiting for jobs") while True: await _job_event.wait() _job_event.clear() + logger.info("Worker woke up — %d job(s) in queue", len(_job_queue)) while _job_queue: - job_id, command = _job_queue.pop(0) + job_id, command, model_name = _job_queue.pop(0) row = job_store.get(job_id) if not row or row["status"] != JobStatus.QUEUED.value: + logger.info("Worker: skipping job %s (status=%s)", job_id, row["status"] if row else "not found") continue - await _run_job(job_id, command) + await _run_job(job_id, command, model_name) + logger.info("Worker: queue drained — waiting for more jobs") @router.get("/") async def read_root(): @@ -347,6 +464,24 @@ def build_table(title: str, jobs: list[dict]) -> str: completed = job_store.list(JobStatus.COMPLETED) + job_store.list(JobStatus.CANCELLED) completed.reverse() + brev_section = "" + if _brev_instance is not None: + instance_name = html.escape(_brev_instance._instance_name) + instance_type = html.escape(_brev_instance._instance_type) + model = html.escape(_brev_instance._current_model or "none") + status = "running" if _brev_instance.is_alive else "stopped" + brev_section = f"""

Brev Instance

+ + + +
InstanceTypeStatusCurrent Model
{instance_name}{instance_type}{status}{model}
""" + else: + brev_section = """

Brev Instance

+ + + +
InstanceTypeStatusCurrent Model
No active instance
""" + html_page = f""" @@ -361,6 +496,7 @@ def build_table(title: str, jobs: list[dict]) -> str:

Job Queue

+{brev_section} {build_table("Running", running)} {build_table("Queued", queued)} {build_table("Completed", completed)} @@ -416,16 +552,16 @@ async def create_job(req: CreateJobRequest): except Exception as e: raise HTTPException(status_code=400, detail=str(e)) - # Build the CLI comand command = build_cli_command(req=req) - # Start the job job_id = str(uuid.uuid4()) + logger.info("Job %s: created (name=%s, agent=%s, dataset=%s, model=%s)", job_id, req.job_name, req.agent.value, req.dataset, req.model_name) job_store.insert(job_id, req.job_name, req.agent.value, req.dataset, req.model_name, command) - _job_queue.append((job_id, command)) + _job_queue.append((job_id, command, req.model_name)) + _job_queue.sort(key=lambda item: item[2]) _job_event.set() + logger.info("Job %s: queued (position %d of %d)", job_id, len(_job_queue), len(_job_queue)) - # Return a success response return CreateJobResponse(message="Job created.", job_id=job_id, job_name=req.job_name, command=command) @router.get("/jobs", response_model=list[JobResponse]) @@ -453,19 +589,21 @@ async def delete_job(job_id: str): if job_row["status"] in (JobStatus.COMPLETED, JobStatus.FAILED, JobStatus.CANCELLED): raise HTTPException(status_code=400, detail=f"Job already {job_row['status']}") - # Remove from queue if still waiting - for i, (qid, _) in enumerate(_job_queue): + for i, (qid, _, _model) in enumerate(_job_queue): if qid == job_id: _job_queue.pop(i) job_store.update_status(job_id, JobStatus.CANCELLED) + logger.info("Job %s: cancelled (was queued)", job_id) return {"message": "Job cancelled", "job_id": job_id} - # Cancel the actively running job if _active_job and _active_job[0] == job_id: job_store.update_status(job_id, JobStatus.CANCELLING) _active_job[1].cancel() + logger.info("Job %s: cancellation requested (was running)", job_id) return {"message": "Job cancelling", "job_id": job_id} + job_store.update_status(job_id, JobStatus.CANCELLED) + logger.info("Job %s: cancelled (not in queue or active — likely in transition)", job_id) return {"message": "Job cancelled", "job_id": job_id} app.include_router(router) diff --git a/src/coding_agent_bench/brev.py b/src/coding_agent_bench/brev.py new file mode 100644 index 0000000..421e05d --- /dev/null +++ b/src/coding_agent_bench/brev.py @@ -0,0 +1,417 @@ +from __future__ import annotations + +import asyncio +import logging +import os +import subprocess +import urllib.request +from dataclasses import dataclass + +logger = logging.getLogger(__name__) + +BREV_LOCAL_PORT = 9000 +BREV_REMOTE_PORT = 8000 +BREV_INSTANCE_NAME = "coding-agent-bench" +BREV_INSTANCE_TYPE = "dmz.h100x2.pcie" + + +@dataclass +class ModelConfig: + container_name: str + docker_command: str + + +MODEL_CONFIGS: dict[str, ModelConfig] = { + "RedHatAI/Qwen3.6-27B-FP8": ModelConfig( + container_name="qwen3.6-27b", + docker_command="""docker run --runtime nvidia --gpus all \ + -v ~/.cache/huggingface:/root/.cache/huggingface \ + --env "HF_TOKEN=$HF_TOKEN" \ + -p 8000:8000 \ + --ipc=host \ + vllm/vllm-openai:v0.24.0 \ + --model RedHatAI/Qwen3.6-27B-FP8 \ + --dtype auto \ + --max-model-len 131072 \ + --trust-remote-code \ + --tensor-parallel-size 2 \ + --gpu-memory-utilization 0.9 \ + --async-scheduling \ + --enable-chunked-prefill \ + --enable-prefix-caching \ + --kv-cache-dtype fp8 \ + --enable-auto-tool-choice \ + --reasoning-parser qwen3 \ + --tool-call-parser qwen3_coder \ + --default-chat-template-kwargs '{"enable_thinking": true}' +""" + ), + "RedHatAI/gemma-4-31B-it-FP8-block": ModelConfig( + container_name="gemma-4-31b-it", + docker_command="""docker run --runtime nvidia --gpus all \ + -v ~/.cache/huggingface:/root/.cache/huggingface \ + --env "HF_TOKEN=$HF_TOKEN" \ + -p 8000:8000 \ + --ipc=host \ + vllm/vllm-openai:v0.24.0 \ + --model RedHatAI/gemma-4-31B-it-FP8-block \ + --dtype auto \ + --max-model-len 131072 \ + --trust-remote-code \ + --tensor-parallel-size 2 \ + --gpu-memory-utilization 0.9 \ + --async-scheduling \ + --enable-chunked-prefill \ + --enable-prefix-caching \ + --kv-cache-dtype fp8 \ + --enable-auto-tool-choice \ + --reasoning-parser gemma4 \ + --tool-call-parser gemma4 \ + --default-chat-template-kwargs '{"enable_thinking": true}' +""" + ), + "RedHatAI/gpt-oss-120b": ModelConfig( + container_name="gpt-oss-120b", + docker_command="""docker run --runtime nvidia --gpus all \ + -v ~/.cache/huggingface:/root/.cache/huggingface \ + --env "HF_TOKEN=$HF_TOKEN" \ + -p 8000:8000 \ + --ipc=host \ + vllm/vllm-openai:v0.24.0 \ + --model RedHatAI/gpt-oss-120b \ + --dtype auto \ + --kv-cache-dtype fp8 \ + --tensor-parallel-size 2 \ + --gpu-memory-utilization 0.9 \ + --async-scheduling \ + --enable-chunked-prefill \ + --enable-prefix-caching \ + --enable-auto-tool-choice \ + --tool-call-parser openai +""", + ), + "RedHatAI/NVIDIA-Nemotron-3-Super-120B-A12B-NVFP4": ModelConfig( + container_name="nemotron-3-super-120b", + docker_command="""docker run --runtime nvidia --gpus all \ + -v ~/.cache/huggingface:/root/.cache/huggingface \ + --env "HF_TOKEN=$HF_TOKEN" \ + -p 8000:8000 \ + --ipc=host \ + vllm/vllm-openai:v0.24.0 \ + --model RedHatAI/NVIDIA-Nemotron-3-Super-120B-A12B-NVFP4 \ + --dtype auto \ + --tensor-parallel-size 2 \ + --gpu-memory-utilization 0.9 \ + --async-scheduling \ + --enable-chunked-prefill \ + --enable-prefix-caching \ + --enable-auto-tool-choice \ + --reasoning-parser nemotron_v3 \ + --tool-call-parser qwen3_coder +""" + ), + "RedHatAI/Mistral-Small-4-119B-2603-NVFP4": ModelConfig( + container_name="mistral-small-4-119b", + docker_command="""docker run --runtime nvidia --gpus all \ + -v ~/.cache/huggingface:/root/.cache/huggingface \ + --env "HF_TOKEN=$HF_TOKEN" \ + -p 8000:8000 \ + --ipc=host \ + vllm/vllm-openai:v0.24.0 \ + --model RedHatAI/Mistral-Small-4-119B-2603-NVFP4 \ + --dtype auto \ + --max-model-len 131072 \ + --trust-remote-code \ + --tensor-parallel-size 2 \ + --gpu-memory-utilization 0.9 \ + --async-scheduling \ + --enable-chunked-prefill \ + --enable-prefix-caching \ + --kv-cache-dtype auto \ + --enable-auto-tool-choice \ + --reasoning-parser mistral \ + --tool-call-parser mistral \ + --default-chat-template-kwargs '{"reasoning_effort": "high"}' \ + --limit-mm-per-prompt '{"image": 0}' +""" + ) +} + + +class BrevInstance: + def __init__( + self, + instance_name: str = BREV_INSTANCE_NAME, + instance_type: str = BREV_INSTANCE_TYPE, + local_port: int = BREV_LOCAL_PORT, + remote_port: int = BREV_REMOTE_PORT, + ): + self._instance_name = instance_name + self._instance_type = instance_type + self._local_port = local_port + self._remote_port = remote_port + self._port_forward_process: asyncio.subprocess.Process | None = None + self._running = False + self._logged_in = False + self._current_model: str | None = None + + @property + def is_alive(self) -> bool: + return self._running + + @property + def server_url(self) -> str: + return "http://brev-model-service" + + async def _run_brev( + self, + args: list[str], + check: bool = True, + timeout_sec: int = 1800, + ) -> tuple[str, str]: + cmd = ["brev", *args] + logger.info("Running: %s", " ".join(cmd)) + + process = await asyncio.create_subprocess_exec( + *cmd, + stdout=asyncio.subprocess.PIPE, + stderr=asyncio.subprocess.PIPE, + ) + + try: + stdout_bytes, stderr_bytes = await asyncio.wait_for( + process.communicate(), timeout=timeout_sec + ) + except asyncio.TimeoutError: + process.terminate() + try: + await asyncio.wait_for(process.communicate(), timeout=10) + except asyncio.TimeoutError: + process.kill() + await process.communicate() + raise RuntimeError( + f"brev command timed out after {timeout_sec}s: {' '.join(cmd)}" + ) + + stdout = stdout_bytes.decode(errors="replace") if stdout_bytes else "" + stderr = stderr_bytes.decode(errors="replace") if stderr_bytes else "" + rc = process.returncode or 0 + + if check and rc != 0: + raise RuntimeError( + f"brev command failed (rc={rc}): {' '.join(cmd)}\n" + f"stdout: {stdout}\nstderr: {stderr}" + ) + + logger.info("brev stdout: %s", stdout.strip()) + if stderr.strip(): + logger.info("brev stderr: %s", stderr.strip()) + + return stdout, stderr + + async def _login(self) -> None: + if self._logged_in: + return + token = os.environ.get("BREV_TOKEN") + if not token: + raise RuntimeError("BREV_TOKEN environment variable is not set") + await self._run_brev(["login", "--token", token]) + self._logged_in = True + + async def ensure_running(self) -> None: + if self._running: + return + + await self._login() + + logger.info("Creating Brev instance %s", self._instance_name) + await self._run_brev( + ["create", self._instance_name, "--type", self._instance_type], + timeout_sec=1800, + ) + self._running = True + + logger.info("Starting port-forward %d:%d", self._local_port, self._remote_port) + self._port_forward_process = await asyncio.create_subprocess_exec( + "brev", "port-forward", self._instance_name, + "--port", f"{self._local_port}:{self._remote_port}", + "--host", "0.0.0.0", + stdout=asyncio.subprocess.PIPE, + stderr=asyncio.subprocess.PIPE, + ) + # Give port-forward a moment to bind + await asyncio.sleep(3) + + if self._port_forward_process.returncode is not None: + stdout = await self._port_forward_process.stdout.read() + stderr = await self._port_forward_process.stderr.read() + raise RuntimeError( + f"brev port-forward exited immediately " + f"(rc={self._port_forward_process.returncode})\n" + f"stdout: {stdout.decode(errors='replace')}\n" + f"stderr: {stderr.decode(errors='replace')}" + ) + + async def _wait_for_health( + self, + timeout_sec: int = 1800, + poll_interval: int = 15, + ) -> None: + health_url = f"http://localhost:{self._local_port}/health" + logger.info("Waiting for model health at %s", health_url) + + for elapsed in range(0, timeout_sec, poll_interval): + try: + req = urllib.request.Request(health_url, method="GET") + with urllib.request.urlopen(req, timeout=10) as resp: + if resp.status == 200: + logger.info("Model healthy after %ds", elapsed) + return + except Exception: + pass + + if elapsed % 60 == 0: + logger.info("Model not ready yet (%ds elapsed)", elapsed) + await asyncio.sleep(poll_interval) + + raise RuntimeError( + f"Model health check timed out after {timeout_sec}s" + ) + + async def start_model(self, model_name: str) -> None: + config = MODEL_CONFIGS.get(model_name) + if config is None: + raise ValueError( + f"No Brev model config for '{model_name}'. " + f"Available: {', '.join(MODEL_CONFIGS.keys()) or '(none)'}" + ) + + logger.info("Starting model container %s", config.container_name) + await self._run_brev( + ["exec", self._instance_name, "--host", config.docker_command], + timeout_sec=60, + ) + self._current_model = model_name + + await self._wait_for_health() + + async def stop_model(self, model_name: str) -> None: + config = MODEL_CONFIGS.get(model_name) + if config is None: + return + + logger.info("Stopping model container %s", config.container_name) + await self._run_brev( + [ + "exec", self._instance_name, "--host", + f"docker container rm -f {config.container_name}", + ], + check=False, + timeout_sec=60, + ) + self._current_model = None + + async def destroy(self) -> None: + logger.info("Destroying Brev instance %s", self._instance_name) + + if self._port_forward_process is not None: + self._port_forward_process.terminate() + try: + await asyncio.wait_for( + self._port_forward_process.communicate(), timeout=10 + ) + except asyncio.TimeoutError: + self._port_forward_process.kill() + await self._port_forward_process.communicate() + self._port_forward_process = None + + # Retry deletion — Brev VMs are expensive, we must not leave them running + last_err: Exception | None = None + for attempt in range(3): + try: + await self._run_brev( + ["delete", self._instance_name], + check=True, + timeout_sec=1200, + ) + self._running = False + self._current_model = None + return + except Exception as e: + last_err = e + logger.warning( + "brev delete attempt %d/3 failed: %s", attempt + 1, e + ) + if attempt < 2: + await asyncio.sleep(5) + + logger.error( + "All brev delete attempts failed for %s: %s", + self._instance_name, last_err, + ) + raise last_err # type: ignore[misc] + + def destroy_sync(self) -> None: + """Synchronous best-effort destroy for use in signal handlers and atexit. + Blocks the calling thread. Does not raise.""" + logger.info("Sync-destroying Brev instance %s", self._instance_name) + + if self._port_forward_process is not None: + try: + self._port_forward_process.kill() + except Exception: + pass + + for attempt in range(3): + try: + subprocess.run( + ["brev", "delete", self._instance_name], + capture_output=True, + timeout=1200, + check=True, + ) + logger.info("Sync brev delete succeeded") + self._running = False + return + except Exception as e: + logger.warning( + "Sync brev delete attempt %d/3 failed: %s", attempt + 1, e + ) + + logger.error( + "All sync brev delete attempts failed for %s", self._instance_name + ) + self._running = False + + @classmethod + def cleanup_orphaned(cls, instance_name: str = BREV_INSTANCE_NAME) -> None: + """Delete a Brev instance by name if it exists. Call at startup to + clean up after a previous crash. Synchronous and best-effort.""" + token = os.environ.get("BREV_TOKEN") + if not token: + return + + try: + subprocess.run( + ["brev", "login", "--token", token], + capture_output=True, + timeout=30, + check=True, + ) + except Exception as e: + logger.warning("brev login failed during orphan cleanup: %s", e) + return + + try: + subprocess.run( + ["brev", "delete", instance_name], + capture_output=True, + timeout=1200, + check=False, + ) + logger.info( + "Orphan cleanup: attempted delete of instance '%s'", + instance_name, + ) + except Exception as e: + logger.warning("Orphan cleanup failed for '%s': %s", instance_name, e)