diff --git a/omnitool/omnibox/Dockerfile b/omnitool/omnibox/Dockerfile index b464714a..dc67536b 100644 --- a/omnitool/omnibox/Dockerfile +++ b/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" diff --git a/omnitool/omniboxes/client/client.py b/omnitool/omniboxes/client/client.py new file mode 100644 index 00000000..8c42b4ff --- /dev/null +++ b/omnitool/omniboxes/client/client.py @@ -0,0 +1,399 @@ +import requests +from PIL import Image +import io +import ipywidgets +import json +import time + +def _pyscript(commands): + script = ";".join(commands) + return f'python -c "{script}"' + +def _moveTo(x, y): + return { + 'command': _pyscript([ + 'import pyautogui', + f'pyautogui.moveTo({x}, {y})' + ]) + } + +def _click(): + return { + 'command': _pyscript([ + 'import pyautogui', + f'pyautogui.click()' + ]) + } + +def _rightClick(): + return { + 'command': _pyscript([ + 'import pyautogui', + f'pyautogui.rightClick()' + ]) + } + +def _doubleClick(): + return { + 'command': _pyscript([ + 'import pyautogui', + f'pyautogui.doubleClick()' + ]) + } + +def _position(ignore_by_mock = True): + return { + 'command': _pyscript([ + 'import pyautogui', + 'import json', + 'p = pyautogui.position()', + "print(json.dumps({'x': p.x, 'y': p.y}))", + ]), + 'ignore_by_mock': ignore_by_mock + } + +def _screensize(ignore_by_mock = True): + return { + 'command': _pyscript([ + 'import pyautogui', + 'import json', + 'sz = pyautogui.size()', + "print(json.dumps({'width': sz.width, 'height': sz.height}))", + ]), + 'ignore_by_mock': ignore_by_mock + } + + +class InstanceClient: + def __init__(self, host = 'localhost', port = 5000): + self.host = host + self.port = port + self.output = None + + def screenshot(self): + data = requests.get(f'http://{self.host}:{self.port}/screenshot') + image_data = io.BytesIO(data.content) + return Image.open(image_data) + + def execute(self, command): + return requests.post(f'http://{self.host}:{self.port}/execute', json = command) + + def do_and_show(self, command, waitTime): + response = self.execute(command) + if response.status_code != 200: + raise Exception(f"Error: {response.status_code} - {response.text}") + time.sleep(waitTime) + self._display(response.json()['output']) + + def _display(self, data = None): + with self.output: + self.output.clear_output(wait = True) + display(self.screenshot()) + print(data) + + def position(self): + return json.loads(self.execute(_position()).json()['output']) + + def screensize(self): + return json.loads(self.execute(_screensize()).json()['output']) + + def ui(self): + self.output = ipywidgets.Output() + self._display() + + screensize = self.screensize() + position = self.position() + + waitTime = ipywidgets.FloatLogSlider(base=2, value = 1, min = -3, max = 5, step = 1, description = 'wait (s)') + + x = ipywidgets.IntSlider(min = 0, max = screensize.get('width', 1), value = position.get('x', 0), description = 'X') + y = ipywidgets.IntSlider(min = 0, max = screensize.get('height', 1), value = position.get('y', 0), description = 'Y') + + click = ipywidgets.Button(description = 'Click') + rightClick = ipywidgets.Button(description = 'Right Click') + doubleClick = ipywidgets.Button(description = 'Double Click') + + x.observe(lambda v: self.do_and_show(_moveTo(x.value, y.value), waitTime.value), names='value') + y.observe(lambda v: self.do_and_show(_moveTo(x.value, y.value), waitTime.value), names='value') + + click.on_click(lambda v: self.do_and_show(_click(), waitTime.value)) + rightClick.on_click(lambda v: self.do_and_show(_rightClick(), waitTime.value)) + doubleClick.on_click(lambda v: self.do_and_show(_doubleClick(), waitTime.value)) + + cmd = ipywidgets.Textarea(description = 'Commands', layout=ipywidgets.Layout(width='50%')) + shell = ipywidgets.Checkbox(description = 'Shell', value = False) + python = ipywidgets.Checkbox(description = 'Python', value = True) + submit = ipywidgets.Button(description = 'Submit') + submit.on_click(lambda v: self.do_and_show({ + 'command': _pyscript(cmd.value.split('\n')) if python.value else '&&'.join(cmd.value.split('\n')), + 'shell': shell.value + }, waitTime.value)) + display(ipywidgets.VBox([ + waitTime, + ipywidgets.HBox([x, y]), + ipywidgets.HBox([click, rightClick, doubleClick]), + ipywidgets.HBox([cmd, ipywidgets.VBox([shell, python]), submit]), + self.output])) + + +class NodeClient: + def __init__(self, host = 'localhost', port = 8000): + self.host = host + self.port = port + self.output = None + self.instance_id = None + + def get_instance(self): + data = requests.post(f'http://{self.host}:{self.port}/get').json() + return data.get('instance_id', None) + + def get_instances_info(self): + data = requests.get(f'http://{self.host}:{self.port}/info') + return data.json() + + def reset_instance(self, instance_id): + data = requests.post(f'http://{self.host}:{self.port}/reset', params = {'instance_id': instance_id}) + if data.status_code != 200: + raise Exception(f"Error: {data.status_code} - {data.text}") + return data.json() + + def screenshot(self, instance_id): + data = requests.get(f'http://{self.host}:{self.port}/screenshot', params = {'instance_id': instance_id}) + image_data = io.BytesIO(data.content) + return Image.open(image_data) + + def execute(self, instance_id, command): + return requests.post( + f'http://{self.host}:{self.port}/execute', + params={'instance_id': instance_id}, + json = command) + + def do_and_show(self, instance_id, command, waitTime): + response = self.execute(instance_id, command) + if response.status_code != 200: + raise Exception(f"Error: {response.status_code} - {response.text}") + time.sleep(waitTime) + self._display(instance_id, response.json()['output']) + + def _display(self, instance_id, data = None): + with self.output: + self.output.clear_output(wait = True) + display(self.screenshot(instance_id)) + print(data) + + def position(self, instance_id): + return json.loads(self.execute(instance_id, _position()).json()['output']) + + def screensize(self, instance_id): + return json.loads(self.execute(instance_id, _screensize()).json()['output']) + + def ui(self): + self.output = ipywidgets.Output() + + get_button = ipywidgets.Button(description = 'Get Instance') + release_button = ipywidgets.Button(description = 'Release Instance') + instances = ipywidgets.RadioButtons(description = 'Instances', options = self.get_instances_info()['in_use'], layout=ipywidgets.Layout(width='50%')) + self.instance_id = instances.value + def on_add_click(b): + instance_id = self.get_instance() + if instance_id: + instances.options = list(instances.options) + [instance_id] + else: + print('No instance available') + + def on_release_click(b): + self.reset_instance(instances.value) + options = list(instances.options) + options.remove(instances.value) + instances.options = options + + get_button.on_click(on_add_click) + release_button.on_click(on_release_click) + + waitTime = ipywidgets.FloatLogSlider(base=2, value = 1, min = -3, max = 5, step = 1, description = 'wait (s)') + + screensize = {} + position = {} + if self.instance_id: + screensize = self.screensize(self.instance_id) + position = self.position(self.instance_id) + self._display(self.instance_id) + + x = ipywidgets.IntSlider(min = 0, max = screensize.get('width', 1), value = position.get('x', 0), description = 'X') + y = ipywidgets.IntSlider(min = 0, max = screensize.get('height', 1), value = position.get('y', 0), description = 'Y') + + click = ipywidgets.Button(description = 'Click') + rightClick = ipywidgets.Button(description = 'Right Click') + doubleClick = ipywidgets.Button(description = 'Double Click') + + x.observe(lambda v: self.do_and_show(self.instance_id, _moveTo(x.value, y.value), waitTime.value), names='value') + y.observe(lambda v: self.do_and_show(self.instance_id, _moveTo(x.value, y.value), waitTime.value), names='value') + + click.on_click(lambda v: self.do_and_show(self.instance_id, _click(), waitTime.value)) + rightClick.on_click(lambda v: self.do_and_show(self.instance_id, _rightClick(), waitTime.value)) + doubleClick.on_click(lambda v: self.do_and_show(self.instance_id, _doubleClick(), waitTime.value)) + + cmd = ipywidgets.Textarea(description = 'Commands', layout=ipywidgets.Layout(width='50%')) + shell = ipywidgets.Checkbox(description = 'Shell', value = False) + python = ipywidgets.Checkbox(description = 'Python', value = True) + submit = ipywidgets.Button(description = 'Submit') + submit.on_click(lambda v: self.do_and_show(self.instance_id, { + 'command': _pyscript(cmd.value.split('\n')) if python.value else '&&'.join(cmd.value.split('\n')), + 'shell': shell.value + }, waitTime.value)) + + def on_instance_change(_): + self.instance_id = instances.value + screensize = self.screensize(self.instance_id) + x.max = screensize.get('width', 1) + y.max = screensize.get('height', 1) + position = self.position(self.instance_id) + x.value = position.get('x', 0) + y.value = position.get('y', 0) + self._display(self.instance_id) + + instances.observe(on_instance_change, names='value') + + display(ipywidgets.VBox([ + ipywidgets.HBox([instances, ipywidgets.VBox([get_button, release_button])]), + waitTime, + ipywidgets.HBox([x, y]), + ipywidgets.HBox([click, rightClick, doubleClick]), + ipywidgets.HBox([cmd, ipywidgets.VBox([shell, python]), submit]), + self.output])) + + +class MasterClient: + def __init__(self, host = 'localhost', port = 7000): + self.host = host + self.port = port + self.output = None + self.instance = None + + def probe(self, instance: dict): + return requests.get(f'http://{self.host}:{self.port}/probe', params=instance).json() + + def get_instance(self): + url = f'http://{self.host}:{self.port}/get' + response = requests.post(url) + if response.status_code != 200: + raise Exception(f"Error: {response.status_code} - {response.text}") + return response.json() + + def get_info(self): + url = f'http://{self.host}:{self.port}/info' + response = requests.get(url) + if response.status_code != 200: + raise Exception(f"Error: {response.status_code} - {response.text}") + return response.json() + + def get_options(self): + info = self.get_info()['nodes'] + print(type(info)) + return [json.dumps({'instance_id': instance, 'node': node['hash']}) for node in info for instance in node['instances']] + + def reset_instance(self, instance: dict): + data = requests.post(f'http://{self.host}:{self.port}/reset', params = instance) + if data.status_code != 200: + raise Exception(f"Error: {data.status_code} - {data.text}") + return data.json() + + def screenshot(self, instance: dict): + data = requests.get(f'http://{self.host}:{self.port}/screenshot', params = instance) + image_data = io.BytesIO(data.content) + return Image.open(image_data) + + def execute(self, instance: dict, command): + return requests.post(f'http://{self.host}:{self.port}/execute', json = dict(command, **instance)) + + def do_and_show(self, instance, command, waitTime): + response = self.execute(instance, command) + if response.status_code != 200: + raise Exception(f"Error: {response.status_code} - {response.text}") + time.sleep(waitTime) + self._display(instance, response.json()['output']) + + def _display(self, instance, data = None): + with self.output: + self.output.clear_output(wait = True) + display(self.screenshot(instance)) + print(data) + + def position(self, instance): + return json.loads(self.execute(instance, _position()).json()['output']) + + def screensize(self, instance): + return json.loads(self.execute(instance, _screensize()).json()['output']) + + def ui(self): + self.output = ipywidgets.Output() + + get_button = ipywidgets.Button(description = 'Get Instance') + release_button = ipywidgets.Button(description = 'Release Instance') + instances = ipywidgets.RadioButtons(description = 'Instances', options = self.get_options(), layout=ipywidgets.Layout(width='50%')) + self.instance = json.loads(instances.value) if instances.value else {} + def on_add_click(b): + new_instance = json.dumps(self.get_instance()) + instances.options = list(instances.options) + [new_instance] + + def on_release_click(b): + self.reset_instance(json.loads(instances.value)) + options = list(instances.options) + options.remove(instances.value) + instances.options = options + + get_button.on_click(on_add_click) + release_button.on_click(on_release_click) + + waitTime = ipywidgets.FloatLogSlider(base=2, value = 1, min = -3, max = 5, step = 1, description = 'wait (s)') + + screensize = {} + position = {} + if self.instance: + screensize = self.screensize(self.instance) + position = self.position(self.instance) + self._display(self.instance) + + x = ipywidgets.IntSlider(min = 0, max = screensize.get('width', 1), value = position.get('x', 0), description = 'X') + y = ipywidgets.IntSlider(min = 0, max = screensize.get('height', 1), value = position.get('y', 0), description = 'Y') + + click = ipywidgets.Button(description = 'Click') + rightClick = ipywidgets.Button(description = 'Right Click') + doubleClick = ipywidgets.Button(description = 'Double Click') + + x.observe(lambda v: self.do_and_show(self.instance, _moveTo(x.value, y.value), waitTime.value), names='value') + y.observe(lambda v: self.do_and_show(self.instance, _moveTo(x.value, y.value), waitTime.value), names='value') + + click.on_click(lambda v: self.do_and_show(self.instance, _click(), waitTime.value)) + rightClick.on_click(lambda v: self.do_and_show(self.instance, _rightClick(), waitTime.value)) + doubleClick.on_click(lambda v: self.do_and_show(self.instance, _doubleClick(), waitTime.value)) + + cmd = ipywidgets.Textarea(description = 'Commands', layout=ipywidgets.Layout(width='50%')) + shell = ipywidgets.Checkbox(description = 'Shell', value = False) + python = ipywidgets.Checkbox(description = 'Python', value = True) + submit = ipywidgets.Button(description = 'Submit') + submit.on_click(lambda v: self.do_and_show(self.instance, { + 'command': _pyscript(cmd.value.split('\n')) if python.value else '&&'.join(cmd.value.split('\n')), + 'shell': shell.value + }, waitTime.value)) + + def on_instance_change(_): + self.instance = json.loads(instances.value) + screensize = self.screensize(self.instance) + x.max = screensize.get('width', 1) + y.max = screensize.get('height', 1) + position = self.position(self.instance) + x.value = position.get('x', 0) + y.value = position.get('y', 0) + self._display(self.instance) + + instances.observe(on_instance_change, names='value') + + display(ipywidgets.VBox([ + ipywidgets.HBox([instances, ipywidgets.VBox([get_button, release_button])]), + waitTime, + ipywidgets.HBox([x, y]), + ipywidgets.HBox([click, rightClick, doubleClick]), + ipywidgets.HBox([cmd, ipywidgets.VBox([shell, python]), submit]), + self.output])) + diff --git a/omnitool/omniboxes/client/notebook.ipynb b/omnitool/omniboxes/client/notebook.ipynb new file mode 100644 index 00000000..972a3832 --- /dev/null +++ b/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 +} diff --git a/omnitool/omniboxes/client/requirements.txt b/omnitool/omniboxes/client/requirements.txt new file mode 100644 index 00000000..dbfd4e08 --- /dev/null +++ b/omnitool/omniboxes/client/requirements.txt @@ -0,0 +1,3 @@ +pillow +ipywidgets +requests \ No newline at end of file diff --git a/omnitool/omniboxes/master/__init__.py b/omnitool/omniboxes/master/__init__.py new file mode 100644 index 00000000..e69de29b diff --git a/omnitool/omniboxes/master/logging_utils.py b/omnitool/omniboxes/master/logging_utils.py new file mode 100644 index 00000000..ee77ff99 --- /dev/null +++ b/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") \ No newline at end of file diff --git a/omnitool/omniboxes/master/node_manager.py b/omnitool/omniboxes/master/node_manager.py new file mode 100644 index 00000000..d4030320 --- /dev/null +++ b/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 + diff --git a/omnitool/omniboxes/master/server.py b/omnitool/omniboxes/master/server.py new file mode 100644 index 00000000..aec69e07 --- /dev/null +++ b/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) \ No newline at end of file diff --git a/omnitool/omniboxes/node/__init__.py b/omnitool/omniboxes/node/__init__.py new file mode 100644 index 00000000..e69de29b diff --git a/omnitool/omniboxes/node/instance.py b/omnitool/omniboxes/node/instance.py new file mode 100644 index 00000000..c5692bab --- /dev/null +++ b/omnitool/omniboxes/node/instance.py @@ -0,0 +1,164 @@ +import os +import subprocess +import time +import requests +from pathlib import Path +from typing import Dict, Any +from logging_utils import default_logger + + +class IInstance: + def create(self): + raise NotImplementedError("Subclasses should implement this!") + + def start(self): + raise NotImplementedError("Subclasses should implement this!") + + def stop(self): + raise NotImplementedError("Subclasses should implement this!") + + def delete(self): + raise NotImplementedError("Subclasses should implement this!") + + def flask_url(self): + raise NotImplementedError("Subclasses should implement this!") + + def is_ready(self): + raise NotImplementedError("Subclasses should implement this!") + + def reset(self): + raise NotImplementedError("Subclasses should implement this!") + + def reset_soft(self): + raise NotImplementedError("Subclasses should implement this!") + + +_compose_template = """ +networks: + omnibox-network: + name: omnibox-network-{instance} + +services: + windows: + image: windows-local + container_name: omni-windows-{instance} + networks: + - omnibox-network + privileged: true + environment: + RAM_SIZE: "2G" + CPU_CORES: "4" + DISK_SIZE: "11G" + devices: + - /dev/kvm + - /dev/net/tun + cap_add: + - NET_ADMIN + ports: + - {web_port}:8006 # Web Viewer access + - {control_port}:5000 # Computer control server + volumes: + - {omniboxes_path}/common/win11iso/custom.iso:/custom.iso + - {omniboxes_path}/common/win11setup/firstboot:/oem + - {omniboxes_path}/common/win11setup/setupscripts:/data + - {omniboxes_path}/omnibox-{instance}:/storage +""" + + +class Instance(IInstance): + def __init__(self, root_path = None, instance_num = 0, logger = None): + self.root_path = Path(root_path or os.path.dirname(__file__)).resolve() + self.instance_num = instance_num + self.config_path = self.root_path / f'{instance_num}.yml' + self.logger = logger or default_logger() + with open(self.config_path, mode='w') as temp_file: + compose_content = _compose_template.format( + instance = self.instance_num, + web_port = 8006 + self.instance_num, + control_port = 5000 + self.instance_num, + omniboxes_path = self.root_path + ) + temp_file.write(compose_content) + + def path(self): + return f"{self.root_path}/omnibox-{self.instance_num}/" + + def _execute(self, command): + self.logger.info(f'Running: {" ".join(command)}') + subprocess.run(command, check=True, stdout=subprocess.DEVNULL, stderr=subprocess.DEVNULL) + + def create(self): + if not os.path.exists(self.path()): + os.makedirs(f"omnibox-{self.instance_num}") + subprocess.run(["cp", "-r", f"{str(self.root_path)}/common/win11storage/.", self.path()]) + + try: + self._execute(["docker", "compose", "-f", str(self.config_path), "-p", f"omnibox-{self.instance_num}", "up", "-d"]) + self.logger.info(f"Instance {self.instance_num} launched successfully!") + except subprocess.CalledProcessError as e: + self.logger.error(f"Error launching instance {self.instance_num}: {e}") + + def start(self): + try: + self._execute(["docker", "compose", "-f", str(self.config_path), "-p", f"omnibox-{self.instance_num}", "start"]) + self.logger.info(f"Instance {self.instance_num} started successfully!") + except subprocess.CalledProcessError as e: + self.logger.error(f"Error starting instance {self.instance_num}: {e}") + + def stop(self): + try: + self._execute(["docker", "compose", "-f", str(self.config_path), "-p", f"omnibox-{self.instance_num}", "stop"]) + self.logger.info(f"Instance {self.instance_num} stopped successfully!") + except subprocess.CalledProcessError as e: + self.logger.error(f"Error stopping instance {self.instance_num}: {e}") + + def delete(self): + try: + self._execute(["docker", "compose", "-f", str(self.config_path), "-p", f"omnibox-{self.instance_num}", "down"]) + self.logger.info(f"Instance {self.instance_num} deleted successfully!") + except subprocess.CalledProcessError as e: + self.logger.error(f"Error deleting instance {self.instance_num}: {e}") + + def flask_url(self): + return f"http://localhost:{5000 + self.instance_num}" + + def is_ready(self): + try: + response = requests.get(self.flask_url() + "/probe", timeout=1) + return response.status_code == 200 + except requests.RequestException as e: + return False + + def reset(self): + self.delete() + if os.path.exists(self.path()): + subprocess.run(["sudo", "rm", "-rf", f"{self.root_path}/omnibox-{self.instance_num}"], check=True) + self.create() + + def reset_soft(self): + self.stop() + if os.path.exists(self.path()): + subprocess.run(["sudo", "rm", "-rf", f"{self.root_path}/omnibox-{self.instance_num}"], check=True) + self.start() + + +def reset_with_callback(instance, callback): + try: + instance.reset() + while not instance.is_ready(): + time.sleep(1) + callback(instance) + except Exception as e: + instance.logger.error(f"Error in reset worker: {str(e)}") + + +# got insufficient perf with this method +def reset_soft_with_callback(instance, callback): + try: + instance.reset_soft() + while not instance.is_ready(): + time.sleep(1) + callback(instance) + except Exception as e: + instance.logger.error(f"Error in reset worker: {str(e)}") + diff --git a/omnitool/omniboxes/node/instance_client.py b/omnitool/omniboxes/node/instance_client.py new file mode 100644 index 00000000..a8e94c2d --- /dev/null +++ b/omnitool/omniboxes/node/instance_client.py @@ -0,0 +1,39 @@ +from typing import Dict, Any +from fastapi.responses import JSONResponse +from fastapi import FastAPI, HTTPException, Response +import requests + +class InstanceClient: + def __init__(self, url): + self.url = url + + def execute(self, command_data: Dict[str, Any]): + """Forward execute command to the Flask server in the specified instance""" + try: + response = requests.post(f"{self.url}/execute", json=command_data, timeout=5) + return JSONResponse(content=response.json(), status_code=response.status_code) + except requests.RequestException as e: + raise HTTPException(status_code=500, detail=f"Error communicating with theinstance {self.url}: {str(e)}") + + def screenshot(self): + """Forward screenshot request to the Flask server in the specified instance""" + try: + response = requests.get(f"{self.url}/screenshot", timeout=5) + + if response.status_code == 200: + return Response(content=response.content, media_type="image/png") + else: + return JSONResponse( + content={"status": "error", "message": f"Failed to get screenshot: {response.text}"}, + status_code=response.status_code + ) + except requests.RequestException as e: + raise HTTPException(status_code=500, detail=f"Error communicating with the instance {self.url}: {str(e)}") + + def probe(self): + """Forward probe request to the Flask server in the specified instance""" + try: + response = requests.get(f"{self.url}/probe", timeout=1) + return JSONResponse(content=response.json(), status_code=response.status_code) + except requests.RequestException as e: + raise HTTPException(status_code=500, detail=f"Error communicating with the instance {self.url}: {str(e)}") diff --git a/omnitool/omniboxes/node/instance_manager.py b/omnitool/omniboxes/node/instance_manager.py new file mode 100644 index 00000000..33e08832 --- /dev/null +++ b/omnitool/omniboxes/node/instance_manager.py @@ -0,0 +1,65 @@ +import concurrent +from instance import Instance, reset_with_callback +from tqdm import tqdm +import time +import uuid +from typing import Dict, Any +from pathlib import Path +import os +from logging_utils import default_logger + +class InstanceManager: + def __init__(self, instance_factory = None, path = None, capacity: int = 2, logger = None): + self.capacity = capacity + self.instance_factory = instance_factory or Instance + self.available_instances = {} # instance_num to instance + self.in_use = {} # key = instance_uuid, value = instance_id + self.logger = logger or default_logger() + self.reset_workers = capacity + self.reset_executor = concurrent.futures.ThreadPoolExecutor(max_workers=self.reset_workers) + self.path = Path(path or os.path.dirname(__file__)).resolve() + + for i in range(self.capacity): + self.reset_executor.submit(reset_with_callback, self.instance_factory(self.path, instance_num=i, logger=self.logger), self.instance_reset_callback) + + # Wait for all instances to be initialized + with tqdm(total=self.capacity, desc="Initializing instances") as pbar: + last_count = 0 + while len(self.available_instances) < self.capacity: + current_count = len(self.available_instances) + if current_count > last_count: + pbar.update(current_count - last_count) + last_count = current_count + time.sleep(0.1) + pbar.update(self.capacity - last_count) + + def instance_reset_callback(self, instance): + self.logger.info(f"Instance {instance.instance_num} is ready") + self.available_instances[instance.instance_num] = instance + + def shutdown(self): + if self.reset_executor: + self.reset_executor.shutdown(wait=True, cancel_futures=False) + + def start(self): + if not self.available_instances: + return None + + instance = self.available_instances.pop(list(self.available_instances.keys())[0]) + instance_uuid = str(uuid.uuid4()) + self.in_use[instance_uuid] = instance + return instance_uuid + + def reset(self, instance_uuid: str): + self.logger.info(f"Resetting instance {instance_uuid}") + if instance_uuid not in self.in_use: + return False + + instance = self.in_use.pop(instance_uuid) + self.logger.info(f"Resetting instance {instance_uuid}: {instance.instance_num} ") + self.reset_executor.submit(reset_with_callback, instance, self.instance_reset_callback) + return True + + def get(self, uuid: str): + return self.in_use.get(uuid, None) + \ No newline at end of file diff --git a/omnitool/omniboxes/node/logging_utils.py b/omnitool/omniboxes/node/logging_utils.py new file mode 100644 index 00000000..2008f488 --- /dev/null +++ b/omnitool/omniboxes/node/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-node") \ No newline at end of file diff --git a/omnitool/omniboxes/node/mock_instance.py b/omnitool/omniboxes/node/mock_instance.py new file mode 100644 index 00000000..46c532ae --- /dev/null +++ b/omnitool/omniboxes/node/mock_instance.py @@ -0,0 +1,116 @@ +import logging +from flask import Flask, request, jsonify, send_file +from PIL import Image, ImageDraw, ImageFont +from io import BytesIO +from instance import IInstance +from threading import Thread +from logging_utils import default_logger + + +class MockApp: + def __init__(self, log_file=None): + """Create and configure the Flask app""" + if log_file: + logging.basicConfig(filename=log_file, level=logging.DEBUG, filemode='w') + logger = logging.getLogger('werkzeug') + + self.app = Flask(__name__) + self.command_history = [] + + @self.app.route('/probe', methods=['GET']) + def probe_endpoint(): + return jsonify({"status": "Probe successful", "message": "Service is operational"}), 200 + + @self.app.route('/execute', methods=['POST']) + def execute_command(): + data = request.json + command = data.get('command', "") + if not data.get('ignore_by_mock', False): + self.command_history.append(str(command)) + return jsonify({ + 'status': 'success', + 'output': '{}', + 'error': '', + 'returncode': 0 + }) + + @self.app.route('/screenshot', methods=['GET']) + def screenshot(): + width, height = 800, 600 + image = Image.new('RGB', (width, height), color='white') + draw = ImageDraw.Draw(image) + + try: + font = ImageFont.truetype("DejaVuSans.ttf", 14) + except IOError: + font = ImageFont.load_default() + + draw.text((10, 10), "Mock Instance - Command History", fill="black", font=font) + draw.line([(10, 40), (width-10, 40)], fill="black", width=1) + + y_position = 50 + max_commands = min(len(self.command_history), 20) # Limit to last 20 commands + + if max_commands == 0: + draw.text((20, y_position), "No commands executed yet", fill="gray", font=font) + else: + for i, cmd in enumerate(self.command_history[-max_commands:]): + timestamp = f"{i+1}." + draw.text((20, y_position), timestamp, fill="blue", font=font) + draw.text((50, y_position), cmd, fill="black", font=font) + y_position += 25 + + # If we're running out of space + if y_position > height - 30: + remaining = len(self.command_history) - (i + 1) + if remaining > 0: + draw.text((20, y_position), f"... {remaining} more commands not shown", + fill="red", font=font) + break + + img_io = BytesIO() + image.save(img_io, 'PNG') + img_io.seek(0) + return send_file(img_io, mimetype='image/png') + + + +class MockInstance(IInstance): + def __init__(self, root_path = None, instance_num = 0, logger = None, base_control_port = 5000): + print(f'Creating mock instance: {instance_num}') + self.root_path = root_path + self.instance_num = instance_num + self.app = MockApp() + self.logger = logger or default_logger() + self.base_control_port = base_control_port + p = Thread(target=self.app.app.run, args=('0.0.0.0', self.control_port), daemon=True) + p.start() + + @property + def control_port(self): + return self.base_control_port + self.instance_num + + def create(self): + self.logger.info(f"Creating dummy instance {self.instance_num}") + + def start(self): + self.logger.info(f"Starting dummy instance {self.instance_num}") + + def stop(self): + self.logger.info(f"Stopping dummy instance {self.instance_num}") + + def delete(self): + self.logger.info(f"Deleting dummy instance {self.instance_num}") + + def flask_url(self): + return f"http://localhost:{self.control_port}" + + def is_ready(self): + return True + + def reset(self): + self.app.command_history = [] + self.logger.info(f"Resetting dummy instance {self.instance_num}") + + def reset_soft(self): + self.logger.info(f"Soft resetting dummy instance {self.instance_num}") diff --git a/omnitool/omniboxes/node/requirements.txt b/omnitool/omniboxes/node/requirements.txt new file mode 100644 index 00000000..4f1aa032 --- /dev/null +++ b/omnitool/omniboxes/node/requirements.txt @@ -0,0 +1,4 @@ +fastapi[standard] +requests +tqdm +flask diff --git a/omnitool/omniboxes/node/server.py b/omnitool/omniboxes/node/server.py new file mode 100644 index 00000000..c8bc2e09 --- /dev/null +++ b/omnitool/omniboxes/node/server.py @@ -0,0 +1,93 @@ +from fastapi import FastAPI, HTTPException +import requests +import uvicorn +from typing import Dict, Any +from contextlib import asynccontextmanager +import argparse +from instance_manager import InstanceManager +from instance_client import InstanceClient +from instance import Instance +from mock_instance import MockInstance +from pathlib import Path + +parser = argparse.ArgumentParser(description="OmniBox Host") +parser.add_argument("--port", type=int, default=8000, help="Port to run the server on") +parser.add_argument('--path', type=str, default='../run', help="Path to the instance directory. Expected to contain prepared common subfolder") +parser.add_argument('--base_control_port', type=int, default=5000, help="Base control port offset for the instances (for testing with mock instances)") +parser.add_argument('--mock', action='store_true', help="Use mock instances") +args = parser.parse_args() + +Path(args.path).mkdir(parents=True, exist_ok=True) +mock_instance_factory = lambda root_path, instance_num, logger: MockInstance(root_path, instance_num, logger, args.base_control_port) +instance_factory = mock_instance_factory if args.mock else Instance +instance_manager = InstanceManager(instance_factory = instance_factory, path = args.path) + +@asynccontextmanager +async def lifespan(app: FastAPI): + # Startup code goes here + yield + instance_manager.shutdown() + + +app = FastAPI(lifespan=lifespan) + +@app.post("/get") +def get_instance(): + """Get an available instance from the pool""" + instance_id = instance_manager.start() + if not instance_id: + raise HTTPException(status_code=503, detail="No instances available") + return {"instance_id": instance_id} + +@app.post("/reset") +def reset_instance(instance_id: str): + """Reset an instance to its initial state and make it available again""" + if instance_manager.reset(instance_id): + return {"status": "success", "message": f"UUID {instance_id} for instance has been queued for reset"} + raise HTTPException(status_code=400, detail=f"Invalid instance UUID: {instance_id}") + +@app.get("/probe") +async def probe_instance(instance_id: str): + """Forward probe request to the Flask server in the specified instance""" + instance= instance_manager.get(instance_id) + if not instance: + raise HTTPException(status_code=400, detail=f"Invalid instance UUID: {instance_id}") + try: + return InstanceClient(instance.flask_url()).probe() + except requests.RequestException as e: + raise HTTPException(status_code=500, detail=f"Error communicating with UUID {instance_id} for instance {instance.instance_num}: {str(e)}") + +@app.get("/screenshot") +async def get_instance_screenshot(instance_id: str): + """Forward screenshot request to the Flask server in the specified instance""" + instance = instance_manager.get(instance_id) + if not instance: + raise HTTPException(status_code=400, detail=f"Invalid instance UUID: {instance_id}") + try: + return InstanceClient(instance.flask_url()).screenshot() + except requests.RequestException as e: + raise HTTPException(status_code=500, detail=f"Error communicating with UUID {instance_id} for instance {instance.instance_num}: {str(e)}") + +@app.post("/execute") +async def execute_instance_command(instance_id: str, command_data: Dict[str, Any]): + """Forward execute command to the Flask server in the specified instance""" + instance= instance_manager.get(instance_id) + if not instance: + raise HTTPException(status_code=400, detail=f"Invalid instance UUID: {instance_id}") + try: + return InstanceClient(instance.flask_url()).execute(command_data) + except requests.RequestException as e: + raise HTTPException(status_code=500, detail=f"Error communicating with instance {instance.instance_num}: {str(e)}") + +@app.get("/info") +def get_available_instances(): + """Get an available instance from the pool""" + return { + 'available': len(instance_manager.available_instances), + 'capacity': instance_manager.capacity, + 'in_use': list(instance_manager.in_use.keys())} + + +if __name__ == "__main__": + uvicorn.run(app, host="0.0.0.0", port=args.port) + diff --git a/omnitool/omniboxes/v0/manageall.py b/omnitool/omniboxes/v0/manageall.py new file mode 100644 index 00000000..4ebf8e86 --- /dev/null +++ b/omnitool/omniboxes/v0/manageall.py @@ -0,0 +1,121 @@ +import os +from fastapi import FastAPI, HTTPException, Response +from fastapi.responses import JSONResponse, StreamingResponse +import requests +import uvicorn +import concurrent.futures +from typing import List, Dict, Any, Optional +from manageinstance import ( + reset_instance_with_callback, flask_url +) +import time +import uuid +from tqdm import tqdm +from contextlib import asynccontextmanager + +@asynccontextmanager +async def lifespan(app: FastAPI): + # Startup code goes here + yield + # Shutdown code goes here + if reset_executor: + reset_executor.shutdown(wait=True, cancel_futures=False) + +app = FastAPI(lifespan=lifespan) + +INSTANCES = 10 +available_instances = [] +in_use = {} # key = instance_uuid, value = instance_id + +RESET_WORKERS = 10 +reset_executor = concurrent.futures.ThreadPoolExecutor(max_workers=RESET_WORKERS) + +def instance_reset_callback(instance_id): + if instance_id not in available_instances: + available_instances.append(instance_id) + +for i in range(INSTANCES): + reset_executor.submit(reset_instance_with_callback, i, instance_reset_callback) + +# Wait for all instances to be initialized +with tqdm(total=INSTANCES, desc="Initializing instances") as pbar: + last_count = 0 + while len(available_instances) < INSTANCES: + current_count = len(available_instances) + if current_count > last_count: + pbar.update(current_count - last_count) + last_count = current_count + time.sleep(0.1) + pbar.update(INSTANCES - last_count) + +@app.get("/getinstance") +def get_instance(): + """Get an available instance from the pool""" + if not available_instances: + raise HTTPException(status_code=503, detail="No instances available") + + instance_id = available_instances.pop(0) + instance_uuid = str(uuid.uuid4()) + in_use[instance_uuid] = instance_id + return {"instance_uuid": instance_uuid} + +@app.post("/resetinstance/{instance_uuid}") +def reset_instance(instance_uuid: str): + """Reset an instance to its initial state and make it available again""" + if instance_uuid not in in_use: + raise HTTPException(status_code=400, detail=f"Invalid instance UUID: {instance_uuid}") + + instance_id = in_use[instance_uuid] + reset_executor.submit(reset_instance_with_callback, instance_id, instance_reset_callback) + + return {"status": "success", "message": f"UUID {instance_uuid} for instance {instance_id} has been queued for reset"} + +@app.post("/executeinstance/{instance_uuid}/execute") +async def execute_instance_command(instance_uuid: str, command_data: Dict[str, Any]): + """Forward execute command to the Flask server in the specified instance""" + if instance_uuid not in in_use: + raise HTTPException(status_code=400, detail=f"Invalid instance UUID: {instance_uuid}") + + instance_id = in_use[instance_uuid] + try: + response = requests.post(flask_url(instance_id), json=command_data, timeout=5) + return JSONResponse(content=response.json(), status_code=response.status_code) + except requests.RequestException as e: + raise HTTPException(status_code=500, detail=f"Error communicating with instance {instance_id}: {str(e)}") + +@app.get("/executeinstance/{instance_uuid}/screenshot") +async def get_instance_screenshot(instance_uuid: str): + """Forward screenshot request to the Flask server in the specified instance""" + if instance_uuid not in in_use: + raise HTTPException(status_code=400, detail=f"Invalid instance UUID: {instance_uuid}") + + instance_id = in_use[instance_uuid] + try: + response = requests.get(flask_url(instance_id) + "/screenshot", timeout=5) + + if response.status_code == 200: + return Response(content=response.content, media_type="image/png") + else: + return JSONResponse( + content={"status": "error", "message": f"Failed to get screenshot: {response.text}"}, + status_code=response.status_code + ) + except requests.RequestException as e: + raise HTTPException(status_code=500, detail=f"Error communicating with UUID {instance_uuid} for instance {instance_id}: {str(e)}") + +@app.get("/executeinstance/{instance_uuid}/probe") +async def probe_instance(instance_uuid: str): + """Forward probe request to the Flask server in the specified instance""" + if instance_uuid not in in_use: + raise HTTPException(status_code=400, detail=f"Invalid instance UUID: {instance_uuid}") + + instance_id = in_use[instance_uuid] + try: + response = requests.get(flask_url(instance_id) + "/probe", timeout=1) + return JSONResponse(content=response.json(), status_code=response.status_code) + except requests.RequestException as e: + raise HTTPException(status_code=500, detail=f"Error communicating with UUID {instance_uuid} for instance {instance_id}: {str(e)}") + +if __name__ == "__main__": + uvicorn.run(app, host="0.0.0.0", port=8000) + diff --git a/omnitool/omniboxes/v0/manageinstance.py b/omnitool/omniboxes/v0/manageinstance.py new file mode 100644 index 00000000..859f3205 --- /dev/null +++ b/omnitool/omniboxes/v0/manageinstance.py @@ -0,0 +1,164 @@ +import os +import subprocess +import tempfile +import threading +import queue +import shutil +import time +import requests +import concurrent.futures + +# Base compose file template +compose_template = """ +networks: + omnibox-network: + name: omnibox-network-{instance} + +services: + windows: + image: windows-local + container_name: omni-windows-{instance} + networks: + - omnibox-network + privileged: true + environment: + RAM_SIZE: "2G" + CPU_CORES: "4" + DISK_SIZE: "11G" + devices: + - /dev/kvm + - /dev/net/tun + cap_add: + - NET_ADMIN + ports: + - {web_port}:8006 # Web Viewer access + - {control_port}:5000 # Computer control server + volumes: + - {omniboxes_path}/common/win11iso/custom.iso:/custom.iso + - {omniboxes_path}/common/win11setup/firstboot:/oem + - {omniboxes_path}/common/win11setup/setupscripts:/data + - {omniboxes_path}/omnibox-{instance}:/storage +""" +omniboxes_path = os.path.dirname(__file__) +def get_compose_path(instance_num): + with tempfile.NamedTemporaryFile(mode='w', suffix='.yml', delete=False) as temp_file: + compose_content = compose_template.format( + instance = instance_num, + web_port = 8006 + instance_num, + control_port = 5000 + instance_num, + omniboxes_path = omniboxes_path + ) + temp_file.write(compose_content) + return temp_file.name + +def create_instance(instance_num): + if not os.path.exists(f"omnibox-{instance_num}"): + os.makedirs(f"omnibox-{instance_num}") + subprocess.run(["cp", "-r", f"{omniboxes_path}/common/win11storage/.", f"{omniboxes_path}/omnibox-{instance_num}/"]) + + temp_file_path = get_compose_path(instance_num) + try: + subprocess.run( + ["docker", "compose", "-f", temp_file_path, "-p", f"omnibox-{instance_num}", "up", "-d"], + check=True, + stdout=subprocess.DEVNULL, # Suppress standard output + stderr=subprocess.DEVNULL # Suppress error output + ) + print(f"Instance {instance_num} launched successfully!") + except subprocess.CalledProcessError as e: + print(f"Error launching instance {instance_num}: {e}") + finally: + os.unlink(temp_file_path) # Clean up the temporary file + +def start_instance(instance_num): + temp_file_path = get_compose_path(instance_num) + try: + subprocess.run( + ["docker", "compose", "-f", temp_file_path, "-p", f"omnibox-{instance_num}", "start"], + check=True, + stdout=subprocess.DEVNULL, + stderr=subprocess.DEVNULL + ) + print(f"Instance {instance_num} started successfully!") + except subprocess.CalledProcessError as e: + print(f"Error starting instance {instance_num}: {e}") + finally: + os.unlink(temp_file_path) + +def stop_instance(instance_num): + temp_file_path = get_compose_path(instance_num) + try: + subprocess.run( + ["docker", "compose", "-f", temp_file_path, "-p", f"omnibox-{instance_num}", "stop"], + check=True, + stdout=subprocess.DEVNULL, + stderr=subprocess.DEVNULL + ) + print(f"Instance {instance_num} stopped successfully!") + except subprocess.CalledProcessError as e: + print(f"Error stopping instance {instance_num}: {e}") + finally: + os.unlink(temp_file_path) + +def delete_instance(instance_num): + temp_file_path = get_compose_path(instance_num) + try: + subprocess.run( + ["docker", "compose", "-f", temp_file_path, "-p", f"omnibox-{instance_num}", "down"], + check=True, + stdout=subprocess.DEVNULL, + stderr=subprocess.DEVNULL + ) + print(f"Instance {instance_num} deleted successfully!") + except subprocess.CalledProcessError as e: + print(f"Error deleting instance {instance_num}: {e}") + finally: + os.unlink(temp_file_path) + +def flask_url(instance_id: int): + return f"http://localhost:{5000 + instance_id}" + +def instance_ready(instance_id: int): + try: + response = requests.get(flask_url(instance_id) + "/probe", timeout=1) + return response.status_code == 200 + except requests.RequestException as e: + return False + +def reset_instance(instance_id): + delete_instance(instance_id) + if os.path.exists(f"omnibox-{instance_id}"): + subprocess.run(["sudo", "rm", "-rf", f"{omniboxes_path}/omnibox-{instance_id}"], check=True) + create_instance(instance_id) + +def reset_instance_soft(instance_id): + stop_instance(instance_id) + if os.path.exists(f"omnibox-{instance_id}"): + subprocess.run(["sudo", "rm", "-rf", f"{omniboxes_path}/omnibox-{instance_id}"], check=True) + start_instance(instance_id) + +def reset_instance_with_callback(instance_id, callback): + try: + reset_instance(instance_id) + while not instance_ready(instance_id): + time.sleep(1) + callback(instance_id) + except Exception as e: + print(f"Error in reset worker: {str(e)}") + +# got insufficient perf with this method +def reset_instance_soft_with_callback(instance_id, callback): + try: + reset_instance_soft(instance_id) + while not instance_ready(instance_id): + time.sleep(1) + callback(instance_id) + except Exception as e: + print(f"Error in reset worker: {str(e)}") + +# num_instances = 3 +# for i in range(num_instances): + # create_instance(i) + # stop_instance(i) + # start_instance(i) + # delete_instance(i)