lzy/sci-gui-agent-benchmark
1
0
Fork 0
Code Issues Pull Requests Actions Packages Projects Releases Wiki Activity

Refactoring VMware Integration and Implementing AWS Support (#44)

Browse Source 
* Initailize aws support

* Add README for the VM server

* Refactor OSWorld for supporting more cloud services.

* Initialize vmware and aws implementation v1, waiting for verification

* Initlize files for azure, gcp and virtualbox support

* Debug on the VMware provider

* Fix on aws interface mapping

* Fix instance type

* Refactor

* Clean

* hk region; debug

* Fix lock

* Remove print

* Remove key_name requirements when allocating aws vm

* Clean README

---------

Co-authored-by: XinyuanWangCS <xywang626@gmail.com>
This commit is contained in:
Tianbao Xie
2024-06-15 20:52:29 +08:00
committed by GitHub
parent c121869219
commit fffa8f8da6
31 changed files with 847 additions and 302 deletions
Download Patch File Download Diff File Expand all files Collapse all files

23
desktop_env/providers/vmware/INSTALL_VMWARE.md Normal file
View File

@@ -0,0 +1,23 @@
## 💾 Installation of VMware Workstation Pro
---
1. Download VMware Workstation Pro from the [official website](https://www.vmware.com/products/workstation-pro/workstation-pro-evaluation.html). The version we are using is 17.5.1. For systems with Apple chips, you should install [VMware Fusion](https://www.vmware.com/go/getfusion).
2. Install VMware Workstation
- **[On Linux](https://docs.vmware.com/en/VMware-Workstation-Pro/17/com.vmware.ws.using.doc/GUID-1F5B1F14-A586-4A56-83FA-2E7D8333D5CA.html):** Run the following command in your terminal, where `xxxx-xxxxxxx` represents the version number and internal version number.
```
sudo sh VMware-Workstation-xxxx-xxxxxxx.architecture.bundle --console
```
- **[On Windows](https://docs.vmware.com/en/VMware-Workstation-Pro/17/com.vmware.ws.using.doc/GUID-F5A7B3CB-9141-458B-A256-E0C3EA805AAA.html):** Ensure that you're logged in as either the Administrator user or as a user who belongs to the local Administrators group. If you're logging in to a domain, make sure your domain account has local administrator privileges. Proceed by double-clicking the `VMware-workstation-xxxx-xxxxxxx.exe` file. Be aware that you might need to reboot your host system to finalize the installation.
- **[For systems with Apple chips](https://docs.vmware.com/en/VMware-Fusion/13/com.vmware.fusion.using.doc/GUID-ACC3A019-93D3-442C-A34E-F7755DF6733B.html):** Double-click the `VMware-Fusion-xxxx-xxxxxxx.dmg` file to open it. In the Finder window that appears, double-click the 'Install Fusion' icon. When prompted, enter your administrator username and password.
> **Note:** You need to fill the activation key during the installation process when prompted.
3. Verify the successful installation by running the following:
```
vmrun -T ws list
```
If the installation along with the environment variable set is successful, you will see the message showing the current running virtual machines.

0
desktop_env/providers/vmware/__init__.py Normal file
View File

379
desktop_env/providers/vmware/manager.py Normal file
View File

@@ -0,0 +1,379 @@
import os
import platform
import random
import re
import threading
from filelock import FileLock
import uuid
import zipfile
from time import sleep
import shutil
import psutil
import subprocess
import requests
from tqdm import tqdm
import logging
from desktop_env.providers.base import VMManager
logger = logging.getLogger("desktopenv.providers.vmware.VMwareVMManager")
logger.setLevel(logging.INFO)
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 = '.vmware_vms'
VMS_DIR = "./vmware_vm_data"
update_lock = threading.Lock()
def generate_new_vm_name(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):
"""Update the VMX file with the new VM name and other parameters, so that the VM can be started successfully without conflict with the original VM."""
with update_lock:
dir_path, vmx_file = os.path.split(vmx_path)
def _generate_mac_address():
# VMware MAC address range starts with 00:0c:29
mac = [0x00, 0x0c, 0x29,
random.randint(0x00, 0x7f),
random.randint(0x00, 0xff),
random.randint(0x00, 0xff)]
return ':'.join(map(lambda x: "%02x" % x, mac))
# Backup the original file
with open(vmx_path, 'r') as file:
original_content = file.read()
# Generate new values
new_uuid_bios = str(uuid.uuid4())
new_uuid_location = str(uuid.uuid4())
new_mac_address = _generate_mac_address()
new_vmci_id = str(random.randint(-2147483648, 2147483647)) # Random 32-bit integer
# Update the content
updated_content = re.sub(r'displayName = ".*?"', f'displayName = "{target_vm_name}"', original_content)
updated_content = re.sub(r'uuid.bios = ".*?"', f'uuid.bios = "{new_uuid_bios}"', updated_content)
updated_content = re.sub(r'uuid.location = ".*?"', f'uuid.location = "{new_uuid_location}"', updated_content)
updated_content = re.sub(r'ethernet0.generatedAddress = ".*?"',
f'ethernet0.generatedAddress = "{new_mac_address}"',
updated_content)
updated_content = re.sub(r'vmci0.id = ".*?"', f'vmci0.id = "{new_vmci_id}"', updated_content)
# Write the updated content back to the file
with open(vmx_path, 'w') as file:
file.write(updated_content)
logger.info(".vmx file updated successfully.")
vmx_file_base_name = os.path.splitext(vmx_file)[0]
assert vmx_file == "Ubuntu.vmx", "The VMX file should be named 'Ubuntu.vmx'."
files_to_rename = ['vmx', 'nvram', 'vmsd', 'vmxf']
for ext in files_to_rename:
original_file = os.path.join(dir_path, f"{vmx_file_base_name}.{ext}")
target_file = os.path.join(dir_path, f"{target_vm_name}.{ext}")
os.rename(original_file, target_file)
# Update the dir_path to the target vm_name, only replace the last character
# Split the path into parts up to the last folder
path_parts = dir_path.rstrip(os.sep).split(os.sep)
path_parts[-1] = target_vm_name
target_dir_path = os.sep.join(path_parts)
os.rename(dir_path, target_dir_path)
logger.info("VM files renamed successfully.")
def _install_vm(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
elif platform.machine().lower() in ['amd64', 'x86_64']:
url = UBUNTU_X86_URL
else:
raise Exception("Unsupported platform or architecture")
# Download the virtual machine image
logger.info("Downloading the virtual machine image...")
downloaded_size = 0
while True:
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)
headers["Range"] = f"bytes={downloaded_size}-"
with requests.get(url, headers=headers, stream=True) as response:
if response.status_code == 416:
# This means the range was not satisfiable, possibly the file was fully downloaded
logger.info("Fully downloaded or the file sized changed.")
break
response.raise_for_status()
total_size = int(response.headers.get('content-length', 0))
with open(downloaded_file_path, "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:
logger.error(f"Download error: {e}")
sleep(1) # Wait for 1 second before retrying
logger.error("Retrying...")
else:
logger.info("Download succeeds.")
break # Download completed successfully
# Unzip the downloaded file
logger.info("Unzipping the downloaded file...☕️")
with zipfile.ZipFile(downloaded_file_path, 'r') as zip_ref:
zip_ref.extractall(os.path.join(vms_dir, vm_name))
logger.info("Files have been successfully extracted to the directory: " + str(os.path.join(vms_dir, vm_name)))
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(vms_dir, vm_name, original_vm_name, original_vm_name + ".vmx"), vm_name)
else:
logger.info(f"Virtual machine exists: {vm_path}")
# 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
def start_vm(vm_path, max_retries=20):
command = f'vmrun {get_vmrun_type()} start "{vm_path}" nogui'
for attempt in range(max_retries):
result = subprocess.run(command, shell=True, text=True, capture_output=True, encoding="utf-8")
if result.returncode == 0:
logger.info("Virtual machine started.")
return True
else:
if "Error" in result.stderr:
logger.error(f"Attempt {attempt + 1} failed with specific error: {result.stderr}")
else:
logger.error(f"Attempt {attempt + 1} failed: {result.stderr}")
if attempt == max_retries - 1:
logger.error("Maximum retry attempts reached, failed to start the virtual machine.")
return False
if not start_vm(vm_path):
raise ValueError("Error encountered during installation, please rerun the code for retrying.")
def get_vm_ip(vm_path, max_retries=20):
command = f'vmrun {get_vmrun_type()} getGuestIPAddress "{vm_path}" -wait'
for attempt in range(max_retries):
result = subprocess.run(command, shell=True, text=True, capture_output=True, encoding="utf-8")
if result.returncode == 0:
return result.stdout.strip()
else:
if "Error" in result.stderr:
logger.error(f"Attempt {attempt + 1} failed with specific error: {result.stderr}")
else:
logger.error(f"Attempt {attempt + 1} failed: {result.stderr}")
if attempt == max_retries - 1:
logger.error("Maximum retry attempts reached, failed to get the IP of virtual machine.")
return None
vm_ip = get_vm_ip(vm_path)
if not vm_ip:
raise ValueError("Error encountered during installation, please rerun the code for retrying.")
# 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=(10, 10))
if response.status_code == 200:
return True
except Exception as e:
logger.error(f"Error: {e}")
logger.error(f"Type: {type(e).__name__}")
logger.error(f"Error detail: {str(e)}")
sleep(2)
return False
# Try downloading the screenshot until successful
while not download_screenshot(vm_ip):
logger.info("Check whether the virtual machine is ready...")
logger.info("Virtual machine is ready. Start to make a snapshot on the virtual machine. It would take a while...")
def create_vm_snapshot(vm_path, max_retries=20):
command = f'vmrun {get_vmrun_type()} snapshot "{vm_path}" "init_state"'
for attempt in range(max_retries):
result = subprocess.run(command, shell=True, text=True, capture_output=True, encoding="utf-8")
if result.returncode == 0:
logger.info("Snapshot created.")
return True
else:
if "Error" in result.stderr:
logger.error(f"Attempt {attempt + 1} failed with specific error: {result.stderr}")
else:
logger.error(f"Attempt {attempt + 1} failed: {result.stderr}")
if attempt == max_retries - 1:
logger.error("Maximum retry attempts reached, failed to create snapshot.")
return False
# Create a snapshot of the virtual machine
if create_vm_snapshot(vm_path, max_retries=MAX_RETRY_TIMES):
return vm_path
else:
raise ValueError("Error encountered during installation, please rerun the code for retrying.")
class VMwareVMManager(VMManager):
def __init__(self, registry_path=REGISTRY_PATH):
self.registry_path = registry_path
self.lock = FileLock(".vmware_lck", timeout=10)
self.initialize_registry()
def initialize_registry(self):
with self.lock: # Locking during initialization
if not os.path.exists(self.registry_path):
with open(self.registry_path, 'w') as file:
file.write('')
def add_vm(self, vm_path, region=None):
assert region in [None, 'local'], "For VMware provider, the region should be neither None or 'local'."
with self.lock:
with open(self.registry_path, 'r') as file:
lines = file.readlines()
new_lines = lines + [f'{vm_path}|free\n']
with open(self.registry_path, 'w') as file:
file.writelines(new_lines)
def occupy_vm(self, vm_path, pid, region=None):
assert region in [None, 'local'], "For VMware provider, the region should be neither None or 'local'."
with self.lock:
new_lines = []
with open(self.registry_path, 'r') as file:
lines = file.readlines()
for line in lines:
registered_vm_path, _ = line.strip().split('|')
if registered_vm_path == vm_path:
new_lines.append(f'{registered_vm_path}|{pid}\n')
else:
new_lines.append(line)
with open(self.registry_path, 'w') as file:
file.writelines(new_lines)
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('|')
if not os.path.exists(vm_path):
logger.info(f"VM {vm_path} not found, releasing it.")
new_lines.append(f'{vm_path}|free\n')
continue
vm_paths.append(vm_path)
if pid_str == "free":
new_lines.append(line)
continue
if int(pid_str) in active_pids:
new_lines.append(line)
else:
new_lines.append(f'{vm_path}|free\n')
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
os.makedirs(vms_dir, exist_ok=True)
vm_names = os.listdir(vms_dir)
for vm_name in vm_names:
# skip the downloaded .zip file
if vm_name == DOWNLOADED_FILE_NAME:
continue
# Skip the .DS_Store file on macOS
if vm_name == ".DS_Store":
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_free_vms(self):
with self.lock: # Lock when reading the registry
free_vms = []
with open(self.registry_path, 'r') as file:
lines = file.readlines()
for line in lines:
vm_path, pid_str = line.strip().split('|')
if pid_str == "free":
free_vms.append((vm_path, pid_str))
return free_vms
def get_vm_path(self, region=None):
assert region in [None, 'local'], "For VMware provider, the region should be neither None or 'local'."
self.check_and_clean(vms_dir=VMS_DIR)
free_vms_paths = self.list_free_vms()
if len(free_vms_paths) == 0:
# No free virtual machine available, generate a new one
logger.info("No free virtual machine available. Generating a new one, which would take a while...☕")
new_vm_name = generate_new_vm_name(vms_dir=VMS_DIR)
new_vm_path = _install_vm(new_vm_name, vms_dir=VMS_DIR,
downloaded_file_name=DOWNLOADED_FILE_NAME)
self.add_vm(new_vm_path)
self.occupy_vm(new_vm_path, os.getpid())
return new_vm_path
else:
# Choose the first free virtual machine
chosen_vm_path = free_vms_paths[0][0]
self.occupy_vm(chosen_vm_path, os.getpid())
return chosen_vm_path

90
desktop_env/providers/vmware/provider.py Normal file
View File

@@ -0,0 +1,90 @@
import logging
import platform
import subprocess
import time
import os
from desktop_env.providers.base import Provider
logger = logging.getLogger("desktopenv.providers.vmware.VMwareProvider")
logger.setLevel(logging.INFO)
WAIT_TIME = 3
def get_vmrun_type(return_list=False):
if platform.system() == 'Windows' or platform.system() == 'Linux':
if return_list:
return ['-T', 'ws']
else:
return '-T ws'
elif platform.system() == 'Darwin': # Darwin is the system name for macOS
if return_list:
return ['-T', 'fusion']
else:
return '-T fusion'
else:
raise Exception("Unsupported operating system")
class VMwareProvider(Provider):
@staticmethod
def _execute_command(command: list):
result = subprocess.run(command, stdout=subprocess.PIPE, stderr=subprocess.PIPE, timeout=60, text=True,
encoding="utf-8")
if result.returncode != 0:
raise Exception("\033[91m" + result.stdout + result.stderr + "\033[0m")
return result.stdout.strip()
def start_emulator(self, path_to_vm: str, headless: bool):
print("Starting VMware VM...")
logger.info("Starting VMware VM...")
while True:
try:
output = subprocess.check_output(f"vmrun {get_vmrun_type()} list", shell=True, stderr=subprocess.STDOUT)
output = output.decode()
output = output.splitlines()
normalized_path_to_vm = os.path.abspath(os.path.normpath(path_to_vm))
if any(os.path.abspath(os.path.normpath(line)) == normalized_path_to_vm for line in output):
logger.info("VM is running.")
break
else:
logger.info("Starting VM...")
VMwareProvider._execute_command(["vmrun"] + get_vmrun_type(return_list=True) + ["start", path_to_vm]) if not headless else \
VMwareProvider._execute_command(
["vmrun"] + get_vmrun_type(return_list=True) + ["start", path_to_vm, "nogui"])
time.sleep(WAIT_TIME)
except subprocess.CalledProcessError as e:
logger.error(f"Error executing command: {e.output.decode().strip()}")
def get_ip_address(self, path_to_vm: str) -> str:
logger.info("Getting VMware VM IP address...")
while True:
try:
output = VMwareProvider._execute_command(
["vmrun"] + get_vmrun_type(return_list=True) + ["getGuestIPAddress", path_to_vm, "-wait"]
)
logger.info(f"VMware VM IP address: {output}")
return output
except Exception as e:
logger.error(e)
time.sleep(WAIT_TIME)
logger.info("Retrying to get VMware VM IP address...")
def save_state(self, path_to_vm: str, snapshot_name: str):
logger.info("Saving VMware VM state...")
VMwareProvider._execute_command(["vmrun"] + get_vmrun_type(return_list=True) + ["snapshot", path_to_vm, snapshot_name])
time.sleep(WAIT_TIME) # Wait for the VM to save
def revert_to_snapshot(self, path_to_vm: str, snapshot_name: str):
logger.info(f"Reverting VMware VM to snapshot: {snapshot_name}...")
VMwareProvider._execute_command(["vmrun"] + get_vmrun_type(return_list=True) + ["revertToSnapshot", path_to_vm, snapshot_name])
time.sleep(WAIT_TIME) # Wait for the VM to revert
return path_to_vm
def stop_emulator(self, path_to_vm: str):
logger.info("Stopping VMware VM...")
VMwareProvider._execute_command(["vmrun"] + get_vmrun_type(return_list=True) + ["stop", path_to_vm])
time.sleep(WAIT_TIME) # Wait for the VM to stop