potassium 0.0.9

The potassium package is a flask-like HTTP server for serving large AI models

Potassium

Potassium is 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/FastAPI
• Bake in best practices for handling large, GPU-bound ML models
• Provide a set of primitives common in ML serving, such as:
  • POST request handlers
  • Websocket / streaming connections
  • Async handlers w/ webhooks
• Maintain a standard interface, to allow the code and models to compile to specialized hardware (ideally on Banana 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 model

The fastest way to get up and running is to use the Banana CLI, which downloads and runs your first model.

Here's a demo video

1. Install the CLI with pip

pip3 install banana-cli

This downloads boilerplate for your potassium app, and automatically installs potassium into the venv.

2. Create a new project directory with

banana init my-app
cd my-app

3. Start the hot-reload dev server

banana dev

4. Call your API (from a separate terminal)

curl -X POST -H "Content-Type: application/json" -d '{"prompt": "Hello I am a [MASK] model."}' http://localhost:8000/

---

Or do it yourself:

1. Install the potassium package

pip3 install potassium

Create a python file called app.py containing:

from potassium import Potassium, Request, Response
from transformers import pipeline
import torch
import time

app = Potassium("my_app")

# @app.init runs at startup, and initializes the app's context
@app.init
def init():
    device = 0 if torch.cuda.is_available() else -1
    model = pipeline('fill-mask', model='bert-base-uncased', device=device)
   
    context = {
        "model": model,
        "hello": "world"
    }

    return context

# @app.handler is an http post handler running for every call
@app.handler()
def handler(context: dict, request: Request) -> Response:
    
    prompt = request.json.get("prompt")
    model = context.get("model")
    outputs = model(prompt)

    return Response(
        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 torch

Start the server with:

python3 app.py

Test the running server with:

curl -X POST -H "Content-Type: application/json" -d '{"prompt": "Hello I am a [MASK] model."}' http://localhost:8000

---

Documentation

potassium.Potassium

from potassium import Potassium

app = Potassium("server")

This instantiates your HTTP app, similar to popular frameworks like Flask

---

@app.init

@app.init
def init():
    device = 0 if torch.cuda.is_available() else -1
    model = pipeline('fill-mask', model='bert-base-uncased', device=device)

    return {
        "model": model
    }

The @app.init decorated function runs once on server startup, and is used to load any reuseable, heavy objects such as:

• Your AI model, loaded to GPU
• Tokenizers
• Precalculated embeddings

The return value is a dictionary which saves to the app's context, and is used later in the handler functions.

There may only be one @app.init function.

---

@app.handler()

@app.handler("/", gpu=True)
def handler(context: dict, request: Request) -> Response:
    
    prompt = request.json.get("prompt")
    model = context.get("model")
    outputs = model(prompt)

    return Response(
        json = {"outputs": outputs}, 
        status=200
    )

The @app.handler decorated 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.handler functions as you'd like, with unique API routes.

The gpu=True argument allows the handler to access the prewarmed context value, and runs the handler as blocking. While the handler is running, potassium will reject any other gpu=True handlers with a 423 Locked error, to ensure that there are no multithreading issues with CUDA. If set to false, the handler may be called at any time, but the context provided will be None. gpu defaults to True.

---

@app.background(path="/background", gpu=True)

@app.background("/background")
def handler(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})

    return

The @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 a send_webhook() helper function for POSTing data onward to a url, or you may add your own custom upload/pipeline code.

When invoked, the client immediately returns a {"success": true} message.

You may configure as many @app.background functions as you'd like, with unique API routes.

The gpu=True argument allows the background handler to access the prewarmed context value, and runs the child background thread as blocking. While the child thread is running, potassium will reject any other gpu=True handlers with a 423 Locked error, to ensure that there are no multithreading issues with CUDA. If set to false, the handler may be called at any time, but the context provided will be None. gpu defaults to True.

---

app.serve()

app.serve runs the server, and is a blocking operation.

---

Store

Potassium includes a key-value storage primative, to help users persist data between calls. This is often used to implement patterns such as an async-worker queue where one background task runs inference and saves the result to the Store, with another handler set to gpu=False built to fetch the result.

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 storage

• Note: not encouraged on Banana serverless or multi-replica environments, as data is stored only on the single replica

from 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")