From 70d3887126d0ea3ee5f09d2e4d64732e0a4b7a19 Mon Sep 17 00:00:00 2001
From: Jacob Tomlinson <jacobtomlinson@users.noreply.github.com>
Date: Mon, 16 Sep 2024 11:41:13 +0100
Subject: [PATCH] Make AzureVMCluster subnet configurable (#402)

---
 dask_cloudprovider/azure/azurevm.py   | 24 ++++++++++++++++++------
 dask_cloudprovider/cloudprovider.yaml |  1 +
 2 files changed, 19 insertions(+), 6 deletions(-)

diff --git a/dask_cloudprovider/azure/azurevm.py b/dask_cloudprovider/azure/azurevm.py
index 8b2529e8..0c0ca95b 100644
--- a/dask_cloudprovider/azure/azurevm.py
+++ b/dask_cloudprovider/azure/azurevm.py
@@ -36,6 +36,7 @@ def __init__(
         *args,
         location: str = None,
         vnet: str = None,
+        subnet: str = None,
         public_ingress: bool = None,
         security_group: str = None,
         vm_size: str = None,
@@ -61,6 +62,8 @@ def __init__(
         self.extra_bootstrap = extra_bootstrap
         self.admin_username = "dask"
         self.admin_password = str(uuid.uuid4())[:32]
+        self.vnet = vnet
+        self.subnet = subnet
         self.security_group = security_group
         self.nic_name = f"dask-cloudprovider-nic-{str(uuid.uuid4())[:8]}"
         self.docker_image = docker_image
@@ -75,18 +78,20 @@ def __init__(
         self.extra_vm_options = extra_vm_options or {}
 
     async def create_vm(self):
-        [subnet_info, *_] = await self.cluster.call_async(
-            self.cluster.network_client.subnets.list,
-            self.cluster.resource_group,
-            self.cluster.vnet,
-        )
+        if not self.subnet:
+            [subnet_info, *_] = await self.cluster.call_async(
+                self.cluster.network_client.subnets.list,
+                self.cluster.resource_group,
+                self.vnet,
+            )
+            self.subnet = subnet_info.id
 
         nic_parameters = {
             "location": self.location,
             "ip_configurations": [
                 {
                     "name": self.nic_name,
-                    "subnet": {"id": subnet_info.id},
+                    "subnet": {"id": self.subnet},
                 }
             ],
             "networkSecurityGroup": {
@@ -276,6 +281,9 @@ class AzureVMCluster(VMCluster):
         The resource group to create components in. List your resource groups with ``az group list``.
     vnet: str
         The vnet to attach VM network interfaces to. List your vnets with ``az network vnet list``.
+    subnet: str (optional)
+        The vnet subnet to attach VM network interfaces to.
+        If omitted it will automatically use the first subnet in your vnet.
     security_group: str
         The security group to apply to your VMs.
         This must allow ports 8786-8787 from wherever you are running this from.
@@ -472,6 +480,7 @@ def __init__(
         location: str = None,
         resource_group: str = None,
         vnet: str = None,
+        subnet: str = None,
         security_group: str = None,
         public_ingress: bool = None,
         vm_size: str = None,
@@ -518,6 +527,7 @@ def __init__(
         self.vnet = self.config.get("azurevm.vnet", override_with=vnet)
         if self.vnet is None:
             raise ConfigError("You must configure a vnet")
+        self.subnet = self.config.get("azurevm.subnet", override_with=subnet)
         self.security_group = self.config.get(
             "azurevm.security_group", override_with=security_group
         )
@@ -569,6 +579,8 @@ def __init__(
         self.options = {
             "cluster": self,
             "config": self.config,
+            "subnet": self.subnet,
+            "vnet": self.vnet,
             "security_group": self.security_group,
             "location": self.location,
             "vm_image": self.vm_image,
diff --git a/dask_cloudprovider/cloudprovider.yaml b/dask_cloudprovider/cloudprovider.yaml
index 7ee5c1da..2e65be81 100755
--- a/dask_cloudprovider/cloudprovider.yaml
+++ b/dask_cloudprovider/cloudprovider.yaml
@@ -68,6 +68,7 @@ cloudprovider:
     subscription_id: null # The Azure subscription ID for the cluster
     azurevm:
       vnet: null # Azure Virtual Network to launch VMs in
+      subnet: null # Azure Virtual Network subnet to launch VMs in
       security_group: null # Network security group to allow 8786 and 8787
       public_ingress: true # Assign a public IP address to the scheduler
       vm_size: "Standard_DS1_v2" # Azure VM size to use for scheduler and workers