delete kernel from kernel_registry when its cleaned

lablup · May 29, 2024 · 09893dc · 09893dc
1 parent fc1a8e4
commit 09893dc
Showing 1 changed file with 4 additions and 0 deletions.
diff --git a/src/ai/backend/agent/agent.py b/src/ai/backend/agent/agent.py
@@ -1084,6 +1084,10 @@ async def _handle_clean_event(self, ev: ContainerLifecycleEvent) -> None:
                         await kernel_obj.close()
                 finally:
                     self.terminating_kernels.discard(ev.kernel_id)
+                    try:
+                        del self.kernel_registry[ev.kernel_id]
+                    except KeyError:
+                        pass
                     if restart_tracker := self.restarting_kernels.get(ev.kernel_id, None):
                         restart_tracker.destroy_event.set()
                     else: