Remove calls to self.cluster.call_async() (#456)

* Add call_async method to VMInterface

* Switch to VMInterface.call_async

Author: jacobtomlinson

dask_cloudprovider/azure/azurevm.py

@@ -79,7 +79,7 @@ def __init__(
 
     async def create_vm(self):
         if not self.subnet:
-            [subnet_info, *_] = await self.cluster.call_async(
+            [subnet_info, *_] = await self.call_async(
                 self.cluster.network_client.subnets.list,
                 self.cluster.resource_group,
                 self.vnet,
@@ -96,7 +96,7 @@ async def create_vm(self):
             ],
             "networkSecurityGroup": {
                 "id": (
-                    await self.cluster.call_async(
+                    await self.call_async(
                         self.cluster.network_client.network_security_groups.get,
                         self.cluster.resource_group,
                         self.security_group,
@@ -107,7 +107,7 @@ async def create_vm(self):
             "tags": self.cluster.get_tags(),
         }
         if self.public_ingress:
-            self.public_ip = await self.cluster.call_async(
+            self.public_ip = await self.call_async(
                 self.cluster.network_client.public_ip_addresses.begin_create_or_update(
                     self.cluster.resource_group,
                     self.nic_name,
@@ -124,7 +124,7 @@ async def create_vm(self):
                 "id": self.public_ip.id
             }
             self.cluster._log("Assigned public IP")
-        self.nic = await self.cluster.call_async(
+        self.nic = await self.call_async(
             self.cluster.network_client.network_interfaces.begin_create_or_update(
                 self.cluster.resource_group,
                 self.nic_name,
@@ -198,17 +198,17 @@ async def create_vm(self):
             self.cluster._log(
                 f"with parameters\n{json.dumps(vm_parameters, sort_keys=True, indent=2)}"
             )
-        await self.cluster.call_async(
+        await self.call_async(
             self.cluster.compute_client.virtual_machines.begin_create_or_update(
                 self.cluster.resource_group, self.name, vm_parameters
             ).wait
         )
-        self.vm = await self.cluster.call_async(
+        self.vm = await self.call_async(
             self.cluster.compute_client.virtual_machines.get,
             self.cluster.resource_group,
             self.name,
         )
-        self.nic = await self.cluster.call_async(
+        self.nic = await self.call_async(
             self.cluster.network_client.network_interfaces.get,
             self.cluster.resource_group,
             self.nic.name,
@@ -222,32 +222,32 @@ async def create_vm(self):
             return private_ip_address, None
 
     async def destroy_vm(self):
-        await self.cluster.call_async(
+        await self.call_async(
             self.cluster.compute_client.virtual_machines.begin_delete(
                 self.cluster.resource_group, self.name
             ).wait
         )
         self.cluster._log(f"Terminated VM {self.name}")
-        for disk in await self.cluster.call_async(
+        for disk in await self.call_async(
             self.cluster.compute_client.disks.list_by_resource_group,
             self.cluster.resource_group,
         ):
             if self.name in disk.name:
-                await self.cluster.call_async(
+                await self.call_async(
                     self.cluster.compute_client.disks.begin_delete(
                         self.cluster.resource_group,
                         disk.name,
                     ).wait
                 )
         self.cluster._log(f"Removed disks for VM {self.name}")
-        await self.cluster.call_async(
+        await self.call_async(
             self.cluster.network_client.network_interfaces.begin_delete(
                 self.cluster.resource_group, self.nic.name
             ).wait
         )
         self.cluster._log("Deleted network interface")
         if self.public_ingress:
-            await self.cluster.call_async(
+            await self.call_async(
                 self.cluster.network_client.public_ip_addresses.begin_delete(
                     self.cluster.resource_group, self.public_ip.name
                 ).wait

dask_cloudprovider/digitalocean/droplet.py

@@ -56,20 +56,20 @@ async def create_vm(self):
             backups=False,
             user_data=self.cluster.render_process_cloud_init(self),
         )
-        await self.cluster.call_async(self.droplet.create)
+        await self.call_async(self.droplet.create)
         for action in self.droplet.get_actions():
             while action.status != "completed":
                 action.load()
                 await asyncio.sleep(0.1)
         while self.droplet.ip_address is None:
-            await self.cluster.call_async(self.droplet.load)
+            await self.call_async(self.droplet.load)
             await asyncio.sleep(0.1)
         self.cluster._log(f"Created droplet {self.name}")
 
         return self.droplet.ip_address, None
 
     async def destroy_vm(self):
-        await self.cluster.call_async(self.droplet.destroy)
+        await self.call_async(self.droplet.destroy)
         self.cluster._log(f"Terminated droplet {self.name}")
 
 

dask_cloudprovider/gcp/instances.py

@@ -206,7 +206,7 @@ async def create_vm(self):
         self.gcp_config = self.create_gcp_config()
 
         try:
-            inst = await self.cluster.call_async(
+            inst = await self.call_async(
                 self.cluster.compute.instances()
                 .insert(project=self.projectid, zone=self.zone, body=self.gcp_config)
                 .execute
@@ -232,7 +232,7 @@ async def create_vm(self):
 
     async def get_internal_ip(self):
         return (
-            await self.cluster.call_async(
+            await self.call_async(
                 self.cluster.compute.instances()
                 .list(
                     project=self.projectid, zone=self.zone, filter=f"name={self.name}"
@@ -243,7 +243,7 @@ async def get_internal_ip(self):
 
     async def get_external_ip(self):
         return (
-            await self.cluster.call_async(
+            await self.call_async(
                 self.cluster.compute.instances()
                 .list(
                     project=self.projectid, zone=self.zone, filter=f"name={self.name}"
@@ -253,7 +253,7 @@ async def get_external_ip(self):
         )["items"][0]["networkInterfaces"][0]["accessConfigs"][0]["natIP"]
 
     async def update_status(self):
-        d = await self.cluster.call_async(
+        d = await self.call_async(
             self.cluster.compute.instances()
             .list(project=self.projectid, zone=self.zone, filter=f"name={self.name}")
             .execute
@@ -276,7 +276,7 @@ def expand_source_image(self, source_image):
 
     async def close(self):
         self.cluster._log(f"Closing Instance: {self.name}")
-        await self.cluster.call_async(
+        await self.call_async(
             self.cluster.compute.instances()
             .delete(project=self.projectid, zone=self.zone, instance=self.name)
             .execute

dask_cloudprovider/generic/vmcluster.py

@@ -63,6 +63,10 @@ async def close(self):
         await self.destroy_vm()
         await super().close()
 
+    async def call_async(self, f, *args, **kwargs):
+        """Run a blocking function in a thread as a coroutine."""
+        return await self.call_async(f, *args, **kwargs)
+
 
 class SchedulerMixin(object):
     """A mixin for Schedulers."""

dask_cloudprovider/hetzner/vserver.py

@@ -51,7 +51,7 @@ def __init__(
         self.docker_image = docker_image
 
     async def create_vm(self):
-        await self.cluster.call_async(
+        await self.call_async(
             self.client.servers.create,
             server_type=self.server_type,
             image=self.image,
@@ -62,14 +62,14 @@ async def create_vm(self):
         self.server = self.client.servers.get_by_name(self.name)
         for action in self.server.get_actions():
             while action.status != Action.STATUS_SUCCESS:
-                await self.cluster.call_async(action.reload)
+                await self.call_async(action.reload)
                 await asyncio.sleep(0.1)
         self.cluster._log(f"Created Hetzner vServer {self.name}")
 
         return self.server.public_net.ipv4.ip, None
 
     async def destroy_vm(self):
-        await self.cluster.call_async(self.client.servers.delete, server=self.server)
+        await self.call_async(self.client.servers.delete, server=self.server)
         self.cluster._log(f"Terminated vServer {self.name}")
 
 

dask_cloudprovider/openstack/instances.py

@@ -111,13 +111,13 @@ async def create_and_assign_floating_ip(self, conn):
         """Create and assign a floating IP to the instance."""
         try:
             # Create a floating IP
-            floating_ip = await self.cluster.call_async(
+            floating_ip = await self.call_async(
                 conn.network.create_ip,
                 floating_network_id=self.config["external_network_id"],
             )
 
             # Find the first port of the instance
-            ports = await self.cluster.call_async(
+            ports = await self.call_async(
                 conn.network.ports,
                 device_id=self.instance.id
             )
@@ -126,7 +126,7 @@ async def create_and_assign_floating_ip(self, conn):
                 raise RuntimeError(f"No network ports found for instance {self.instance.id}")
 
             # Assign the floating IP to the instance's port
-            await self.cluster.call_async(
+            await self.call_async(
                 conn.network.update_ip,
                 floating_ip,
                 port_id=ports[0].id
@@ -170,7 +170,7 @@ async def destroy_vm(self):
         try:
             instance = conn.compute.get_server(self.instance.id)
             if instance:
-                await self.cluster.call_async(conn.compute.delete_server, instance.id)
+                await self.call_async(conn.compute.delete_server, instance.id)
                 self.cluster._log(f"Terminated instance {self.name}")
             else:
                 self.cluster._log(f"Instance {self.name} not found or already deleted.")