Merge pull request #3 from cudoventures/create-vm-pool

Create vm pool

dev/create-vm-pool.py

@@ -0,0 +1,40 @@
+import sys
+import os
+from time import sleep
+
+sys.path.append(os.path.abspath(os.path.join(os.path.dirname(__file__), '..', 'helpers')))
+
+import cudo_api
+
+from cudo_compute import Disk, CreateVMBody
+from cudo_compute.rest import ApiException
+
+
+vm_name = 'launch-lot'
+id = 1
+ids = []
+api = cudo_api.virtual_machines()
+try:
+    for i in range(12):
+        vm_id = vm_name + str(id)
+        disk = Disk(storage_class="STORAGE_CLASS_NETWORK", size_gib=100,
+                    id="my-disk-id-" + vm_id)
+
+        request = CreateVMBody(vm_id=vm_id, machine_type="intel-broadwell",
+                               data_center_id="gb-bournemouth-1", boot_disk_image_id='ubuntu-nvidia-docker',
+                               memory_gib=4, vcpus=2, gpus=0, gpu_model="", boot_disk=disk)
+
+        vm = api.create_vm(cudo_api.project_id_throwable(), request)
+        print(vm)
+        ids.append(vm_id)
+        id += 1
+
+except ApiException as e:
+    print(e)
+
+sleep(80)
+for del_id in ids:
+    res = api.terminate_vm(cudo_api.project_id_throwable(), del_id)
+    print(res)
+
+print("done")

helpers/cudo_api.py

@@ -1,6 +1,16 @@
+from asyncio import timeout
+
 import cudo_compute as cudo
 import os
+from time import sleep
 import importlib.metadata
+from concurrent.futures import ThreadPoolExecutor
+from queue import Queue
+import atexit
+import threading
+
+from cudo_compute.models.create_vm_response import CreateVMResponse
+from cudo_compute.models.vm import VM
 
 home = os.path.expanduser("~")
 
@@ -63,45 +73,32 @@ def project_id_throwable():
 
 
 # APIs
+c, err = client()
+if err:
+    raise Exception(err)
+
+
 def api_keys():
-    c, err = client()
-    if err:
-        raise Exception(err)
     return cudo.APIKeysApi(c)
 
 
 def disks():
-    c, err = client()
-    if err:
-        raise Exception(err)
     return cudo.DisksApi(c)
 
 
 def networks():
-    c, err = client()
-    if err:
-        raise Exception(err)
     return cudo.NetworksApi(c)
 
 
 def object_storage():
-    c, err = client()
-    if err:
-        raise Exception(err)
     return cudo.ObjectStorageApi(c)
 
 
 def permissions():
-    c, err = client()
-    if err:
-        raise Exception(err)
     return cudo.PermissionsApi(c)
 
 
 def projects():
-    c, err = client()
-    if err:
-        raise Exception(err)
     return cudo.ProjectsApi(c)
 
 
@@ -113,21 +110,83 @@ def ssh_keys():
 
 
 def search():
-    c, err = client()
-    if err:
-        raise Exception(err)
     return cudo.SearchApi(c)
 
 
 def user():
-    c, err = client()
-    if err:
-        raise Exception(err)
     return cudo.UserApi(c)
 
 
-def virtual_machines():
-    c, err = client()
-    if err:
-        raise Exception(err)
+def legacy_virtual_machines():
     return cudo.VirtualMachinesApi(c)
+
+
+class PooledVirtualMachinesApi(cudo.VirtualMachinesApi):
+    def __init__(self, api_client=None):
+        self.task_queue = None
+        self.max_workers = 5
+        self.shutdown_event = threading.Event()
+        self.workers_active = False
+        self.executor = None
+        atexit.register(self.stop_workers)
+        super().__init__(api_client)
+
+    def create_vm(self, project_id, create_vm_body, **kwargs):
+        self.start_queue()
+        self.task_queue.put((project_id, create_vm_body))
+        self.start_workers()
+        return CreateVMResponse(id=create_vm_body.vm_id, vm=VM())
+
+    def worker(self):
+        while self.workers_active:
+            if not self.task_queue:
+                break
+            req = self.task_queue.get(timeout=1)
+            create_vm_body = None
+            try:
+                project, create_vm_body = req
+                vm = super().create_vm(project, create_vm_body)
+                print(f"Created VM: {vm.to_dict()}")
+                wait = True
+                while wait:
+                    res = self.get_vm(project, create_vm_body.vm_id)
+                    if (res.vm.state == 'ACTIVE' or res.vm.state == 'FAILED' or res.vm.state == 'STOPPED'
+                            or res.vm.state == 'SUSPENDED' or res.vm.state == 'DELETED'):
+                        wait = False
+                    else:
+                        sleep(5)
+            except Exception as e:
+                if create_vm_body:
+                    print(f"Error creating VM: {create_vm_body.vm_id} {e}")
+                else:
+                    print(f"Error creating VM: {e}")
+
+            self.task_queue.task_done()
+
+    def start_queue(self):
+        if not self.task_queue:
+            self.task_queue = Queue()
+
+    def start_workers(self):
+        if not self.workers_active:
+            self.workers_active = True
+            self.executor = ThreadPoolExecutor(max_workers=self.max_workers)
+
+            for _ in range(self.max_workers):
+                self.executor.submit(self.worker)
+
+    def stop_workers(self):
+        if not self.shutdown_event.is_set():
+            try:
+                self.workers_active = False
+                self.shutdown_event.set()
+
+                self.executor.shutdown(wait=False)
+
+            except Exception as e:
+                print(f"Error shutting down: {e}")
+
+pool = PooledVirtualMachinesApi(c)
+
+def virtual_machines():
+    return pool

pyproject.toml

@@ -4,7 +4,7 @@ build-backend = "hatchling.build"
 
 [project]
 name = "cudo-compute"
-version = "0.2.0"
+version = "0.3.0"
 authors = [
     { name = "Cudo Ventures", email = "[email protected]" },
 ]