cube-vm-backend 0.1.0__tar.gz

cube_vm_backend-0.1.0/PKG-INFO

@@ -0,0 +1,10 @@
+Metadata-Version: 2.3
+Name: cube-vm-backend
+Version: 0.1.0
+Summary: VM backend implementations for CUBE desktop-automation benchmarks
+Requires-Dist: cube-standard
+Requires-Dist: requests>=2.28
+Requires-Dist: tqdm>=4.60
+Requires-Dist: pytest>=8.0.0 ; extra == 'dev'
+Requires-Python: >=3.12
+Provides-Extra: dev

cube_vm_backend-0.1.0/pyproject.toml

@@ -0,0 +1,36 @@
+[project]
+name = "cube-vm-backend"
+version = "0.1.0"
+description = "VM backend implementations for CUBE desktop-automation benchmarks"
+requires-python = ">=3.12"
+dependencies = [
+    "cube-standard",
+    "requests>=2.28",
+    "tqdm>=4.60",
+]
+
+[project.optional-dependencies]
+dev = [
+    "pytest>=8.0.0",
+]
+
+[build-system]
+requires = ["uv_build>=0.6.0,<0.7.0"]
+build-backend = "uv_build"
+
+[tool.uv.build-backend]
+module-name = "cube_vm_backend"
+
+[tool.ruff]
+fix = true
+line-length = 120
+indent-width = 4
+
+[tool.ruff.format]
+quote-style = "double"
+indent-style = "space"
+skip-magic-trailing-comma = false
+line-ending = "auto"
+
+[tool.ruff.lint]
+extend-select = ["I"]

cube_vm_backend-0.1.0/src/cube_vm_backend/__init__.py

@@ -0,0 +1,11 @@
+"""cube-vm-backend: VM backend implementations for CUBE desktop-automation benchmarks."""
+
+from cube_vm_backend.local import LocalQEMUVM, LocalQEMUVMBackend
+from cube_vm_backend.qemu_manager import QEMUConfig, QEMUManager
+
+__all__ = [
+    "LocalQEMUVM",
+    "LocalQEMUVMBackend",
+    "QEMUConfig",
+    "QEMUManager",
+]

cube_vm_backend-0.1.0/src/cube_vm_backend/local.py

