Skip to content

Commit

Permalink
Add Stream API Example (#95)
Browse files Browse the repository at this point in the history
* Add Stream API Example

* [pre-commit.ci] auto fixes from pre-commit.com hooks

for more information, see https://pre-commit.ci

* Modify example so it uses the requests.post input

---------

Co-authored-by: pre-commit-ci[bot] <66853113+pre-commit-ci[bot]@users.noreply.github.com>
  • Loading branch information
rasbt and pre-commit-ci[bot] authored May 20, 2024
1 parent bfa9219 commit 3c1c2d3
Show file tree
Hide file tree
Showing 2 changed files with 36 additions and 3 deletions.
6 changes: 3 additions & 3 deletions README.md
Original file line number Diff line number Diff line change
Expand Up @@ -8,9 +8,9 @@
**High-throughput serving engine for AI models**

<pre>
✅ Batching ✅ Streaming ✅ Auto-GPU, multi-GPU
✅ Multi-modal ✅ PyTorch/JAX/TF ✅ Full control
✅ Auth ✅ Built on Fast API
✅ Batching ✅ Streaming ✅ Auto-GPU, multi-GPU
✅ Multi-modal ✅ PyTorch/JAX/TF ✅ Full control
✅ Auth ✅ Built on Fast API
</pre>


Expand Down
33 changes: 33 additions & 0 deletions src/litserve/examples/simple_example.py
Original file line number Diff line number Diff line change
Expand Up @@ -68,3 +68,36 @@ def predict(self, x):
def encode_response(self, output):
# float will take the output value directly onto CPU memory
return {"output": float(output)}


class SimpleStreamAPI(ls.LitAPI):
"""
Run as:
```
server = ls.LitServer(SimpleStreamAPI(), stream=True)
server.run(port=8000)
```
Then, in a new Python session, retrieve the responses as follows:
```
import requests
url = "http://127.0.0.1:8000/predict"
resp = requests.post(url, json={"input": "Hello world"}, headers=None, stream=True)
for line in resp.iter_content(5000):
if line:
print(line.decode("utf-8"))
```
"""

def setup(self, device) -> None:
self.model = lambda x, y: f"{x}: {y}"

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

def predict(self, x):
for i in range(3):
yield self.model(i, x.encode("utf-8").decode())

def encode_response(self, output_stream):
for output in output_stream:
yield {"output": output}

0 comments on commit 3c1c2d3

Please sign in to comment.