Skip to content
Open
Show file tree
Hide file tree
Changes from all commits
Commits
File filter

Filter by extension

Filter by extension

Conversations
Failed to load comments.
Loading
Jump to
Jump to file
Failed to load files.
Loading
Diff view
Diff view
10 changes: 5 additions & 5 deletions src/server/worker_registry.py
Original file line number Diff line number Diff line change
Expand Up @@ -97,13 +97,13 @@ def _commit_completed_packet(
) -> bool:
"""Complete the request future. Return True if the worker should exit."""
if completed.error is not None:
logger.error(
f"[{model_name}] Inference failed, triggering model unload..."
)
# Resolve the caller and keep the worker/model alive; unloading here runs the
# pipeline destructor, which can touch a corrupted device context after certain
# errors and SIGABRT the whole process, taking every other loaded model with it.
logger.error(f"[{model_name}] Inference failed: {completed.error}")
if packet.result_future is not None and not packet.result_future.done():
packet.result_future.set_exception(completed.error)
asyncio.create_task(registry.register_unload(model_name))
return True
return False
if packet.result_future is not None and not packet.result_future.done():
packet.result_future.set_result(completed)
return False
Expand Down
9 changes: 6 additions & 3 deletions tests/unit/test_worker_registry_unit.py
Original file line number Diff line number Diff line change
Expand Up @@ -310,7 +310,7 @@ async def transcribe(self, gen_config):
assert packet.segments is None


def test_commit_treats_error_field_not_error_prefix() -> None:
def test_commit_resolves_error_without_unloading_the_model() -> None:
class DummyRegistry:
def __init__(self):
self.unloaded = []
Expand Down Expand Up @@ -342,11 +342,14 @@ async def _run():
result_future=err_future,
error=RuntimeError("gpu oom"),
)
assert worker_module._commit_completed_packet(err, err, "m", registry) is True
# On error the worker resolves the caller's future and keeps the model loaded;
# unloading here runs the pipeline destructor, which can SIGABRT the process if
# the device context is corrupted (see queue_worker_llm docstring).
assert worker_module._commit_completed_packet(err, err, "m", registry) is False
with pytest.raises(RuntimeError, match="gpu oom"):
err_future.result()
await asyncio.sleep(0)
assert registry.unloaded == ["m"]
assert registry.unloaded == []

asyncio.run(_run())

Expand Down