wrap_litserve_start() hangs on requests when worker setup() is slow

[Mystic98]

🐛 Bug

wrap_litserve_start() is a pytest utility that starts the LitServe server in a context manager for testing. It launches inference worker processes but yields immediately without waiting for them to finish setup() (model loading, etc.).

This causes TestClient.post() to hang indefinitely because the worker hasn't started listening on the request queue yet. The request gets enqueued but never processed.

To Reproduce

import litserve as ls
from fastapi.testclient import TestClient


class SlowSetupAPI(ls.LitAPI):
    def setup(self, device):
        import time
        time.sleep(5)  # simulate slow model loading
        self.model = lambda x: x

    def decode_request(self, request):
        return request["input"]

    def predict(self, x):
        return self.model(x)

    def encode_response(self, output):
        return {"output": output}


def test_predict():
    api = SlowSetupAPI()
    server = ls.LitServer(api, accelerator="cpu")

    with ls.utils.wrap_litserve_start(server):
        client = TestClient(server.app)
        # This hangs forever because the worker is still in setup()
        response = client.post("/predict", json={"input": "hello"})

    assert response.status_code == 200

Code sample

The root cause is in litserve/utils.py. wrap_litserve_start calls launch_inference_worker() but never waits for workers_setup_status to reach READY, unlike LitServer.run() which has this loop:

# From server.py run() — this wait is missing from wrap_litserve_start
while not all(v == WorkerSetupStatus.READY for v in self.workers_setup_status.values()):
    if any(v == WorkerSetupStatus.ERROR for v in self.workers_setup_status.values()):
        raise RuntimeError("One or more workers failed to start. Shutting down LitServe")
    time.sleep(0.05)

Expected behavior

wrap_litserve_start should wait for all workers to be ready before yielding, matching the behavior of LitServer.run(). This would ensure TestClient requests are processed immediately without hanging.

Environment

• LitServe Version: 0.2.17
• OS: Linux
• Python version: 3.12
• How you installed: pip install litserve

Reproduction setup with uv

uv init --python 3.12
uv add "litserve==0.2.17" pytest pytest-timeout httpx
uv run pytest test_bug.py -v --timeout=30

Additional context

Related issues: #263 (setup not awaited), #663 (health check returns 200 when workers not ready), #660 (stale worker status).

Workaround:

import time

with ls.utils.wrap_litserve_start(server):
    while not all(v == "ready" for v in server.workers_setup_status.values()):
        time.sleep(0.5)
    with TestClient(server.app) as client:
        response = client.post("/predict", json={"input": "hello"})

Suggested fix: Add the worker readiness wait loop to wrap_litserve_start before yielding:

@contextmanager
def wrap_litserve_start(server: "LitServer", worker_monitor: bool = False):
    # ... existing setup code ...

    # Wait for all workers to be ready (same as LitServer.run)
    while not all(
        v == WorkerSetupStatus.READY for v in server.workers_setup_status.values()
    ):
        if any(v == WorkerSetupStatus.ERROR for v in server.workers_setup_status.values()):
            raise RuntimeError("One or more workers failed to start")
        time.sleep(0.05)

    try:
        yield server
    finally:
        # ... existing teardown code ...

[hesam-oxe]

Great analysis @Mystic98!

I can implement the suggested fix — just the worker readiness wait loop
in wrap_litserve_start().

/claim