Fix issues with stopping VMs

src/aleph/vm/models.py

@@ -90,6 +90,10 @@
             else self.systemd_manager.is_service_active(self.controller_service)
         )
 
+    @property
+    def is_stopping(self) -> bool:
+        return bool(self.times.stopping_at and not self.times.stopped_at)
+
     @property
     def is_program(self):
         return isinstance(self.message, ProgramContent)
@@ -289,8 +293,8 @@
             if self.times.stopped_at is not None:
                 logger.debug(f"VM={self.vm.vm_id} already stopped")
                 return
-            await self.all_runs_complete()
             self.times.stopping_at = datetime.now(tz=timezone.utc)
+            await self.all_runs_complete()
             await self.record_usage()
             await self.vm.teardown()
             self.times.stopped_at = datetime.now(tz=timezone.utc)

src/aleph/vm/orchestrator/run.py

@@ -117,7 +117,7 @@
     Execute the code corresponding to the 'code id' in the path.
     """
 
-    execution: Optional[VmExecution] = await pool.get_running_vm(vm_hash=vm_hash)
+    execution: Optional[VmExecution] = pool.get_running_vm(vm_hash=vm_hash)
 
     if not execution:
         execution = await create_vm_execution_or_raise_http_error(vm_hash=vm_hash, pool=pool)
@@ -214,7 +214,7 @@
     Execute code in response to an event.
     """
 
-    execution: Optional[VmExecution] = await pool.get_running_vm(vm_hash=vm_hash)
+    execution: Optional[VmExecution] = pool.get_running_vm(vm_hash=vm_hash)
 
     if not execution:
         execution = await create_vm_execution_or_raise_http_error(vm_hash=vm_hash, pool=pool)
@@ -260,7 +260,7 @@
 
 
 async def start_persistent_vm(vm_hash: ItemHash, pubsub: Optional[PubSub], pool: VmPool) -> VmExecution:
-    execution: Optional[VmExecution] = await pool.get_running_vm(vm_hash=vm_hash)
+    execution: Optional[VmExecution] = pool.get_running_vm(vm_hash=vm_hash)
 
     if not execution:
         logger.info(f"Starting persistent virtual machine with id: {vm_hash}")
@@ -280,7 +280,7 @@
 
 async def stop_persistent_vm(vm_hash: ItemHash, pool: VmPool) -> Optional[VmExecution]:
     logger.info(f"Stopping persistent VM {vm_hash}")
-    execution = await pool.get_running_vm(vm_hash)
+    execution = pool.get_running_vm(vm_hash)
 
     if execution:
         await execution.stop()

src/aleph/vm/pool.py

@@ -89,9 +89,10 @@
 
         # Check if an execution is already present for this VM, then return it.
         # Do not `await` in this section.
-        try:
-            return self.executions[vm_hash]
-        except KeyError:
+        current_execution = self.get_running_vm(vm_hash)
+        if current_execution:
+            return current_execution
+        else:
             execution = VmExecution(
                 vm_hash=vm_hash,
                 message=message,
@@ -161,10 +162,10 @@
                 msg = "No available value for vm_id."
                 raise ValueError(msg)
 
-    async def get_running_vm(self, vm_hash: ItemHash) -> Optional[VmExecution]:
+    def get_running_vm(self, vm_hash: ItemHash) -> Optional[VmExecution]:
         """Return a running VM or None. Disables the VM expiration task."""
         execution = self.executions.get(vm_hash)
-        if execution and execution.is_running:
+        if execution and execution.is_running and not execution.is_stopping:
             execution.cancel_expiration()
             return execution
         else: