Support Docker VM manager and provider (#75)

* Add docker provider framework

* Update VM download link

* Add stop container

* Update docker manager & provider

* Update

* Update

* Update provider

desktop_env/providers/docker/manager.py

@@ -0,0 +1,123 @@
+import os
+import platform
+import random
+import re
+
+import threading
+from filelock import FileLock
+import uuid
+import zipfile
+
+from time import sleep
+import shutil
+import psutil
+import subprocess
+import requests
+from tqdm import tqdm
+import docker
+
+import logging
+
+from desktop_env.providers.base import VMManager
+
+logger = logging.getLogger("desktopenv.providers.vmware.VMwareVMManager")
+logger.setLevel(logging.INFO)
+
+MAX_RETRY_TIMES = 10
+RETRY_INTERVAL = 5
+
+UBUNTU_X86_URL = "https://huggingface.co/datasets/xlangai/ubuntu_osworld/resolve/main/Ubuntu.qcow2"
+VMS_DIR = "./docker_vm_data"
+
+# Determine the platform and CPU architecture to decide the correct VM image to download
+# if platform.system() == 'Darwin':  # macOS
+#     # if os.uname().machine == 'arm64':  # Apple Silicon
+#     URL = UBUNTU_ARM_URL
+# # else:
+# #     url = UBUNTU_X86_URL
+# elif platform.machine().lower() in ['amd64', 'x86_64']:
+#     URL = UBUNTU_X86_URL
+# else:
+#     raise Exception("Unsupported platform or architecture")
+URL = UBUNTU_X86_URL
+
+DOWNLOADED_FILE_NAME = URL.split('/')[-1]
+
+if platform.system() == 'Windows':
+    docker_path = r"C:\Program Files\Docker\Docker"
+    os.environ["PATH"] += os.pathsep + docker_path
+
+def _download_vm(vms_dir: str):
+    # Download the virtual machine image
+    logger.info("Downloading the virtual machine image...")
+    downloaded_size = 0
+
+    URL = UBUNTU_X86_URL
+    DOWNLOADED_FILE_NAME = URL.split('/')[-1]
+    downloaded_file_name = DOWNLOADED_FILE_NAME
+
+    os.makedirs(vms_dir, exist_ok=True)
+
+    while True:
+        downloaded_file_path = os.path.join(vms_dir, downloaded_file_name)
+        headers = {}
+        if os.path.exists(downloaded_file_path):
+            downloaded_size = os.path.getsize(downloaded_file_path)
+            headers["Range"] = f"bytes={downloaded_size}-"
+
+        with requests.get(URL, headers=headers, stream=True) as response:
+            if response.status_code == 416:
+                # This means the range was not satisfiable, possibly the file was fully downloaded
+                logger.info("Fully downloaded or the file size changed.")
+                break
+
+            response.raise_for_status()
+            total_size = int(response.headers.get('content-length', 0))
+
+            with open(downloaded_file_path, "ab") as file, tqdm(
+                    desc="Progress",
+                    total=total_size,
+                    unit='iB',
+                    unit_scale=True,
+                    unit_divisor=1024,
+                    initial=downloaded_size,
+                    ascii=True
+            ) as progress_bar:
+                try:
+                    for data in response.iter_content(chunk_size=1024):
+                        size = file.write(data)
+                        progress_bar.update(size)
+                except (requests.exceptions.RequestException, IOError) as e:
+                    logger.error(f"Download error: {e}")
+                    sleep(RETRY_INTERVAL)
+                    logger.error("Retrying...")
+                else:
+                    logger.info("Download succeeds.")
+                    break  # Download completed successfully
+
+class DockerVMManager(VMManager):
+    def __init__(self, registry_path=""):
+        pass
+
+    def add_vm(self, vm_path):
+        pass
+
+    def check_and_clean(self):
+        pass
+
+    def delete_vm(self, vm_path):
+        pass
+
+    def initialize_registry(self):
+        pass
+
+    def list_free_vms(self):
+        return os.path.join(VMS_DIR, DOWNLOADED_FILE_NAME)
+
+    def occupy_vm(self, vm_path):
+        pass
+
+    def get_vm_path(self, os_type, region):
+        if not os.path.exists(os.path.join(VMS_DIR, DOWNLOADED_FILE_NAME)):
+            _download_vm(VMS_DIR)
+        return os.path.join(VMS_DIR, DOWNLOADED_FILE_NAME)

desktop_env/providers/docker/provider.py

@@ -0,0 +1,67 @@
+import logging
+import os
+import platform
+import subprocess
+import time
+import docker
+import psutil
+import requests
+ 
+from desktop_env.providers.base import Provider
+ 
+logger = logging.getLogger("desktopenv.providers.vmware.VMwareProvider")
+logger.setLevel(logging.INFO)
+ 
+WAIT_TIME = 3
+RETRY_INTERVAL = 1
+ 
+class DockerProvider(Provider):
+    def __init__(self, region: str):
+        self.client = docker.from_env()
+        self.vnc_port = self._get_available_port(8006)
+        self.server_port = self._get_available_port(5000)
+        # self.remote_debugging_port = self._get_available_port(1337)
+        self.chromium_port = self._get_available_port(9222)
+        self.environment = {"DISK_SIZE": "32G", "RAM_SIZE": "4G", "CPU_CORES": "4"} # Modify if needed
+ 
+    @staticmethod
+    def _get_available_port(port: int):
+        while port < 65354:
+            if port not in [conn.laddr.port for conn in psutil.net_connections()]:
+                return port
+            port += 1
+ 
+    def start_emulator(self, path_to_vm: str, headless: bool, os_type: str):
+        logger.info(f"Occupying ports: {self.vnc_port}, {self.server_port}, {self.chromium_port}")
+        self.container = self.client.containers.run("happysixd/osworld-docker", environment=self.environment, cap_add=["NET_ADMIN"], devices=["/dev/kvm"], volumes={os.path.abspath(path_to_vm): {"bind": "/Ubuntu.qcow2", "mode": "ro"}}, ports={8006: self.vnc_port, 5000: self.server_port, 9222: self.chromium_port}, detach=True)
+        def download_screenshot(ip, port):
+            url = f"http://{ip}:{port}/screenshot"
+            try:
+                # max trey times 1, max timeout 1
+                response = requests.get(url, timeout=(10, 10))
+                if response.status_code == 200:
+                    return True
+            except Exception as e:
+                time.sleep(RETRY_INTERVAL)
+            return False
+ 
+        # Try downloading the screenshot until successful
+        while not download_screenshot("localhost", self.server_port):
+            logger.info("Check whether the virtual machine is ready...")
+ 
+    def get_ip_address(self, path_to_vm: str) -> str:
+        return f"localhost:{self.server_port}:{self.chromium_port}:{self.vnc_port}"
+ 
+    def save_state(self, path_to_vm: str, snapshot_name: str):
+        raise NotImplementedError("Not available for Docker.")
+ 
+    def revert_to_snapshot(self, path_to_vm: str, snapshot_name: str):
+        pass
+ 
+    def stop_emulator(self, path_to_vm: str):
+        logger.info("Stopping VM...")
+        self.container.stop()
+        self.container.remove()
+        time.sleep(WAIT_TIME)
+ 
+# docker run -it --rm -e "DISK_SIZE=64G" -e "RAM_SIZE=8G" -e "CPU_CORES=8" --volume /home/$USER/osworld/docker_vm_data/Ubuntu.qcow2:/Ubuntu.qcow2:ro --cap-add NET_ADMIN --device /dev/kvm -p 8008:8006 -p 5002:5000 happysixd/osworld-docker