Reset ctx_pid before logging on prediction cancel or failure

replicate · Dec 2, 2024 · cdf9985 · cdf9985
1 parent 6211adb
commit cdf9985
Showing 1 changed file with 3 additions and 2 deletions.
diff --git a/python/coglet/file_runner.py b/python/coglet/file_runner.py
@@ -170,18 +170,19 @@ async def _predict(self, pid: str, req: Dict[str, Any]) -> None:
                         self._respond(pid, resp)
             else:
                 resp['output'] = await self.runner.predict(req['input'])
-            self.ctx_pid.set(None)
             resp['status'] = 'succeeded'
+            self.ctx_pid.set(None)
             self.logger.info('prediction completed: id=%s', pid)
         except asyncio.CancelledError:
             resp['status'] = 'canceled'
+            self.ctx_pid.set(None)
             self.logger.error('prediction canceled: id=%s', pid)
         except Exception as e:
             resp['error'] = str(e)
             resp['status'] = 'failed'
+            self.ctx_pid.set(None)
             self.logger.error('prediction failed: id=%s %s', pid, e)
         finally:
-            self.ctx_pid.set(None)
             resp['completed_at'] = util.now_iso()
         self._respond(pid, resp)