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:
@@ -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
|
||||
|
||||
Reference in New Issue
Block a user