@@ -0,0 +1,139 @@
+"""LocalQEMUVMBackend — VMBackend implementation using QEMU/KVM on the local host.
+
+Uses a read-only base qcow2 image + per-task copy-on-write overlays.
+Reset strategy: delete overlay + reboot (ResetIsolation.RESTART, ~30s).
+
+Image:
+    path_to_vm must point to an existing qcow2 base image.
+    Benchmarks that need auto-download (e.g. OSWorld) subclass this backend
+    and override ensure_resource() to fetch the image before launch.
+
+Port forwarding:
+    SLIRP user-mode networking — no root or bridge required.
+    Each VM gets unique host ports forwarded to guest ports 5000/9222/8006/8080.
+"""
+
+from __future__ import annotations
+
+import logging
+from pathlib import Path
+
+from cube.vm import VM, VMBackend, VMConfig
+
+from cube_vm_backend.qemu_manager import QEMUConfig, QEMUManager
+
+logger = logging.getLogger(__name__)
+
+# Default cache location — can be overridden via cache_dir field
+_DEFAULT_CACHE_DIR = Path.home() / ".cube" / "vm_data"
+
+
+class LocalQEMUVM(VM):
+    """Runtime handle to a QEMU/KVM VM managed by LocalQEMUVMBackend.
+
+    Not serializable. The caller owns the lifecycle via stop().
+    """
+
+    def __init__(self, manager: QEMUManager) -> None:
+        self._manager = manager
+
+    @property
+    def endpoint(self) -> str:
+        """Base URL of the in-VM HTTP agent: ``http://localhost:<port>``."""
+        return f"http://localhost:{self._manager.server_port}"
+
+    @property
+    def chromium_port(self) -> int:
+        """Host port forwarded to guest Chromium DevTools (9222)."""
+        return self._manager.chromium_port
+
+    @property
+    def vlc_port(self) -> int:
+        """Host port forwarded to guest VLC HTTP (8080)."""
+        return self._manager.vlc_port
+
+    @property
+    def server_port(self) -> int:
+        """Host port forwarded to guest Flask agent (5000)."""
+        return self._manager.server_port
+
+    def restore_snapshot(self, name: str) -> None:
+        """Restore the VM to its initial state.
+
+        Implementation: delete overlay + create fresh overlay + reboot QEMU.
+        This provides RESTART isolation (~30s). The ``name`` argument is
+        accepted for API compatibility but ignored — only one snapshot state
+        (the base image) is available with the overlay strategy.
+        """
+        logger.info("Restoring VM snapshot '%s' (overlay reset)", name)
+        self._manager.reset()
+
+    def stop(self) -> None:
+        """Shut down the VM and clean up overlay and socket files."""
+        self._manager.stop()
+
+
+class LocalQEMUVMBackend(VMBackend):
+    """VMBackend that runs QEMU/KVM directly on the local host.
+
+    Attributes
+    ----------
+    cache_dir : str
+        Directory for the base qcow2 image and per-task overlays.
+    path_to_vm : str | None
+        Explicit path to a pre-existing qcow2 base image.
+        If None, the image is auto-downloaded to cache_dir on first use.
+    headless : bool
+        Suppress the graphical display (default True).
+    memory : str
+        RAM allocation passed to QEMU -m (e.g. "4G").
+    cpus : int
+        Number of vCPUs.
+    """
+
+    cache_dir: str = str(_DEFAULT_CACHE_DIR)
+    path_to_vm: str | None = None
+    headless: bool = True
+    memory: str = "4G"
+    cpus: int = 4
+
+    def ensure_resource(self, config: VMConfig) -> None:
+        """Validate or prepare the base qcow2 image before launch.
+
+        Override in subclasses to add auto-download behaviour (e.g. OSWorld).
+        Default implementation raises if path_to_vm is not set.
+        """
+        if self.path_to_vm is None:
+            raise ValueError(
+                "path_to_vm must be set on LocalQEMUVMBackend. "
+                "Provide a path to an existing qcow2 image, or use a subclass "
+                "that handles image acquisition (e.g. OSWorldQEMUVMBackend)."
+            )
+
+    def launch(self, config: VMConfig) -> LocalQEMUVM:
+        """Ensure image exists, then start a QEMU VM and return a live handle.
+
+        Blocks until the in-VM HTTP agent is reachable.
+        """
+        self.ensure_resource(config)
+
+        vm_dir = Path(self.cache_dir)
+        base_image = Path(self.path_to_vm)  # type: ignore[arg-type]
+
+        qemu_config = QEMUConfig(
+            base_image=base_image,
+            overlay_dir=vm_dir / "overlays",
+            memory=self.memory,
+            cpus=self.cpus,
+            headless=self.headless,
+            screen_width=config.screen_size[0],
+            screen_height=config.screen_size[1],
+        )
+        manager = QEMUManager(qemu_config)
+        manager.start()
+        logger.info("VM launched at endpoint http://localhost:%d", manager.server_port)
+        return LocalQEMUVM(manager)
+
+    def close(self) -> None:
+        """No-op at the backend level — each VM is stopped individually via vm.stop()."""
+        pass

cube_vm_backend-0.1.0/src/cube_vm_backend/qemu_manager.py

