diff --git a/deploy/job-queue-service.yml b/deploy/job-queue-service.yml index 31fd779..bb455b4 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 8443 \ + 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 8443 \ --ssl-certfile /etc/job-queue/tls/tls.crt \ --ssl-keyfile /etc/job-queue/tls/tls.key resources: diff --git a/src/coding_agent_bench/api.py b/src/coding_agent_bench/api.py index 9d24d93..642ffe6 100644 --- a/src/coding_agent_bench/api.py +++ b/src/coding_agent_bench/api.py @@ -45,10 +45,11 @@ class QueuedJob(NamedTuple): command: list[str] server_url: str model_name: str + adopt_existing: bool = False _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 @@ -89,13 +90,17 @@ def _worker_server_url_errors( class JobStatus(str, Enum): QUEUED = "queued" RUNNING = "running" + COMPLETING = "completing" COMPLETED = "completed" + FAILING = "failing" FAILED = "failed" CANCELLING = "cancelling" CANCELLED = "cancelled" 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 @@ -162,6 +167,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()) @@ -225,6 +233,29 @@ 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.""" + 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.""" + 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: @@ -421,21 +452,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) @@ -457,7 +494,7 @@ async def _verify_api_key(key: str = Depends(_api_key_header)) -> str: async def lifespan(_app: FastAPI) -> AsyncIterator[None]: """Start queue workers and cleanly stop them with the FastAPI application.""" global _shutting_down, _nebius - job_store.mark_orphaned() + _shutting_down = False # Initialize Nebius orchestrator if enabled background_tasks: list[asyncio.Task] = [] @@ -482,6 +519,9 @@ async def lifespan(_app: FastAPI) -> AsyncIterator[None]: background_tasks.append(asyncio.create_task(_nebius.idle_cleanup_loop())) logger.info("Nebius orchestrator initialized") + 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 @@ -495,13 +535,9 @@ 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") + # 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) @@ -522,7 +558,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) @@ -567,16 +603,133 @@ async def _best_effort_cleanup(oj: OpenshiftJob, signal: bool = False) -> str | return "; ".join(errors) if errors else None +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: + """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: + """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: + """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 without blocking recovery forever.""" + assert _nebius is not None + 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} " + 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: + """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(): + 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"], + adopt_existing=True, + )) + if _job_queue: + 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], server_url: str | None = None, managed_endpoint: bool = False, openrouter: bool = False, + adopt_existing: bool = False, ): """Validate, run, and monitor an OpenShift Job.""" - global _active_job - if server_url: server_url_errors = _worker_server_url_errors( server_url, @@ -590,90 +743,93 @@ async def _run_job( ) return - oj = OpenshiftJob(job_name=job_id) - task = asyncio.current_task() - assert task is not None - _active_job = (job_id, task, oj) + oj = OpenshiftJob(job_name=job_id, clean_legacy_pods=adopt_existing) 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]) + 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() 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 job_store.get(job_id)["status"] == JobStatus.QUEUED.value: + job_store.update_status(job_id, JobStatus.RUNNING) + 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} " + 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", []) + 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 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) + raise 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) - - finally: - _active_job = None + await _retry_terminal_job(job_id, oj, JobStatus.FAILED, error=error) def _reorder_queue_for_nebius(): @@ -689,7 +845,8 @@ def _reorder_queue_for_nebius(): current_model = states[0].current_model def _sort_key(item: QueuedJob): - """Rank a queued job by how efficiently it can reuse Nebius capacity.""" + if item.adopt_existing: + return -1 gpu = _parse_nebius_url(item.server_url) if gpu is None: return 0 # non-nebius, keep in place @@ -702,68 +859,167 @@ 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.""" + job_id, command, server_url, model_name, adopt_existing = queued + oj = OpenshiftJob(job_name=job_id, clean_legacy_pods=adopt_existing) + + model_config: ModelConfig | None = None + nebius_instance_name: str | None = None + try: + nebius_gpu_config = _parse_nebius_url(server_url) + job_server_url: str | None = server_url + managed_endpoint = False + 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 nebius_gpu_config is not None and not _nebius: + job_store.update_status( + job_id, + JobStatus.FAILED, + error="Nebius is not enabled on this server", + ) + return + + if adopt_existing: + for attempt in range(1, CLEANUP_MAX_ATTEMPTS + 1): + try: + existing = await oj._get_job() + break + 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: + 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 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) + is_resume = len(command) == 3 and command[0] == "sh" and command[1] == "-c" + if is_resume: + 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] + job_server_url = real_url + managed_endpoint = True + 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, + server_url=None if adopt_existing else job_server_url, + managed_endpoint=managed_endpoint, + 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 _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) + + 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, + JobStatus.COMPLETING.value, + JobStatus.FAILING.value, + JobStatus.CANCELLING.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) - job_server_url = server_url - managed_endpoint = False - if nebius_gpu_config is not None and not _nebius: - job_store.update_status( - job_id, - JobStatus.FAILED, - error="Nebius is not enabled on this server", - ) + if not adopt_existing and row["status"] != JobStatus.QUEUED.value: continue - 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] - job_server_url = real_url - managed_endpoint = True - 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, - server_url=job_server_url, - managed_endpoint=managed_endpoint, - openrouter=is_openrouter(server_url), - ) - - if nebius_instance_name and _nebius: - await _nebius.mark_job_completed(nebius_instance_name) + processing_task = asyncio.create_task(_process_queued_job(QueuedJob(job_id, command, server_url, model_name, adopt_existing))) + _active_job = (job_id, processing_task) + try: + await processing_task + except asyncio.CancelledError: + if _shutting_down: + raise + 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: + if _active_job and _active_job[0] == job_id: + _active_job = None @router.get("/") async def read_root(): @@ -791,7 +1047,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() @@ -1080,15 +1341,35 @@ 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"] == JobStatus.CANCELLING.value: + return {"message": "Job cancelling", "job_id": job_id} + + 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 + # 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: @@ -1148,6 +1429,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.""" @@ -1193,6 +1488,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 f941f83..009286f 100644 --- a/src/coding_agent_bench/builder.py +++ b/src/coding_agent_bench/builder.py @@ -81,6 +81,12 @@ def _build_command( # Add environment 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 e995bd7..43c9742 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.""" @@ -37,6 +38,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 +53,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": [ @@ -71,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( { @@ -90,6 +95,7 @@ def _job_spec( "kind": "Job", "metadata": {"name": self._pod_name, "labels": {"app": "harbor"}}, "spec": { + "backoffLimit": 0, "template": { "spec": { "restartPolicy": "Never", @@ -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 @@ -189,7 +212,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, ) @@ -207,9 +239,37 @@ 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) + matches_parent = any( + item.get("name") == "HARBOR_PARENT" + and item.get("value") == self._pod_name + for item in env + ) + 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 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, ) 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", + ] diff --git a/tests/test_recovery_cleanup.py b/tests/test_recovery_cleanup.py new file mode 100644 index 0000000..80b7ab9 --- /dev/null +++ b/tests/test_recovery_cleanup.py @@ -0,0 +1,212 @@ +import asyncio + + +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, api): + 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): + from coding_agent_bench import api + + disable_retry_delays(monkeypatch, api) + 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): + from coding_agent_bench import api + + disable_retry_delays(monkeypatch, api) + 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): + from coding_agent_bench import api + + disable_retry_delays(monkeypatch, api) + + 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 + + +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"] + + +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