Skip to content

Commit

Permalink
create vm pool
Browse files Browse the repository at this point in the history
  • Loading branch information
sean committed Oct 3, 2024
1 parent 4fa5269 commit b9ccb77
Showing 1 changed file with 19 additions and 19 deletions.
38 changes: 19 additions & 19 deletions helpers/cudo_api.py
Original file line number Diff line number Diff line change
Expand Up @@ -130,7 +130,7 @@ def user():
raise Exception(err)
return cudo.UserApi(c)

class ExtendedVirtualMachinesApi(cudo.VirtualMachinesApi):
class PooledVirtualMachinesApi(cudo.VirtualMachinesApi):
def __init__(self, api_client=None):
max_workers=2 #TODO make larger
self.executor = ThreadPoolExecutor(max_workers=max_workers)
Expand All @@ -139,65 +139,65 @@ def __init__(self, api_client=None):
self.shutdown_event = threading.Event()
self.workers_active = False

# Register the shutdown method to be called at exit
atexit.register(self.shutdown)

super().__init__(api_client)
# Additional initialization if needed

def start_workers(self):

if not self.workers_active:
self.workers_active = True
self.shutdown_event.clear()
for _ in range(self.max_workers):
self.executor.submit(self.worker)
print("Workers started.") #TODO

def stop_workers(self):
if self.workers_active:
print("stopping workers") #TODO
self.workers_active = False
for _ in range(self.max_workers):
self.task_queue.put(None)
print("Workers stopped.") #TODO
print("workers stopped") #TODO

def worker(self):
while True:
print("worker getting tasks")
req = self.task_queue.get()
if req is None:
break
project, create_vm_body, kwargs = req
super().create_vm(project, create_vm_body)
print("worker processing task")
try:
project, create_vm_body = req
vm = super().create_vm(project, create_vm_body)
print(f"Created VM: {vm.to_dict()}")
except Exception as e:
print(f"Error creating VM: {create_vm_body.vm_id} {e}")

self.task_queue.task_done()

# Check if the task queue is empty and call shutdown if it is
if self.task_queue.empty():
self.shutdown()

def create_vm(self, project_id, create_vm_body, **kwargs):
print("Adding VM...") #TODO
print("create vm") #TODO
self.task_queue.put((project_id, create_vm_body))
self.start_workers()

# Custom implementation or additional logic
print("Creating VM with custom logic")
# return
return {"id":create_vm_body.vm_id}

def shutdown(self):
if not self.shutdown_event.is_set():
print("shutting down...") # TODO

self.shutdown_event.set()
# Wait for all tasks to be processed
self.task_queue.join()
print("All VMs added. Waiting for all tasks to complete...")

# Stop worker threads
self.stop_workers()

self.executor.shutdown(wait=True)
print("Executor shutdown complete.")
print("shutdown") # TODO


def virtual_machines():
c, err = client()
if err:
raise Exception(err)
return ExtendedVirtualMachinesApi(c)
return PooledVirtualMachinesApi(c) #cudo.VirtualMachinesApi(c)

0 comments on commit b9ccb77

Please sign in to comment.