mig

Author: runzhen

src/aks-preview/azcli_aks_live_test/configs/ext_matrix_default.json

@@ -24,7 +24,8 @@
             "test_aks_nodepool_add_with_gpu_instance_profile",
             "test_aks_gpu_driver_type",
             "test_aks_nodepool_add_with_enable_managed_gpu",
-            "test_aks_nodepool_update_with_enable_managed_gpu"
+            "test_aks_nodepool_update_with_enable_managed_gpu",
+            "test_aks_nodepool_add_with_gpu_mig_strategy"
         ],
         "pod ip allocation mode static block, missing feature registration": [
             "test_aks_create_with_pod_ip_allocation_mode_static_block"

src/aks-preview/azext_aks_preview/agentpool_decorator.py

@@ -1395,6 +1395,7 @@ def set_up_gpu_mig_strategy(self, agentpool: AgentPool) -> AgentPool:
             if agentpool.gpu_profile.nvidia is None:
                 agentpool.gpu_profile.nvidia = self.models.NvidiaGPUProfile()  # pylint: disable=no-member
             agentpool.gpu_profile.nvidia.mig_strategy = gpu_mig_strategy
+            agentpool.gpu_profile.driver = CONST_GPU_DRIVER_INSTALL
         return agentpool
 
     def set_up_pod_ip_allocation_mode(self, agentpool: AgentPool) -> AgentPool:
@@ -1778,6 +1779,7 @@ def update_gpu_mig_strategy(self, agentpool: AgentPool) -> AgentPool:
             if agentpool.gpu_profile.nvidia is None:
                 agentpool.gpu_profile.nvidia = self.models.NvidiaGPUProfile()  # pylint: disable=no-member
             agentpool.gpu_profile.nvidia.mig_strategy = gpu_mig_strategy
+            agentpool.gpu_profile.driver = CONST_GPU_DRIVER_INSTALL
         return agentpool
 
     def update_artifact_streaming(self, agentpool: AgentPool) -> AgentPool:

src/aks-preview/azext_aks_preview/tests/latest/test_aks_commands.py

@@ -5515,6 +5515,64 @@ def test_aks_nodepool_add_with_gpu_instance_profile(
             checks=[self.is_empty()],
         )
 
+    @live_only()
+    @AllowLargeResponse()
+    @AKSCustomResourceGroupPreparer(
+        random_name_length=17, name_prefix="clitest", location="westus3"
+    )
+    def test_aks_nodepool_add_with_gpu_mig_strategy(
+        self, resource_group, resource_group_location
+    ):
+        aks_name = self.create_random_name("cliakstest", 16)
+        node_pool_name = self.create_random_name("c", 6)
+        node_pool_name_second = self.create_random_name("c", 6)
+        self.kwargs.update(
+            {
+                "resource_group": resource_group,
+                "name": aks_name,
+                "node_pool_name": node_pool_name,
+                "node_pool_name_second": node_pool_name_second,
+                "ssh_key_value": self.generate_ssh_keys(),
+            }
+        )
+
+        create_cmd = (
+            "aks create --resource-group={resource_group} --name={name} "
+            "--nodepool-name {node_pool_name} -c 1 "
+            "--ssh-key-value={ssh_key_value}"
+        )
+        self.cmd(
+            create_cmd,
+            checks=[
+                self.check("provisioningState", "Succeeded"),
+            ],
+        )
+
+        # nodepool add with gpu-mig-strategy
+        self.cmd(
+            "aks nodepool add "
+            "--resource-group={resource_group} "
+            "--cluster-name={name} "
+            "--name={node_pool_name_second} "
+            "--enable-managed-gpu=true "
+            "--gpu-instance-profile=MIG3g "
+            "--gpu-mig-strategy=Single "
+            "-c 1 "
+            "--aks-custom-headers UseGPUDedicatedVHD=true "
+            "--node-vm-size=Standard_NC24ads_A100_v4",
+            checks=[
+                self.check("provisioningState", "Succeeded"),
+                self.check("gpuInstanceProfile", "MIG3g"),
+                self.check("gpuProfile.nvidia.migStrategy", "Single"),
+            ],
+        )
+
+        # delete
+        self.cmd(
+            "aks delete -g {resource_group} -n {name} --yes --no-wait",
+            checks=[self.is_empty()],
+        )
+
     @live_only()  # live only due to workspace is not mocked correctly and role assignment is not mocked
     @AllowLargeResponse()
     @AKSCustomResourceGroupPreparer(
@@ -6907,7 +6965,7 @@ def test_aks_nodepool_add_with_enable_managed_gpu(
         self.cmd(
             "aks nodepool add --resource-group={resource_group} --cluster-name={name} --name={node_pool_name} "
             "--node-vm-size={node_vm_size} --node-count 1 "
-            " --enable-managed-gpu",
+            "--enable-managed-gpu=true ",
             checks=[
                 self.check("provisioningState", "Succeeded"),
                 self.check("gpuProfile.driver", "Install"),
@@ -16644,7 +16702,7 @@ def test_aks_nodepool_update_with_enable_managed_gpu(
             "--resource-group={resource_group} "
             "--cluster-name={name} "
             "--name={node_pool_name} "
-            "--enable-managed-gpu",
+            "--enable-managed-gpu=true ",
             checks=[
                 self.check("provisioningState", "Succeeded"),
                 self.check("gpuProfile.driver", "Install"),