Avoid deprecated method to get Azure subscription ID

dask_cloudprovider/azure/azurevm.py

@@ -1,6 +1,7 @@
 import base64
 import json
 import uuid
+from typing import Optional
 
 import dask
 from dask_cloudprovider.config import ClusterConfig
@@ -11,9 +12,9 @@
     WorkerMixin,
 )
 from dask_cloudprovider.exceptions import ConfigError
+from dask_cloudprovider.azure.utils import _get_default_subscription
 
 try:
-    from azure.common.credentials import get_cli_profile
     from azure.mgmt.network import NetworkManagementClient
     from azure.mgmt.compute import ComputeManagementClient
     from azure.identity import DefaultAzureCredential
@@ -266,7 +267,7 @@ class AzureVMCluster(VMCluster):
     security_group: str
         The security group to apply to your VMs.
         This must allow ports 8786-8787 from wherever you are running this from.
-        List your security greoups with ``az network nsg list``.
+        List your security groups with ``az network nsg list``.
     public_ingress: bool
         Assign a public IP address to the scheduler. Default ``True``.
     vm_size: str
@@ -336,7 +337,10 @@ class AzureVMCluster(VMCluster):
         All three fields "name", "publisher", "product" must be passed in the dictionary if set. For e.g.
 
         ``{"name": "ngc-base-version-21-02-2", "publisher": "nvidia","product": "ngc_azure_17_11"}``
-
+    subscription_id: str (optional)
+        The ID of the Azure Subscription to create the virtual machines in. If not specified, then
+        dask-cloudprovider will attempt to use the configured default for the Azure CLI. List your
+        subscriptions with ``az account list``.
 
     Examples
     --------
@@ -464,6 +468,7 @@ def __init__(
         docker_image=None,
         debug: bool = False,
         marketplace_plan: dict = {},
+        subscription_id: Optional[str] = None,
         **kwargs,
     ):
         self.config = ClusterConfig(dask.config.get("cloudprovider.azure", {}))
@@ -480,7 +485,13 @@ def __init__(
         self.public_ingress = self.config.get(
             "azurevm.public_ingress", override_with=public_ingress
         )
-        self.subscription_id = get_cli_profile().get_subscription_id()
+        # We prefer code > Dask config > Azure CLI
+        self.subscription_id = self.config.get(
+            "subscription_id", override_with=subscription_id
+        )
+        if self.subscription_id is None:
+            self.subscription_id = _get_default_subscription()
+
         self.credentials = DefaultAzureCredential()
         self.compute_client = ComputeManagementClient(
             self.credentials, self.subscription_id

dask_cloudprovider/azure/tests/test_azurevm.py

@@ -3,26 +3,13 @@
 import dask
 
 azure_compute = pytest.importorskip("azure.mgmt.compute")
-from azure.common.credentials import get_azure_cli_credentials
 
 from dask_cloudprovider.azure import AzureVMCluster
 from dask.distributed import Client
 from distributed.core import Status
 
 
 def skip_without_credentials(func):
-    try:
-        get_azure_cli_credentials()
-    except FileNotFoundError:
-        return pytest.mark.skip(
-            reason="""
-        You must configure your Azure credentials to run this test.
-
-            $ az login
-
-        """
-        )(func)
-
     rg = dask.config.get("cloudprovider.azure.resource_group", None)
     vnet = dask.config.get("cloudprovider.azure.azurevm.vnet", None)
     security_group = dask.config.get("cloudprovider.azure.azurevm.security_group", None)
@@ -33,9 +20,9 @@ def skip_without_credentials(func):
         You must configure your Azure resource group and vnet to run this test.
 
             $ export DASK_CLOUDPROVIDER__AZURE__LOCATION="<LOCATION>"
-            $ export DASK_CLOUDPROVIDER__AZURE__AZUREVM__RESOURCE_GROUP="<RESOURCE GROUP>"
+            $ export DASK_CLOUDPROVIDER__AZURE__RESOURCE_GROUP="<RESOURCE GROUP>"
             $ export DASK_CLOUDPROVIDER__AZURE__AZUREVM__VNET="<VNET>"
-            $ export DASK_CLOUDPROVIDER__AZURE__AZUREVM__SECURITY_GROUP="<SECUROTY GROUP>"
+            $ export DASK_CLOUDPROVIDER__AZURE__AZUREVM__SECURITY_GROUP="<SECURITY GROUP>"
 
         """
         )(func)

dask_cloudprovider/azure/utils.py

@@ -1,5 +1,7 @@
 import asyncio
 import datetime
+import json
+import subprocess
 import logging
 
 import aiohttp
@@ -14,6 +16,21 @@
 )
 
 
+def _get_default_subscription() -> str:
+    """
+    Get the default Azure subscription ID, as configured by the Azure CLI.
+    """
+    out = subprocess.check_output(["az", "account", "list", "--query", "[?isDefault]"])
+    accounts = json.loads(out)
+    if accounts:
+        subscription_id = accounts[0]["id"]
+        return subscription_id
+    raise ValueError(
+        "Could not find a default subscription. "
+        "Run 'az account set' to set a default subscription."
+    )
+
+
 class AzurePreemptibleWorkerPlugin(WorkerPlugin):
     """A worker plugin for azure spot instances
 

dask_cloudprovider/cloudprovider.yaml

@@ -64,6 +64,7 @@ cloudprovider:
   azure:
     location: null # The Azure location to launch your cluster
     resource_group: null # The Azure resource group for the cluster
+    subscription_id: null # The Azure subscription ID for the cluster
     azurevm:
       vnet: null # Azure Virtual Network to launch VMs in
       security_group: null # Network security group to allow 8786 and 8787

doc/source/azure.rst

@@ -93,6 +93,32 @@ or specific IP.
 
 Again take note of this security group name for later.
 
+Dask Configuration
+^^^^^^^^^^^^^^^^^^
+
+You'll provide the names or IDs of the Azure resources when you create a :class:`AzureVMCluster`. You can specify
+these values manually, or use Dask's `configuration system <https://docs.dask.org/en/stable/configuration.html>`_
+system. For example, the ``resource_group`` value can be specified using an environment variable:
+
+.. code-block:: console
+
+   $ export DASK_CLOUDPROVIDER__AZURE__RESOURCE_GROUP="<resource group name>"
+   $ python
+
+Or you can set it in a YAML configuration file.
+
+.. code-block:: yaml
+
+   cloudprovider:
+     azure:
+       resource_group: "<resource group name>"
+       azurevm:
+        vnet: "<vnet name>"
+
+Note that the options controlling the VMs are under the `cloudprovider.azure.azurevm` key.
+
+See :doc:`config` for more.
+
 AzureVM
 -------
 

setup.py

@@ -8,10 +8,8 @@
 extras_require = {
     "aws": ["aiobotocore>=0.10.2"],
     "azure": [
-        "azure-mgmt-compute>=18.0.0,<19",
-        "azure-mgmt-network>=16.0.0,<17",
-        "azure-cli-core>=2.15.1,<2.21.0",
-        "msrestazure",
+        "azure-mgmt-compute>=18.0.0",
+        "azure-mgmt-network>=16.0.0",
         "azure-identity",
     ],
     "digitalocean": ["python-digitalocean>=1.15.0"],