From bead06afc8504b984e64e80f981a52b856029350 Mon Sep 17 00:00:00 2001 From: Timothyxxx <384084775@qq.com> Date: Fri, 10 May 2024 16:24:07 +0800 Subject: [PATCH] Improve the logic in auto-installation; Add auto-remove on the failed files --- desktop_env/envs/__init__.py | 88 ++++++++++++++++++++---------------- 1 file changed, 48 insertions(+), 40 deletions(-) diff --git a/desktop_env/envs/__init__.py b/desktop_env/envs/__init__.py index 133812e..c787e66 100644 --- a/desktop_env/envs/__init__.py +++ b/desktop_env/envs/__init__.py @@ -7,28 +7,25 @@ import threading import uuid import zipfile from time import sleep -import logging - -logger = logging.getLogger("desktopenv.envinit") - +import shutil import psutil import requests from tqdm import tqdm -__version__ = "0.1.6" +__version__ = "0.1.8" MAX_RETRY_TIMES = 10 UBUNTU_ARM_URL = "https://huggingface.co/datasets/xlangai/ubuntu_arm/resolve/main/Ubuntu.zip" UBUNTU_X86_URL = "https://huggingface.co/datasets/xlangai/ubuntu_x86/resolve/main/Ubuntu.zip" +DOWNLOADED_FILE_NAME = "Ubuntu.zip" REGISTRY_PATH = '.vms' -REGISTRY_IDX_PATH = ".vms_idx" +VMS_DIR = "./vm_data" update_lock = threading.Lock() class VirtualMachineManager: - def __init__(self, registry_path=REGISTRY_PATH, registry_idx_path=REGISTRY_IDX_PATH): + def __init__(self, registry_path=REGISTRY_PATH): self.registry_path = registry_path - self.registry_idx_path = registry_idx_path self.lock = threading.Lock() self.initialize_registry() @@ -37,9 +34,6 @@ class VirtualMachineManager: if not os.path.exists(self.registry_path): with open(self.registry_path, 'w') as file: file.write('') - if not os.path.exists(self.registry_idx_path): - with open(self.registry_idx_path, 'w') as file: - file.write('0') def add_vm(self, vm_path): with self.lock: @@ -77,14 +71,18 @@ class VirtualMachineManager: with open(self.registry_path, 'w') as file: file.writelines(new_lines) - def check_and_clean(self): - with self.lock: # Lock when cleaning up the registry + def check_and_clean(self, vms_dir): + with self.lock: # Lock when cleaning up the registry and vms_dir + + # Check and clean on the running vms, detect the released ones and mark then as 'free' active_pids = {p.pid for p in psutil.process_iter()} new_lines = [] + vm_paths = [] with open(self.registry_path, 'r') as file: lines = file.readlines() for line in lines: vm_path, pid_str = line.strip().split('|') + vm_paths.append(vm_path) if pid_str == "free": new_lines.append(line) continue @@ -96,6 +94,19 @@ class VirtualMachineManager: with open(self.registry_path, 'w') as file: file.writelines(new_lines) + # Check and clean on the files inside vms_dir, delete the unregistered ones + vm_names = os.listdir(vms_dir) + for vm_name in vm_names: + # skip the downloaded .zip file + if vm_name == DOWNLOADED_FILE_NAME: + continue + flag = True + for vm_path in vm_paths: + if vm_name + ".vmx" in vm_path: + flag = False + if flag: + shutil.rmtree(os.path.join(vms_dir, vm_name)) + def list_vms(self): with self.lock: # Lock when reading the registry all_vms = [] @@ -117,17 +128,15 @@ class VirtualMachineManager: free_vms.append((vm_path, pid_str)) return free_vms - def generate_new_vm_name(self): - with self.lock: # Lock when generating a new path - with open(self.registry_idx_path, 'r') as file: - idx = int(file.read()) - - new_name = f"Ubuntu{idx}" - - with open(self.registry_idx_path, 'w') as file: - file.write(str(idx + 1)) - - return new_name + def generate_new_vm_name(self, vms_dir): + registry_idx = 0 + while True: + attempted_new_name = f"Ubuntu{registry_idx}" + if os.path.exists( + os.path.join(vms_dir, attempted_new_name, attempted_new_name, attempted_new_name + ".vmx")): + registry_idx += 1 + else: + return attempted_new_name def _update_vm(vmx_path, target_vm_name): @@ -188,17 +197,16 @@ def _update_vm(vmx_path, target_vm_name): print("VM files renamed successfully.") -def _install_virtual_machine(vm_name, working_dir="./vm_data", downloaded_file_name="Ubuntu.zip", original_vm_name="Ubuntu"): - os.makedirs(working_dir, exist_ok=True) +def _install_virtual_machine(vm_name, vms_dir, downloaded_file_name, original_vm_name="Ubuntu"): + os.makedirs(vms_dir, exist_ok=True) + def __download_and_unzip_vm(): # 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 - url = UBUNTU_ARM_URL # a workout as most new macs are using apple silicon and they are frequently misrecognized as x86_64 - logger.warning("Your platform is temporarily considered to be ARM. If this is a mistake, please manually replace the downloaded VM under ./vm_data with the VM of x86 version from %s", UBUNTU_X86_URL) + # 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: @@ -209,7 +217,7 @@ def _install_virtual_machine(vm_name, working_dir="./vm_data", downloaded_file_n downloaded_size = 0 while True: - downloaded_file_path = os.path.join(working_dir, downloaded_file_name) + 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) @@ -248,15 +256,15 @@ def _install_virtual_machine(vm_name, working_dir="./vm_data", downloaded_file_n # Unzip the downloaded file print("Unzipping the downloaded file...☕️") with zipfile.ZipFile(downloaded_file_path, 'r') as zip_ref: - zip_ref.extractall(os.path.join(working_dir, vm_name)) - print("Files have been successfully extracted to the directory:", os.path.join(working_dir, vm_name)) + zip_ref.extractall(os.path.join(vms_dir, vm_name)) + print("Files have been successfully extracted to the directory:", os.path.join(vms_dir, vm_name)) - vm_path = os.path.join(working_dir, vm_name, vm_name, vm_name + ".vmx") + vm_path = os.path.join(vms_dir, vm_name, vm_name, vm_name + ".vmx") # Execute the function to download and unzip the VM, and update the vm metadata if not os.path.exists(vm_path): __download_and_unzip_vm() - _update_vm(os.path.join(working_dir, vm_name, original_vm_name, original_vm_name + ".vmx"), vm_name) + _update_vm(os.path.join(vms_dir, vm_name, original_vm_name, original_vm_name + ".vmx"), vm_name) else: print(f"Virtual machine exists: {vm_path}") @@ -326,13 +334,13 @@ def _install_virtual_machine(vm_name, working_dir="./vm_data", downloaded_file_n def _get_vm_path(): vm_manager = VirtualMachineManager(REGISTRY_PATH) - vm_manager.check_and_clean() + vm_manager.check_and_clean(vms_dir=VMS_DIR) free_vms_paths = vm_manager.list_free_vms() if len(free_vms_paths) == 0: # No free virtual machine available, generate a new one print("No free virtual machine available. Generating a new one, which would take a while...☕") - new_vm_name = vm_manager.generate_new_vm_name() - new_vm_path = _install_virtual_machine(new_vm_name) + new_vm_name = vm_manager.generate_new_vm_name(vms_dir=VMS_DIR) + new_vm_path = _install_virtual_machine(new_vm_name, vms_dir=VMS_DIR, downloaded_file_name=DOWNLOADED_FILE_NAME) vm_manager.add_vm(new_vm_path) vm_manager.occupy_vm(new_vm_path, os.getpid()) return new_vm_path