@@ -0,0 +1,398 @@
+"""QEMU/KVM VM lifecycle manager.
+
+Uses SLIRP user-mode networking with port forwarding (no root / bridge required).
+Snapshot strategy: read-only base qcow2 + per-instance copy-on-write overlay
+(reset = stop QEMU, delete overlay, create new overlay, start QEMU).
+Communicates with running QEMU via QMP Unix socket for clean shutdown.
+"""
+
+import json
+import logging
+import os
+import socket
+import subprocess
+import tempfile
+import time
+from pathlib import Path
+from typing import Optional
+
+import requests
+
+logger = logging.getLogger(__name__)
+
+_VM_READY_TIMEOUT = 300  # seconds to wait for VM screenshot endpoint
+_VM_READY_POLL_INTERVAL = 2  # seconds between readiness polls
+_QMP_TIMEOUT = 10  # seconds to wait for QMP socket
+
+
+class QEMUConfig:
+    """Configuration for a QEMU/KVM virtual machine.
+
+    Parameters
+    ----------
+    base_image : Path
+        Path to the read-only base qcow2 disk image.
+    overlay_dir : Path
+        Directory where per-instance overlay qcow2 files are created.
+    memory : str
+        RAM allocation passed to QEMU ``-m`` flag (e.g. ``"4G"``).
+    cpus : int
+        Number of vCPUs (``-smp``).
+    headless : bool
+        If True, suppress the graphical display (``-display none``).
+    screen_width : int
+        Horizontal resolution injected into the guest via kernel cmdline (informational).
+    screen_height : int
+        Vertical resolution (informational).
+    enable_kvm : bool
+        Automatically enable KVM hardware acceleration if ``/dev/kvm`` is present.
+    """
+
+    def __init__(
+        self,
+        base_image: Path,
+        overlay_dir: Path,
+        memory: str = "4G",
+        cpus: int = 4,
+        headless: bool = True,
+        screen_width: int = 1920,
+        screen_height: int = 1080,
+        enable_kvm: bool = True,
+    ) -> None:
+        self.base_image = Path(base_image)
+        self.overlay_dir = Path(overlay_dir)
+        self.memory = memory
+        self.cpus = cpus
+        self.headless = headless
+        self.screen_width = screen_width
+        self.screen_height = screen_height
+        self.enable_kvm = enable_kvm
+
+
+class QEMUManager:
+    """Manages the full lifecycle of an OSWorld QEMU/KVM virtual machine.
+
+    Usage::
+
+        manager = QEMUManager(config)
+        manager.start()          # boot VM (allocates ports, creates overlay)
+        manager.reset()          # restore initial state (stop → new overlay → boot)
+        manager.stop()           # shut down VM cleanly
+
+    After ``start()`` the following properties are available:
+
+    - :attr:`server_port`   — host port forwarded to guest :5000  (Flask agent)
+    - :attr:`chromium_port` — host port forwarded to guest :9222  (Chromium DevTools)
+    - :attr:`vnc_port`      — host port forwarded to guest :8006  (VNC/noVNC)
+    - :attr:`vlc_port`      — host port forwarded to guest :8080  (VLC HTTP)
+
+    Parameters
+    ----------
+    config : QEMUConfig
+        VM configuration (image paths, memory, CPU, display settings).
+    """
+
+    def __init__(self, config: QEMUConfig) -> None:
+        self.config = config
+        self._process: Optional[subprocess.Popen] = None
+        self._pid_file: Optional[Path] = None
+        self._overlay_path: Optional[Path] = None
+        self._qmp_socket: Optional[Path] = None
+
+        self._server_port: Optional[int] = None
+        self._chromium_port: Optional[int] = None
+        self._vnc_port: Optional[int] = None
+        self._vlc_port: Optional[int] = None
+
+    # ------------------------------------------------------------------
+    # Public properties
+    # ------------------------------------------------------------------
+
+    @property
+    def server_port(self) -> int:
+        if self._server_port is None:
+            raise RuntimeError("VM not started — call start() first")
+        return self._server_port
+
+    @property
+    def chromium_port(self) -> int:
+        if self._chromium_port is None:
+            raise RuntimeError("VM not started — call start() first")
+        return self._chromium_port
+
+    @property
+    def vnc_port(self) -> int:
+        if self._vnc_port is None:
+            raise RuntimeError("VM not started — call start() first")
+        return self._vnc_port
+
+    @property
+    def vlc_port(self) -> int:
+        if self._vlc_port is None:
+            raise RuntimeError("VM not started — call start() first")
+        return self._vlc_port
+
+    # ------------------------------------------------------------------
+    # Lifecycle
+    # ------------------------------------------------------------------
+
+    def start(self) -> None:
+        """Allocate ports, create overlay, launch QEMU, and wait for VM readiness."""
+        self.config.overlay_dir.mkdir(parents=True, exist_ok=True)
+
+        self._server_port = _reserve_free_port(5000)
+        self._chromium_port = _reserve_free_port(9222)
+        self._vnc_port = _reserve_free_port(8006)
+        self._vlc_port = _reserve_free_port(8080)
+
+        self._overlay_path = self.config.overlay_dir / f"overlay_{self._server_port}.qcow2"
+        self._qmp_socket = Path(tempfile.gettempdir()) / f"qemu_qmp_{self._server_port}.sock"
+        self._pid_file = Path(tempfile.gettempdir()) / f"qemu_{self._server_port}.pid"
+
+        self._create_overlay()
+        self._launch_qemu()
+        self._wait_for_ready()
+
+    def reset(self) -> None:
+        """Restore the VM to its initial state.
+
+        Stops the running instance, deletes the overlay, creates a fresh one,
+        and boots the VM again.
+        """
+        self._stop_qemu()
+        if self._overlay_path and self._overlay_path.exists():
+            self._overlay_path.unlink()
+        self._create_overlay()
+        self._launch_qemu()
+        self._wait_for_ready()
+
+    def stop(self) -> None:
+        """Shut down the VM and clean up overlay and socket files."""
+        self._stop_qemu()
+        for path in (self._overlay_path, self._qmp_socket, self._pid_file):
+            if path and path.exists():
+                try:
+                    path.unlink()
+                except OSError:
+                    pass
+        # Release port reservations so other workers can reuse these ports
+        for port in (self._server_port, self._chromium_port, self._vnc_port, self._vlc_port):
+            _release_port_reservation(port)
+
+    # ------------------------------------------------------------------
+    # Internal helpers
+    # ------------------------------------------------------------------
+
+    def _create_overlay(self) -> None:
+        """Create a fresh qcow2 overlay on top of the read-only base image."""
+        if self._overlay_path.exists():
+            logger.warning("Stale overlay found at %s — removing before recreating", self._overlay_path)
+            self._overlay_path.unlink()
+        cmd = [
+            "qemu-img",
+            "create",
+            "-f",
+            "qcow2",
+            "-b",
+            str(self.config.base_image),
+            "-F",
+            "qcow2",
+            str(self._overlay_path),
+        ]
+        logger.info("Creating overlay: %s", " ".join(cmd))
+        result = subprocess.run(cmd, capture_output=True, text=True)
+        if result.returncode != 0:
+            raise RuntimeError(f"qemu-img create failed (exit {result.returncode}):\n{result.stderr}")
+
+    def _launch_qemu(self) -> None:
+        """Build QEMU command and launch as a background subprocess."""
+        for path in (self._qmp_socket, self._pid_file):
+            if path and path.exists():
+                logger.warning("Removing stale file before launch: %s", path)
+                path.unlink()
+
+        qemu_cmd = ["qemu-system-x86_64"]
+
+        # KVM hardware acceleration
+        if self.config.enable_kvm and os.path.exists("/dev/kvm"):
+            qemu_cmd += ["-enable-kvm"]
+            logger.info("KVM acceleration enabled")
+        else:
+            logger.warning("KVM not available — running without hardware acceleration (slow)")
+
+        # Machine resources
+        qemu_cmd += ["-m", self.config.memory, "-smp", str(self.config.cpus)]
+
+        # Disk
+        qemu_cmd += ["-drive", f"file={self._overlay_path},format=qcow2,if=virtio"]
+
+        # Network: SLIRP user-mode with port forwarding
+        hostfwds = ",".join(
+            [
+                f"hostfwd=tcp::{self._server_port}-:5000",
+                f"hostfwd=tcp::{self._chromium_port}-:9222",
+                f"hostfwd=tcp::{self._vnc_port}-:8006",
+                f"hostfwd=tcp::{self._vlc_port}-:8080",
+            ]
+        )
+        qemu_cmd += [
+            "-netdev",
+            f"user,id=net0,{hostfwds}",
+            "-device",
+            "virtio-net-pci,netdev=net0",
+        ]
+
+        # Display
+        if self.config.headless:
+            qemu_cmd += ["-display", "none"]
+        else:
+            qemu_cmd += ["-vga", "virtio"]
+
+        # QMP control socket
+        qemu_cmd += ["-qmp", f"unix:{self._qmp_socket},server,nowait"]
+
+        # PID file and daemonize
+        qemu_cmd += ["-pidfile", str(self._pid_file), "-daemonize"]
+
+        logger.info("Starting QEMU: %s", " ".join(qemu_cmd))
+        result = subprocess.run(qemu_cmd, capture_output=True, text=True)
+        if result.returncode != 0:
+            raise RuntimeError(
+                f"QEMU failed to start (exit {result.returncode}):\nstdout: {result.stdout}\nstderr: {result.stderr}"
+            )
+        logger.info("QEMU launched (server_port=%d)", self._server_port)
+
+    def _wait_for_ready(self) -> None:
+        """Poll the guest's /screenshot endpoint until the VM is ready."""
+        deadline = time.time() + _VM_READY_TIMEOUT
+        url = f"http://localhost:{self._server_port}/screenshot"
+        logger.info("Waiting for VM to be ready at %s ...", url)
+        while time.time() < deadline:
+            try:
+                resp = requests.get(url, timeout=(5, 5))
+                if resp.status_code == 200:
+                    logger.info("VM is ready")
+                    return
+            except Exception:
+                pass
+            logger.info("VM not ready yet, retrying in %ds...", _VM_READY_POLL_INTERVAL)
+            time.sleep(_VM_READY_POLL_INTERVAL)
+        raise TimeoutError(f"VM failed to become ready within {_VM_READY_TIMEOUT}s")
+
+    def _stop_qemu(self) -> None:
+        """Send QMP 'quit' to the running QEMU instance and wait for it to exit."""
+        if self._qmp_socket and self._qmp_socket.exists():
+            try:
+                _qmp_quit(str(self._qmp_socket))
+                logger.info("Sent QMP quit")
+            except Exception as exc:
+                logger.warning("QMP quit failed (%s), falling back to SIGTERM", exc)
+                self._kill_by_pid()
+        elif self._pid_file and self._pid_file.exists():
+            self._kill_by_pid()
+
+    def _kill_by_pid(self) -> None:
+        """Terminate the QEMU process via the PID file."""
+        if not self._pid_file or not self._pid_file.exists():
+            return
+        try:
+            pid = int(self._pid_file.read_text().strip())
+            os.kill(pid, 15)  # SIGTERM
+            logger.info("Sent SIGTERM to QEMU pid %d", pid)
+            # Wait up to 5s for clean exit, then SIGKILL
+            for _ in range(10):
+                time.sleep(0.5)
+                try:
+                    os.kill(pid, 0)  # check if still alive
+                except ProcessLookupError:
+                    return  # process exited cleanly
+            os.kill(pid, 9)  # SIGKILL
+            logger.warning("QEMU pid %d did not exit after SIGTERM — sent SIGKILL", pid)
+        except Exception as exc:
+            logger.warning("Failed to kill QEMU by pid: %s", exc)
+
+
+# ------------------------------------------------------------------
+# QMP communication
+# ------------------------------------------------------------------
+
+
+def _qmp_quit(socket_path: str) -> None:
+    """Connect to a QEMU QMP Unix socket and send the 'quit' command."""
+    with socket.socket(socket.AF_UNIX, socket.SOCK_STREAM) as sock:
+        sock.settimeout(_QMP_TIMEOUT)
+        sock.connect(socket_path)
+
+        # Read the QMP greeting
+        greeting = _qmp_recv(sock)
+        logger.debug("QMP greeting: %s", greeting)
+
+        # Negotiate capabilities
+        sock.sendall(json.dumps({"execute": "qmp_capabilities"}).encode())
+        _qmp_recv(sock)
+
+        # Send quit
+        sock.sendall(json.dumps({"execute": "quit"}).encode())
+        try:
+            _qmp_recv(sock)
+        except (ConnectionResetError, OSError):
+            pass  # Connection closed after quit — expected
+
+
+def _qmp_recv(sock: socket.socket) -> dict:
+    """Read a complete JSON object from a QMP socket."""
+    data = b""
+    while True:
+        chunk = sock.recv(4096)
+        if not chunk:
+            break
+        data += chunk
+        try:
+            return json.loads(data.decode())
+        except json.JSONDecodeError:
+            continue
+    return {}
+
+
+# ------------------------------------------------------------------
+# Port allocation
+# ------------------------------------------------------------------
+
+
+def _reserve_free_port(start: int = 5000) -> int:
+    """Atomically find and reserve a free TCP port at or above ``start``.
+
+    Uses O_EXCL flag to create a reservation file — this is race-free even
+    across multiple processes (e.g. parallel Ray workers), unlike bind+release.
+    Call _release_port_reservation() when the port is no longer needed.
+    """
+    tmp = Path(tempfile.gettempdir())
+    for port in range(start, 65354):
+        reservation = tmp / f"qemu_port_{port}.reserved"
+        try:
+            # Atomically create reservation — fails if another process beat us to it
+            fd = os.open(str(reservation), os.O_CREAT | os.O_EXCL | os.O_WRONLY, 0o644)
+            os.close(fd)
+        except FileExistsError:
+            continue  # Another worker reserved this port
+        # Verify the port is actually free on the network
+        with socket.socket(socket.AF_INET, socket.SOCK_STREAM) as s:
+            s.setsockopt(socket.SOL_SOCKET, socket.SO_REUSEADDR, 1)
+            try:
+                s.bind(("localhost", port))
+                return port  # Reserved and free
+            except OSError:
+                reservation.unlink()  # Port in use — release reservation and try next
+                continue
+    raise RuntimeError(f"No free port found starting from {start}")
+
+
+def _release_port_reservation(port: Optional[int]) -> None:
+    """Remove the reservation file for a port allocated by _reserve_free_port."""
+    if port is None:
+        return
+    reservation = Path(tempfile.gettempdir()) / f"qemu_port_{port}.reserved"
+    try:
+        reservation.unlink(missing_ok=True)
+    except OSError:
+        pass