WebSockets

Documentation Index

Fetch the complete documentation index at: https://docs.baseten.co/llms.txt

Use this file to discover all available pages before exploring further.

​

WebSockets in Truss models

1. Initialize your Truss:

truss init websocket-model

2. Replace the predict method in model/model.py with a websocket method. For example:

import fastapi

class Model:
    async def websocket(self, websocket: fastapi.WebSocket):
        try:
            while True:
                message = await websocket.receive_text()
                await websocket.send_text(f"WS obtained: {message}")
        except fastapi.WebSocketDisconnect:
            pass

3. Set runtime.transport.kind=websocket in config.yaml:

...
runtime:
  transport:
    kind: websocket

4. Deploy the model:

truss push

​

Constraints and behavior

• Message exchange runs in a loop until the client disconnects. To close the connection from the server, call websocket.close().
• WebSockets support bidirectional streaming, so you don’t need multiple HTTP round-trips.
• Don’t implement predict, preprocess, or postprocess. Baseten doesn’t call them.
• Baseten accepts the connection for you, so don’t call websocket.accept(). You can close the connection yourself when you’re done; otherwise Baseten closes it after your websocket method returns.

​

Call the model

websocat -H="Authorization: Api-Key $BASETEN_API_KEY" \
   wss://model-{MODEL_ID}.api.baseten.co/environments/production/websocket
Hello # Your input.
WS obtained: Hello # Echoed from model.
# ctrl+c to close connection.

​

WebSockets in Chains

​

Example chainlet

import fastapi
import truss_chains as chains

class Dependency(chains.ChainletBase):
    async def run_remote(self, name: str) -> str:
        return f"Hello from dependency, {name}."

@chains.mark_entrypoint
class WSEntrypoint(chains.ChainletBase):
    def __init__(self, dependency=chains.depends(Dependency)):
        self._dependency = dependency

    async def run_remote(self, websocket: chains.WebSocketProtocol) -> None:
        try:
            while True:
                message = await websocket.receive_text()
                if message == "dep":
                    response = await self._dependency.run_remote("WSEntrypoint")
                else:
                    response = f"You said: {message}"
                await websocket.send_text(response)
        except fastapi.WebSocketDisconnect:
            print("Disconnected.")

​

Constraints and behavior

• Your run_remote signature must use WebSocketProtocol. It mirrors fastapi.WebSocket, except you can’t call accept(). Baseten has already accepted the connection by the time your chainlet runs.
• run_remote accepts no other arguments when using WebSockets.
• The return type must be None. Send any data back to the client through the WebSocket instead.
• WebSockets are only supported on the entrypoint chainlet, not on dependencies.
• Unlike Truss models, Chains don’t require you to set runtime.transport.kind.

​

Call the chain

websocat -H="Authorization: Api-Key $BASETEN_API_KEY" \
   wss://chain-{CHAIN_ID}.api.baseten.co/environments/production/websocket

​

WebSockets with custom servers

​

Configuration

base_image:
  image: bryanzhang2/custom_ws:v0.0.4
docker_server:
  start_command: /app/server
  readiness_endpoint: /health
  liveness_endpoint: /health
  predict_endpoint: /v1/websocket
  server_port: 8081
model_name: custom_ws
runtime:
  transport:
    kind: "websocket"

​

Required fields

• predict_endpoint: The WebSocket endpoint path on your server, for example /v1/websocket or /ws.
• runtime.transport.kind: Must be "websocket".
• start_command: Command that starts your WebSocket server.
• readiness_endpoint: HTTP path for readiness probes.
• liveness_endpoint: HTTP path for liveness probes.

​

Call the model

websocat -H="Authorization: Api-Key $BASETEN_API_KEY" \
   wss://model-{MODEL_ID}.api.baseten.co/environments/production/websocket

​

Deployment and concurrency considerations

​

Scheduling

• Resource utilization: HTTP requests are stateless, so Baseten can rebalance them freely. WebSocket connections stay pinned to a replica for their lifetime and count against that replica’s concurrency target even when idle. Manage connection efficiency on the client side.
• Stateful complexity: WebSocket handlers often hold server-side state, which adds lifecycle work (disconnects, cleanup, reconnection logic).

​

Lifetime guarantees

​

Concurrency changes

​

Promotion

​

Maximum message size

​

Monitoring

​

Inference volume

• HTTP status codes for connections that failed before the WebSocket upgrade completed.
• WebSocket close codes for connections that completed the upgrade and later closed.

​

HTTP status codes

​

WebSocket close codes

​

End-to-end connection duration

​

Connection input and output size

• Connection input size: Bytes sent by the client to the server.
• Connection output size: Bytes sent by the server to the client.