Files
sci-gui-agent-benchmark/setup_vm.py

80 lines
2.9 KiB
Python

import platform
import subprocess
import requests
# Define the path to the virtual machine
VM_PATH = r"Ubuntu\Ubuntu.vmx" # change this to the path of your downloaded virtual machine
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
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
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)
# 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)
# Execute the function to download and unzip the VM
download_and_unzip_vm()
# Determine the platform of the host machine and decide the parameter for vmrun
def get_vmrun_type():
if platform.system() == 'Windows' or platform.system() == 'Linux':
return '-T ws'
elif platform.system() == 'Darwin': # Darwin is the system name for macOS
return '-T fusion'
else:
raise Exception("Unsupported operating system")
# Start the virtual machine
subprocess.run(f'vmrun {get_vmrun_type()} start "{VM_PATH}"', shell=True)
print("Virtual machine started...")
# 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())
vm_ip = get_vm_ip.stdout.strip()
# Function used to check whether the virtual machine is ready
def download_screenshot(ip):
url = f"http://{ip}:5000/screenshot"
try:
# max trey times 1, max timeout 1
response = requests.get(url, timeout=(1, 1))
if response.status_code == 200:
return True
except requests.exceptions.ConnectionError:
pass
return False
# Try downloading the screenshot until successful
while not download_screenshot(vm_ip):
print("Check whether the virtual machine is ready...")
print("Virtual machine is ready. Start to make a snapshot on the virtual machine. It would take a while...")
# Create a snapshot of the virtual machine
subprocess.run(f'vmrun {get_vmrun_type()} snapshot "{VM_PATH}" "init_state"', shell=True)
print("Snapshot created.")