From 1ce28b9adb661201d86bfc33dd2bc4ba99a5b762 Mon Sep 17 00:00:00 2001 From: Hugo Herter Date: Tue, 20 Feb 2024 16:58:11 +0100 Subject: [PATCH] Cleanup: Apply review feedback --- src/aleph/vm/pool.py | 22 +++++++++++----------- 1 file changed, 11 insertions(+), 11 deletions(-) diff --git a/src/aleph/vm/pool.py b/src/aleph/vm/pool.py index 2fcfcd608..3a0d8ed78 100644 --- a/src/aleph/vm/pool.py +++ b/src/aleph/vm/pool.py @@ -1,3 +1,4 @@ +from __future__ import annotations import asyncio import json import logging @@ -128,11 +129,7 @@ async def create_a_vm( self.forget_vm(vm_hash) raise - async def forget_on_stop(stop_event: asyncio.Event): - await stop_event.wait() - self.forget_vm(vm_hash) - - _ = asyncio.create_task(forget_on_stop(stop_event=execution.stop_event)) + self._schedule_forget_on_stop(execution) return execution @@ -202,6 +199,14 @@ def forget_vm(self, vm_hash: ItemHash) -> None: except KeyError: pass + def _schedule_forget_on_stop(self, execution: VmExecution): + """Create a task that will remove the VM from the pool after it stops.""" + async def forget_on_stop(stop_event: asyncio.Event): + await stop_event.wait() + self.forget_vm(execution.vm_hash) + + _ = asyncio.create_task(forget_on_stop(stop_event=execution.stop_event)) + async def _load_persistent_executions(self): """Load persistent executions from the database.""" saved_executions = await get_execution_records() @@ -224,12 +229,7 @@ async def _load_persistent_executions(self): persistent=saved_execution.persistent, ) - # Ensure the VM is forgotten from the pool if it is stopped. - async def forget_on_stop(stop_event: asyncio.Event): - await stop_event.wait() - self.forget_vm(vm_hash) - - _ = asyncio.create_task(forget_on_stop(stop_event=execution.stop_event)) + self._schedule_forget_on_stop(execution) if execution.is_running: # TODO: Improve the way that we re-create running execution