Package: potassium
PotassiumPotassiumis an open source web framework, built to tackle the unique challenges of serving custom models in production.The goal of this project is to:Provide a familiar web framework similar to Flask/FastAPIBake in best practices for handling large, GPU-bound ML modelsProvide a set of primitives common in ML serving, such as:POST request handlersWebsocket / streaming connectionsAsync handlers w/ webhooksMaintain a standard interface, to allow the code and models to compile to specialized hardware (ideally onBanana Serverless GPUs😉)Stability Notes:Potassium uses Semantic Versioning, in that major versions imply breaking changes, and v0 implies instability even between minor/patch versions. Be sure to lock your versions, as we're still in v0!Quickstart: Serving a Huggingface BERT modelThe fastest way to get up and running is to use theBanana CLI, which downloads and runs your first model.Here's a demo videoInstall the CLI with pippip3installbanana-cliThis downloads boilerplate for your potassium app, and automatically installs potassium into the venv.Create a new project directory withbananainitmy-appcdmy-appStart the dev server../venv/bin/activate
python3app.pyCall your API (from a separate terminal)curl-XPOST-H"Content-Type: application/json"-d'{"prompt": "Hello I am a [MASK] model."}'http://localhost:8000/Or do it yourself:Install the potassium packagepip3installpotassiumCreate a python file calledapp.pycontaining:frompotassiumimportPotassium,Request,Responsefromtransformersimportpipelineimporttorchimporttimeapp=Potassium("my_app")# @app.init runs at startup, and initializes the app's context@app.initdefinit():device=0iftorch.cuda.is_available()else-1model=pipeline('fill-mask',model='bert-base-uncased',device=device)context={"model":model,"hello":"world"}returncontext# @app.handler is an http post handler running for every call@app.handler()defhandler(context:dict,request:Request)->Response:prompt=request.json.get("prompt")model=context.get("model")outputs=model(prompt)returnResponse(json={"outputs":outputs},status=200)if__name__=="__main__":app.serve()This runs a Huggingface BERT model.For this example, you'll also need to install transformers and torch.pip3 install transformers torchStart the server with:python3app.pyTest the running server with:curl-XPOST-H"Content-Type: application/json"-d'{"prompt": "Hello I am a [MASK] model."}'http://localhost:8000Documentationpotassium.PotassiumfrompotassiumimportPotassiumapp=Potassium("server")This instantiates your HTTP app, similar to popular frameworks likeFlask@app.init@app.initdefinit():device=0iftorch.cuda.is_available()else-1model=pipeline('fill-mask',model='bert-base-uncased',device=device)return{"model":model}The@app.initdecorated function runs once on server startup, and is used to load any reuseable, heavy objects such as:Your AI model, loaded to GPUTokenizersPrecalculated embeddingsThe return value is a dictionary which saves to the app'scontext, and is used later in the handler functions.There may only be one@app.initfunction.@app.handler()@app.handler("/")defhandler(context:dict,request:Request)->Response:prompt=request.json.get("prompt")model=context.get("model")outputs=model(prompt)returnResponse(json={"outputs":outputs},status=200)The@app.handlerdecorated function runs for every http call, and is used to run inference or training workloads against your model(s).You may configure as many@app.handlerfunctions as you'd like, with unique API routes.The context dict passed in is a mutable reference, so you can modify it in-place to persist objects between warm handlers.@app.background(path="/background")@app.background("/background")defhandler(context:dict,request:Request)->Response:prompt=request.json.get("prompt")model=context.get("model")outputs=model(prompt)send_webhook(url="http://localhost:8001",json={"outputs":outputs})returnThe@app.background()decorated function runs a nonblocking job in the background, for tasks where results aren't expected to return clientside. It's on you to forward the data to wherever you please. Potassium supplies asend_webhook()helper function for POSTing data onward to a url, or you may add your own custom upload/pipeline code.When invoked, the server immediately returns a{"success": true}message.You may configure as many@app.backgroundfunctions as you'd like, with unique API routes.The context dict passed in is a mutable reference, so you can modify it in-place to persist objects between warm handlers.app.serve()app.serveruns the server, and is a blocking operation.Pre-warming your appPotassium comes with a built-in endpoint for those cases where you want to "warm up" your app to better control the timing of your inference calls. You don'tneedto call it, since your inference call requiresinit()to have run once on server startup anyway, but this gives you a bit more control.Once your model is warm (i.e., cold boot finished), this endpoint returns a 200. If a cold boot is required, theinit()function is first called while the server starts up, and then a 200 is returned from this endpoint.You don't need any extra code to enable it, it comes out of the box and you can call it at/_k/warmupas either a GET or POST request.StorePotassium includes a key-value storage primative, to help users persist data between calls.Example usage: your own Redis backend (encouraged)from potassium.store import Store, RedisConfig

store = Store(
    backend="redis",
    config = RedisConfig(
        host = "localhost",
        port = 6379
    )
)

# in one handler
store.set("key", "value", ttl=60)

# in another handler
value = store.get("key")Example usage: using local storageNote: not encouraged on Banana serverless or multi-replica environments, as data is stored only on the single replicafrom potassium.store import Store, RedisConfig

store = Store(
    backend="local"
)

# in one handler
store.set("key", "value", ttl=60)

# in another handler
value = store.get("key")