Skip to content

Streaming APIs and Large models

Two things make a large model workable behind a recommendation endpoint: loading it once per replica instead of once per request, and streaming the answer back as it is produced instead of making the caller wait for all of it.

Both are opt-in, and both are configured on the component in your peach.conf.

Initializing States

A component can name an init function alongside its method:

endpoints:
  assistant:
    url: /assistant
    py_env: llm
    components:
      main:
        notebook: notebooks/Endpoints.ipynb
        init: init_assistant     # runs once, when the replica starts
        method: ask_assistant    # runs on every request

init is called once, when the Ray Serve replica starts — not at import time, and not per request. Whatever it returns is kept on the replica and handed to your method as a parameter called model:

def init_assistant():
    # Runs once per replica. Load the expensive things here.
    return {
        "client": SomeModelClient(...),
        "index": load_index(...),
    }


def ask_assistant(model, query: str, size: int = 10):
    # `model` is exactly what init_assistant() returned
    return model["client"].answer(query, model["index"], size)

Two rules follow from how it is wired:

  • Your method has to declare a model parameter to receive the state. Request parameters are matched against your function's signature, so a parameter you do not declare is simply not passed.
  • The state is only as fresh as the replica. It is loaded at startup and then held, so a model retrained afterwards is not picked up on its own — see Drift Detection below.

Per replica, not per endpoint

Each replica runs init for itself, so a model loaded this way is held once per replica and multiplied by your replica count — that is what the component's memory request has to cover. If init raises, the replica never becomes healthy.

The equivalent for a decorator endpoint is init_fn, which works the same way except that it exposes its result as init_state rather than a model parameter — see Decorator endpoints.

Streaming APIs

A model that takes twenty seconds to produce a full answer is unusable if the caller sees nothing until it is done. Streaming sends each piece as it is produced.

Make the method an async generator — async def, and yield each chunk instead of returning a value:

async def ask_assistant(model, query: str):
    async for piece in model["client"].stream(query):
        yield piece + "\n"

The caller then asks for a stream by sending stream with the request, as a query parameter or in the JSON body:

POST /assistant
{"query": "...", "stream": "true"}

When stream is set, the gateway returns the generator's chunks to the caller as they arrive. Anything else is treated as a normal, non-streaming call.

stream has to come from the caller

Putting stream in the component's args: defaults is not enough. The gateway decides whether to stream before those defaults are applied, so a default there is never seen in time. The caller has to send it on the request.

Omitting it is the common mistake, and it surfaces as an error about not being able to pickle an async_generator — the endpoint tried to return your generator as an ordinary value. The message says what to pass.

Three things behave differently once a response streams:

  • No response envelope. A normal endpoint answers inside {"result": ..., "status": "ok"}. A stream does not — the caller receives your chunks exactly as you yield them, so you choose the framing (a newline after each chunk, JSON lines, SSE, whatever your client parses).
  • No fallback. The fallback component covers a primary that fails or returns nothing. A stream has already started responding by the time anything can go wrong, so nothing can be substituted for it. Handle errors inside the generator and yield something your client understands.
  • No request timeout. Non-streaming component calls are bounded by a per-codops timeout, so one slow call cannot pile up and starve the endpoints sharing its replica. Streams are long-lived by design and are deliberately exempt.

Drift Detection of initialized states

State loaded by init is as old as the replica. When the model behind it is retrained — usually by a task on its own schedule — the replicas keep serving the copy they loaded at startup.

Returning drift_detected tells the replica to reload. Return your result as a (items, meta) tuple with the flag in the meta dict:

def ask_assistant(model, query: str):
    items = model["client"].answer(query)
    if model["index"].is_stale():
        return items, {"drift_detected": True}
    return items

The current request is answered first, from the state already loaded; the replica then re-runs init so the next request uses fresh state. Each replica decides for itself, so they refresh as each one observes the drift rather than all at once.

This is the way to pick up a retrained model without a redeploy. Deciding when state is stale is your code's job — a version marker written by the task and compared against the one loaded at startup is the usual approach.

Note

drift_detected is read from the response, so a streaming method — which returns chunks rather than a result — cannot use it.