From 86e922a4a93301e13c78582fa212fd98694eb07a Mon Sep 17 00:00:00 2001 From: Taylor Agarwal Date: Tue, 1 Sep 2026 10:47:50 -0400 Subject: [PATCH 01/11] Keep jobs running across queue restarts --- deploy/job-queue-service.yml | 13 +- src/coding_agent_bench/api.py | 382 ++++++++++++++++++++++++---------- src/coding_agent_bench/job.py | 33 ++- 3 files changed, 310 insertions(+), 118 deletions(-) diff --git a/deploy/job-queue-service.yml b/deploy/job-queue-service.yml index d1cf16e..cc65f10 100644 --- a/deploy/job-queue-service.yml +++ b/deploy/job-queue-service.yml @@ -22,6 +22,9 @@ metadata: name: job-queue app: job-queue spec: + replicas: 1 + strategy: + type: Recreate selector: matchLabels: app: job-queue @@ -33,6 +36,7 @@ spec: app: job-queue component: api spec: + terminationGracePeriodSeconds: 60 serviceAccountName: harbor-orchestrator securityContext: fsGroup: 1001 @@ -45,9 +49,10 @@ spec: mkdir -p ~/.ssh && \ ([ -f /app/data/nebius-ssh-key ] || ssh-keygen -t ed25519 -f /app/data/nebius-ssh-key -N "" -q) && \ chmod 600 /app/data/nebius-ssh-key && \ - NEBIUS_SSH_PUBLIC_KEY_PATH=/app/data/nebius-ssh-key.pub \ - NEBIUS_SSH_PRIVATE_KEY_PATH=/app/data/nebius-ssh-key \ - uv run uvicorn coding_agent_bench.api:app --host 0.0.0.0 --port 8000 + exec env \ + NEBIUS_SSH_PUBLIC_KEY_PATH=/app/data/nebius-ssh-key.pub \ + NEBIUS_SSH_PRIVATE_KEY_PATH=/app/data/nebius-ssh-key \ + uv run uvicorn coding_agent_bench.api:app --host 0.0.0.0 --port 8000 resources: requests: cpu: "1" @@ -109,4 +114,4 @@ spec: kind: Service name: job-queue-service weight: 100 - wildcardPolicy: None \ No newline at end of file + wildcardPolicy: None diff --git a/src/coding_agent_bench/api.py b/src/coding_agent_bench/api.py index f442385..35d9599 100644 --- a/src/coding_agent_bench/api.py +++ b/src/coding_agent_bench/api.py @@ -43,6 +43,7 @@ class QueuedJob(NamedTuple): command: list[str] server_url: str model_name: str + adopt_existing: bool = False _job_queue: list[QueuedJob] = [] _job_event = asyncio.Event() @@ -65,7 +66,9 @@ def _parse_nebius_url(server_url: str) -> str | None: class JobStatus(str, Enum): QUEUED = "queued" RUNNING = "running" + COMPLETING = "completing" COMPLETED = "completed" + FAILING = "failing" FAILED = "failed" CANCELLING = "cancelling" CANCELLED = "cancelled" @@ -137,6 +140,9 @@ async def acquire_instance(self, model_name: str, gpu_config: str) -> tuple[str, # Create the instance if we haven't tracked it yet if instance_name not in self._instances: + if await self._manager.instance_exists(instance_name): + logger.info(f"Deleting untracked nebius instance {instance_name} before reuse") + await self._manager.delete_instance(instance_name) logger.info(f"Creating nebius instance {instance_name} with {gpu_config}") await self._manager.create_instance(instance_name, self._subnet_id, gpu_config) self._instances[instance_name] = NebiusInstanceState(instance_name=instance_name, gpu_config=gpu_config, last_job_completed_at=time.time()) @@ -198,6 +204,20 @@ async def mark_job_completed(self, instance_name: str): state.job_running = False state.last_job_completed_at = time.time() + async def adopt_running_instance(self, model_name: str, gpu_config: str) -> str: + """Restore tracking for the deterministic instance used by a running job.""" + instance_name = self._pick_instance_name() + if not await self._manager.instance_exists(instance_name): + logger.warning(f"Nebius instance {instance_name} is missing for recovered job") + return instance_name + self._instances[instance_name] = NebiusInstanceState( + instance_name=instance_name, + gpu_config=gpu_config, + current_model=model_name, + job_running=True, + ) + return instance_name + async def idle_cleanup_loop(self): """Periodically delete idle instances and evict stale entries.""" while True: @@ -343,21 +363,27 @@ def list(self, status: JobStatus | None = None) -> list[dict]: """List all jobs.""" conn = self._connect() if status: - rows = conn.execute("SELECT * FROM jobs WHERE status = ?", (status.value,)).fetchall() + rows = conn.execute("SELECT * FROM jobs WHERE status = ? ORDER BY rowid", (status.value,)).fetchall() else: - rows = conn.execute("SELECT * FROM jobs").fetchall() + rows = conn.execute("SELECT * FROM jobs ORDER BY rowid").fetchall() conn.close() return [dict(row) for row in rows] - def mark_orphaned(self): - """Mark queued or running jobs as failed on server restart.""" + def list_recoverable(self) -> "list[dict]": + """List non-terminal jobs in their original enqueue order.""" conn = self._connect() - conn.execute( - "UPDATE jobs SET status = ?, error = ? WHERE status IN (?, ?, ?)", - (JobStatus.FAILED.value, "Server restarted", JobStatus.QUEUED.value, JobStatus.RUNNING.value, JobStatus.CANCELLING.value), - ) - conn.commit() + rows = conn.execute( + "SELECT * FROM jobs WHERE status IN (?, ?, ?, ?, ?) ORDER BY rowid", + ( + JobStatus.QUEUED.value, + JobStatus.RUNNING.value, + JobStatus.COMPLETING.value, + JobStatus.FAILING.value, + JobStatus.CANCELLING.value, + ), + ).fetchall() conn.close() + return [dict(row) for row in rows] job_store = JobStore(db_path) @@ -377,7 +403,7 @@ async def _verify_api_key(key: str = Depends(_api_key_header)) -> str: @asynccontextmanager async def lifespan(_app: FastAPI) -> AsyncIterator[None]: global _shutting_down, _nebius - job_store.mark_orphaned() + _shutting_down = False # Initialize Nebius orchestrator if enabled background_tasks: list[asyncio.Task] = [] @@ -402,6 +428,7 @@ async def lifespan(_app: FastAPI) -> AsyncIterator[None]: background_tasks.append(asyncio.create_task(_nebius.idle_cleanup_loop())) logger.info("Nebius orchestrator initialized") + await _restore_jobs() worker_task = asyncio.create_task(_worker()) cleanup_task = asyncio.create_task(_build_pod_cleanup_loop()) yield @@ -415,13 +442,6 @@ async def lifespan(_app: FastAPI) -> AsyncIterator[None]: await task except asyncio.CancelledError: pass - if _nebius is not None: - for name in list(_nebius._instances): - try: - logger.info(f"Shutdown: deleting nebius instance {name}") - await _nebius._manager.delete_instance(name) - except Exception: - logger.exception(f"Failed to delete nebius instance {name} during shutdown") app = FastAPI(lifespan=lifespan) @@ -442,7 +462,7 @@ async def _run_oc(command: list[str], timeout_sec: int = 30) -> str: stdout_bytes, _ = await asyncio.wait_for( process.communicate(), timeout=timeout_sec ) - except asyncio.TimeoutError: + except (asyncio.TimeoutError, asyncio.CancelledError): process.terminate() try: await asyncio.wait_for(process.communicate(), timeout=5) @@ -487,7 +507,111 @@ 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], openrouter: bool = False): +async def _finish_cancellation(job_id: str, oj: OpenshiftJob, signal: bool) -> bool: + """Complete cancellation, leaving it recoverable when cleanup fails.""" + cleanup_err = await _best_effort_cleanup(oj, signal=signal) + if cleanup_err: + job_store.update_status(job_id, JobStatus.CANCELLING, error=f"cleanup failed: {cleanup_err}") + return False + job_store.update_status(job_id, JobStatus.CANCELLED) + return True + + +async def _retry_cancellation(job_id: str, oj: OpenshiftJob) -> None: + """Keep the serial queue blocked until an active workload is removed.""" + while job_store.get(job_id)["status"] == JobStatus.CANCELLING.value: + await asyncio.sleep(5) + try: + existing = await oj._get_job() + await _finish_cancellation(job_id, oj, signal=existing is not None) + except Exception: + logger.exception(f"Cancellation cleanup retry failed for {job_id}") + + +def _terminal_error(error: str | None) -> str | None: + """Remove a prior cleanup suffix before retrying terminal cleanup.""" + if not error or error.startswith("cleanup failed:"): + return None + return error.split("; cleanup failed:", 1)[0] + + +async def _finish_terminal_job( + job_id: str, + oj: OpenshiftJob, + final_status: JobStatus, + error: str | None = None, +) -> bool: + """Delete parent and child workloads before recording a terminal status.""" + pending_status = JobStatus.COMPLETING if final_status == JobStatus.COMPLETED else JobStatus.FAILING + base_error = _terminal_error(error) + job_store.update_status(job_id, pending_status, error=base_error) + cleanup_err = await _best_effort_cleanup(oj) + if cleanup_err: + combined_error = f"cleanup failed: {cleanup_err}" + if base_error: + combined_error = f"{base_error}; {combined_error}" + job_store.update_status(job_id, pending_status, error=combined_error) + return False + job_store.update_status(job_id, final_status, error=base_error) + return True + + +async def _retry_terminal_job( + job_id: str, + oj: OpenshiftJob, + final_status: JobStatus, + error: str | None = None, +) -> None: + """Keep the serial queue blocked until terminal workload cleanup succeeds.""" + while not await _finish_terminal_job(job_id, oj, final_status, error=error): + await asyncio.sleep(5) + + +async def _restore_jobs() -> None: + """Rebuild the in-memory dispatcher state from SQLite and OpenShift.""" + recovered: list[QueuedJob] = [] + waiting: list[QueuedJob] = [] + _job_event.clear() + for row in job_store.list_recoverable(): + job_id = row["job_id"] + queued = QueuedJob( + job_id, + json.loads(row["command"]), + row["server_url"], + row["model_name"], + ) + oj = OpenshiftJob(job_name=job_id) + existing = await oj._get_job() + + if row["status"] in (JobStatus.COMPLETING.value, JobStatus.FAILING.value): + final_status = JobStatus.COMPLETED if row["status"] == JobStatus.COMPLETING.value else JobStatus.FAILED + if not await _finish_terminal_job(job_id, oj, final_status, error=row["error"]): + raise RuntimeError(f"Terminal cleanup is still pending for {job_id}") + elif row["status"] == JobStatus.CANCELLING.value: + if not await _finish_cancellation(job_id, oj, signal=existing is not None): + raise RuntimeError(f"Cancellation cleanup is still pending for {job_id}") + elif existing is not None: + recovered.append(queued._replace(adopt_existing=True)) + elif row["status"] == JobStatus.RUNNING.value: + if not await _finish_terminal_job( + job_id, + oj, + JobStatus.FAILED, + error="OpenShift Job missing after server restart", + ): + raise RuntimeError(f"Missing workload cleanup is still pending for {job_id}") + else: + waiting.append(queued) + + _job_queue.clear() + _job_queue.extend(recovered) + _job_queue.extend(waiting) + if _job_queue: + logger.info(f"Recovered {len(recovered)} active and {len(waiting)} queued jobs") + _job_event.set() + + +async def _run_job(job_id: str, command: list[str], adopt_existing: bool = False, openrouter: bool = False): """Run and monitor an Openshift Job.""" global _active_job @@ -497,81 +621,69 @@ async def _run_job(job_id: str, command: list[str], openrouter: bool = False): _active_job = (job_id, task, oj) try: - is_resume = len(command) == 3 and command[0] == "sh" and command[1] == "-c" - if is_resume: - job_spec = oj._resume_job_spec(command[2]) - else: - job_spec = oj._job_spec(command, openrouter=openrouter) - await oj._run_oc_command( - ["apply", "-f", "-"], - stdin_data=json.dumps(job_spec).encode(), - ) - await oj._wait_for_job_pod_ready() - job_store.update_status(job_id, JobStatus.RUNNING) + if not adopt_existing: + is_resume = len(command) == 3 and command[0] == "sh" and command[1] == "-c" + if is_resume: + job_spec = oj._resume_job_spec(command[2]) + else: + job_spec = oj._job_spec(command, openrouter=openrouter) + await oj._run_oc_command( + ["apply", "-f", "-"], + stdin_data=json.dumps(job_spec).encode(), + ) + job_store.update_status(job_id, JobStatus.RUNNING) + await oj._wait_for_job_pod_ready() + elif job_store.get(job_id)["status"] == JobStatus.QUEUED.value: + job_store.update_status(job_id, JobStatus.RUNNING) consecutive_missing = 0 max_missing = 6 # 6 polls × 5s = 30s before declaring pod gone while True: - stdout, _ = await oj._run_oc_command( - ["get", "pod", f"--selector=job-name={oj._pod_name}", "-o", "json"], - check=False, - ) - if stdout: - pods = json.loads(stdout).get("items", []) - else: - pods = [] - - if not pods: + try: + job = await oj._get_job() + except Exception: + logger.exception(f"Unable to query OpenShift Job for {job_id}; monitoring will retry") + await asyncio.sleep(5) + continue + if job is None: consecutive_missing += 1 if consecutive_missing >= max_missing: - cleanup_err = await _best_effort_cleanup(oj) - error = "Pod vanished (likely deleted externally)" - if cleanup_err: - error += f"; cleanup failed: {cleanup_err}" - job_store.update_status(job_id, JobStatus.FAILED, error=error) + await _retry_terminal_job( + job_id, + oj, + JobStatus.FAILED, + error="OpenShift Job vanished (likely deleted externally)", + ) return await asyncio.sleep(5) continue consecutive_missing = 0 - phase = pods[0].get("status", {}).get("phase", "") - if phase == "Succeeded": - cleanup_err = await _best_effort_cleanup(oj) - job_store.update_status( - job_id, JobStatus.COMPLETED, - error=f"cleanup failed: {cleanup_err}" if cleanup_err else None, - ) + conditions = { + condition.get("type"): condition + for condition in job.get("status", {}).get("conditions", []) + if condition.get("status") == "True" + } + if "Complete" in conditions: + await _retry_terminal_job(job_id, oj, JobStatus.COMPLETED) return - if phase in ("Failed", "Unknown", "Error"): - reason = pods[0].get("status", {}).get("reason", "") - message = pods[0].get("status", {}).get("message", "") - cleanup_err = await _best_effort_cleanup(oj) - 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) + if "Failed" in conditions: + reason = conditions["Failed"].get("reason", "") + message = conditions["Failed"].get("message", "") + error = f"Failed: reason={reason}, message={message}" + await _retry_terminal_job(job_id, oj, JobStatus.FAILED, error=error) return await asyncio.sleep(5) except asyncio.CancelledError: 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) raise - 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) + await _finish_cancellation(job_id, oj, signal=True) except Exception as e: - 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) + await _retry_terminal_job(job_id, oj, JobStatus.FAILED, error=error) finally: _active_job = None @@ -590,6 +702,8 @@ def _reorder_queue_for_nebius(): current_model = states[0].current_model def _sort_key(item: QueuedJob): + if item.adopt_existing: + return -1 gpu = _parse_nebius_url(item.server_url) if gpu is None: return 0 # non-nebius, keep in place @@ -602,51 +716,95 @@ def _sort_key(item: QueuedJob): _job_queue.sort(key=_sort_key) +async def _process_queued_job(queued: QueuedJob) -> None: + """Provision or adopt one queued job without risking the dispatcher task.""" + global _active_job + + job_id, command, server_url, model_name, adopt_existing = queued + task = asyncio.current_task() + assert task is not None + oj = OpenshiftJob(job_name=job_id) + _active_job = (job_id, task, oj) + + model_config: ModelConfig | None = None + nebius_instance_name: str | None = None + try: + nebius_gpu_config = _parse_nebius_url(server_url) + if adopt_existing and nebius_gpu_config is not None and _nebius: + try: + nebius_instance_name = await _nebius.adopt_running_instance(model_name, nebius_gpu_config) + except Exception: + logger.exception(f"Failed to restore Nebius tracking for job {job_id}") + elif nebius_gpu_config is not None and _nebius: + try: + nebius_instance_name, real_url = await _nebius.acquire_instance(model_name, gpu_config=nebius_gpu_config) + is_resume = len(command) == 3 and command[0] == "sh" and command[1] == "-c" + if is_resume: + row = job_store.get(job_id) + job_name = row["job_name"] + orig_name = job_name.removesuffix("--resume") + py_job_dir = f"/app/jobs/{orig_name}" + step = _build_url_replace_shell_step(real_url, py_job_dir) + command = list(command) + command[2] = command[2].replace(" && uv run", f"{step} && uv run", 1) + else: + command = [real_url if arg == server_url else arg for arg in command] + model_config = MODEL_REGISTRY.get(model_name) + await _nebius.mark_job_started(nebius_instance_name) + except Exception as e: + logger.exception(f"Nebius provisioning failed for job {job_id}") + job_store.update_status(job_id, JobStatus.FAILED, error=f"Nebius provisioning failed: {e}") + return + + if "--model-max-len" not in command and model_config is not None: + command += ["--model-max-len", str(model_config.model_max_len)] + + await _run_job(job_id, command, openrouter=is_openrouter(server_url), adopt_existing=adopt_existing) + + if job_store.get(job_id)["status"] == JobStatus.CANCELLING.value: + await _retry_cancellation(job_id, oj) + + if nebius_instance_name and _nebius: + await _nebius.mark_job_completed(nebius_instance_name) + except asyncio.CancelledError: + if _shutting_down: + raise + if not await _finish_cancellation(job_id, oj, signal=True): + await _retry_cancellation(job_id, oj) + finally: + if _active_job and _active_job[0] == job_id: + _active_job = None + + async def _worker(): """Process jobs from the queue one at a time.""" + global _active_job + while True: await _job_event.wait() _job_event.clear() while _job_queue: _reorder_queue_for_nebius() - job_id, command, server_url, model_name = _job_queue.pop(0) + job_id, command, server_url, model_name, adopt_existing = _job_queue.pop(0) row = job_store.get(job_id) - if not row or row["status"] != JobStatus.QUEUED.value: + recoverable_statuses = (JobStatus.QUEUED.value, JobStatus.RUNNING.value) + if not row or (adopt_existing and row["status"] not in recoverable_statuses): continue - - # For nebius jobs: provision instance, swap model, patch the command - model_config: ModelConfig | None = None - nebius_instance_name: str | None = None - nebius_gpu_config = _parse_nebius_url(server_url) - if nebius_gpu_config is not None and _nebius: - try: - nebius_instance_name, real_url = await _nebius.acquire_instance(model_name, gpu_config=nebius_gpu_config) - is_resume = len(command) == 3 and command[0] == "sh" and command[1] == "-c" - if is_resume: - # Inject URL replacement into the resume shell command - job_name = row["job_name"] - orig_name = job_name.removesuffix("--resume") - py_job_dir = f"/app/jobs/{orig_name}" - step = _build_url_replace_shell_step(real_url, py_job_dir) - command = list(command) - command[2] = command[2].replace(" && uv run", f"{step} && uv run", 1) - else: - command = [real_url if arg == server_url else arg for arg in command] - model_config = MODEL_REGISTRY.get(model_name) - await _nebius.mark_job_started(nebius_instance_name) - except Exception as e: - logger.exception(f"Nebius provisioning failed for job {job_id}") - job_store.update_status(job_id, JobStatus.FAILED, error=f"Nebius provisioning failed: {e}") - continue - - # Add "--model-max-len" arg to the command to match served model max len, if not already present - if "--model-max-len" not in command and model_config is not None: - command += ["--model-max-len", str(model_config.model_max_len)] - - await _run_job(job_id, command, openrouter=is_openrouter(server_url)) - - if nebius_instance_name and _nebius: - await _nebius.mark_job_completed(nebius_instance_name) + if not adopt_existing and row["status"] != JobStatus.QUEUED.value: + continue + processing_task = asyncio.create_task(_process_queued_job(QueuedJob(job_id, command, server_url, model_name, adopt_existing))) + _active_job = (job_id, processing_task, OpenshiftJob(job_name=job_id)) + try: + await processing_task + except asyncio.CancelledError: + if _shutting_down: + raise + oj = OpenshiftJob(job_name=job_id) + if not await _finish_cancellation(job_id, oj, signal=False): + await _retry_cancellation(job_id, oj) + finally: + if _active_job and _active_job[0] == job_id: + _active_job = None @router.get("/") async def read_root(): @@ -888,7 +1046,13 @@ async def delete_job(job_id: str): if not job_row: raise HTTPException(status_code=404, detail="Job not found") - if job_row["status"] in (JobStatus.COMPLETED, JobStatus.FAILED, JobStatus.CANCELLED): + if job_row["status"] in ( + JobStatus.COMPLETING, + JobStatus.COMPLETED, + JobStatus.FAILING, + JobStatus.FAILED, + JobStatus.CANCELLED, + ): raise HTTPException(status_code=400, detail=f"Job already {job_row['status']}") # Remove from queue if still waiting diff --git a/src/coding_agent_bench/job.py b/src/coding_agent_bench/job.py index e995bd7..658e85b 100644 --- a/src/coding_agent_bench/job.py +++ b/src/coding_agent_bench/job.py @@ -37,6 +37,7 @@ def _resume_job_spec(self, shell_command: str) -> dict: "kind": "Job", "metadata": {"name": self._pod_name, "labels": {"app": "harbor"}}, "spec": { + "backoffLimit": 0, "template": { "spec": { "restartPolicy": "Never", @@ -51,6 +52,7 @@ def _resume_job_spec(self, shell_command: str) -> dict: "args": [shell_command], "env": [ {"name": "HOME", "value": "/tmp"}, + {"name": "HARBOR_PARENT", "value": self._pod_name}, ], "volumeMounts": [{"name": "jobs", "mountPath": "/app/jobs"}], "envFrom": [ @@ -90,6 +92,7 @@ def _job_spec( "kind": "Job", "metadata": {"name": self._pod_name, "labels": {"app": "harbor"}}, "spec": { + "backoffLimit": 0, "template": { "spec": { "restartPolicy": "Never", @@ -114,6 +117,9 @@ def _job_spec( "envFrom": [ {"secretRef": {"name": "harbor-minio"}} ], + "env": [ + {"name": "HARBOR_PARENT", "value": self._pod_name}, + ], } ], } @@ -144,7 +150,7 @@ async def _run_oc_command( ) else: stdout_bytes, stderr_bytes = await process.communicate(input=stdin_data) - except asyncio.TimeoutError: + except (asyncio.TimeoutError, asyncio.CancelledError) as exc: process.terminate() try: stdout_bytes, stderr_bytes = await asyncio.wait_for( @@ -153,10 +159,9 @@ async def _run_oc_command( except asyncio.TimeoutError: process.kill() stdout_bytes, stderr_bytes = await process.communicate() - raise RuntimeError( - f"oc command timed out after {timeout_sec} seconds: " - f"{' '.join(full_command)}" - ) + if isinstance(exc, asyncio.CancelledError): + raise + raise RuntimeError(f"oc command timed out after {timeout_sec} seconds: {' '.join(full_command)}") stdout = stdout_bytes.decode(errors="replace") if stdout_bytes else None stderr = stderr_bytes.decode(errors="replace") if stderr_bytes else None @@ -172,6 +177,24 @@ async def _run_oc_command( return stdout, stderr + async def _get_job(self) -> dict | None: + """Return the OpenShift Job resource, or None when it does not exist.""" + try: + stdout, _ = await self._run_oc_command( + ["get", f"job/{self._pod_name}", "-o", "json"], + timeout_sec=30, + ) + except RuntimeError as exc: + if "notfound" in str(exc).lower() or "not found" in str(exc).lower(): + return None + raise + if not stdout: + raise RuntimeError(f"oc returned no data for job/{self._pod_name}") + try: + return json.loads(stdout) + except json.JSONDecodeError as exc: + raise RuntimeError(f"oc returned invalid JSON for job/{self._pod_name}") from exc + async def _signal_job_pod(self) -> None: """Send SIGTERM to the harbor process inside the job pod so it can run its own cleanup (stopping task pods via From 898ae95486bc0b8c13cbac42ad08f7e09f78cf36 Mon Sep 17 00:00:00 2001 From: Taylor Agarwal Date: Tue, 1 Sep 2026 11:50:57 -0400 Subject: [PATCH 02/11] Fix persistent queue recovery cleanup --- src/coding_agent_bench/api.py | 206 ++++++++++++++++++++++-------- src/coding_agent_bench/builder.py | 6 + src/coding_agent_bench/job.py | 39 +++++- 3 files changed, 196 insertions(+), 55 deletions(-) diff --git a/src/coding_agent_bench/api.py b/src/coding_agent_bench/api.py index 35d9599..3060cb5 100644 --- a/src/coding_agent_bench/api.py +++ b/src/coding_agent_bench/api.py @@ -208,8 +208,7 @@ async def adopt_running_instance(self, model_name: str, gpu_config: str) -> str: """Restore tracking for the deterministic instance used by a running job.""" instance_name = self._pick_instance_name() if not await self._manager.instance_exists(instance_name): - logger.warning(f"Nebius instance {instance_name} is missing for recovered job") - return instance_name + raise RuntimeError(f"Nebius instance {instance_name} is missing for recovered job") self._instances[instance_name] = NebiusInstanceState( instance_name=instance_name, gpu_config=gpu_config, @@ -218,6 +217,15 @@ async def adopt_running_instance(self, model_name: str, gpu_config: str) -> str: ) return instance_name + async def delete_recovered_instance(self) -> None: + """Delete the deterministic VM left behind by an interrupted terminal transition.""" + async with self._lock: + instance_name = self._pick_instance_name() + if await self._manager.instance_exists(instance_name): + logger.info(f"Deleting recovered nebius instance {instance_name}") + await self._manager.delete_instance(instance_name) + self._instances.pop(instance_name, None) + async def idle_cleanup_loop(self): """Periodically delete idle instances and evict stale entries.""" while True: @@ -428,7 +436,9 @@ async def lifespan(_app: FastAPI) -> AsyncIterator[None]: background_tasks.append(asyncio.create_task(_nebius.idle_cleanup_loop())) logger.info("Nebius orchestrator initialized") - await _restore_jobs() + has_recoverable_nebius = await _restore_jobs() + if _nebius is not None and not has_recoverable_nebius: + background_tasks.append(asyncio.create_task(_delete_recovered_nebius("startup"))) worker_task = asyncio.create_task(_worker()) cleanup_task = asyncio.create_task(_build_pod_cleanup_loop()) yield @@ -567,55 +577,43 @@ async def _retry_terminal_job( await asyncio.sleep(5) -async def _restore_jobs() -> None: - """Rebuild the in-memory dispatcher state from SQLite and OpenShift.""" - recovered: list[QueuedJob] = [] - waiting: list[QueuedJob] = [] +async def _delete_recovered_nebius(job_id: str) -> None: + """Retry deletion so a recovered terminal transition cannot leak its VM.""" + assert _nebius is not None + while True: + try: + await _nebius.delete_recovered_instance() + return + except Exception: + logger.exception(f"Unable to delete recovered Nebius instance for {job_id}; retrying") + await asyncio.sleep(5) + + +async def _restore_jobs() -> bool: + """Rebuild the dispatcher without making startup depend on OpenShift.""" _job_event.clear() + _job_queue.clear() + has_recoverable_nebius = False for row in job_store.list_recoverable(): - job_id = row["job_id"] - queued = QueuedJob( - job_id, + has_recoverable_nebius |= _parse_nebius_url(row["server_url"]) is not None + _job_queue.append(QueuedJob( + row["job_id"], json.loads(row["command"]), row["server_url"], row["model_name"], - ) - oj = OpenshiftJob(job_name=job_id) - existing = await oj._get_job() - - if row["status"] in (JobStatus.COMPLETING.value, JobStatus.FAILING.value): - final_status = JobStatus.COMPLETED if row["status"] == JobStatus.COMPLETING.value else JobStatus.FAILED - if not await _finish_terminal_job(job_id, oj, final_status, error=row["error"]): - raise RuntimeError(f"Terminal cleanup is still pending for {job_id}") - elif row["status"] == JobStatus.CANCELLING.value: - if not await _finish_cancellation(job_id, oj, signal=existing is not None): - raise RuntimeError(f"Cancellation cleanup is still pending for {job_id}") - elif existing is not None: - recovered.append(queued._replace(adopt_existing=True)) - elif row["status"] == JobStatus.RUNNING.value: - if not await _finish_terminal_job( - job_id, - oj, - JobStatus.FAILED, - error="OpenShift Job missing after server restart", - ): - raise RuntimeError(f"Missing workload cleanup is still pending for {job_id}") - else: - waiting.append(queued) - - _job_queue.clear() - _job_queue.extend(recovered) - _job_queue.extend(waiting) + adopt_existing=True, + )) if _job_queue: - logger.info(f"Recovered {len(recovered)} active and {len(waiting)} queued jobs") + logger.info(f"Recovered {len(_job_queue)} non-terminal jobs") _job_event.set() + return has_recoverable_nebius async def _run_job(job_id: str, command: list[str], adopt_existing: bool = False, openrouter: bool = False): """Run and monitor an Openshift Job.""" global _active_job - oj = OpenshiftJob(job_name=job_id) + oj = OpenshiftJob(job_name=job_id, clean_legacy_pods=adopt_existing) task = asyncio.current_task() assert task is not None _active_job = (job_id, task, oj) @@ -633,8 +631,23 @@ async def _run_job(job_id: str, command: list[str], adopt_existing: bool = False ) job_store.update_status(job_id, JobStatus.RUNNING) await oj._wait_for_job_pod_ready() - elif job_store.get(job_id)["status"] == JobStatus.QUEUED.value: - job_store.update_status(job_id, JobStatus.RUNNING) + else: + if job_store.get(job_id)["status"] == JobStatus.QUEUED.value: + job_store.update_status(job_id, JobStatus.RUNNING) + while True: + try: + existing = await oj._get_job() + break + except Exception: + logger.exception(f"Unable to inspect recovered OpenShift Job {job_id}; retrying") + await asyncio.sleep(5) + conditions = { + condition.get("type") + for condition in (existing or {}).get("status", {}).get("conditions", []) + if condition.get("status") == "True" + } + if not conditions.intersection({"Complete", "Failed"}): + await oj._wait_for_job_pod_ready() consecutive_missing = 0 max_missing = 6 # 6 polls × 5s = 30s before declaring pod gone @@ -723,18 +736,63 @@ async def _process_queued_job(queued: QueuedJob) -> None: job_id, command, server_url, model_name, adopt_existing = queued task = asyncio.current_task() assert task is not None - oj = OpenshiftJob(job_name=job_id) + oj = OpenshiftJob(job_name=job_id, clean_legacy_pods=adopt_existing) _active_job = (job_id, task, oj) model_config: ModelConfig | None = None nebius_instance_name: str | None = None try: nebius_gpu_config = _parse_nebius_url(server_url) + row = job_store.get(job_id) + if not row: + return + + if row["status"] in (JobStatus.COMPLETING.value, JobStatus.FAILING.value): + final_status = JobStatus.COMPLETED if row["status"] == JobStatus.COMPLETING.value else JobStatus.FAILED + if nebius_gpu_config is not None and _nebius: + await _delete_recovered_nebius(job_id) + await _retry_terminal_job(job_id, oj, final_status, error=row["error"]) + return + + if row["status"] == JobStatus.CANCELLING.value: + if nebius_gpu_config is not None and _nebius: + await _delete_recovered_nebius(job_id) + await _retry_cancellation(job_id, oj) + return + + if adopt_existing: + while True: + try: + existing = await oj._get_job() + break + except Exception: + logger.exception(f"Unable to reconcile recovered OpenShift Job {job_id}; retrying") + await asyncio.sleep(5) + if existing is None: + if row["status"] == JobStatus.RUNNING.value: + if nebius_gpu_config is not None and _nebius: + await _delete_recovered_nebius(job_id) + await _retry_terminal_job( + job_id, + oj, + JobStatus.FAILED, + error="OpenShift Job missing after server restart", + ) + return + adopt_existing = False + if adopt_existing and nebius_gpu_config is not None and _nebius: try: nebius_instance_name = await _nebius.adopt_running_instance(model_name, nebius_gpu_config) - except Exception: + except Exception as e: logger.exception(f"Failed to restore Nebius tracking for job {job_id}") + await _retry_terminal_job( + job_id, + oj, + JobStatus.FAILED, + error=str(e), + ) + return elif nebius_gpu_config is not None and _nebius: try: nebius_instance_name, real_url = await _nebius.acquire_instance(model_name, gpu_config=nebius_gpu_config) @@ -769,6 +827,8 @@ async def _process_queued_job(queued: QueuedJob) -> None: except asyncio.CancelledError: if _shutting_down: raise + if _parse_nebius_url(server_url) is not None and _nebius: + await _delete_recovered_nebius(job_id) if not await _finish_cancellation(job_id, oj, signal=True): await _retry_cancellation(job_id, oj) finally: @@ -787,19 +847,29 @@ async def _worker(): _reorder_queue_for_nebius() job_id, command, server_url, model_name, adopt_existing = _job_queue.pop(0) row = job_store.get(job_id) - recoverable_statuses = (JobStatus.QUEUED.value, JobStatus.RUNNING.value) + recoverable_statuses = ( + JobStatus.QUEUED.value, + JobStatus.RUNNING.value, + JobStatus.COMPLETING.value, + JobStatus.FAILING.value, + JobStatus.CANCELLING.value, + ) if not row or (adopt_existing and row["status"] not in recoverable_statuses): continue if not adopt_existing and row["status"] != JobStatus.QUEUED.value: continue processing_task = asyncio.create_task(_process_queued_job(QueuedJob(job_id, command, server_url, model_name, adopt_existing))) - _active_job = (job_id, processing_task, OpenshiftJob(job_name=job_id)) + _active_job = ( + job_id, + processing_task, + OpenshiftJob(job_name=job_id, clean_legacy_pods=adopt_existing), + ) try: await processing_task except asyncio.CancelledError: if _shutting_down: raise - oj = OpenshiftJob(job_name=job_id) + oj = OpenshiftJob(job_name=job_id, clean_legacy_pods=adopt_existing) if not await _finish_cancellation(job_id, oj, signal=False): await _retry_cancellation(job_id, oj) finally: @@ -830,7 +900,12 @@ def build_table(title: str, jobs: list[dict]) -> str: rows = f'No jobs' return f"

{title}

{header}{rows}
" - running = job_store.list(JobStatus.RUNNING) + job_store.list(JobStatus.CANCELLING) + running = ( + job_store.list(JobStatus.RUNNING) + + job_store.list(JobStatus.COMPLETING) + + job_store.list(JobStatus.FAILING) + + job_store.list(JobStatus.CANCELLING) + ) queued = job_store.list(JobStatus.QUEUED) completed = job_store.list(JobStatus.COMPLETED) + job_store.list(JobStatus.FAILED) + job_store.list(JobStatus.CANCELLED) completed.reverse() @@ -1046,6 +1121,9 @@ async def delete_job(job_id: str): if not job_row: raise HTTPException(status_code=404, detail="Job not found") + if job_row["status"] == JobStatus.CANCELLING.value: + return {"message": "Job cancelling", "job_id": job_id} + if job_row["status"] in ( JobStatus.COMPLETING, JobStatus.COMPLETED, @@ -1055,12 +1133,23 @@ async def delete_job(job_id: str): ): raise HTTPException(status_code=400, detail=f"Job already {job_row['status']}") - # Remove from queue if still waiting + # Remove from the queue only when the persisted job has never started. for i, queued in enumerate(_job_queue): if queued.job_id == job_id: - _job_queue.pop(i) - job_store.update_status(job_id, JobStatus.CANCELLED) - return {"message": "Job cancelled", "job_id": job_id} + if ( + job_row["status"] == JobStatus.QUEUED.value + and not queued.adopt_existing + ): + _job_queue.pop(i) + job_store.update_status(job_id, JobStatus.CANCELLED) + return {"message": "Job cancelled", "job_id": job_id} + if job_row["status"] in ( + JobStatus.QUEUED.value, + JobStatus.RUNNING.value, + ): + job_store.update_status(job_id, JobStatus.CANCELLING) + _job_event.set() + return {"message": "Job cancelling", "job_id": job_id} # Cancel the actively running job if _active_job and _active_job[0] == job_id: @@ -1120,6 +1209,20 @@ def _build_url_replace_shell_step(server_url: str, py_job_dir: str) -> str: return f" && python3 -c {shlex.quote(replace_script)}" +def _build_parent_env_shell_step(py_job_dir: str) -> str: + """Update a resumed Harbor config so new task pods retain parent ownership.""" + lines = [ + "import json, os", + f"path = {json.dumps(f'{py_job_dir}/config.json')}", + "with open(path) as f: config = json.load(f)", + "kwargs = config.setdefault('environment', {}).setdefault('kwargs', {})", + "env = kwargs.setdefault('persistent_env', {})", + "env['HARBOR_PARENT'] = os.environ['HARBOR_PARENT']", + "with open(path, 'w') as f: json.dump(config, f)", + ] + return f" && python3 -c {shlex.quote(chr(10).join(lines))}" + + @router.post("/jobs/{job_id}/resume") async def resume_job(job_id: str, req: ResumeJobRequest = ResumeJobRequest()): """Resume a completed/failed job by retrying errored tasks via harbor jobs resume.""" @@ -1165,6 +1268,7 @@ async def resume_job(job_id: str, req: ResumeJobRequest = ResumeJobRequest()): shell_command = ( "mc alias set minio http://harbor-minio:9000 $MINIO_ROOT_USER $MINIO_ROOT_PASSWORD" f" && mc cp --recursive minio/results/{shlex.quote(original_job_name)}/ {job_dir}/" + f"{_build_parent_env_shell_step(py_job_dir)}" f"{url_replace_step}" f" && uv run --no-sync --no-cache harbor jobs resume -p {job_dir}{filter_flags}" f" ; mc rm --recursive --force minio/results/{shlex.quote(original_job_name)}/" diff --git a/src/coding_agent_bench/builder.py b/src/coding_agent_bench/builder.py index d79d3e1..6fa75ff 100644 --- a/src/coding_agent_bench/builder.py +++ b/src/coding_agent_bench/builder.py @@ -86,6 +86,12 @@ def _build_command( ] else: args += ["--env", environment] + harbor_parent = os.environ.get("HARBOR_PARENT") + if environment == "openshift" and harbor_parent: + args += [ + "--ek", + f"persistent_env={json.dumps({'HARBOR_PARENT': harbor_parent})}", + ] # Add mounts if mounts is not None: diff --git a/src/coding_agent_bench/job.py b/src/coding_agent_bench/job.py index 658e85b..e563bfb 100644 --- a/src/coding_agent_bench/job.py +++ b/src/coding_agent_bench/job.py @@ -26,9 +26,10 @@ def preflight(cls) -> None: "Please run 'oc login' and try again." ) - def __init__(self, job_name: str): + def __init__(self, job_name: str, clean_legacy_pods: bool = False): self._job_name = job_name self._pod_name = f"coding-agent-bench--{self._job_name}"[:58] + self._clean_legacy_pods = clean_legacy_pods def _resume_job_spec(self, shell_command: str) -> dict: """Build a pod spec for a resume job with a raw shell command.""" @@ -212,7 +213,16 @@ async def _signal_job_pod(self) -> None: return await self._run_oc_command( - ["exec", pod_name, "--", "kill", "-TERM", "1"], + [ + "exec", pod_name, "--", "sh", "-c", + "for f in /proc/[0-9]*/cmdline; do " + "pid=${f#/proc/}; pid=${pid%/cmdline}; " + "[ \"$pid\" = 1 ] && continue; " + "[ \"$pid\" = \"$$\" ] && continue; " + "cmd=$(tr '\\0' ' ' < \"$f\" 2>/dev/null) || continue; " + "case \"$cmd\" in *'harbor run'*|*'harbor jobs resume'*) " + "kill -TERM \"$pid\" 2>/dev/null || true;; esac; done", + ], check=False, ) @@ -230,9 +240,30 @@ async def _signal_job_pod(self) -> None: await asyncio.sleep(2) async def _delete_harbor_pods(self): - """Delete all pods spawned by harbor that are associated with this job.""" + """Delete task pods whose environment identifies this parent Job.""" + stdout, _ = await self._run_oc_command( + ["get", "pods", "--selector=app=harbor,harbor-session", "-o", "json"], + timeout_sec=60, + ) + pods = json.loads(stdout or "{}").get("items", []) + pod_names = [] + for pod in pods: + env = [ + item + for container in pod.get("spec", {}).get("containers", []) + for item in container.get("env", []) + ] + has_parent = any(item.get("name") == "HARBOR_PARENT" for item in env) + if any( + item.get("name") == "HARBOR_PARENT" + and item.get("value") == self._pod_name + for item in env + ) or (self._clean_legacy_pods and not has_parent): + pod_names.append(pod["metadata"]["name"]) + if not pod_names: + return await self._run_oc_command( - ["delete", "pods", f"--selector=harbor-parent={self._pod_name}", "--ignore-not-found"], + ["delete", "pods", *pod_names, "--ignore-not-found"], timeout_sec=60, ) From f0b98cffe75b8a9262053e6d4ede4ae14b5290ce Mon Sep 17 00:00:00 2001 From: Taylor Agarwal Date: Thu, 3 Sep 2026 11:53:41 -0400 Subject: [PATCH 03/11] Fix locking when adopting Nebius instances --- src/coding_agent_bench/api.py | 21 +++++++++++---------- 1 file changed, 11 insertions(+), 10 deletions(-) diff --git a/src/coding_agent_bench/api.py b/src/coding_agent_bench/api.py index 3060cb5..85ba71a 100644 --- a/src/coding_agent_bench/api.py +++ b/src/coding_agent_bench/api.py @@ -206,16 +206,17 @@ async def mark_job_completed(self, instance_name: str): async def adopt_running_instance(self, model_name: str, gpu_config: str) -> str: """Restore tracking for the deterministic instance used by a running job.""" - instance_name = self._pick_instance_name() - if not await self._manager.instance_exists(instance_name): - raise RuntimeError(f"Nebius instance {instance_name} is missing for recovered job") - self._instances[instance_name] = NebiusInstanceState( - instance_name=instance_name, - gpu_config=gpu_config, - current_model=model_name, - job_running=True, - ) - return instance_name + async with self._lock: + instance_name = self._pick_instance_name() + if not await self._manager.instance_exists(instance_name): + raise RuntimeError(f"Nebius instance {instance_name} is missing for recovered job") + self._instances[instance_name] = NebiusInstanceState( + instance_name=instance_name, + gpu_config=gpu_config, + current_model=model_name, + job_running=True, + ) + return instance_name async def delete_recovered_instance(self) -> None: """Delete the deterministic VM left behind by an interrupted terminal transition.""" From 779aaa434fc1399e0da398815294ced40a8053cd Mon Sep 17 00:00:00 2001 From: Taylor Agarwal Date: Thu, 3 Sep 2026 11:54:30 -0400 Subject: [PATCH 04/11] Remove stale active job references --- src/coding_agent_bench/api.py | 24 ++---------------------- 1 file changed, 2 insertions(+), 22 deletions(-) diff --git a/src/coding_agent_bench/api.py b/src/coding_agent_bench/api.py index 85ba71a..45ce391 100644 --- a/src/coding_agent_bench/api.py +++ b/src/coding_agent_bench/api.py @@ -47,7 +47,7 @@ class QueuedJob(NamedTuple): _job_queue: list[QueuedJob] = [] _job_event = asyncio.Event() -_active_job: tuple[str, asyncio.Task, OpenshiftJob] | None = None +_active_job: tuple[str, asyncio.Task] | None = None _shutting_down = False _nebius: "NebiusOrchestrator | None" = None @@ -612,12 +612,7 @@ async def _restore_jobs() -> bool: async def _run_job(job_id: str, command: list[str], adopt_existing: bool = False, openrouter: bool = False): """Run and monitor an Openshift Job.""" - global _active_job - oj = OpenshiftJob(job_name=job_id, clean_legacy_pods=adopt_existing) - task = asyncio.current_task() - assert task is not None - _active_job = (job_id, task, oj) try: if not adopt_existing: @@ -699,9 +694,6 @@ async def _run_job(job_id: str, command: list[str], adopt_existing: bool = False error = str(e) await _retry_terminal_job(job_id, oj, JobStatus.FAILED, error=error) - finally: - _active_job = None - def _reorder_queue_for_nebius(): """Stable-sort the queue so nebius jobs that can reuse the current instance @@ -732,13 +724,8 @@ def _sort_key(item: QueuedJob): async def _process_queued_job(queued: QueuedJob) -> None: """Provision or adopt one queued job without risking the dispatcher task.""" - global _active_job - job_id, command, server_url, model_name, adopt_existing = queued - task = asyncio.current_task() - assert task is not None oj = OpenshiftJob(job_name=job_id, clean_legacy_pods=adopt_existing) - _active_job = (job_id, task, oj) model_config: ModelConfig | None = None nebius_instance_name: str | None = None @@ -832,9 +819,6 @@ async def _process_queued_job(queued: QueuedJob) -> None: await _delete_recovered_nebius(job_id) if not await _finish_cancellation(job_id, oj, signal=True): await _retry_cancellation(job_id, oj) - finally: - if _active_job and _active_job[0] == job_id: - _active_job = None async def _worker(): @@ -860,11 +844,7 @@ async def _worker(): if not adopt_existing and row["status"] != JobStatus.QUEUED.value: continue processing_task = asyncio.create_task(_process_queued_job(QueuedJob(job_id, command, server_url, model_name, adopt_existing))) - _active_job = ( - job_id, - processing_task, - OpenshiftJob(job_name=job_id, clean_legacy_pods=adopt_existing), - ) + _active_job = (job_id, processing_task) try: await processing_task except asyncio.CancelledError: From df263e42e9c25e17181bad3466e5b46be6af82ef Mon Sep 17 00:00:00 2001 From: Taylor Agarwal Date: Thu, 3 Sep 2026 11:55:35 -0400 Subject: [PATCH 05/11] Bound recovery cleanup retries --- src/coding_agent_bench/api.py | 49 +++++++++++++++---- tests/test_recovery_cleanup.py | 87 ++++++++++++++++++++++++++++++++++ 2 files changed, 126 insertions(+), 10 deletions(-) create mode 100644 tests/test_recovery_cleanup.py diff --git a/src/coding_agent_bench/api.py b/src/coding_agent_bench/api.py index 45ce391..142397a 100644 --- a/src/coding_agent_bench/api.py +++ b/src/coding_agent_bench/api.py @@ -75,6 +75,8 @@ class JobStatus(str, Enum): NEBIUS_IDLE_TIMEOUT = int(os.environ.get("NEBIUS_IDLE_TIMEOUT_SECONDS", "600")) +CLEANUP_MAX_ATTEMPTS = int(os.environ.get("CLEANUP_MAX_ATTEMPTS", "120")) +CLEANUP_RETRY_INTERVAL_SECONDS = float(os.environ.get("CLEANUP_RETRY_INTERVAL_SECONDS", "5")) @dataclass @@ -529,14 +531,25 @@ async def _finish_cancellation(job_id: str, oj: OpenshiftJob, signal: bool) -> b async def _retry_cancellation(job_id: str, oj: OpenshiftJob) -> None: - """Keep the serial queue blocked until an active workload is removed.""" - while job_store.get(job_id)["status"] == JobStatus.CANCELLING.value: - await asyncio.sleep(5) + """Retry cancellation cleanup without blocking the serial queue forever.""" + for attempt in range(1, CLEANUP_MAX_ATTEMPTS + 1): + if job_store.get(job_id)["status"] != JobStatus.CANCELLING.value: + return + await asyncio.sleep(CLEANUP_RETRY_INTERVAL_SECONDS) try: existing = await oj._get_job() await _finish_cancellation(job_id, oj, signal=existing is not None) except Exception: logger.exception(f"Cancellation cleanup retry failed for {job_id}") + if job_store.get(job_id)["status"] != JobStatus.CANCELLING.value: + return + logger.warning( + f"Cancellation cleanup attempt {attempt}/{CLEANUP_MAX_ATTEMPTS} failed for {job_id}" + ) + + row = job_store.get(job_id) + logger.error(f"Cancellation cleanup exhausted for {job_id}; advancing the queue") + job_store.update_status(job_id, JobStatus.CANCELLED, error=row["error"]) def _terminal_error(error: str | None) -> str | None: @@ -573,21 +586,37 @@ async def _retry_terminal_job( final_status: JobStatus, error: str | None = None, ) -> None: - """Keep the serial queue blocked until terminal workload cleanup succeeds.""" - while not await _finish_terminal_job(job_id, oj, final_status, error=error): - await asyncio.sleep(5) + """Retry terminal cleanup without blocking the serial queue forever.""" + for attempt in range(1, CLEANUP_MAX_ATTEMPTS + 1): + if await _finish_terminal_job(job_id, oj, final_status, error=error): + return + logger.warning( + f"Terminal cleanup attempt {attempt}/{CLEANUP_MAX_ATTEMPTS} failed for {job_id}" + ) + if attempt < CLEANUP_MAX_ATTEMPTS: + await asyncio.sleep(CLEANUP_RETRY_INTERVAL_SECONDS) + + row = job_store.get(job_id) + logger.error(f"Terminal cleanup exhausted for {job_id}; advancing the queue") + job_store.update_status(job_id, final_status, error=row["error"]) async def _delete_recovered_nebius(job_id: str) -> None: - """Retry deletion so a recovered terminal transition cannot leak its VM.""" + """Retry deletion without blocking recovery forever.""" assert _nebius is not None - while True: + for attempt in range(1, CLEANUP_MAX_ATTEMPTS + 1): try: await _nebius.delete_recovered_instance() return except Exception: - logger.exception(f"Unable to delete recovered Nebius instance for {job_id}; retrying") - await asyncio.sleep(5) + logger.exception( + f"Unable to delete recovered Nebius instance for {job_id} " + f"(attempt {attempt}/{CLEANUP_MAX_ATTEMPTS})" + ) + if attempt < CLEANUP_MAX_ATTEMPTS: + await asyncio.sleep(CLEANUP_RETRY_INTERVAL_SECONDS) + + logger.error(f"Nebius cleanup exhausted for {job_id}; advancing recovery") async def _restore_jobs() -> bool: diff --git a/tests/test_recovery_cleanup.py b/tests/test_recovery_cleanup.py new file mode 100644 index 0000000..b599a1b --- /dev/null +++ b/tests/test_recovery_cleanup.py @@ -0,0 +1,87 @@ +import asyncio + +from coding_agent_bench import api + + +class FakeJobStore: + def __init__(self, status, error="cleanup failed: unavailable"): + self.row = {"status": status.value, "error": error} + + def get(self, _job_id): + return self.row + + def update_status(self, _job_id, status, error=None): + self.row = {"status": status.value, "error": error} + + +def disable_retry_delays(monkeypatch): + async def no_sleep(_seconds): + pass + + monkeypatch.setattr(api, "CLEANUP_MAX_ATTEMPTS", 2) + monkeypatch.setattr(api.asyncio, "sleep", no_sleep) + + +def test_terminal_cleanup_advances_queue_after_retry_limit(monkeypatch): + disable_retry_delays(monkeypatch) + store = FakeJobStore(api.JobStatus.FAILING) + monkeypatch.setattr(api, "job_store", store) + + attempts = 0 + + async def fail_cleanup(*_args, **_kwargs): + nonlocal attempts + attempts += 1 + return False + + monkeypatch.setattr(api, "_finish_terminal_job", fail_cleanup) + + asyncio.run( + api._retry_terminal_job("job-1", object(), api.JobStatus.FAILED, error="failed") + ) + + assert attempts == 2 + assert store.row == { + "status": api.JobStatus.FAILED.value, + "error": "cleanup failed: unavailable", + } + + +def test_cancellation_advances_queue_after_retry_limit(monkeypatch): + disable_retry_delays(monkeypatch) + store = FakeJobStore(api.JobStatus.CANCELLING) + monkeypatch.setattr(api, "job_store", store) + + class Job: + async def _get_job(self): + return {} + + async def fail_cleanup(*_args, **_kwargs): + return False + + monkeypatch.setattr(api, "_finish_cancellation", fail_cleanup) + + asyncio.run(api._retry_cancellation("job-1", Job())) + + assert store.row == { + "status": api.JobStatus.CANCELLED.value, + "error": "cleanup failed: unavailable", + } + + +def test_nebius_cleanup_stops_after_retry_limit(monkeypatch): + disable_retry_delays(monkeypatch) + + class Nebius: + attempts = 0 + + async def delete_recovered_instance(self): + self.attempts += 1 + raise RuntimeError("unavailable") + + nebius = Nebius() + monkeypatch.setattr(api, "_nebius", nebius) + + asyncio.run(api._delete_recovered_nebius("job-1")) + + assert nebius.attempts == 2 From 7d6547c35078b3d6d97b4cdae48f57bce55d80ae Mon Sep 17 00:00:00 2001 From: Taylor Agarwal Date: Thu, 3 Sep 2026 11:56:09 -0400 Subject: [PATCH 06/11] Document Nebius restart cleanup policy --- src/coding_agent_bench/api.py | 3 +++ 1 file changed, 3 insertions(+) diff --git a/src/coding_agent_bench/api.py b/src/coding_agent_bench/api.py index 142397a..2992600 100644 --- a/src/coding_agent_bench/api.py +++ b/src/coding_agent_bench/api.py @@ -455,6 +455,9 @@ async def lifespan(_app: FastAPI) -> AsyncIterator[None]: await task except asyncio.CancelledError: pass + # Do not delete Nebius instances here: OpenShift jobs survive queue restarts + # and the next queue process must be able to adopt their deterministic VM. + # Permanent decommissioning therefore requires external instance cleanup. app = FastAPI(lifespan=lifespan) From 88c63682ea020b688bd44f18c1973987fdc5fa67 Mon Sep 17 00:00:00 2001 From: Taylor Agarwal Date: Thu, 3 Sep 2026 11:58:05 -0400 Subject: [PATCH 07/11] Scope legacy Harbor pod cleanup --- src/coding_agent_bench/job.py | 11 ++++++-- tests/test_job_cleanup.py | 50 +++++++++++++++++++++++++++++++++++ 2 files changed, 59 insertions(+), 2 deletions(-) create mode 100644 tests/test_job_cleanup.py diff --git a/src/coding_agent_bench/job.py b/src/coding_agent_bench/job.py index e563bfb..21ae622 100644 --- a/src/coding_agent_bench/job.py +++ b/src/coding_agent_bench/job.py @@ -254,11 +254,18 @@ async def _delete_harbor_pods(self): for item in container.get("env", []) ] has_parent = any(item.get("name") == "HARBOR_PARENT" for item in env) - if any( + matches_parent = any( item.get("name") == "HARBOR_PARENT" and item.get("value") == self._pod_name for item in env - ) or (self._clean_legacy_pods and not has_parent): + ) + labels = pod.get("metadata", {}).get("labels", {}) + matches_legacy_parent = ( + self._clean_legacy_pods + and not has_parent + and labels.get("harbor-parent") == self._pod_name + ) + if matches_parent or matches_legacy_parent: pod_names.append(pod["metadata"]["name"]) if not pod_names: return diff --git a/tests/test_job_cleanup.py b/tests/test_job_cleanup.py new file mode 100644 index 0000000..9a1af02 --- /dev/null +++ b/tests/test_job_cleanup.py @@ -0,0 +1,50 @@ +import asyncio +import json + +from coding_agent_bench.job import OpenshiftJob + + +def test_legacy_cleanup_only_deletes_pods_owned_by_job(monkeypatch): + job = OpenshiftJob("job-1", clean_legacy_pods=True) + pods = { + "items": [ + { + "metadata": {"name": "current-match"}, + "spec": { + "containers": [ + {"env": [{"name": "HARBOR_PARENT", "value": job._pod_name}]} + ] + }, + }, + { + "metadata": { + "name": "legacy-match", + "labels": {"harbor-parent": job._pod_name}, + }, + "spec": {"containers": [{"env": []}]}, + }, + { + "metadata": {"name": "other-job"}, + "spec": {"containers": [{"env": []}]}, + }, + ] + } + commands = [] + + async def run_oc(command, **_kwargs): + commands.append(command) + if command[0] == "get": + return json.dumps(pods), None + return None, None + + monkeypatch.setattr(job, "_run_oc_command", run_oc) + + asyncio.run(job._delete_harbor_pods()) + + assert commands[1] == [ + "delete", + "pods", + "current-match", + "legacy-match", + "--ignore-not-found", + ] From c30d2eac4a009934fb895f5a747db9672442af0d Mon Sep 17 00:00:00 2001 From: Taylor Agarwal Date: Thu, 3 Sep 2026 11:59:58 -0400 Subject: [PATCH 08/11] Isolate recovery test imports --- tests/test_recovery_cleanup.py | 16 ++++++++++------ 1 file changed, 10 insertions(+), 6 deletions(-) diff --git a/tests/test_recovery_cleanup.py b/tests/test_recovery_cleanup.py index b599a1b..144a12c 100644 --- a/tests/test_recovery_cleanup.py +++ b/tests/test_recovery_cleanup.py @@ -1,7 +1,5 @@ import asyncio -from coding_agent_bench import api - class FakeJobStore: def __init__(self, status, error="cleanup failed: unavailable"): @@ -14,7 +12,7 @@ def update_status(self, _job_id, status, error=None): self.row = {"status": status.value, "error": error} -def disable_retry_delays(monkeypatch): +def disable_retry_delays(monkeypatch, api): async def no_sleep(_seconds): pass @@ -23,7 +21,9 @@ async def no_sleep(_seconds): def test_terminal_cleanup_advances_queue_after_retry_limit(monkeypatch): - disable_retry_delays(monkeypatch) + from coding_agent_bench import api + + disable_retry_delays(monkeypatch, api) store = FakeJobStore(api.JobStatus.FAILING) monkeypatch.setattr(api, "job_store", store) @@ -48,7 +48,9 @@ async def fail_cleanup(*_args, **_kwargs): def test_cancellation_advances_queue_after_retry_limit(monkeypatch): - disable_retry_delays(monkeypatch) + from coding_agent_bench import api + + disable_retry_delays(monkeypatch, api) store = FakeJobStore(api.JobStatus.CANCELLING) monkeypatch.setattr(api, "job_store", store) @@ -70,7 +72,9 @@ async def fail_cleanup(*_args, **_kwargs): def test_nebius_cleanup_stops_after_retry_limit(monkeypatch): - disable_retry_delays(monkeypatch) + from coding_agent_bench import api + + disable_retry_delays(monkeypatch, api) class Nebius: attempts = 0 From b0c6e050f64fe1c16738c74a8e2842ac2e55acc2 Mon Sep 17 00:00:00 2001 From: Taylor Agarwal Date: Thu, 3 Sep 2026 12:00:38 -0400 Subject: [PATCH 09/11] Preserve OpenRouter job environment --- src/coding_agent_bench/job.py | 7 +++---- 1 file changed, 3 insertions(+), 4 deletions(-) diff --git a/src/coding_agent_bench/job.py b/src/coding_agent_bench/job.py index 21ae622..43c9742 100644 --- a/src/coding_agent_bench/job.py +++ b/src/coding_agent_bench/job.py @@ -74,7 +74,9 @@ def _job_spec( ) -> dict: # Only openrouter jobs need the OpenRouter key, so scope the secret to # them rather than exposing it to every job pod. - env: list[dict] = [] + env: list[dict] = [ + {"name": "HARBOR_PARENT", "value": self._pod_name}, + ] if openrouter: env.append( { @@ -118,9 +120,6 @@ def _job_spec( "envFrom": [ {"secretRef": {"name": "harbor-minio"}} ], - "env": [ - {"name": "HARBOR_PARENT", "value": self._pod_name}, - ], } ], } From c662dd876f4ad896a052552e2e80f1f0a4a2b772 Mon Sep 17 00:00:00 2001 From: Taylor Agarwal Date: Thu, 3 Sep 2026 12:59:46 -0400 Subject: [PATCH 10/11] Bound recovered job inspection retries --- src/coding_agent_bench/api.py | 28 +++++++++++---- tests/test_recovery_cleanup.py | 63 ++++++++++++++++++++++++++++++++++ 2 files changed, 84 insertions(+), 7 deletions(-) diff --git a/src/coding_agent_bench/api.py b/src/coding_agent_bench/api.py index 2992600..0d71034 100644 --- a/src/coding_agent_bench/api.py +++ b/src/coding_agent_bench/api.py @@ -662,13 +662,18 @@ async def _run_job(job_id: str, command: list[str], adopt_existing: bool = False else: if job_store.get(job_id)["status"] == JobStatus.QUEUED.value: job_store.update_status(job_id, JobStatus.RUNNING) - while True: + for attempt in range(1, CLEANUP_MAX_ATTEMPTS + 1): try: existing = await oj._get_job() break except Exception: - logger.exception(f"Unable to inspect recovered OpenShift Job {job_id}; retrying") - await asyncio.sleep(5) + logger.exception( + f"Unable to inspect recovered OpenShift Job {job_id} " + f"(attempt {attempt}/{CLEANUP_MAX_ATTEMPTS})" + ) + if attempt == CLEANUP_MAX_ATTEMPTS: + raise + await asyncio.sleep(CLEANUP_RETRY_INTERVAL_SECONDS) conditions = { condition.get("type") for condition in (existing or {}).get("status", {}).get("conditions", []) @@ -781,13 +786,22 @@ async def _process_queued_job(queued: QueuedJob) -> None: return if adopt_existing: - while True: + for attempt in range(1, CLEANUP_MAX_ATTEMPTS + 1): try: existing = await oj._get_job() break - except Exception: - logger.exception(f"Unable to reconcile recovered OpenShift Job {job_id}; retrying") - await asyncio.sleep(5) + except Exception as e: + logger.exception( + f"Unable to reconcile recovered OpenShift Job {job_id} " + f"(attempt {attempt}/{CLEANUP_MAX_ATTEMPTS})" + ) + if attempt < CLEANUP_MAX_ATTEMPTS: + await asyncio.sleep(CLEANUP_RETRY_INTERVAL_SECONDS) + continue + if nebius_gpu_config is not None and _nebius: + await _delete_recovered_nebius(job_id) + await _retry_terminal_job(job_id, oj, JobStatus.FAILED, error=str(e)) + return if existing is None: if row["status"] == JobStatus.RUNNING.value: if nebius_gpu_config is not None and _nebius: diff --git a/tests/test_recovery_cleanup.py b/tests/test_recovery_cleanup.py index 144a12c..0cd3399 100644 --- a/tests/test_recovery_cleanup.py +++ b/tests/test_recovery_cleanup.py @@ -89,3 +89,66 @@ async def delete_recovered_instance(self): asyncio.run(api._delete_recovered_nebius("job-1")) assert nebius.attempts == 2 + + +def test_run_job_stops_recovery_probe_after_retry_limit(monkeypatch): + from coding_agent_bench import api + + disable_retry_delays(monkeypatch, api) + store = FakeJobStore(api.JobStatus.RUNNING) + monkeypatch.setattr(api, "job_store", store) + + class Job: + attempts = 0 + + def __init__(self, **_kwargs): + pass + + async def _get_job(self): + type(self).attempts += 1 + raise RuntimeError("unavailable") + + terminal_errors = [] + + async def finish_terminal(_job_id, _job, _status, error=None): + terminal_errors.append(error) + + monkeypatch.setattr(api, "OpenshiftJob", Job) + monkeypatch.setattr(api, "_retry_terminal_job", finish_terminal) + + asyncio.run(api._run_job("job-1", [], adopt_existing=True)) + + assert Job.attempts == 2 + assert terminal_errors == ["unavailable"] + + +def test_process_queued_job_stops_recovery_probe_after_retry_limit(monkeypatch): + from coding_agent_bench import api + + disable_retry_delays(monkeypatch, api) + store = FakeJobStore(api.JobStatus.RUNNING) + monkeypatch.setattr(api, "job_store", store) + + class Job: + attempts = 0 + + def __init__(self, **_kwargs): + pass + + async def _get_job(self): + type(self).attempts += 1 + raise RuntimeError("unavailable") + + terminal_errors = [] + + async def finish_terminal(_job_id, _job, _status, error=None): + terminal_errors.append(error) + + monkeypatch.setattr(api, "OpenshiftJob", Job) + monkeypatch.setattr(api, "_retry_terminal_job", finish_terminal) + + queued = api.QueuedJob("job-1", [], "https://example.com", "model", True) + asyncio.run(api._process_queued_job(queued)) + + assert Job.attempts == 2 + assert terminal_errors == ["unavailable"] From 98ed20f8a34ace3fa508986d3c4317ba4efa421a Mon Sep 17 00:00:00 2001 From: Taylor Agarwal Date: Thu, 3 Sep 2026 14:51:07 -0400 Subject: [PATCH 11/11] Delete Nebius instance on job cancellation --- src/coding_agent_bench/api.py | 1 + tests/test_recovery_cleanup.py | 58 ++++++++++++++++++++++++++++++++++ 2 files changed, 59 insertions(+) diff --git a/src/coding_agent_bench/api.py b/src/coding_agent_bench/api.py index 0d71034..c1f3a5b 100644 --- a/src/coding_agent_bench/api.py +++ b/src/coding_agent_bench/api.py @@ -726,6 +726,7 @@ async def _run_job(job_id: str, command: list[str], adopt_existing: bool = False if _shutting_down: raise await _finish_cancellation(job_id, oj, signal=True) + raise except Exception as e: error = str(e) diff --git a/tests/test_recovery_cleanup.py b/tests/test_recovery_cleanup.py index 0cd3399..80b7ab9 100644 --- a/tests/test_recovery_cleanup.py +++ b/tests/test_recovery_cleanup.py @@ -152,3 +152,61 @@ async def finish_terminal(_job_id, _job, _status, error=None): assert Job.attempts == 2 assert terminal_errors == ["unavailable"] + + +def test_cancelled_nebius_job_deletes_instance_instead_of_marking_idle(monkeypatch): + from coding_agent_bench import api + + store = FakeJobStore(api.JobStatus.RUNNING) + monkeypatch.setattr(api, "job_store", store) + monkeypatch.setattr(api, "_shutting_down", False) + + existing = {"status": {"conditions": []}} + + class Job: + instances = 0 + + def __init__(self, **_kwargs): + self.instance = type(self).instances + type(self).instances += 1 + self.calls = 0 + + async def _get_job(self): + self.calls += 1 + if self.instance == 1 and self.calls == 2: + store.update_status("job-1", api.JobStatus.CANCELLING) + raise asyncio.CancelledError + return existing + + async def _wait_for_job_pod_ready(self): + pass + + async def _signal_job_pod(self): + pass + + async def _delete_job(self): + pass + + class Nebius: + deleted = 0 + completed = 0 + + async def adopt_running_instance(self, _model_name, _gpu_config): + return "instance-1" + + async def delete_recovered_instance(self): + self.deleted += 1 + + async def mark_job_completed(self, _instance_name): + self.completed += 1 + + nebius = Nebius() + monkeypatch.setattr(api, "OpenshiftJob", Job) + monkeypatch.setattr(api, "_nebius", nebius) + + queued = api.QueuedJob("job-1", [], "nebius-gpu", "model", True) + asyncio.run(api._process_queued_job(queued)) + + assert nebius.deleted == 1 + assert nebius.completed == 0 + assert store.row["status"] == api.JobStatus.CANCELLED.value