Improve the logic in auto-installation; Add auto-remove on the failed files
This commit is contained in:
Timothyxxx
@@ -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
|
||||
|
||||
Reference in New Issue
Block a user