diff --git a/potassium/potassium.py b/potassium/potassium.py index 92c3dc7..2b57e0b 100644 --- a/potassium/potassium.py +++ b/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, }) diff --git a/setup.py b/setup.py index 1020912..9b49edd 100644 --- a/setup.py +++ b/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', diff --git a/tests/test_endpoints.py b/tests/test_endpoints.py index 7ac4418..f54b198 100644 --- a/tests/test_endpoints.py +++ b/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