Fix: Reboot endpoint was not implemented

Author: hoh

src/aleph/vm/orchestrator/supervisor.py

@@ -32,7 +32,13 @@
     status_public_config,
     update_allocations,
 )
-from .views.operator import operate_erase, operate_expire, operate_stop, stream_logs
+from .views.operator import (
+    operate_erase,
+    operate_expire,
+    operate_reboot,
+    operate_stop,
+    stream_logs,
+)
 
 logger = logging.getLogger(__name__)
 
@@ -78,6 +84,7 @@ async def allow_cors_on_endpoint(request: web.Request):
         web.post("/control/machine/{ref}/expire", operate_expire),
         web.post("/control/machine/{ref}/stop", operate_stop),
         web.post("/control/machine/{ref}/erase", operate_erase),
+        web.post("/control/machine/{ref}/reboot", operate_reboot),
         web.options(
             "/control/machine/{ref}/{view:.*}",
             allow_cors_on_endpoint,

src/aleph/vm/orchestrator/views/operator.py

@@ -7,6 +7,7 @@
 from typing import Any, Callable, Coroutine, Literal, Union
 
 import aiohttp.web_exceptions
+import pydantic
 from aiohttp import web
 from aiohttp.web_urldispatcher import UrlMappingMatchInfo
 from aleph_message.exceptions import UnknownHashError
@@ -18,6 +19,7 @@
 from pydantic.main import BaseModel
 
 from aleph.vm.models import VmExecution
+from aleph.vm.orchestrator.run import create_vm_execution
 from aleph.vm.pool import VmPool
 
 logger = logging.getLogger(__name__)
@@ -28,7 +30,7 @@ def is_token_still_valid(timestamp):
     Checks if a token has exprired based on its timestamp
     """
     current_datetime = datetime.now(tz=timezone.utc)
-    target_datetime = datetime.fromisoformat(timestamp, tz=timezone.utc)
+    target_datetime = datetime.fromisoformat(timestamp)
 
     return target_datetime > current_datetime
 
@@ -51,7 +53,7 @@ class SignedPubKeyPayload(BaseModel):
     # alg: Literal["ECDSA"]
     domain: str
     address: str
-    expires: str 
+    expires: str
 
     @property
     def json_web_key(self) -> Jwk:
@@ -219,7 +221,6 @@ def get_execution_or_404(ref: ItemHash, pool: VmPool) -> VmExecution:
 
 @require_jwk_authentication
 async def stream_logs(request: web.Request, authenticated_sender: str) -> web.StreamResponse:
-    # TODO: Add user authentication
     vm_hash = get_itemhash_or_400(request.match_info)
     pool: VmPool = request.app["vm_pool"]
     execution = get_execution_or_404(vm_hash, pool=pool)
@@ -260,7 +261,6 @@ async def operate_expire(request: web.Request, authenticated_sender: str) -> web
     """Stop the virtual machine, smoothly if possible.
 
     A timeout may be specified to delay the action."""
-    # TODO: Add user authentication
     vm_hash = get_itemhash_or_400(request.match_info)
     timeout = float(ItemHash(request.match_info["timeout"]))
     if not 0 < timeout < timedelta(days=10).total_seconds():
@@ -317,9 +317,14 @@ async def operate_reboot(request: web.Request, authenticated_sender: str) -> web
         logger.debug(f"Unauthorized sender {authenticated_sender} for {vm_hash}")
         return web.Response(status=401, body="Unauthorized sender")
 
-    # TODO: implement this endpoint
-    logger.info(f"Rebooting {execution.vm_hash}")
-    return web.Response(status=200, body=f"Rebooted {execution.vm_hash}")
+    if execution.is_running:
+        logger.info(f"Rebooting {execution.vm_hash}")
+        await pool.stop_vm(vm_hash)
+        pool.forget_vm(vm_hash)
+        await create_vm_execution(vm_hash=vm_hash, pool=pool)
+        return web.Response(status=200, body=f"Rebooted VM with ref {vm_hash}")
+    else:
+        return web.Response(status=200, body="Starting VM (was not running) with ref {vm_hash}")
 
 
 @require_jwk_authentication

src/aleph/vm/pool.py

@@ -136,6 +136,15 @@ async def get_running_vm(self, vm_hash: ItemHash) -> Optional[VmExecution]:
         else:
             return None
 
+    async def stop_vm(self, vm_hash: ItemHash) -> Optional[VmExecution]:
+        """Stop a VM."""
+        execution = self.executions.get(vm_hash)
+        if execution:
+            await execution.stop()
+            return execution
+        else:
+            return None
+
     def forget_vm(self, vm_hash: ItemHash) -> None:
         """Remove a VM from the executions pool.