WIP: Add API to update long-running VM allocations

Author: hoh

vm_supervisor/resources.py

@@ -1,5 +1,7 @@
+
 from datetime import datetime, timezone
 from functools import lru_cache
+from typing import Set, Optional
 from typing import Tuple
 
 import cpuinfo
@@ -117,3 +119,9 @@ async def about_system_usage(request: web.Request):
     return web.json_response(
         text=usage.json(exclude_none=True),
     )
+
+
+class Allocation(BaseModel):
+    long_running_vms: Set[str]
+    on_demand_vms: Optional[Set[str]] = None
+    jobs: Optional[Set] = None

vm_supervisor/run.py

@@ -216,3 +216,31 @@ async def run_code_on_event(vm_hash: VmHash, event, pubsub: PubSub):
             execution.stop_after_timeout(timeout=settings.REUSE_TIMEOUT)
         else:
             await execution.stop()
+
+
+async def start_long_running(vm_hash: VmHash, pubsub: PubSub) -> VmExecution:
+    execution: Optional[VmExecution] = await pool.get_running_vm(vm_hash=vm_hash)
+
+    if not execution:
+        execution = await create_vm_execution(vm_hash=vm_hash)\
+
+    # If the VM was already running in lambda mode, it should not expire
+    # as long as it is also scheduled as long-running
+    execution.marked_as_long_running = True
+    execution.cancel_expiration()
+
+    await execution.becomes_ready()
+
+    # if settings.WATCH_FOR_UPDATES:
+    #     # FIXME: Is this added for every request ?
+    #     execution.start_watching_for_updates(pubsub=request.app["pubsub"])
+
+    return execution
+
+
+async def stop_long_running(vm_hash: VmHash) -> Optional[VmExecution]:
+    logger.info(f"Stopping long running {vm_hash}")
+    execution = await pool.get_running_vm(vm_hash)
+    if execution:
+        await execution.stop()
+    return execution

vm_supervisor/supervisor.py

@@ -26,6 +26,7 @@
     status_check_fastapi,
     about_execution_records,
     status_check_version,
+    update_allocations,
 )
 
 logger = logging.getLogger(__name__)
@@ -53,6 +54,7 @@ async def server_version_middleware(
         web.get("/about/executions/records", about_execution_records),
         web.get("/about/usage/system", about_system_usage),
         web.get("/about/config", about_config),
+        web.post("/control/allocations", update_allocations),
         web.get("/status/check/fastapi", status_check_fastapi),
         web.get("/status/check/version", status_check_version),
         web.route("*", "/vm/{ref}{suffix:.*}", run_code_from_path),

vm_supervisor/views.py

@@ -3,18 +3,20 @@
 import os.path
 from string import Template
 from typing import Awaitable, Optional
-from packaging.version import Version, InvalidVersion
 
 import aiodns
 import aiohttp
 from aiohttp import web
 from aiohttp.web_exceptions import HTTPNotFound
+from packaging.version import Version, InvalidVersion
+from pydantic import ValidationError
 
 from . import status, __version__
 from .conf import settings
 from .metrics import get_execution_records
 from .models import VmHash
-from .run import run_code_on_request, pool
+from .resources import Allocation
+from .run import run_code_on_request, pool, start_long_running
 from .utils import b32_to_b16, get_ref_from_dns, dumps_for_json
 
 logger = logging.getLogger(__name__)
@@ -166,3 +168,34 @@ async def status_check_version(request: web.Request):
         )
     else:
         return web.HTTPForbidden(text=f"Outdated: version {current} < {reference}")
+
+
+async def update_allocations(request: web.Request):
+    # TODO: Add some form of authentication
+    try:
+        data = await request.json()
+        allocation = Allocation(**data)
+    except ValidationError as error:
+        return web.json_response(
+            data=error.json(), status=web.HTTPBadRequest.status_code
+        )
+
+    pubsub = request.app["pubsub"]
+
+    for vm_hash in allocation.long_running_vms:
+        vm_hash = VmHash(vm_hash)
+        logger.info(f"Starting long running VM {vm_hash}")
+        await start_long_running(vm_hash, pubsub)
+
+    for execution in pool.get_long_running_executions():
+        if execution.vm_hash not in allocation.long_running_vms:
+            logger.info(f"Stopping long running VM {execution.vm_hash}")
+            await execution.stop()
+            execution.marked_as_long_running = False
+
+    if allocation.on_demand_vms:
+        logger.info("Not supported yet: 'allocation.on_demand_vms'")
+    if allocation.jobs:
+        logger.info("Not supported yet: 'allocation.on_demand_vms'")
+
+    return web.json_response(data={"success": True})