Package: litserve
litserveLightweight inference server for AI/ML models based on FastAPI.litserve supports:serving models across multiple GPUsfull flexibility in the definition of request and responseautomatic schema validationhandling of timeouts and disconnectionsAPI key authenticationWhile being fairly capable and fast, the server is extremely simple and hackable.For example, adding batching and streaming logic should be relatively straightforward
and we'll probably add those directly as part of the library in the near future.Installpipinstallgit+https://github.com/Lightning-AI/litserve.gitUseCreating a simple API serverfromlitserveimportLitAPIclassSimpleLitAPI(LitAPI):defsetup(self,device):"""Setup the model so it can be called in `predict`."""self.model=lambdax:x**2defdecode_request(self,request):"""Convert the request payload to your model input."""returnrequest["input"]defpredict(self,x):"""Run the model on the input and return the output."""returnself.model(x)defencode_response(self,output):"""Convert the model output to a response payload."""return{"output":output}Now instantiate the API and start the server on an accelerator:fromlitserveimportLitServerapi=SimpleLitAPI()server=LitServer(api,accelerator="cpu")server.run(port=8000)The server expects the client to send aPOSTto the/predictURL with a JSON payload.
The way the payload is structured is up to the implementation of theLitAPIsubclass.Once the server starts it generates an example client you can use like this:pythonclient.pythat simply posts a sample request to the server:response=requests.post("http://127.0.0.1:8000/predict",json={"input":4.0})Pydantic modelsYou can also define request and response asPydantic models,
with the advantage that the API will automatically validate the request.frompydanticimportBaseModelclassPredictRequest(BaseModel):input:floatclassPredictResponse(BaseModel):output:floatclassSimpleLitAPI2(LitAPI):defsetup(self,device):self.model=lambdax:x**2defdecode_request(self,request:PredictRequest)->float:returnrequest.inputdefpredict(self,x):returnself.model(x)defencode_response(self,output:float)->PredictResponse:returnPredictResponse(output=output)if__name__=="__main__":api=SimpleLitAPI()server=LitServer(api,accelerator="cpu")server.run(port=8888)Serving on GPULitServerhas the ability to coordinate serving from multiple GPUs.For example, running the API server on a 4-GPU machine, with a PyTorch model served by each GPU:fromfastapiimportRequest,ResponsefromlitserveimportLitAPI,LitServerimporttorchimporttorch.nnasnnclassLinear(nn.Module):def__init__(self):super().__init__()self.linear=nn.Linear(1,1)self.linear.weight.data.fill_(2.0)self.linear.bias.data.fill_(1.0)defforward(self,x):returnself.linear(x)classSimpleLitAPI(LitAPI):defsetup(self,device):# move the model to the correct device# keep track of the device for moving data accordinglyself.model=Linear().to(device)self.device=devicedefdecode_request(self,request:Request):# get the input and create a 1D tensor on the correct devicecontent=request["input"]returntorch.tensor([content],device=self.device)defpredict(self,x):# the model expects a batch dimension, so create itreturnself.model(x[None,:])defencode_response(self,output)->Response:# float will take the output value directly onto CPU memoryreturn{"output":float(output)}if__name__=="__main__":# accelerator="cuda", devices=4 will lead to 4 workers serving the# model from "cuda:0", "cuda:1", "cuda:2", "cuda:3" respectivelyserver=LitServer(SimpleLitAPI(),accelerator="cuda",devices=4)server.run(port=8000)Thedevicesvariable can also be an array specifying what device id to
run the model on:server=LitServer(SimpleLitAPI(),accelerator="cuda",devices=[0,3])Last, you can run multiple copies of the same model from the same device,
if the model is small. The following will load two copies of the model on
each of the 4 GPUs:server=LitServer(SimpleLitAPI(),accelerator="cuda",devices=4,workers_per_device=2)Timeouts and disconnectionsThe server will remove a queued request if the client requesting it disconnects.You can configure a timeout (in seconds) after which clients will receive a504HTTP
response (Gateway Timeout) indicating that their request has timed out.For example, this is how you can configure the server with a timeout of 30 seconds per response.server=LitServer(SimpleLitAPI(),accelerator="cuda",devices=4,timeout=30)This is useful to avoid requests queuing up beyond the ability of the server to respond.Using API key authenticationIn order to secure the API behind an API key, just define the env var when
starting the serverLIT_SERVER_API_KEY=supersecretkeypythonmain.pyClients are expected to auth with the same API key set in theX-API-KeyHTTP header.Licenselitserve is released under theApache 2.0license.
See LICENSE file for details.