68 lines
2.8 KiB
Python
68 lines
2.8 KiB
Python
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}")
|
|
print(path_to_vm)
|
|
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": "/System.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 |