Update setup_vm script

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