AWS marketplace: Add streaming and bidirectional stream support for model package images

Service: marketplace · 2026-10-29 · Documentation low

File: marketplace/latest/userguide/ml-model-package-images.md

Summary

Adds a new section documenting how to make an AWS Marketplace SageMaker AI inference container support streaming responses (InvokeEndpointWithResponseStream) and bidirectional streams (InvokeEndpointWithBidirectionalStream). Covers container contract requirements (keep /ping and /invocations, serve HTTP and WebSocket on port 8080, add the com.amazonaws.sagemaker.capabilities.bidirectional-streaming Docker label, answer WebSocket ping frames, plan for 30-minute sessions), a FastAPI web server example replacing Flask, a serve script, a Dockerfile with the capability label, and local test steps (docker inspect label check, curl streaming test, websockets client test).

Security assessment

The change is purely functional feature documentation for container streaming/bidirectional-stream capabilities (Docker label, port 8080, WebSocket ping/pong, session duration, FastAPI example). It does not address a vulnerability, credential handling, encryption, authentication/authorization, or any incident; the only operational constraints (port binding, capability label, ping/pong keepalive) are protocol/contract requirements rather than security hardening.

Evidence

+  * **For bidirectional streams, add the Docker label** `com.amazonaws.sagemaker.capabilities.bidirectional-streaming=true`. SageMaker AI reads the label when a buyer deploys your model. Without it, buyers can't open bidirectional streams to the endpoint. Streaming responses don't need a label.

Diff

