Fix thread lock in AWS, VirtualBox, and VMware (#49)

* 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

* Add Azure provider

* hk region; debug

* Fix lock

* Remove print

* Remove key_name requirements when allocating aws vm

* Clean README

* Fix reset

* Fix bugs

* Add VirtualBox and Azure providers

* Add VirtualBox OVF link

* Raise exception on macOS host

* Init RAEDME for VBox

* Update VirtualBox VM download link

* Update requirements and setup.py; Improve robustness on Windows

* Fix network adapter

* Go through on Windows machine

* Add default adapter option

* Fix minor error

* Change resolution before creating snapshot

* Fix small error

* Change default provider option

* Fix thread lock

* Refactor for more smooth VMware support

---------

Co-authored-by: Timothyxxx <384084775@qq.com>
Co-authored-by: XinyuanWangCS <xywang626@gmail.com>
Co-authored-by: Tianbao Xie <47296835+Timothyxxx@users.noreply.github.com>
This commit is contained in:
HappySix
2024-06-20 23:54:50 +08:00
committed by GitHub
parent 16c3defe20
commit df70b11cf6
5 changed files with 171 additions and 63 deletions

View File

@@ -300,7 +300,7 @@ def _install_vm(vm_name, vms_dir, downloaded_file_name, original_vm_name="Ubuntu
class VirtualBoxVMManager(VMManager):
def __init__(self, registry_path=REGISTRY_PATH):
self.registry_path = registry_path
self.lock = FileLock(LOCK_FILE_NAME, timeout=10)
self.lock = FileLock(LOCK_FILE_NAME, timeout=60)
self.initialize_registry()
def initialize_registry(self):
@@ -309,7 +309,14 @@ class VirtualBoxVMManager(VMManager):
with open(self.registry_path, 'w') as file:
file.write('')
def add_vm(self, vm_path, region=None):
def add_vm(self, vm_path, lock_needed=True):
if lock_needed:
with self.lock:
self._add_vm(vm_path)
else:
self._add_vm(vm_path)
def _add_vm(self, vm_path, region=None):
assert region in [None, 'local'], "For VirtualBox provider, the region should be neither None or 'local'."
with self.lock:
with open(self.registry_path, 'r') as file:
@@ -318,7 +325,14 @@ class VirtualBoxVMManager(VMManager):
with open(self.registry_path, 'w') as file:
file.writelines(new_lines)
def occupy_vm(self, vm_path, pid, region=None):
def occupy_vm(self, vm_path, pid, lock_needed=True):
if lock_needed:
with self.lock:
self._occupy_vm(vm_path, pid)
else:
self._occupy_vm(vm_path, pid)
def _occupy_vm(self, vm_path, pid, region=None):
assert region in [None, 'local'], "For VirtualBox provider, the region should be neither None or 'local'."
with self.lock:
new_lines = []
@@ -333,7 +347,24 @@ class VirtualBoxVMManager(VMManager):
with open(self.registry_path, 'w') as file:
file.writelines(new_lines)
def check_and_clean(self, vms_dir):
def delete_vm(self, vm_path, lock_needed=True):
if lock_needed:
with self.lock:
self._delete_vm(vm_path)
else:
self._delete_vm(vm_path)
def _delete_vm(self, vm_path):
raise NotImplementedError
def check_and_clean(self, vms_dir, lock_needed=True):
if lock_needed:
with self.lock:
self._check_and_clean(vms_dir)
else:
self._check_and_clean(vms_dir)
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()}
@@ -379,7 +410,14 @@ class VirtualBoxVMManager(VMManager):
if flag:
shutil.rmtree(os.path.join(vms_dir, vm_name))
def list_free_vms(self):
def list_free_vms(self, lock_needed=True):
if lock_needed:
with self.lock:
return self._list_free_vms()
else:
return self._list_free_vms()
def _list_free_vms(self):
with self.lock: # Lock when reading the registry
free_vms = []
with open(self.registry_path, 'r') as file:
@@ -391,20 +429,30 @@ class VirtualBoxVMManager(VMManager):
return free_vms
def get_vm_path(self, region=None):
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
with self.lock:
if not VirtualBoxVMManager.checked_and_cleaned:
VirtualBoxVMManager.checked_and_cleaned = True
self._check_and_clean(vms_dir=VMS_DIR)
allocation_needed = False
with self.lock:
free_vms_paths = self._list_free_vms()
if len(free_vms_paths) == 0:
# No free virtual machine available, generate a new one
allocation_needed = True
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
if allocation_needed:
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,
bridged_adapter_name=region)
self.add_vm(new_vm_path)
self.occupy_vm(new_vm_path, os.getpid())
downloaded_file_name=DOWNLOADED_FILE_NAME,
bridged_adapter_name=region)
with self.lock:
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