Spaces:
Running on Zero
Running on Zero
| """Concurrency and lifecycle guarantees for the inference queue.""" | |
| from __future__ import annotations | |
| import asyncio | |
| import threading | |
| import time | |
| from contextvars import ContextVar | |
| import pytest | |
| from core.errors import GatewayError, QueueCapacityError | |
| from core.queue import InferenceQueue | |
| async def test_queue_executes_only_one_operation_at_a_time() -> None: | |
| queue = InferenceQueue(capacity=4, timeout_seconds=5) | |
| active = 0 | |
| maximum_active = 0 | |
| lock = threading.Lock() | |
| def operation(value: int) -> int: | |
| nonlocal active, maximum_active | |
| with lock: | |
| active += 1 | |
| maximum_active = max(maximum_active, active) | |
| time.sleep(0.02) | |
| with lock: | |
| active -= 1 | |
| return value | |
| results = await asyncio.gather( | |
| *(queue.submit(str(index), "test", lambda index=index: operation(index)) for index in range(3)) | |
| ) | |
| await queue.stop() | |
| assert results == [0, 1, 2] | |
| assert maximum_active == 1 | |
| async def test_queue_timeout_cancels_job_before_inference() -> None: | |
| queue = InferenceQueue(capacity=2, timeout_seconds=0.01) | |
| queued_job_ran = False | |
| async def submit_first() -> None: | |
| await queue.submit("first", "slow", lambda: time.sleep(0.04)) | |
| def queued_operation() -> None: | |
| nonlocal queued_job_ran | |
| queued_job_ran = True | |
| first = asyncio.create_task(submit_first()) | |
| await asyncio.sleep(0.005) | |
| with pytest.raises(GatewayError, match="timed out") as captured: | |
| await queue.submit("second", "queued", queued_operation) | |
| assert captured.value.status_code == 504 | |
| assert captured.value.code == "queue_timeout" | |
| record = queue.job("second") | |
| assert record is not None | |
| assert record["status"] == "timeout" | |
| await first | |
| await asyncio.sleep(0) | |
| assert queued_job_ran is False | |
| assert queue.active is False | |
| await queue.stop() | |
| async def test_cancellation_waits_for_started_operation() -> None: | |
| queue = InferenceQueue(capacity=2, timeout_seconds=1) | |
| started = threading.Event() | |
| finished = threading.Event() | |
| output = [] | |
| def operation() -> int: | |
| started.set() | |
| time.sleep(0.04) | |
| output.append("complete") | |
| finished.set() | |
| return 7 | |
| task = asyncio.create_task(queue.submit("request", "slow", operation)) | |
| assert await asyncio.to_thread(started.wait, 0.1) | |
| task.cancel() | |
| with pytest.raises(asyncio.CancelledError): | |
| await task | |
| assert finished.is_set() | |
| assert output == ["complete"] | |
| assert queue.job("request")["status"] == "succeeded" | |
| result = await queue.submit("next", "test", lambda: 42) | |
| assert result == 42 | |
| await queue.stop() | |
| async def test_cancellation_before_start_marks_job_cancelled_and_does_not_run() -> None: | |
| queue = InferenceQueue(capacity=2, timeout_seconds=1) | |
| blocker_started = threading.Event() | |
| blocker_finished = threading.Event() | |
| queued_ran = False | |
| def blocker() -> None: | |
| blocker_started.set() | |
| time.sleep(0.05) | |
| blocker_finished.set() | |
| def queued_operation() -> None: | |
| nonlocal queued_ran | |
| queued_ran = True | |
| first = asyncio.create_task(queue.submit("blocker", "slow", blocker)) | |
| assert await asyncio.to_thread(blocker_started.wait, 0.1) | |
| second = asyncio.create_task(queue.submit("cancelled", "queued", queued_operation)) | |
| await asyncio.sleep(0.005) | |
| second.cancel() | |
| with pytest.raises(asyncio.CancelledError): | |
| await second | |
| await first | |
| assert blocker_finished.is_set() | |
| assert queued_ran is False | |
| record = queue.job("cancelled") | |
| assert record is not None | |
| assert record["status"] == "cancelled" | |
| assert record["finished_at"] is not None | |
| await queue.stop() | |
| async def test_queue_preserves_request_context_for_worker() -> None: | |
| queue = InferenceQueue(capacity=1, timeout_seconds=5) | |
| request_context: ContextVar[str] = ContextVar("test_request", default="missing") | |
| token = request_context.set("request-identity") | |
| try: | |
| result = await queue.submit("request", "context", request_context.get) | |
| finally: | |
| request_context.reset(token) | |
| await queue.stop() | |
| assert result == "request-identity" | |
| async def test_duplicate_request_id_is_rejected_without_overwriting_history() -> None: | |
| queue = InferenceQueue(capacity=2, timeout_seconds=1, history_size=10) | |
| assert await queue.submit("duplicate", "first", lambda: 1) == 1 | |
| with pytest.raises(GatewayError) as captured: | |
| await queue.submit("duplicate", "second", lambda: 2) | |
| assert captured.value.status_code == 409 | |
| assert captured.value.code == "request_id_conflict" | |
| record = queue.job("duplicate") | |
| assert record is not None | |
| assert record["label"] == "first" | |
| assert record["status"] == "succeeded" | |
| await queue.stop() | |
| async def test_queue_history_is_bounded_and_newest_first() -> None: | |
| queue = InferenceQueue(capacity=2, timeout_seconds=1, history_size=2) | |
| await queue.submit("one", "first", lambda: 1) | |
| await queue.submit("two", "second", lambda: 2) | |
| await queue.submit("three", "third", lambda: 3) | |
| records = queue.jobs() | |
| assert [record["request_id"] for record in records] == ["three", "two"] | |
| assert queue.job("one") is None | |
| await queue.stop() | |
| async def test_operation_failure_is_recorded() -> None: | |
| queue = InferenceQueue(capacity=2, timeout_seconds=1) | |
| def fail() -> None: | |
| raise RuntimeError("model failed") | |
| with pytest.raises(RuntimeError, match="model failed"): | |
| await queue.submit("failed", "failure", fail) | |
| record = queue.job("failed") | |
| assert record is not None | |
| assert record["status"] == "failed" | |
| assert record["error_code"] == "inference_failed" | |
| assert record["error_message"] == "model failed" | |
| assert record["finished_at"] is not None | |
| await queue.stop() | |
| async def test_queue_full_rejects_without_running_the_rejected_operation() -> None: | |
| queue = InferenceQueue(capacity=1, timeout_seconds=1) | |
| started = threading.Event() | |
| release = threading.Event() | |
| rejected_ran = False | |
| def blocker() -> None: | |
| started.set() | |
| release.wait(0.2) | |
| def rejected_operation() -> None: | |
| nonlocal rejected_ran | |
| rejected_ran = True | |
| first = asyncio.create_task(queue.submit("running", "blocker", blocker)) | |
| assert await asyncio.to_thread(started.wait, 0.1) | |
| queued = asyncio.create_task(queue.submit("queued", "waiting", lambda: 2)) | |
| await asyncio.sleep(0.005) | |
| with pytest.raises(QueueCapacityError): | |
| await queue.submit("rejected", "overflow", rejected_operation) | |
| rejected_record = queue.job("rejected") | |
| assert rejected_record is not None | |
| assert rejected_record["status"] == "rejected" | |
| assert rejected_record["error_code"] == "queue_full" | |
| release.set() | |
| assert await first is None | |
| assert await queued == 2 | |
| assert rejected_ran is False | |
| await queue.stop() | |
| async def test_shutdown_cancels_queued_jobs_but_waits_for_active_operation() -> None: | |
| queue = InferenceQueue(capacity=2, timeout_seconds=1) | |
| started = threading.Event() | |
| release = threading.Event() | |
| def active_operation() -> int: | |
| started.set() | |
| release.wait(0.2) | |
| return 1 | |
| active_task = asyncio.create_task(queue.submit("active", "active", active_operation)) | |
| assert await asyncio.to_thread(started.wait, 0.1) | |
| queued_task = asyncio.create_task(queue.submit("queued", "queued", lambda: 2)) | |
| await asyncio.sleep(0.005) | |
| stop_task = asyncio.create_task(queue.stop(timeout_seconds=0.5)) | |
| await asyncio.sleep(0.005) | |
| queued_record = queue.job("queued") | |
| assert queued_record is not None | |
| assert queued_record["status"] == "cancelled" | |
| release.set() | |
| assert await active_task == 1 | |
| assert await stop_task is True | |
| with pytest.raises(asyncio.CancelledError): | |
| await queued_task | |
| async def test_shutdown_timeout_resolves_active_request_and_records_timeout() -> None: | |
| queue = InferenceQueue(capacity=1, timeout_seconds=1) | |
| started = threading.Event() | |
| release = threading.Event() | |
| def active_operation() -> int: | |
| started.set() | |
| release.wait(0.5) | |
| return 9 | |
| submit_task = asyncio.create_task(queue.submit("shutdown-active", "slow", active_operation)) | |
| assert await asyncio.to_thread(started.wait, 0.1) | |
| assert await queue.stop(timeout_seconds=0.01) is False | |
| with pytest.raises(GatewayError) as captured: | |
| await asyncio.wait_for(submit_task, timeout=0.2) | |
| assert captured.value.status_code == 503 | |
| assert captured.value.code == "shutdown_timeout" | |
| record = queue.job("shutdown-active") | |
| assert record is not None | |
| assert record["status"] == "timeout" | |
| assert record["error_code"] == "shutdown_timeout" | |
| assert record["finished_at"] is not None | |
| # The Python/native worker thread is not force-killed by asyncio cancellation. | |
| release.set() | |
| async def test_cancellation_remains_primary_when_started_operation_fails() -> None: | |
| queue = InferenceQueue(capacity=1, timeout_seconds=1) | |
| started = threading.Event() | |
| release = threading.Event() | |
| def fail_after_release() -> int: | |
| started.set() | |
| release.wait(0.5) | |
| raise RuntimeError("model failed after cancellation") | |
| task = asyncio.create_task( | |
| queue.submit("cancel-then-fail", "failure", fail_after_release) | |
| ) | |
| assert await asyncio.to_thread(started.wait, 0.1) | |
| # Let submit resume from the start signal and enter its result wait. | |
| await asyncio.sleep(0.01) | |
| task.cancel() | |
| release.set() | |
| with pytest.raises(asyncio.CancelledError): | |
| await task | |
| record = queue.job("cancel-then-fail") | |
| assert record is not None | |
| assert record["status"] == "failed" | |
| assert record["error_code"] == "inference_failed" | |
| assert record["error_message"] == "model failed after cancellation" | |
| assert record["finished_at"] is not None | |
| await queue.stop() | |
| async def test_shutdown_timeout_remains_failed_on_repeated_stop() -> None: | |
| queue = InferenceQueue(capacity=1, timeout_seconds=1) | |
| started = threading.Event() | |
| release = threading.Event() | |
| def active_operation() -> int: | |
| started.set() | |
| release.wait(0.5) | |
| return 1 | |
| submit_task = asyncio.create_task( | |
| queue.submit("repeat-stop", "slow", active_operation) | |
| ) | |
| assert await asyncio.to_thread(started.wait, 0.1) | |
| assert await queue.stop(timeout_seconds=0.01) is False | |
| assert await queue.stop(timeout_seconds=0.01) is False | |
| with pytest.raises(GatewayError) as captured: | |
| await asyncio.wait_for(submit_task, timeout=0.2) | |
| assert captured.value.code == "shutdown_timeout" | |
| record = queue.job("repeat-stop") | |
| assert record is not None | |
| assert record["status"] == "timeout" | |
| release.set() | |
| async def test_cancelled_caller_during_shutdown_timeout_keeps_consistent_state() -> None: | |
| queue = InferenceQueue(capacity=1, timeout_seconds=1) | |
| started = threading.Event() | |
| release = threading.Event() | |
| def active_operation() -> int: | |
| started.set() | |
| release.wait(0.5) | |
| return 5 | |
| task = asyncio.create_task( | |
| queue.submit("cancel-during-stop", "slow", active_operation) | |
| ) | |
| assert await asyncio.to_thread(started.wait, 0.1) | |
| await asyncio.sleep(0.01) | |
| task.cancel() | |
| stop_task = asyncio.create_task(queue.stop(timeout_seconds=0.01)) | |
| assert await stop_task is False | |
| with pytest.raises(asyncio.CancelledError): | |
| await asyncio.wait_for(task, timeout=0.2) | |
| record = queue.job("cancel-during-stop") | |
| assert record is not None | |
| assert record["status"] == "timeout" | |
| assert record["error_code"] == "shutdown_timeout" | |
| assert record["finished_at"] is not None | |
| release.set() | |
| async def test_submit_racing_shutdown_cannot_enqueue_behind_worker_sentinel() -> None: | |
| queue = InferenceQueue(capacity=1, timeout_seconds=1) | |
| entered = asyncio.Event() | |
| async def submit_job() -> None: | |
| entered.set() | |
| await queue.submit("submit-stop-race", "race", lambda: 1) | |
| # Hold admission while both tasks queue for the lifecycle lock. Whichever | |
| # acquires it first, shutdown must either cancel the job or reject admission. | |
| await queue._lifecycle_lock.acquire() | |
| submit_task = asyncio.create_task(submit_job()) | |
| await entered.wait() | |
| stop_task = asyncio.create_task(queue.stop(timeout_seconds=0.2)) | |
| await asyncio.sleep(0) | |
| queue._lifecycle_lock.release() | |
| stop_result = await asyncio.wait_for(stop_task, timeout=0.5) | |
| assert stop_result is True | |
| try: | |
| await asyncio.wait_for(submit_task, timeout=0.5) | |
| except (asyncio.CancelledError, GatewayError): | |
| pass | |
| assert queue.stopping is True | |
| assert queue.depth == 0 | |
| async def test_concurrent_stop_calls_enqueue_only_one_sentinel() -> None: | |
| queue = InferenceQueue(capacity=2, timeout_seconds=1) | |
| started = threading.Event() | |
| release = threading.Event() | |
| def operation() -> int: | |
| started.set() | |
| release.wait(0.5) | |
| return 3 | |
| submit_task = asyncio.create_task(queue.submit("concurrent-stop", "slow", operation)) | |
| assert await asyncio.to_thread(started.wait, 0.1) | |
| first_stop = asyncio.create_task(queue.stop(timeout_seconds=0.5)) | |
| second_stop = asyncio.create_task(queue.stop(timeout_seconds=0.5)) | |
| await asyncio.sleep(0.01) | |
| release.set() | |
| assert await asyncio.wait_for(submit_task, timeout=0.5) == 3 | |
| results = await asyncio.gather(first_stop, second_stop) | |
| assert results == [True, True] | |
| assert queue.depth == 0 | |
| assert queue.stopping is True | |