diff --git a/marketplace/latest/userguide/ml-model-package-images.md b/marketplace/latest/userguide/ml-model-package-images.md
index bc28046..86afbda 100644
--- a//marketplace/latest/userguide/ml-model-package-images.md
+++ b//marketplace/latest/userguide/ml-model-package-images.md
@@ -7 +7 @@
-OverviewCreate an inference image for model packages
+OverviewCreate an inference image for model packagesSupport streaming responses and bidirectional streams
@@ -18,0 +19,2 @@ An Amazon SageMaker AI model package is a pre-trained model that makes predictio
+  * Support streaming responses and bidirectional streams
+
@@ -91 +93 @@ For the inference image to be compatible with SageMaker AI, the Docker image mus
-The following walkthrough uses the Docker CLI in a development environment using a Linux Ubuntu distribution. 
+The following walkthrough uses the Docker CLI in a development environment using a Linux Ubuntu distribution. To support streaming responses or bidirectional streams, see Support streaming responses and bidirectional streams.
@@ -330,0 +333,147 @@ When your inference image is ready and tested, you can continue to [Uploading yo
+## Support streaming responses and bidirectional streams
+
+Buyers can call your endpoint in three ways. The inference image in the previous section supports real-time requests. To support streaming responses and bidirectional streams, make the changes in this section.
+
+Ways to call your endpoint How the buyer calls your endpoint | SageMaker AI API | What your container does  
+---|---|---  
+Real-time request | `InvokeEndpoint` | Returns the complete response to a POST request at `/invocations`.  
+Streaming response | `InvokeEndpointWithResponseStream` | Returns the response to a POST request at `/invocations` in parts as your model produces them, using HTTP chunked transfer encoding.  
+Bidirectional stream | `InvokeEndpointWithBidirectionalStream` | Accepts a WebSocket connection at `/invocations-bidirectional-stream` and exchanges messages with the buyer in both directions until the session ends.  
+  
+Your container must meet these requirements:
+
+  * **Keep`/ping` and `/invocations`.** SageMaker AI checks `/ping` on every endpoint. AWS Marketplace validates your model package with a batch transform job that calls `/invocations`.
+
+  * **Serve HTTP and WebSocket connections on port 8080** from the same web server.
+
+  * **For bidirectional streams, add the Docker label** `com.amazonaws.sagemaker.capabilities.bidirectional-streaming=true`. SageMaker AI reads the label when a buyer deploys your model. Without it, buyers can't open bidirectional streams to the endpoint. Streaming responses don't need a label.
+
+  * **Answer WebSocket ping frames with pong frames.** SageMaker AI sends a ping every 60 seconds and closes the connection after 5 missed pongs. Most WebSocket servers, including Uvicorn, answer pings for you.
+
+  * **Plan for sessions of up to 30 minutes.** SageMaker AI ends a bidirectional stream after 30 minutes, and the buyer's application opens a new one. Store session state with the connection, so a new stream can start cleanly.
+
+
+
+
+For the full contract, see [Container contract to support bidirectional streaming capabilities](https://docs.aws.amazon.com/sagemaker/latest/dg/your-algorithms-inference-code.html#your-algorithms-inference-algo-bidi) in the _Amazon SageMaker AI Developer Guide_.
+
+### Create the web server script for streaming
+
+The following example replaces the Flask server from the previous section with FastAPI, because Flask doesn't support WebSocket connections. One server handles all three ways of calling the endpoint. As before, the example has no inference logic: it returns the input in uppercase.
+
+**`./web_app_serve.py`**
+    
+    
+    from fastapi import FastAPI, Request, Response, WebSocket, WebSocketDisconnect
+    from fastapi.responses import StreamingResponse
+    
+    app = FastAPI()
+    
+    
+    # Health check, used by every invocation type
+    @app.get("/ping")
+    async def ping():
+        return Response(status_code=200)
+    
+    
+    # Real-time requests, streaming responses, and batch transform
+    @app.post("/invocations")
+    async def invocations(request: Request):
+        input_str = (await request.body()).decode("utf8")
+    
+        async def generate():
+            # Replace this loop with your model code. Yield each part as soon as it's ready.
+            for word in input_str.split():
+                yield word.upper() + "\n"
+    
+        return StreamingResponse(generate(), media_type="text/plain")
+    
+    
+    # Bidirectional streams
+    @app.websocket("/invocations-bidirectional-stream")
+    async def bidirectional_stream(websocket: WebSocket):
+        await websocket.accept()
+        try:
+            while True:
+                message = await websocket.receive_text()
+                # Replace this line with your model code.
+                await websocket.send_text(message.upper())
+        except WebSocketDisconnect:
+            pass
+
+`/invocations` streams every response. Real-time requests and batch transform jobs receive the parts as one complete response, so you don't need separate code for them.
+
+### Create the script for the streaming container run
+
+**`./serve`**
+    
+    
+    #!/bin/bash
+    
+    # Run the web server on port 8080 for SageMaker
+    exec uvicorn web_app_serve:app --app-dir / --host 0.0.0.0 --port 8080
+
+### Create the `Dockerfile` for streaming
+
+`**./Dockerfile**`
+    
+    
+    FROM python:3.12-slim
+    
+    # Tell SageMaker that this container supports bidirectional streams
+    LABEL com.amazonaws.sagemaker.capabilities.bidirectional-streaming=true
+    
+    # Install the web server
+    RUN pip install --no-cache-dir fastapi "uvicorn[standard]"
+    
+    # Add the web server script and the script that SageMaker runs
+    COPY web_app_serve.py /web_app_serve.py
+    COPY serve /opt/program/serve
+    RUN chmod 755 /opt/program/serve
+    ENV PATH=/opt/program:${PATH}
+
+`uvicorn[standard]` includes the WebSocket library that Uvicorn needs. If you don't support bidirectional streams, leave out the `LABEL` line.
+
+### Test streaming locally
+
+Build and run the image as described in Step 2: Build and testing the image locally, then run the following tests.
+
+###### To test streaming locally
+
+  1. Confirm that the image has the label. The command returns `true`.
+    
+        sudo docker inspect --format '{{ index .Config.Labels "com.amazonaws.sagemaker.capabilities.bidirectional-streaming" }}' my-inference-image
+
+  2. Test a streaming response. The `--no-buffer` option prints each part as it arrives.
+    
+        curl --no-buffer --request POST --data "hello world" http://127.0.0.1:8080/invocations
+
+The command prints `HELLO` and `WORLD` on separate lines.
+
+  3. Test a bidirectional stream with the following Python script. Install the client library with `pip install websockets`.
+
+**`./test_bidi.py`**
+    
+        import asyncio
+    
+    import websockets
+    
+    
+    async def main():
+        async with websockets.connect("ws://127.0.0.1:8080/invocations-bidirectional-stream") as ws:
+            for text in ["hello", "world"]:
+                await ws.send(text)
+                print(await ws.recv())
+    
+    
+    asyncio.run(main())
+    
+        python test_bidi.py
+
+The script prints `HELLO` and `WORLD`.
+
+
+
+
+If your product uses inference pricing, your container also reports usage for streaming responses and bidirectional streams. For more information, see [Reporting usage for inference pricing](./ml-inference-metering.html).
+
@@ -339 +488 @@ Packaging your code into images
-Creating algorithm images
+Reporting usage for inference pricing