diff --git a/setup_vm.py b/setup_vm.py index 5164a18..fafdd15 100644 --- a/setup_vm.py +++ b/setup_vm.py @@ -1,35 +1,76 @@ +import os import platform import subprocess - import requests +from tqdm import tqdm +import zipfile +from time import sleep + +import socket # Define the path to the virtual machine VM_PATH = r"Ubuntu\Ubuntu.vmx" # change this to the path of your downloaded virtual machine +MAX_RETRY_TIMES = 10 def download_and_unzip_vm(): # Determine the platform and CPU architecture to decide the correct VM image to download - if platform.system() == 'Darwin' and platform.machine() == 'arm64': # macOS with Apple Silicon + if platform.machine() == 'arm64': # macOS with Apple Silicon url = "https://huggingface.co/datasets/xlangai/ubuntu_arm/resolve/main/Ubuntu.zip" - elif platform.system() in ['Linux', - 'Windows'] and platform.machine() == 'x86_64': # Linux or Windows with x86_64 CPU + elif platform.machine().lower() == 'amd64': url = "https://huggingface.co/datasets/xlangai/ubuntu_x86/resolve/main/Ubuntu.zip" else: raise Exception("Unsupported platform or architecture") # Download the virtual machine image print("Downloading the virtual machine image...") - subprocess.run(['wget', url], stdout=subprocess.PIPE, stderr=subprocess.STDOUT, text=True, - universal_newlines=True) + filename = "Ubuntu.zip" + downloaded_size = 0 + + while True: + headers = {} + if os.path.exists(filename): + downloaded_size = os.path.getsize(filename) + headers["Range"] = f"bytes={downloaded_size}-" + + with requests.get(url, headers=headers, stream=True) as response: + response.raise_for_status() + total_size = int(response.headers.get('content-length', 0)) + + with open(filename, "ab") as file, tqdm( + desc="Progress", + total=total_size, + unit='iB', + unit_scale=True, + unit_divisor=1024, + initial=downloaded_size, + ascii=True + ) as progress_bar: + try: + for data in response.iter_content(chunk_size=1024): + size = file.write(data) + progress_bar.update(size) + except (requests.exceptions.RequestException, IOError) as e: + print(f"Download error: {e}") + sleep(1) # Wait for 1 second before retrying + print("Retrying...") + else: + print("Download succeeds.") + break # Download completed successfully # Unzip the downloaded file - print("Unzipping the downloaded file...") - subprocess.run(['unzip', 'Ubuntu.zip'], stdout=subprocess.PIPE, stderr=subprocess.STDOUT, text=True, - universal_newlines=True) + print("Unzipping the downloaded file...☕️") + current_directory = os.getcwd() + with zipfile.ZipFile('Ubuntu.zip', 'r') as zip_ref: + zip_ref.extractall(current_directory) + print("Files have been successfully extracted to the current working directory:", current_directory) # Execute the function to download and unzip the VM -download_and_unzip_vm() +if not os.path.exists(VM_PATH): + download_and_unzip_vm() +else: + print(f"Virtual machine exists: {VM_PATH}") # Determine the platform of the host machine and decide the parameter for vmrun @@ -44,17 +85,37 @@ def get_vmrun_type(): # Start the virtual machine subprocess.run(f'vmrun {get_vmrun_type()} start "{VM_PATH}"', shell=True) -print("Virtual machine started...") +print("Starting virtual machine...") # Get the IP address of the virtual machine -get_vm_ip = subprocess.run(f'vmrun {get_vmrun_type()} getGuestIPAddress "{VM_PATH}" -wait', shell=True, - capture_output=True, - text=True) -print("Virtual machine IP address:", get_vm_ip.stdout.strip()) +for i in range(MAX_RETRY_TIMES): + get_vm_ip = subprocess.run(f'vmrun {get_vmrun_type()} getGuestIPAddress "{VM_PATH}" -wait', shell=True, + capture_output=True, + text=True) + if "Error" in get_vm_ip.stdout: + print("Retry on getting IP") + continue + print("Virtual machine IP address:", get_vm_ip.stdout.strip()) + break vm_ip = get_vm_ip.stdout.strip() +def is_url_accessible(url, timeout=1): + try: + response = requests.head(url, timeout=timeout) + return response.status_code == 200 + except requests.exceptions.RequestException: + return False + + +print("--------------------------------") +url = f"http://{vm_ip}:5000/screenshot" +ckeck_url = is_url_accessible(url) +print(f"check url: {url} | is accessible: {ckeck_url}") +print("--------------------------------") + + # Function used to check whether the virtual machine is ready def download_screenshot(ip): url = f"http://{ip}:5000/screenshot" @@ -63,8 +124,11 @@ def download_screenshot(ip): response = requests.get(url, timeout=(1, 1)) if response.status_code == 200: return True - except requests.exceptions.ConnectionError: - pass + except Exception as e: + print(f"Error: {e}") + print(f"Type: {type(e).__name__}") + print(f"Error detail: {str(e)}") + sleep(2) return False