omni boxes

omnitool/omnibox/Dockerfile

@@ -1,7 +1,7 @@
 ARG VERSION_ARG="latest"
 FROM scratch AS build-amd64
 
-COPY --from=qemux/qemu-docker:6.08 / /
+COPY --from=qemux/qemu:6.18 / /
 
 ARG DEBCONF_NOWARNINGS="yes"
 ARG DEBIAN_FRONTEND="noninteractive"

omnitool/omniboxes/client/notebook.ipynb

@@ -0,0 +1,84 @@
+{
+ "cells": [
+  {
+   "cell_type": "markdown",
+   "id": "578457c0",
+   "metadata": {},
+   "source": [
+    "# Test instance"
+   ]
+  },
+  {
+   "cell_type": "code",
+   "execution_count": null,
+   "id": "ae7b2936",
+   "metadata": {},
+   "outputs": [],
+   "source": [
+    "import client\n",
+    "client.InstanceClient(port = 5000).ui()"
+   ]
+  },
+  {
+   "cell_type": "markdown",
+   "id": "741598f1",
+   "metadata": {},
+   "source": [
+    "# Test node"
+   ]
+  },
+  {
+   "cell_type": "code",
+   "execution_count": null,
+   "id": "eca33f8c",
+   "metadata": {},
+   "outputs": [],
+   "source": [
+    "import client\n",
+    "c = client.NodeClient(port = 8000)\n",
+    "c.ui()"
+   ]
+  },
+  {
+   "cell_type": "markdown",
+   "id": "c2b1810b",
+   "metadata": {},
+   "source": [
+    "# Test master node"
+   ]
+  },
+  {
+   "cell_type": "code",
+   "execution_count": null,
+   "id": "bf6188a0",
+   "metadata": {},
+   "outputs": [],
+   "source": [
+    "import client\n",
+    "c = client.MasterClient(port = 7000)\n",
+    "c.ui()"
+   ]
+  }
+ ],
+ "metadata": {
+  "kernelspec": {
+   "display_name": "omnibox",
+   "language": "python",
+   "name": "python3"
+  },
+  "language_info": {
+   "codemirror_mode": {
+    "name": "ipython",
+    "version": 3
+   },
+   "file_extension": ".py",
+   "mimetype": "text/x-python",
+   "name": "python",
+   "nbconvert_exporter": "python",
+   "pygments_lexer": "ipython3",
+   "version": "3.12.9"
+  }
+ },
+ "nbformat": 4,
+ "nbformat_minor": 5
+}

omnitool/omniboxes/client/requirements.txt

@@ -0,0 +1,3 @@
+pillow
+ipywidgets
+requests

omnitool/omniboxes/master/logging_utils.py

@@ -0,0 +1,8 @@
+import logging
+
+def default_logger():
+    logging.basicConfig(
+        level=logging.INFO,
+        format="%(asctime)s - %(name)s - %(levelname)s - %(message)s",
+    )
+    return logging.getLogger("omnibox-master")

omnitool/omniboxes/master/node_manager.py

@@ -0,0 +1,104 @@
+import asyncio
+import httpx
+import uvicorn
+import logging
+from fastapi import FastAPI, HTTPException, status, Response
+from pydantic import BaseModel, Field
+from typing import Dict, List, Optional, Any
+import logging_utils
+
+
+class NodeRegistration(BaseModel):
+    url: str
+
+
+class NodeStatus(BaseModel):
+    url: str
+    hash: str
+    healthy: bool
+    capacity: int
+    available: int
+    instances: List[str]
+
+    @staticmethod
+    def failed(url):
+        return NodeStatus(url=url, hash=url_hash(url), healthy=False, capacity=0, available=0, instances=[])
+
+
+def url_hash(url: str) -> str:
+    return url
+
+
+def _default_logger():
+    logging.basicConfig(
+        level=logging.INFO,
+        format="%(asctime)s - %(name)s - %(levelname)s - %(message)s",
+    )
+    return logging.getLogger("omnibox-master")
+
+
+class NodeManager:
+    def __init__(self, logger = None, update_timeout: int = 10):
+        self._nodes = {}
+        self._node_info = {}
+        self.logger = logger or _default_logger()
+        self.update_timeout = update_timeout
+
+    async def _get_status(self, node_url):
+        try:
+            async with httpx.AsyncClient(timeout=5.0) as client:
+                response = await client.get(f"{node_url}/info")
+                if response.status_code != 200:
+                    return NodeStatus.failed(node_url)
+                data = response.json()
+                return NodeStatus(
+                    url = node_url,
+                    hash = url_hash(node_url),
+                    healthy = True,
+                    capacity = data.get("capacity", 0),
+                    available = data.get("available", 0),
+                    instances = data.get("in_use", [])
+                )
+        except Exception as e:
+            self.logger.warning(f"Node {node_url} health check failed: {str(e)}")
+            return NodeStatus.failed(node_url)
+
+    async def update_statuses(self):
+        for node_url in self._nodes.keys():
+            self._node_info[node_url] = await self._get_status(node_url)
+
+    async def update_statuses_worker(self):
+        while True:
+            await self.update_statuses()
+            await asyncio.sleep(self.update_timeout)
+
+    async def register_node(self, node: NodeRegistration):
+        self._nodes[url_hash(node.url)] = node
+        self._node_info[url_hash(node.url)] = await self._get_status(node.url)            
+
+    async def unregister_node(self, node_url: str):
+        """Unregister a worker node from the master"""
+        node_hash = url_hash(node_url)
+        if node_hash not in self._nodes:
+            self.logger.warning(f"Node {node_url} not found for unregistration")
+            return False
+
+        self._nodes.pop(node_hash, None)
+        self._node_info.pop(node_hash, None)
+        return True
+
+    def get_best_node(self) -> Optional[str]:
+        result = None
+        available = 0
+        for node in self._node_info.values():
+            if node.healthy and node.available > available:
+                result = node
+                available = node.available
+        return result
+
+    def get_node(self, hash: str) -> Optional[NodeStatus]:
+        return self._node_info.get(hash, None)
+
+    def node_info(self):
+        return self._node_info
+

