From 0a21e24a81411414651d58740b817b2c8fe600a1 Mon Sep 17 00:00:00 2001
From: FredWuCZ <fredwucz@outlook.com>
Date: Sat, 28 Sep 2024 16:57:58 +0800
Subject: [PATCH] Update provider

---
 desktop_env/providers/docker/provider.py | 49 ++++++++++++++++--------
 1 file changed, 32 insertions(+), 17 deletions(-)

diff --git a/desktop_env/providers/docker/provider.py b/desktop_env/providers/docker/provider.py
index 6dcca46..9ceab84 100644
--- a/desktop_env/providers/docker/provider.py
+++ b/desktop_env/providers/docker/provider.py
@@ -5,14 +5,16 @@ 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()
@@ -20,33 +22,46 @@ class DockerProvider(Provider):
         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": "64G", "RAM_SIZE": "8G", "CPU_CORES": "8"} # Modify if needed
-
+        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('qemux/qemu-docker', environment={"DISK_SIZE": "64G", "RAM_SIZE": "6G", "CPU_CORES": "8"}, volumes={"/Users/happysix/Programs/HKUNLP/Qemu/Ubuntu.qcow2": {"bind": "/Ubuntu.qcow2", "mode": "ro"}, "/Users/happysix/Programs/HKUNLP/Qemu/snapshot.qcow2": {"bind": "/boot.qcow2", "mode": "rw"}}, cap_add=["NET_ADMIN"], ports={8006: self.vnc_port, 5000: self.server_port}, detach=True)
-        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}, detach=True)
-        time.sleep(20)
-
+        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(WAIT_TIME)
+        self.container.stop()
         self.container.remove()
-
-# docker run -it --rm -e "DISK_SIZE=64G" -e "RAM_SIZE=8G" -e "CPU_CORES=8" --volume C:\Users\admin\Documents\OSWorld\docker_vm_data\Ubuntu.qcow2:/Ubuntu.qcow2:ro --cap-add NET_ADMIN --device /dev/kvm -p 8006:8006 -p 5000:5000 osworld-docker
\ No newline at end of file
+        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
\ No newline at end of file