File size: 5,122 Bytes
09d4173 | 1 2 3 4 5 6 7 8 9 10 11 12 13 14 15 16 17 18 19 20 21 22 23 24 25 26 27 28 29 30 31 32 33 34 35 36 37 38 39 40 41 42 43 44 45 46 47 48 49 50 51 52 53 54 55 56 57 58 59 60 61 62 63 64 65 66 67 68 69 70 71 72 73 74 75 76 77 78 79 80 81 82 83 84 85 86 87 88 89 90 91 92 93 94 95 96 97 98 99 100 101 102 103 104 105 106 107 108 109 110 111 112 113 114 115 116 117 118 119 120 121 122 123 124 125 126 127 128 129 130 131 132 133 134 135 136 137 138 139 140 141 142 143 144 145 146 147 148 149 150 151 | """A completed score must not wait forever for a cancelled disconnect poll."""
import asyncio
import pytest
from jev_adapter.backend import AdapterError, ScoringResult
from jev_adapter.protocol import SystemOneRequest
from jev_adapter.server import create_app
class ControlledBackend:
model = "test-model"
def __init__(self):
self.release = asyncio.Event()
self.entered = asyncio.Event()
self.cancelled = asyncio.Event()
def labels(self, count):
return tuple("AB"[:count]), tuple(range(count))
async def evaluate(self, prompt, images, labels, token_ids, assistant_prefix):
self.entered.set()
try:
await self.release.wait()
except asyncio.CancelledError:
self.cancelled.set()
raise
return ScoringResult((-0.1, -2.0), 10)
class CancellationSwallowingRequest:
"""Reproduce a request probe whose internal cancellation scope wins a race.
CancelledError can be suppressed within a probe, leaving its caller alive.
A real Request.is_disconnected uses a self-cancelling AnyIO CancelScope;
deliberately control that boundary here rather than rely on scheduler luck.
"""
def __init__(self):
self.entered = asyncio.Event()
self.swallowed = asyncio.Event()
self.probe_task = None
self.calls = 0
async def is_disconnected(self):
self.probe_task = asyncio.current_task()
self.calls += 1
self.entered.set()
try:
await asyncio.Future()
except asyncio.CancelledError:
self.swallowed.set()
return False
def endpoint_and_body(backend):
app = create_app(backend)
endpoint = next(
route.endpoint
for route in app.routes
if getattr(route, "path", None) == "/v1/systemone"
)
body = SystemOneRequest.model_validate(
{
"model": backend.model,
"state": "A billing question.",
"questions": {
"route": {
"type": "choice",
"instructions": "Choose the team.",
"criteria": {"billing": None, "technical": None},
}
},
}
)
return endpoint, body
@pytest.mark.asyncio
async def test_completed_score_returns_when_disconnect_probe_swallows_cancel():
backend = ControlledBackend()
endpoint, body = endpoint_and_body(backend)
request = CancellationSwallowingRequest()
response = asyncio.create_task(endpoint(body, request))
try:
await asyncio.wait_for(request.entered.wait(), timeout=1)
await asyncio.wait_for(backend.entered.wait(), timeout=1)
backend.release.set()
await asyncio.wait_for(request.swallowed.wait(), timeout=1)
done, _ = await asyncio.wait({response}, timeout=0.1)
assert response in done, "completed inference is stuck cleaning up its watcher"
result = response.result()
assert result["answers"]["route"]["choice"] == "billing"
assert request.calls == 1
assert request.probe_task.done()
finally:
# Do not let a deliberately cancellation-resistant fake leak after red.
request.is_disconnected = disconnected_now
if request.probe_task is not None:
request.probe_task.cancel()
response.cancel()
await asyncio.gather(response, return_exceptions=True)
async def disconnected_now():
return True
@pytest.mark.asyncio
async def test_client_disconnect_still_cancels_unfinished_inference():
backend = ControlledBackend()
endpoint, body = endpoint_and_body(backend)
class DisconnectedRequest:
async def is_disconnected(self):
await backend.entered.wait()
return True
with pytest.raises(AdapterError) as error:
await asyncio.wait_for(endpoint(body, DisconnectedRequest()), timeout=1)
assert error.value.code == "client_disconnected"
assert error.value.status == 499
assert backend.cancelled.is_set()
@pytest.mark.asyncio
async def test_cancelled_handler_joins_work_and_cancellation_resistant_watcher():
backend = ControlledBackend()
endpoint, body = endpoint_and_body(backend)
request = CancellationSwallowingRequest()
response = asyncio.create_task(endpoint(body, request))
try:
await asyncio.wait_for(request.entered.wait(), timeout=1)
await asyncio.wait_for(backend.entered.wait(), timeout=1)
response.cancel()
done, _ = await asyncio.wait({response}, timeout=0.1)
assert response in done, "cancelled handler did not release its child tasks"
with pytest.raises(asyncio.CancelledError):
response.result()
assert backend.cancelled.is_set()
assert request.swallowed.is_set()
assert request.probe_task.done()
finally:
request.is_disconnected = disconnected_now
if request.probe_task is not None:
request.probe_task.cancel()
response.cancel()
await asyncio.gather(response, return_exceptions=True)
|