omnitool/omniboxes/master/server.py

@@ -0,0 +1,147 @@
+import asyncio
+import httpx
+import uvicorn
+import logging
+from fastapi import FastAPI, HTTPException, status, Response
+from pydantic import BaseModel, Field
+from typing import Dict, List, Optional, Any
+from contextlib import asynccontextmanager
+from fastapi.responses import JSONResponse
+from node_manager import NodeManager, NodeRegistration
+import requests
+import argparse
+
+# Configure logging
+logging.basicConfig(
+    level=logging.INFO,
+    format="%(asctime)s - %(name)s - %(levelname)s - %(message)s",
+)
+logger = logging.getLogger("omnibox-master")
+
+node_manager = NodeManager(logger = logger)
+
+@asynccontextmanager
+async def lifespan(app: FastAPI):
+    await node_manager.register_node(NodeRegistration(url = "http://localhost:8000"))
+#    await node_manager.register_node(NodeRegistration(url = "http://localhost:8001"))
+    tasks = asyncio.create_task(node_manager.update_statuses_worker())
+    yield
+    # Shutdown code goes here
+
+app = FastAPI(
+    title="OmniBox Master Node",
+    description="Manages redirection of instance operations to worker nodes",
+    lifespan = lifespan)
+
+@app.post("/get")
+def create_instance():
+    """Get an available new instance from less occupied node"""
+    node = node_manager.get_best_node()
+    if node is None:
+        raise HTTPException(
+            status_code=status.HTTP_503_SERVICE_UNAVAILABLE,
+            detail="No available nodes with capacity to create new instance"
+        )
+
+    data = requests.post(f'{node.url}/get').json()
+    if 'instance_id' in data:
+        return {
+            'instance_id': data['instance_id'],
+            'node': node.hash,
+        }
+    raise HTTPException(
+        status_code=status.HTTP_503_SERVICE_UNAVAILABLE,
+        detail="No available nodes with capacity to create new instance"
+    )
+
+@app.post("/reset")
+async def reset(instance_id: str, node: str):
+    """Reset an existing instance by delegating to the worker node that hosts it"""
+    node_info = node_manager.get_node(node)
+    if node_info is None:
+        raise HTTPException(
+            status_code=status.HTTP_400_BAD_REQUEST,
+            detail=f"Node {node} is not found"
+        )
+    response = requests.post(f'{node_info.url}/reset', params={"instance_id": instance_id})
+    return JSONResponse(content=response.json(), status_code=response.status_code)
+
+
+@app.get("/probe")
+async def probe(instance_id: str, node: str):
+    """Probe the instance by delegating to the worker node that hosts it"""
+    node_info = node_manager.get_node(node)
+    if node_info is None:
+        raise HTTPException(
+            status_code=status.HTTP_400_BAD_REQUEST,
+            detail=f"Node {node} is not found"
+        )
+    response = requests.get(f'{node_info.url}/probe', params={"instance_id": instance_id})
+    return JSONResponse(content=response.json(), status_code=response.status_code)
+
+@app.get("/screenshot")
+async def screenshot(instance_id: str, node: str):
+    """Make a screenshot of an existing instance by delegating to the worker node that hosts it"""
+    node_info = node_manager.get_node(node)
+    if node_info is None:
+        raise HTTPException(
+            status_code=status.HTTP_400_BAD_REQUEST,
+            detail=f"Node {node} is not found"
+        )
+    response = requests.get(f'{node_info.url}/screenshot', params={"instance_id": instance_id})   
+    if response.status_code == 200:
+        return Response(content=response.content, media_type="image/png")
+
+    return JSONResponse(
+        content={"status": "error", "message": f"Failed to get screenshot: {response.text}"},
+        status_code=response.status_code
+    )
+
+@app.post("/execute")
+async def execute(command_data: Dict[str, Any]):
+    """Forward execute command to the Flask server in the specified instance"""
+    node = command_data.pop('node')
+    instance_id = command_data.pop('instance_id')
+    node_info = node_manager.get_node(node)
+    if node_info is None:
+        raise HTTPException(
+            status_code=status.HTTP_400_BAD_REQUEST,
+            detail=f"Node {node} is not found"
+        )
+    response = requests.post(f'{node_info.url}/execute',
+                            params={"instance_id": instance_id},
+                            json = command_data)   
+    if response.status_code == 200:
+        return JSONResponse(content=response.json(), status_code=response.status_code)
+
+    return JSONResponse(
+        content={"status": "error", "message": f"Failed to execute command: {response.text}"},
+        status_code=response.status_code
+    )
+
+@app.get("/info")
+def get_info():
+    node_info = node_manager.node_info()
+    return JSONResponse(
+        content={
+            "nodes": [
+                {
+                    "url": node.url,
+                    "hash": node.hash,
+                    "healthy": node.healthy,
+                    "capacity": node.capacity,
+                    "available": node.available,
+                    "instances": node.instances
+                }
+                for node in node_info.values()
+            ]
+        },
+        status_code=status.HTTP_200_OK
+    )
+
+
+if __name__ == "__main__":
+    parser = argparse.ArgumentParser(description="OmniBox Host")
+    parser.add_argument("--port", type=int, default=7000, help="Port to run the server on")
+    args = parser.parse_args()
+    uvicorn.run(app, host="0.0.0.0", port=args.port)