Warmup calls (#36)

Dockerfile

@@ -1,5 +1,14 @@
 FROM python:3.8-slim-buster
 
+ENV NODE_MAJOR=20
+RUN apt-get update && \
+    apt-get install -y ca-certificates curl gnupg && \
+    mkdir -p /etc/apt/keyrings && \
+    curl -fsSL https://deb.nodesource.com/gpgkey/nodesource-repo.gpg.key | gpg --dearmor -o /etc/apt/keyrings/nodesource.gpg && \
+    echo "deb [signed-by=/etc/apt/keyrings/nodesource.gpg] https://deb.nodesource.com/node_$NODE_MAJOR.x nodistro main" | tee /etc/apt/sources.list.d/nodesource.list && \
+    apt-get update && \
+    apt-get install nodejs -y
+
 WORKDIR /potassium
 
 RUN pip install pyright pytest

README.md

@@ -212,6 +212,14 @@ The context dict passed in is a mutable reference, so you can modify it in-place
 
 `app.serve` runs the server, and is a blocking operation.
 
+---
+## Pre-warming your app
+
+Potassium 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't *need* to call it, since your inference call requires `init()` 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, the `init()` 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/warmup` as either a GET or POST request.
 
 ---
 

potassium/potassium.py

@@ -234,7 +234,17 @@ def handle(path):
 
             endpoint = self._endpoints[route]
             return self._handle_generic(endpoint, request)
+
+        @flask_app.route('/_k/warmup', methods=["POST"])
+        def warm():
+            res = make_response({
+                "warm": True,
+            })
+            res.status_code = 200
+            res.headers['X-Endpoint-Type'] = "warmup"
+            return res
 
+        @flask_app.route('/_k/status', methods=["GET"])
         @flask_app.route('/__status__', methods=["GET"])
         def status():
             idle_time = 0

setup.py

@@ -7,7 +7,7 @@
 setup(
     name='potassium',
     packages=['potassium'],
-    version='0.2.1',
+    version='0.3.0',
     license='Apache License 2.0',
     # Give a short description about your library
     description='The potassium package is a flask-like HTTP server for serving large AI models',

tests/test_endpoints.py

@@ -183,3 +183,19 @@ def wait_for_background_task():
     assert order_of_execution_queue.get() == "send_background_task"
     assert order_of_execution_queue.get() == "background_task_completed"
 
+def test_warmup():
+    app = potassium.Potassium("my_app")
+
+    @app.init
+    def init():
+        return {}
+
+    @app.handler()
+    def handler(context: dict, request: potassium.Request) -> potassium.Response:
+        raise Exception("should not be called")
+
+    client = app.test_client()
+
+    res = client.post("/_k/warmup", json={})
+    assert res.status_code == 200
+    assert res.json == {"warm": True}