diff --git a/dev/create-vm-pool.py b/dev/create-vm-pool.py new file mode 100644 index 0000000..f550140 --- /dev/null +++ b/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") diff --git a/helpers/cudo_api.py b/helpers/cudo_api.py index aeb3cb5..c11b51c 100644 --- a/helpers/cudo_api.py +++ b/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 diff --git a/pyproject.toml b/pyproject.toml index ff69b37..a188310 100644 --- a/pyproject.toml +++ b/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 = "dev@cudoventures.com" }, ]