small change

bananaml · Sep 13, 2023 · 8a26304 · 8a26304
1 parent 71391f8
commit 8a26304
Show file tree

Hide file tree

Showing 2 changed files with 5 additions and 3 deletions.
diff --git a/potassium/potassium.py b/potassium/potassium.py
@@ -167,6 +167,7 @@ def _handle_generic(self, endpoint, flask_request):
                 res.status_code = 500
                 res.headers['X-Endpoint-Type'] = endpoint.type
             self._idle_start_time = time.time()
+            self._inference_start_time = 0
             self._gpu_lock.release()
         elif endpoint.type == "background":
             req = Request(
@@ -185,6 +186,7 @@ def task(endpoint, lock, req):
                         self._background_task_cv.notify_all()
 
                         self._idle_start_time = time.time()
+                        self._inference_start_time = 0
                         lock.release()
 
             thread = Thread(target=task, args=(endpoint, self._gpu_lock, req))
@@ -228,11 +230,11 @@ def status():
             idle_time = 0
             gpu_available = not self._gpu_lock.locked()
 
-            if self._inference_start_time != 0:
+            if self._inference_start_time != 0:            
                 inference_time = int((time.time() - self._inference_start_time)*1000)
             else:
                 inference_time = 0
-
+            
             if gpu_available:
                 idle_time = int((time.time() - self._idle_start_time)*1000)
 

diff --git a/tests/test_endpoints.py b/tests/test_endpoints.py
@@ -133,7 +133,7 @@ def background(context: dict, request: potassium.Request):
     assert res.json["gpu_available"] == True
     assert res.json["sequence_number"] == 1
     assert res.json["idle_time"] > 0
-    assert res.json["inference_time"] > 0
+    assert res.json["inference_time"] == 0
 
 def test_wait_for_background_task():
     app = potassium.Potassium("my_app")