Improve the logic in auto-installation; Add auto-remove on the failed files

desktop_env/envs/__init__.py

@@ -7,28 +7,25 @@ import threading
 import uuid
 import zipfile
 from time import sleep
-import logging
+import shutil
-
-logger = logging.getLogger("desktopenv.envinit")
-
 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):
+    def check_and_clean(self, vms_dir):
-        with self.lock:  # Lock when cleaning up the registry
+        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):
+    def generate_new_vm_name(self, vms_dir):
-        with self.lock:  # Lock when generating a new path
+        registry_idx = 0
-            with open(self.registry_idx_path, 'r') as file:
+        while True:
-                idx = int(file.read())
+            attempted_new_name = f"Ubuntu{registry_idx}"
-
+            if os.path.exists(
-            new_name = f"Ubuntu{idx}"
+                    os.path.join(vms_dir, attempted_new_name, attempted_new_name, attempted_new_name + ".vmx")):
-
+                registry_idx += 1
-            with open(self.registry_idx_path, 'w') as file:
+            else:
-                file.write(str(idx + 1))
+                return attempted_new_name
-
-        return 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"):
+def _install_virtual_machine(vm_name, vms_dir, downloaded_file_name, original_vm_name="Ubuntu"):
-    os.makedirs(working_dir, exist_ok=True)
+    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
+            # if os.uname().machine == 'arm64':  # Apple Silicon
-                #url = UBUNTU_ARM_URL
+            url = UBUNTU_ARM_URL
-            #else:
+        # else:
-                #url = UBUNTU_X86_URL
+        #     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)
         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))
+            zip_ref.extractall(os.path.join(vms_dir, vm_name))
-        print("Files have been successfully extracted to the directory:", os.path.join(working_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_name = vm_manager.generate_new_vm_name(vms_dir=VMS_DIR)
-        new_vm_path = _install_virtual_machine(new_vm_name)
+        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