LitServe icon indicating copy to clipboard operation
LitServe copied to clipboard

[WIP] dynamic stream loop selection

Open aniketmaurya opened this issue 1 year ago • 0 comments

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

aniketmaurya avatar May 22 '24 00:05 aniketmaurya