Update provider

This commit is contained in:
FredWuCZ
2024-09-28 16:57:58 +08:00
parent 2b4b6d8fdf
commit 0a21e24a81

View File

@@ -5,6 +5,7 @@ import subprocess
import time import time
import docker import docker
import psutil import psutil
import requests
from desktop_env.providers.base import Provider from desktop_env.providers.base import Provider
@@ -12,6 +13,7 @@ logger = logging.getLogger("desktopenv.providers.vmware.VMwareProvider")
logger.setLevel(logging.INFO) logger.setLevel(logging.INFO)
WAIT_TIME = 3 WAIT_TIME = 3
RETRY_INTERVAL = 1
class DockerProvider(Provider): class DockerProvider(Provider):
def __init__(self, region: str): def __init__(self, region: str):
@@ -20,7 +22,7 @@ class DockerProvider(Provider):
self.server_port = self._get_available_port(5000) self.server_port = self._get_available_port(5000)
# self.remote_debugging_port = self._get_available_port(1337) # self.remote_debugging_port = self._get_available_port(1337)
self.chromium_port = self._get_available_port(9222) 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 @staticmethod
def _get_available_port(port: int): def _get_available_port(port: int):
@@ -31,9 +33,21 @@ class DockerProvider(Provider):
def start_emulator(self, path_to_vm: str, headless: bool, os_type: str): 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}") 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, 9222: self.chromium_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) def download_screenshot(ip, port):
time.sleep(20) 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: def get_ip_address(self, path_to_vm: str) -> str:
return f"localhost:{self.server_port}:{self.chromium_port}:{self.vnc_port}" return f"localhost:{self.server_port}:{self.chromium_port}:{self.vnc_port}"
@@ -46,7 +60,8 @@ class DockerProvider(Provider):
def stop_emulator(self, path_to_vm: str): def stop_emulator(self, path_to_vm: str):
logger.info("Stopping VM...") logger.info("Stopping VM...")
self.container.stop(WAIT_TIME) self.container.stop()
self.container.remove() self.container.remove()
time.sleep(WAIT_TIME)
# 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 # 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