Skip to content

Commit

Permalink
fix warmup deadlock in cherry (#41)
Browse files Browse the repository at this point in the history
* fix warmup deadlock in cherry

* add tests
  • Loading branch information
Peddle authored Oct 25, 2023
1 parent b6962c3 commit 7fdb6aa
Show file tree
Hide file tree
Showing 3 changed files with 10 additions and 1 deletion.
3 changes: 3 additions & 0 deletions potassium/potassium.py
Original file line number Diff line number Diff line change
Expand Up @@ -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)
Expand All @@ -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,
})
Expand Down
2 changes: 1 addition & 1 deletion setup.py
Original file line number Diff line number Diff line change
Expand Up @@ -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',
Expand Down
6 changes: 6 additions & 0 deletions tests/test_endpoints.py
Original file line number Diff line number Diff line change
Expand Up @@ -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

0 comments on commit 7fdb6aa

Please sign in to comment.