Fixes 404 deadlock (#40)

* fixes the 404 deadlock draft

* cleaner fix

* rebase main and increment version

potassium/potassium.py

@@ -48,6 +48,7 @@ def __init__(self, name):
         self._gpu_lock = Lock()
         self._background_task_cv = Condition()
         self._sequence_number = 0
+        self._sequence_number_lock = Lock()
         self._idle_start_time = 0
         self._last_inference_start_time = None
         self._flask_app = self._create_flask_app()
@@ -140,7 +141,6 @@ def _handle_generic(self, endpoint, flask_request):
         # potassium rejects if lock already in use
         try:
             self._gpu_lock.acquire(blocking=False)
-            self._sequence_number += 1
         except:
             res = make_response()
             res.status_code = 423
@@ -223,6 +223,8 @@ def _create_flask_app(self):
         @flask_app.route('/', defaults={'path': ''}, methods=["POST"])
         @flask_app.route('/<path:path>', methods=["POST"])
         def handle(path):
+            with self._sequence_number_lock:
+                self._sequence_number += 1
             route = "/" + path
             if route not in self._endpoints:
                 abort(404)

setup.py

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

@@ -143,6 +143,15 @@ def background(context: dict, request: potassium.Request):
     assert res.json["idle_time"] > 0
     assert res.json["inference_time"] == 0
 
+    res = client.post("/this_path_does_not_exist", json={})
+    assert res.status_code == 404
+    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"] == 2
+
+
 def test_wait_for_background_task():
     app = potassium.Potassium("my_app")