LitServe
LitServe copied to clipboard
[WIP] dynamic stream loop selection
fixes #99
- Select streaming or non-streaming loop based on client's
request["stream"]parameter.
Server
# server.py
import logging
import litserve as ls
logging.basicConfig(level=logging.INFO)
# STEP 1: DEFINE YOUR MODEL API
class SimpleLitAPI(ls.LitAPI):
def setup(self, device):
# Setup the model so it can be called in `predict`.
self.model = lambda x: x**2
def decode_request(self, request):
# Convert the request payload to your model input.
input = request["input"]
stream = request.get("stream", False)
return input, stream
def predict(self, x):
input, stream = x
def gen():
yield self.model(input)
if stream:
return gen()
return self.model(input)
def encode_response(self, output):
# YEILD or return based on output
...
# STEP 2: START THE SERVER
if __name__ == "__main__":
api = SimpleLitAPI()
server = ls.LitServer(api, accelerator="auto")
server.run(port=8000)
Client
import requests
response = requests.post("http://127.0.0.1:8000/predict", json={"input": 4.0, "stream": False})
print(f"Status: {response.status_code}\nResponse:\n {response.text}")
EDIT: Updated to remove the iter_* API