From 4ae448bfd3fecfe6693306bf3c10fa3d77ab5674 Mon Sep 17 00:00:00 2001 From: Hannah Smith Date: Wed, 15 Jul 2026 14:35:03 -0400 Subject: [PATCH 1/2] =?UTF-8?q?=E2=9C=A8=20Add=20concurrent=20job=20queue?= =?UTF-8?q?=20support?= MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit The job queue now runs up to MAX_CONCURRENT_JOBS (default 3) benchmark jobs in parallel instead of one at a time. This allows running benchmarks for multiple models simultaneously when each targets a different vLLM endpoint on separate Brev instances. --- src/coding_agent_bench/api.py | 131 +++++++++++++++++++--------------- 1 file changed, 73 insertions(+), 58 deletions(-) diff --git a/src/coding_agent_bench/api.py b/src/coding_agent_bench/api.py index 52c10f5..4a009fb 100644 --- a/src/coding_agent_bench/api.py +++ b/src/coding_agent_bench/api.py @@ -23,9 +23,12 @@ _job_queue: list[tuple[str, list[str]]] = [] _job_event = asyncio.Event() -_active_job: tuple[str, asyncio.Task, OpenshiftJob] | None = None +_active_jobs: dict[str, tuple[asyncio.Task, OpenshiftJob]] = {} _shutting_down = False +MAX_CONCURRENT_JOBS = int(os.environ.get("MAX_CONCURRENT_JOBS", "3")) +_job_semaphore = asyncio.Semaphore(MAX_CONCURRENT_JOBS) + db_path = Path(os.environ.get("JOB_STORE_PATH", "jobs.db")) class JobStatus(str, Enum): @@ -171,7 +174,10 @@ async def lifespan(_app: FastAPI) -> AsyncIterator[None]: _shutting_down = True worker_task.cancel() cleanup_task.cancel() - for task in (worker_task, cleanup_task): + for _, (task, _oj) in list(_active_jobs.items()): + task.cancel() + all_tasks = [worker_task, cleanup_task] + [t for t, _ in _active_jobs.values()] + for task in all_tasks: try: await task except asyncio.CancelledError: @@ -240,74 +246,82 @@ async def _best_effort_cleanup(oj: OpenshiftJob, signal: bool = False) -> str | async def _run_job(job_id: str, command: list[str]): """Run and monitor an Openshift Job.""" - global _active_job oj = OpenshiftJob(job_name=job_id) task = asyncio.current_task() assert task is not None - _active_job = (job_id, task, oj) + _active_jobs[job_id] = (task, oj) try: - job_spec = oj._job_spec(command) - 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) - - 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", []) - if pods: - 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, - ) - 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) - return - await asyncio.sleep(5) + async with _job_semaphore: + try: + job_spec = oj._job_spec(command) + 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) + + 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", []) + if pods: + 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, + ) + 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) + 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) + + 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) 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) + job_store.update_status(job_id, JobStatus.FAILED, error="Server shut down") 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) - - 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) + job_store.update_status(job_id, JobStatus.CANCELLED) finally: - _active_job = None + _active_jobs.pop(job_id, None) + _job_event.set() async def _worker(): - """Process jobs from the queue one at a time.""" + """Process jobs from the queue, up to MAX_CONCURRENT_JOBS at a time.""" while True: await _job_event.wait() _job_event.clear() @@ -316,7 +330,7 @@ async def _worker(): row = job_store.get(job_id) if not row or row["status"] != JobStatus.QUEUED.value: continue - await _run_job(job_id, command) + asyncio.create_task(_run_job(job_id, command)) @router.get("/") async def read_root(): @@ -461,9 +475,10 @@ async def delete_job(job_id: str): return {"message": "Job cancelled", "job_id": job_id} # Cancel the actively running job - if _active_job and _active_job[0] == job_id: + active = _active_jobs.get(job_id) + if active: job_store.update_status(job_id, JobStatus.CANCELLING) - _active_job[1].cancel() + active[0].cancel() return {"message": "Job cancelling", "job_id": job_id} return {"message": "Job cancelled", "job_id": job_id} From 5f5fe45d2580cb7f0913ca85e650c7238ca0b324 Mon Sep 17 00:00:00 2001 From: Hannah Smith Date: Wed, 15 Jul 2026 14:55:07 -0400 Subject: [PATCH 2/2] =?UTF-8?q?=F0=9F=90=9B=20Address=20CodeRabbit=20revie?= =?UTF-8?q?w=20feedback?= MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit - Validate MAX_CONCURRENT_JOBS env var at startup (reject non-int and < 1) - Pass signal=True during shutdown cleanup to avoid orphaned task pods - Track background tasks in a set to prevent GC and satisfy RUF006 --- src/coding_agent_bench/api.py | 13 ++++++++++--- 1 file changed, 10 insertions(+), 3 deletions(-) diff --git a/src/coding_agent_bench/api.py b/src/coding_agent_bench/api.py index 4a009fb..d842d22 100644 --- a/src/coding_agent_bench/api.py +++ b/src/coding_agent_bench/api.py @@ -24,9 +24,15 @@ _job_queue: list[tuple[str, list[str]]] = [] _job_event = asyncio.Event() _active_jobs: dict[str, tuple[asyncio.Task, OpenshiftJob]] = {} +_background_tasks: set[asyncio.Task] = set() _shutting_down = False -MAX_CONCURRENT_JOBS = int(os.environ.get("MAX_CONCURRENT_JOBS", "3")) +try: + MAX_CONCURRENT_JOBS = int(os.environ.get("MAX_CONCURRENT_JOBS", "3")) +except ValueError as exc: + raise RuntimeError("MAX_CONCURRENT_JOBS must be an integer") from exc +if MAX_CONCURRENT_JOBS < 1: + raise RuntimeError("MAX_CONCURRENT_JOBS must be at least 1") _job_semaphore = asyncio.Semaphore(MAX_CONCURRENT_JOBS) db_path = Path(os.environ.get("JOB_STORE_PATH", "jobs.db")) @@ -292,7 +298,7 @@ async def _run_job(job_id: str, command: list[str]): except asyncio.CancelledError: if _shutting_down: - cleanup_err = await _best_effort_cleanup(oj) + cleanup_err = await _best_effort_cleanup(oj, signal=True) error = "Server shut down" if cleanup_err: error += f"; cleanup failed: {cleanup_err}" @@ -330,7 +336,8 @@ async def _worker(): row = job_store.get(job_id) if not row or row["status"] != JobStatus.QUEUED.value: continue - asyncio.create_task(_run_job(job_id, command)) + _background_tasks.add(task := asyncio.create_task(_run_job(job_id, command))) + task.add_done_callback(_background_tasks.discard) @router.get("/") async def read_root():