Skip to content

fix: Correctly register and deregister terminating kernels #2157

New issue

Have a question about this project? Sign up for a free GitHub account to open an issue and contact its maintainers and the community.

Sign up for GitHub

By clicking “Sign up for GitHub”, you agree to our terms of service and privacy statement. We’ll occasionally send you account related emails.

Already on GitHub? Sign in to your account

Closed
fregataa wants to merge 4 commits into from
Closed

fix: Correctly register and deregister terminating kernels #2157

Show file tree
Hide file tree
Changes from all commits
Commits
Show all changes
4 commits
Select commit Hold shift + click to select a range
c184aa0
fix: Correctly register and deregister terminating kernels
fregataa May 17, 2024
4a6c251
add news fragment
fregataa May 17, 2024
85a3df8
register terminating_kernels right before destroy_kernel
fregataa May 17, 2024
a4ff916
some comments
fregataa May 17, 2024
File filter

Filter by extension

Filter by extension
Conversations
Failed to load comments.
Loading
Jump to
Jump to file
Failed to load files.
Loading
Diff view
Diff view
1 change: 1 addition & 0 deletions changes/2157.fix.md
View file
Open in desktop
Original file line number Diff line number Diff line change
@@ -0,0 +1 @@
Correctly register and deregister Agent's terminating kernels
8 changes: 6 additions & 2 deletions src/ai/backend/agent/agent.py
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -997,7 +997,6 @@ async def _handle_destroy_event(self, ev: ContainerLifecycleEvent) -> None:
assert current_task is not None
if ev.kernel_id not in self._ongoing_destruction_tasks:
self._ongoing_destruction_tasks[ev.kernel_id] = current_task
self.terminating_kernels.add(ev.kernel_id)
async with self.registry_lock:
kernel_obj = self.kernel_registry.get(ev.kernel_id)
if kernel_obj is None:
Expand All @@ -1023,6 +1022,7 @@ async def _handle_destroy_event(self, ev: ContainerLifecycleEvent) -> None:
if kernel_obj.runner is not None:
await kernel_obj.runner.close()
kernel_obj.clean_event = ev.done_future
self.terminating_kernels.add(ev.kernel_id)
try:
await self.destroy_kernel(ev.kernel_id, ev.container_id)
except Exception as e:
Expand All @@ -1042,6 +1042,10 @@ async def _handle_destroy_event(self, ev: ContainerLifecycleEvent) -> None:
done_future=ev.done_future,
),
)
else:
# Items in `terminating_kernels` are deleted only in _handle_clean_event()
# Should delete the kernel_id that will not be cleaned
self.terminating_kernels.discard(ev.kernel_id)
except asyncio.CancelledError:
pass
except Exception:
Expand Down Expand Up @@ -1304,7 +1308,7 @@ async def sync_container_lifecycles(self, interval: float) -> None:
LifecycleEvent.CLEAN,
KernelLifecycleEventReason.SELF_TERMINATED,
)
# Check if: there are containers not spawned by me.
# Check if: there are containers not spawned by me or already deleted from my registry.
for kernel_id in alive_kernels.keys() - known_kernels.keys():
if kernel_id in self.restarting_kernels:
continue
Expand Down