litserve 的认证机制简单说明

litserve 已经内置了认证能力可以通过token（环境变量模式）以及通过自定义请求头

模式

• 内部token 模式

LIT_SERVER_API_KEY=A_SECRET_KEY python server.py

访问

import requests

# Define the API key in the request headers
headers = {"X-API-Key": "A_SECRET_KEY"}

response = requests.post("http://127.0.0.1:8000/predict", json={"input": 4.0}, headers=headers)
print(response.json())

• 自定义请求头模式

from fastapi import Depends, HTTPException
from fastapi.security import HTTPAuthorizationCredentials, HTTPBearer

from litserve import LitAPI, LitServer

class SimpleAuthedLitAPI(LitAPI):
    def setup(self, device):
        self.model = lambda x: x ** 2

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

    def predict(self, x):
        return self.model(x)

    def encode_response(self, output):
        return {"output": output}

    def authorize(self, auth: HTTPAuthorizationCredentials = Depends(HTTPBearer())):
        if auth.scheme != "Bearer" or auth.credentials != "1234":
            raise HTTPException(status_code=401, detail="Bad token")

if __name__ == '__main__':
    api = SimpleAuthedLitAPI()
    server = LitServer(api)
    server.run()

请求

import requests

# Define authorization in the request headers
headers = {"Authorization": "Bearer 1234"}

response = requests.post("http://127.0.0.1:8000/predict", json={"input": 4.0}, headers=headers)
print(response.json())

内部处理是基于了自定义解析以及fastapi 的依赖注入

self.app.add_api_route(
    endpoint,
    stream_predict if stream else predict,
    methods=methods,
    dependencies=[Depends(self.setup_auth())],
)

def setup_auth(self):
    if hasattr(self.lit_api, "authorize") and callable(self.lit_api.authorize):
        return self.lit_api.authorize
    if LIT_SERVER_API_KEY:
        return api_key_auth
    return no_auth

说明

litserve 的认证能力还是比较简单的，但是我们可以自己灵活的进行扩展

参考资料

src/litserve/server.py

https://lightning.ai/docs/litserve/features/authentication