Added support for using Marketplace VM images with plans in AzureVMCl…

…uster (#294)

* Optionally add plan information to vm_parameters in Azure if using a Marketplace VM Image

* Modified the variable name and added better explanation

* Removed flag and add marketplace_plan dict so that we can use dask_config as well. Added check to see if correct plan values are set. Override vm_image values from plan information.

* Linting fixes

* Add missing dependency

Co-authored-by: Jacob Tomlinson <jtomlinson@nvidia.com>

dask_cloudprovider/azure/azurevm.py

@@ -43,6 +43,7 @@ def __init__(
         env_vars: dict = {},
         bootstrap: bool = None,
         auto_shutdown: bool = None,
+        marketplace_plan: dict = {},
         **kwargs,
     ):
         super().__init__(*args, **kwargs)
@@ -64,6 +65,7 @@ def __init__(
         self.disk_size = disk_size
         self.auto_shutdown = auto_shutdown
         self.env_vars = env_vars
+        self.marketplace_plan = marketplace_plan
 
     async def create_vm(self):
         [subnet_info, *_] = await self.cluster.call_async(
@@ -155,6 +157,23 @@ async def create_vm(self):
             },
             "tags": self.cluster.get_tags(),
         }
+
+        if self.marketplace_plan:
+            # Ref: https://docs.microsoft.com/en-us/rest/api/compute/virtual-machines/create-or-update#create-a-vm-with-a-marketplace-image-plan. # noqa
+            # Creating a marketplace VM with a plan will override default vm_image values.
+            vm_parameters["plan"] = self.marketplace_plan
+            vm_parameters["storage_profile"]["image_reference"][
+                "sku"
+            ] = self.marketplace_plan["name"]
+            vm_parameters["storage_profile"]["image_reference"][
+                "publisher"
+            ] = self.marketplace_plan["publisher"]
+            vm_parameters["storage_profile"]["image_reference"][
+                "offer"
+            ] = self.marketplace_plan["product"]
+            vm_parameters["storage_profile"]["image_reference"]["version"] = "latest"
+            self.cluster._log("Using Marketplace VM image with a Plan")
+
         self.cluster._log("Creating VM")
         if self.cluster.debug:
             self.cluster._log(
@@ -304,6 +323,14 @@ class AzureVMCluster(VMCluster):
         be created automatically. Default is ``True``.
     debug: bool, optional
         More information will be printed when constructing clusters to enable debugging.
+    marketplace_plan: dict (optional)
+        Plan information dict necessary for creating a virtual machine from Azure Marketplace image or
+        a custom image sourced from a Marketplace image with a plan. Default is {}.
+
+        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"}``
+
 
     Examples
     --------
@@ -430,6 +457,7 @@ def __init__(
         auto_shutdown: bool = None,
         docker_image=None,
         debug: bool = False,
+        marketplace_plan: dict = {},
         **kwargs,
     ):
         self.config = dask.config.get("cloudprovider.azure.azurevm", {})
@@ -504,6 +532,18 @@ def __init__(
         )
         self.docker_image = docker_image or self.config.get("docker_image")
         self.debug = debug
+        self.marketplace_plan = marketplace_plan or self.config.get("marketplace_plan")
+        if self.marketplace_plan:
+            # Check that self.marketplace_plan contains the right options with values
+            if not all(
+                self.marketplace_plan.get(item, "") != ""
+                for item in ["name", "publisher", "product"]
+            ):
+                raise ConfigError(
+                    """To create a virtual machine from Marketplace image or a custom image sourced
+                from a Marketplace image with a plan, all 3 fields 'name', 'publisher' and 'product' must be passed."""
+                )
+
         self.options = {
             "cluster": self,
             "config": self.config,
@@ -515,6 +555,7 @@ def __init__(
             "bootstrap": self.bootstrap,
             "auto_shutdown": self.auto_shutdown,
             "docker_image": self.docker_image,
+            "marketplace_plan": self.marketplace_plan,
         }
         self.scheduler_options = {
             "vm_size": self.scheduler_vm_size,

dask_cloudprovider/cloudprovider.yaml

@@ -70,6 +70,10 @@ cloudprovider:
         version: "latest"
       bootstrap: true # It is assumed that the VHD does not have Docker and needs bootstrapping. Set this to false if using a custom VHD with Docker already installed.
       auto_shutdown: true # Shutdown instances automatically if the scheduler or worker services time out.
+      marketplace_plan:  null # This needs to be passed in if the user wants to use a Marketplace VM with a plan.
+        # name: "ngc-base-version-21-02-2"
+        # publisher: "nvidia"
+        # product: "ngc_azure_17_11"
 
   digitalocean:
     token: null # API token for interacting with the Digital Ocean API

setup.py

@@ -11,6 +11,7 @@
         "azure-mgmt-compute>=18.0.0",
         "azure-mgmt-network>=16.0.0",
         "azure-cli-core>=2.15.1",
+        "msrestazure",
     ],
     "digitalocean": ["python-digitalocean>=1.15.0"],
     "gcp": ["google-api-python-client>=1.12.5", "google-auth>=1.23.0"],