cleaner fix

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
@@ -218,18 +218,11 @@ 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:
-                try:
-                    self._gpu_lock.acquire(blocking=False)
-                    self._sequence_number += 1
-                    self._gpu_lock.release()
-                    return make_response("Route not found", 404)
-                except:
-                    res = make_response()
-                    res.status_code = 423
-                    return res
-
+                abort(404)
 
             endpoint = self._endpoints[route]
             return self._handle_generic(endpoint, request)

tests/test_endpoints.py

@@ -135,6 +135,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")