From 905dd76213e8dd3b8c2a5f6dccd663b58d62e459 Mon Sep 17 00:00:00 2001 From: Sebastien Tardif Date: Sat, 29 Aug 2026 11:20:30 -0700 Subject: [PATCH] fix(queue): use a thread lock for JobQueue across event loops Space UI calls asyncio.run(queue.submit()) on a new loop while the eval worker holds the same lock on its own loop. asyncio.Lock is loop-bound, so public submit can hang or raise. Use threading.Lock and release it before Hub uploads. Signed-off-by: Sebastien Tardif --- clawbench/queue.py | 116 +++++++++++++++++++++++-------------------- tests/test_queue.py | 117 ++++++++++++++++++++++++++++++++++++++++++++ 2 files changed, 179 insertions(+), 54 deletions(-) diff --git a/clawbench/queue.py b/clawbench/queue.py index 2c0a664..8da7fdb 100644 --- a/clawbench/queue.py +++ b/clawbench/queue.py @@ -17,6 +17,7 @@ import logging import os import tempfile +import threading from enum import Enum from pathlib import Path @@ -103,7 +104,9 @@ class JobQueue: def __init__(self) -> None: self._jobs: dict[str, Job] = {} - self._lock = asyncio.Lock() + # Gradio calls asyncio.run(queue.submit()) on a fresh loop while the + # worker keeps its own loop. asyncio.Lock is loop-bound and can hang. + self._lock = threading.Lock() self._dataset_repo = resolve_dataset_repo(HF_TOKEN) LOCAL_QUEUE_DIR.mkdir(parents=True, exist_ok=True) self._load_local() @@ -179,7 +182,7 @@ def _save_local(self) -> None: async def submit(self, request: SubmissionRequest) -> Job: """Submit a new evaluation job.""" import uuid - async with self._lock: + with self._lock: max_runs = _env_int("CLAWBENCH_MAX_RUNS_PER_SUBMISSION", 3, minimum=1, maximum=100) if request.runs_per_task > max_runs: raise ValueError( @@ -231,9 +234,9 @@ async def submit(self, request: SubmissionRequest) -> Job: ) self._jobs[job.job_id] = job self._save_local() - await self._sync_to_hub() - logger.info("Job %s submitted for model %s", job.job_id, request.model) - return job + await self._sync_to_hub() + logger.info("Job %s submitted for model %s", job.job_id, request.model) + return job async def get_status(self, job_id: str) -> Job | None: return self._jobs.get(job_id) @@ -249,7 +252,7 @@ async def claim_pending(self, limit: int = 1) -> list[Job]: """Atomically claim up to ``limit`` pending jobs for evaluation.""" if limit <= 0: return [] - async with self._lock: + with self._lock: claimed: list[Job] = [] pending = sorted( (job for job in self._jobs.values() if job.status == JobStatus.PENDING), @@ -271,8 +274,9 @@ async def claim_pending(self, limit: int = 1) -> list[Job]: claimed.append(job) if claimed: self._save_local() - await self._sync_to_hub() - return claimed + if claimed: + await self._sync_to_hub() + return claimed async def update_progress( self, @@ -283,7 +287,7 @@ async def update_progress( current_run_total: int | None, progress_message: str | None, ) -> None: - async with self._lock: + with self._lock: job = self._jobs.get(job_id) if not job or job.status != JobStatus.EVALUATING: return @@ -293,13 +297,13 @@ async def update_progress( job.current_run_total = current_run_total job.progress_message = progress_message self._save_local() - await self._sync_to_hub() + await self._sync_to_hub() async def reclaim_stale_jobs(self, stale_after_seconds: int) -> list[Job]: """Return evaluating jobs to pending when their heartbeat is stale.""" if stale_after_seconds <= 0: return [] - async with self._lock: + with self._lock: reclaimed: list[Job] = [] cutoff = datetime.datetime.now(datetime.timezone.utc) - datetime.timedelta(seconds=stale_after_seconds) now_iso = _now_iso() @@ -327,59 +331,63 @@ async def reclaim_stale_jobs(self, stale_after_seconds: int) -> list[Job]: reclaimed.append(job) if reclaimed: self._save_local() - await self._sync_to_hub() logger.warning("Reclaimed %d stale evaluating jobs", len(reclaimed)) - return reclaimed + if reclaimed: + await self._sync_to_hub() + return reclaimed async def mark_evaluating(self, job_id: str) -> None: - async with self._lock: + with self._lock: job = self._jobs.get(job_id) - if job: - job.status = JobStatus.EVALUATING - now_iso = _now_iso() - if job.started_at is None: - job.attempt_count += 1 - job.started_at = now_iso - job.last_progress_at = now_iso - job.finished_at = None - job.error = None - job.result_id = None - job.current_task_id = None - job.current_run_index = None - job.current_run_total = None - job.progress_message = "Queued for evaluation" - self._save_local() - await self._sync_to_hub() + if not job: + return + job.status = JobStatus.EVALUATING + now_iso = _now_iso() + if job.started_at is None: + job.attempt_count += 1 + job.started_at = now_iso + job.last_progress_at = now_iso + job.finished_at = None + job.error = None + job.result_id = None + job.current_task_id = None + job.current_run_index = None + job.current_run_total = None + job.progress_message = "Queued for evaluation" + self._save_local() + await self._sync_to_hub() async def mark_finished(self, job_id: str, result_id: str) -> None: - async with self._lock: + with self._lock: job = self._jobs.get(job_id) - if job: - job.status = JobStatus.FINISHED - job.finished_at = _now_iso() - job.last_progress_at = job.finished_at - job.result_id = result_id - job.current_task_id = None - job.current_run_index = None - job.current_run_total = None - job.progress_message = "Finished" - self._save_local() - await self._sync_to_hub() + if not job: + return + job.status = JobStatus.FINISHED + job.finished_at = _now_iso() + job.last_progress_at = job.finished_at + job.result_id = result_id + job.current_task_id = None + job.current_run_index = None + job.current_run_total = None + job.progress_message = "Finished" + self._save_local() + await self._sync_to_hub() async def mark_failed(self, job_id: str, error: str) -> None: - async with self._lock: + with self._lock: job = self._jobs.get(job_id) - if job: - job.status = JobStatus.FAILED - job.finished_at = _now_iso() - job.last_progress_at = job.finished_at - job.error = error - job.current_task_id = None - job.current_run_index = None - job.current_run_total = None - job.progress_message = "Failed" - self._save_local() - await self._sync_to_hub() + if not job: + return + job.status = JobStatus.FAILED + job.finished_at = _now_iso() + job.last_progress_at = job.finished_at + job.error = error + job.current_task_id = None + job.current_run_index = None + job.current_run_total = None + job.progress_message = "Failed" + self._save_local() + await self._sync_to_hub() async def _sync_to_hub(self) -> None: """Push queue state to HF Dataset for persistence across restarts.""" diff --git a/tests/test_queue.py b/tests/test_queue.py index 57641e3..0b4f674 100644 --- a/tests/test_queue.py +++ b/tests/test_queue.py @@ -1,5 +1,7 @@ +import asyncio import datetime import json +import threading import pytest @@ -462,3 +464,118 @@ async def fake_sync() -> None: assert fresh_job.current_task_id == "t1-bugfix-discount" assert save_calls == ["saved"] assert sync_calls == ["synced"] + + +async def _hold_queue_lock(queue: JobQueue, held: threading.Event, release: threading.Event) -> None: + """Hold JobQueue._lock whether it is asyncio.Lock or threading.Lock.""" + lock = queue._lock + if hasattr(lock, "__aenter__"): + async with lock: + held.set() + await asyncio.to_thread(release.wait, 5) + return + with lock: + held.set() + await asyncio.to_thread(release.wait, 5) + + +def test_submit_from_second_event_loop_while_worker_holds_lock(tmp_path, monkeypatch): + """Space UI uses asyncio.run() on a new loop while EvalWorker holds the queue lock.""" + monkeypatch.setattr(queue_module, "LOCAL_QUEUE_DIR", tmp_path) + monkeypatch.setattr(queue_module, "HF_TOKEN", "") + queue = JobQueue() + + held = threading.Event() + release = threading.Event() + worker_errors: list[BaseException] = [] + + def worker() -> None: + try: + asyncio.run(_hold_queue_lock(queue, held, release)) + except BaseException as exc: + worker_errors.append(exc) + + worker_thread = threading.Thread(target=worker, daemon=True) + worker_thread.start() + assert held.wait(timeout=2), "worker never acquired the queue lock" + + outcome: dict[str, Job | BaseException] = {} + + def submit_from_other_loop() -> None: + try: + outcome["job"] = asyncio.run( + queue.submit(SubmissionRequest(model="anthropic/claude-sonnet-4-6", submitter="space-ui")) + ) + except BaseException as exc: + outcome["error"] = exc + + submit_thread = threading.Thread(target=submit_from_other_loop, daemon=True) + submit_thread.start() + # Worker still holds the lock. Release after submit is waiting so a + # thread lock can proceed; an asyncio.Lock bound to the worker loop + # deadlocks or raises instead. + release.set() + submit_thread.join(timeout=2) + worker_thread.join(timeout=2) + + assert worker_errors == [] + assert not submit_thread.is_alive(), "submit deadlocked across event loops" + assert "error" not in outcome, outcome.get("error") + job = outcome["job"] + assert isinstance(job, Job) + assert job.status == JobStatus.PENDING + assert job.job_id in queue._jobs + + +def test_submit_does_not_block_on_hub_sync_from_second_loop(tmp_path, monkeypatch): + """Hub uploads must not keep the queue lock, or Space submit waits on HF.""" + monkeypatch.setattr(queue_module, "LOCAL_QUEUE_DIR", tmp_path) + monkeypatch.setattr(queue_module, "HF_TOKEN", "") + queue = JobQueue() + + sync_started = threading.Event() + release_sync = threading.Event() + + async def blocking_sync() -> None: + if sync_started.is_set(): + return + sync_started.set() + await asyncio.to_thread(release_sync.wait, 5) + + monkeypatch.setattr(queue, "_sync_to_hub", blocking_sync) + + first_errors: list[BaseException] = [] + + def first_submit() -> None: + try: + asyncio.run(queue.submit(SubmissionRequest(model="anthropic/claude-sonnet-4-6", submitter="worker"))) + except BaseException as exc: + first_errors.append(exc) + + first_thread = threading.Thread(target=first_submit, daemon=True) + first_thread.start() + assert sync_started.wait(timeout=2), "first submit never reached hub sync" + + outcome: dict[str, Job | BaseException] = {} + + def second_submit() -> None: + try: + outcome["job"] = asyncio.run( + queue.submit(SubmissionRequest(model="huggingface/Qwen/Qwen3-32B", submitter="space-ui")) + ) + except BaseException as exc: + outcome["error"] = exc + + second_thread = threading.Thread(target=second_submit, daemon=True) + second_thread.start() + second_thread.join(timeout=2) + release_sync.set() + first_thread.join(timeout=2) + + assert first_errors == [] + assert not second_thread.is_alive(), "submit waited on hub upload lock" + assert "error" not in outcome, outcome.get("error") + job = outcome["job"] + assert isinstance(job, Job) + assert job.request.model == "huggingface/Qwen/Qwen3-32B" + assert job.status == JobStatus.PENDING