fix warmup deadlock in cherry

potassium/potassium.py

@@ -225,6 +225,7 @@ def _create_flask_app(self):
         def handle(path):
             with self._sequence_number_lock:
                 self._sequence_number += 1
+
             route = "/" + path
             if route not in self._endpoints:
                 abort(404)
@@ -234,6 +235,8 @@ def handle(path):
 
         @flask_app.route('/_k/warmup', methods=["POST"])
         def warm():
+            with self._sequence_number_lock:
+                self._sequence_number += 1
             res = make_response({
                 "warm": True,
             })

setup.py

@@ -7,7 +7,7 @@
 setup(
     name='potassium',
     packages=['potassium'],
-    version='0.3.1',
+    version='0.3.2',
     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

@@ -208,3 +208,9 @@ def handler(context: dict, request: potassium.Request) -> potassium.Response:
     res = client.post("/_k/warmup", json={})
     assert res.status_code == 200
     assert res.json == {"warm": True}
+
+    res = client.get("/__status__", json={})
+    assert res.status_code == 200
+    assert res.json is not None
+    assert res.json["gpu_available"] == True
+    assert res.json["sequence_number"] == 1