99 lines
3.3 KiB
Python
99 lines
3.3 KiB
Python
|
from __future__ import annotations
|
||
|
|
|
||
|
|
import asyncio
|
||
|
|
from concurrent.futures import ThreadPoolExecutor
|
||
|
|
from time import sleep
|
||
|
|
from typing import Any
|
||
|
|
|
||
|
|
import pytest
|
||
|
|
|
||
|
|
from cmvr_edge_ai.workers import (
|
||
|
|
PersistentProcessWorker,
|
||
|
|
ProcessWorkerError,
|
||
|
|
RemoteWorkerError,
|
||
|
|
)
|
||
|
|
|
||
|
|
|
||
|
|
def _worker_handler(payload: Any) -> Any:
|
||
|
|
if payload == "explode":
|
||
|
|
raise RuntimeError("intentional remote failure")
|
||
|
|
if payload == "slow":
|
||
|
|
sleep(0.3)
|
||
|
|
return "late-result"
|
||
|
|
return payload * 2
|
||
|
|
|
||
|
|
|
||
|
|
def _worker_factory(): # type: ignore[no-untyped-def]
|
||
|
|
"""Top-level factory so the spawn start method can import and pickle it."""
|
||
|
|
|
||
|
|
return _worker_handler
|
||
|
|
|
||
|
|
|
||
|
|
async def _keep_restricted_event_loop_responsive() -> None:
|
||
|
|
"""Provide wakeups where sandboxed self-pipe notifications are unavailable."""
|
||
|
|
|
||
|
|
while True:
|
||
|
|
await asyncio.sleep(0.01)
|
||
|
|
|
||
|
|
|
||
|
|
def test_spawn_worker_returns_results_and_surfaces_remote_errors() -> None:
|
||
|
|
async def exercise(): # type: ignore[no-untyped-def]
|
||
|
|
loop = asyncio.get_running_loop()
|
||
|
|
host_executor = ThreadPoolExecutor(max_workers=2)
|
||
|
|
loop.set_default_executor(host_executor)
|
||
|
|
ticker = asyncio.create_task(_keep_restricted_event_loop_responsive())
|
||
|
|
worker = PersistentProcessWorker(
|
||
|
|
_worker_factory,
|
||
|
|
start_method="spawn",
|
||
|
|
queue_capacity=1,
|
||
|
|
)
|
||
|
|
try:
|
||
|
|
await worker.start()
|
||
|
|
assert await worker.submit(21, timeout_s=10) == 42
|
||
|
|
with pytest.raises(RemoteWorkerError) as captured:
|
||
|
|
await worker.submit("explode", timeout_s=10)
|
||
|
|
assert captured.value.remote_type == "RuntimeError"
|
||
|
|
assert "intentional remote failure" in str(captured.value)
|
||
|
|
finally:
|
||
|
|
await worker.stop(timeout_s=5)
|
||
|
|
ticker.cancel()
|
||
|
|
await asyncio.gather(ticker, return_exceptions=True)
|
||
|
|
host_executor.shutdown(wait=True, cancel_futures=True)
|
||
|
|
loop._default_executor = None # type: ignore[attr-defined]
|
||
|
|
return worker
|
||
|
|
|
||
|
|
worker = asyncio.run(exercise())
|
||
|
|
|
||
|
|
assert worker.is_alive is False
|
||
|
|
|
||
|
|
|
||
|
|
def test_timed_out_request_poisons_worker_and_rejects_later_submits() -> None:
|
||
|
|
async def exercise(): # type: ignore[no-untyped-def]
|
||
|
|
loop = asyncio.get_running_loop()
|
||
|
|
host_executor = ThreadPoolExecutor(max_workers=2)
|
||
|
|
loop.set_default_executor(host_executor)
|
||
|
|
ticker = asyncio.create_task(_keep_restricted_event_loop_responsive())
|
||
|
|
worker = PersistentProcessWorker(
|
||
|
|
_worker_factory,
|
||
|
|
start_method="spawn",
|
||
|
|
queue_capacity=1,
|
||
|
|
)
|
||
|
|
try:
|
||
|
|
await worker.start()
|
||
|
|
with pytest.raises(TimeoutError, match="timed out"):
|
||
|
|
await worker.submit("slow", timeout_s=0.05)
|
||
|
|
assert worker.is_alive is True
|
||
|
|
with pytest.raises(ProcessWorkerError, match="worker is poisoned"):
|
||
|
|
await worker.submit(22, timeout_s=1)
|
||
|
|
finally:
|
||
|
|
await worker.stop(timeout_s=2)
|
||
|
|
ticker.cancel()
|
||
|
|
await asyncio.gather(ticker, return_exceptions=True)
|
||
|
|
host_executor.shutdown(wait=True, cancel_futures=True)
|
||
|
|
loop._default_executor = None # type: ignore[attr-defined]
|
||
|
|
return worker
|
||
|
|
|
||
|
|
worker = asyncio.run(exercise())
|
||
|
|
|
||
|
|
assert worker.is_alive is False
|