{AKS} Add MIG profiles to aks-preview

src/aks-preview/HISTORY.rst

@@ -12,6 +12,10 @@ To release a new version, please select a new version number (usually plus 1 to
 Pending
 +++++++
 
+22.0.0b11
+++++++++++
+* `az aks nodepool add/update`: Add `--gpu-mig-profiles` to configure an ordered list of NVIDIA MIG partition profiles.
+
 22.0.0b10
 +++++++++
 * Bump API version to 2026-07-02-preview.

src/aks-preview/azext_aks_preview/_help.py

@@ -2632,6 +2632,10 @@
         - name: --gpu-mig-strategy
           type: string
           short-summary: Specify the MIG (Multi-Instance GPU) strategy for managed MIG support. Valid values are "Single" and "Mixed". When not specified, managed MIG is disabled.
+        - name: --gpu-mig-profiles
+          type: string array
+          short-summary: Space-separated ordered list of MIG partition profiles. Valid values are "MIG1g", "MIG2g", "MIG3g", "MIG4g", and "MIG7g".
+          long-summary: Duplicate profiles are allowed and meaningful. "Single" accepts one profile; "Mixed" accepts an ordered profile layout. This option cannot be used with --gpu-instance-profile. The service validates the VM size and profile layout.
         - name: --ssh-access
           type: string
           short-summary: Configure SSH setting for the node pool. Use "disabled" to disable SSH access, "localuser" to enable SSH access using private key.
@@ -2718,6 +2722,8 @@
           text: az aks nodepool add -g MyResourceGroup -n managedsystem1 --cluster-name MyManagedCluster --mode ManagedSystem
         - name: Create a node pool with blue-green upgrade strategy and default parameters
           text: az aks nodepool add -g MyResourceGroup -n nodepool1 --cluster-name MyManagedCluster --upgrade-strategy BlueGreen
+        - name: Create a node pool with an ordered mixed MIG profile layout
+          text: az aks nodepool add -g MyResourceGroup -n nodepool1 --cluster-name MyManagedCluster --enable-managed-gpu --gpu-mig-strategy Mixed --gpu-mig-profiles MIG1g MIG2g MIG2g MIG2g
 """
 
 helps['aks nodepool scale'] = """
@@ -2797,6 +2803,10 @@
     short-summary: Update a node pool properties.
     long-summary: Update a node pool to enable/disable cluster-autoscaler or change min-count or max-count.  When called with no optional arguments this attempts to move the node pool to its goal state without changing the current node pool configuration. This can be used to move out of a non succeeded state.
     parameters:
+        - name: --gpu-mig-profiles
+          type: string array
+          short-summary: Space-separated ordered list of MIG partition profiles. Valid values are "MIG1g", "MIG2g", "MIG3g", "MIG4g", and "MIG7g".
+          long-summary: Duplicate profiles are allowed and meaningful. The service rejects changes to this immutable property. This option cannot be used with --gpu-instance-profile.
         - name: --enable-cluster-autoscaler -e
           type: bool
           short-summary: Enable cluster autoscaler. For VMSS pools, enables autoscaler on the pool. For VirtualMachines pools, converts all manual scale profiles to autoscale profiles using the same min/max counts.

src/aks-preview/azext_aks_preview/_params.py

@@ -2516,6 +2516,13 @@ def load_arguments(self, _):
             is_preview=True,
             help="Specify the GPU Multi-Instance GPU (MIG) strategy. Allowed values: Single, Mixed.",
         )
+        c.argument(
+            "gpu_mig_profiles",
+            arg_type=get_enum_type(gpu_instance_profiles),
+            nargs="+",
+            is_preview=True,
+            help="Space-separated ordered list of GPU MIG profiles. Allowed values: MIG1g, MIG2g, MIG3g, MIG4g, MIG7g.",
+        )
         # in creation scenario, use "localuser" as default
         c.argument(
             'ssh_access',
@@ -2714,6 +2721,13 @@ def load_arguments(self, _):
             is_preview=True,
             help="Specify the GPU Multi-Instance GPU (MIG) strategy. Allowed values: Single, Mixed.",
         )
+        c.argument(
+            "gpu_mig_profiles",
+            arg_type=get_enum_type(gpu_instance_profiles),
+            nargs="+",
+            is_preview=True,
+            help="Space-separated ordered list of GPU MIG profiles. Allowed values: MIG1g, MIG2g, MIG3g, MIG4g, MIG7g.",
+        )
         # prepared image specification
         c.argument(
             'prepared_image_specification_id',

src/aks-preview/azext_aks_preview/agentpool_decorator.py

@@ -858,6 +858,23 @@ def get_gpu_mig_strategy(self) -> Union[str, None]:
 
         return gpu_mig_strategy
 
+    def get_gpu_mig_profiles(self) -> Union[List[str], None]:
+        """Obtain the value of gpu_mig_profiles.
+        :return: list of strings or None
+        """
+        gpu_mig_profiles = self.raw_param.get("gpu_mig_profiles")
+
+        if self.decorator_mode == DecoratorMode.CREATE:
+            if (
+                self.agentpool and
+                self.agentpool.gpu_profile is not None and
+                self.agentpool.gpu_profile.nvidia is not None and
+                self.agentpool.gpu_profile.nvidia.mig_profiles is not None
+            ):
+                gpu_mig_profiles = self.agentpool.gpu_profile.nvidia.mig_profiles
+
+        return gpu_mig_profiles
+
     def get_enable_secure_boot(self) -> bool:
         """Obtain the value of enable_secure_boot.
         :return: bool
@@ -1595,6 +1612,20 @@ def set_up_gpu_mig_strategy(self, agentpool: AgentPool) -> AgentPool:
             agentpool.gpu_profile.driver = CONST_GPU_DRIVER_INSTALL
         return agentpool
 
+    def set_up_gpu_mig_profiles(self, agentpool: AgentPool) -> AgentPool:
+        """Set up gpu mig profiles for the AgentPool object."""
+        self._ensure_agentpool(agentpool)
+
+        gpu_mig_profiles = self.context.get_gpu_mig_profiles()
+        if gpu_mig_profiles is not None:
+            if agentpool.gpu_profile is None:
+                agentpool.gpu_profile = self.models.GPUProfile()  # pylint: disable=no-member
+            if agentpool.gpu_profile.nvidia is None:
+                agentpool.gpu_profile.nvidia = self.models.NvidiaGPUProfile()  # pylint: disable=no-member
+            agentpool.gpu_profile.nvidia.mig_profiles = gpu_mig_profiles
+            agentpool.gpu_profile.driver = CONST_GPU_DRIVER_INSTALL
+        return agentpool
+
     def set_up_pod_ip_allocation_mode(self, agentpool: AgentPool) -> AgentPool:
         """Set up pod ip allocation mode for the AgentPool object."""
         self._ensure_agentpool(agentpool)
@@ -1796,6 +1827,8 @@ def construct_agentpool_profile_preview(self) -> AgentPool:
         agentpool = self.set_up_driver_type(agentpool)
         # set up gpu_mig_strategy
         agentpool = self.set_up_gpu_mig_strategy(agentpool)
+        # set up gpu_mig_profiles
+        agentpool = self.set_up_gpu_mig_profiles(agentpool)
         # set up agentpool ssh access
         agentpool = self.set_up_ssh_access(agentpool)
         # set up agentpool pod ip allocation mode
@@ -2013,6 +2046,20 @@ def update_gpu_mig_strategy(self, agentpool: AgentPool) -> AgentPool:
             agentpool.gpu_profile.driver = CONST_GPU_DRIVER_INSTALL
         return agentpool
 
+    def update_gpu_mig_profiles(self, agentpool: AgentPool) -> AgentPool:
+        """Update gpu mig profiles for the AgentPool object."""
+        self._ensure_agentpool(agentpool)
+
+        gpu_mig_profiles = self.context.get_gpu_mig_profiles()
+        if gpu_mig_profiles is not None:
+            if agentpool.gpu_profile is None:
+                agentpool.gpu_profile = self.models.GPUProfile()  # pylint: disable=no-member
+            if agentpool.gpu_profile.nvidia is None:
+                agentpool.gpu_profile.nvidia = self.models.NvidiaGPUProfile()  # pylint: disable=no-member
+            agentpool.gpu_profile.nvidia.mig_profiles = gpu_mig_profiles
+            agentpool.gpu_profile.driver = CONST_GPU_DRIVER_INSTALL
+        return agentpool
+
     def update_artifact_streaming(self, agentpool: AgentPool) -> AgentPool:
         """Update artifact streaming property for the AgentPool object.
         :return: the AgentPool object
@@ -2289,6 +2336,9 @@ def update_agentpool_profile_preview(self, agentpools: List[AgentPool] = None) -
         # update gpu mig strategy
         agentpool = self.update_gpu_mig_strategy(agentpool)
 
+        # update gpu mig profiles
+        agentpool = self.update_gpu_mig_profiles(agentpool)
+
         # update crg id
         agentpool = self.update_crg(agentpool)
 

src/aks-preview/azext_aks_preview/custom.py

@@ -2325,6 +2325,7 @@ def aks_agentpool_add(
     gpu_driver=None,
     driver_type=None,
     gpu_mig_strategy=None,
+    gpu_mig_profiles=None,
     ssh_access=CONST_SSH_ACCESS_LOCALUSER,
     # trusted launch
     enable_secure_boot=False,
@@ -2425,6 +2426,7 @@ def aks_agentpool_update(
     zones=None,
     gpu_driver=None,
     gpu_mig_strategy=None,
+    gpu_mig_profiles=None,
     # crg
     crg_id=None,
     # prepared image specification

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

@@ -631,6 +631,35 @@ def common_get_gpu_mig_strategy(self):
         ctx_0.attach_agentpool(agentpool_0)
         self.assertEqual(ctx_0.get_gpu_mig_strategy(), "Single")
 
+    def common_get_gpu_mig_profiles(self):
+        profiles = ["MIG1g", "MIG2g", "MIG2g", "MIG2g"]
+        ctx = AKSPreviewAgentPoolContext(
+            self.cmd,
+            AKSAgentPoolParamDict({"gpu_mig_profiles": None}),
+            self.models,
+            DecoratorMode.CREATE,
+            self.agentpool_decorator_mode,
+        )
+        self.assertIsNone(ctx.get_gpu_mig_profiles())
+
+        agentpool = self.create_initialized_agentpool_instance(
+            gpu_profile=self.models.GPUProfile(
+                nvidia=self.models.NvidiaGPUProfile(mig_profiles=profiles)
+            )
+        )
+        ctx.attach_agentpool(agentpool)
+        self.assertEqual(ctx.get_gpu_mig_profiles(), profiles)
+
+        requested_profiles = ["MIG7g", "MIG1g"]
+        ctx = AKSPreviewAgentPoolContext(
+            self.cmd,
+            AKSAgentPoolParamDict({"gpu_mig_profiles": requested_profiles}),
+            self.models,
+            DecoratorMode.CREATE,
+            self.agentpool_decorator_mode,
+        )
+        self.assertEqual(ctx.get_gpu_mig_profiles(), requested_profiles)
+
     def common_get_os_sku(self):
         # default
         ctx_1 = AKSPreviewAgentPoolContext(
@@ -1471,6 +1500,9 @@ def test_get_driver_type(self):
     def test_get_gpu_mig_strategy(self):
         self.common_get_gpu_mig_strategy()
 
+    def test_get_gpu_mig_profiles(self):
+        self.common_get_gpu_mig_profiles()
+
     def test_get_enable_secure_boot(self):
         self.common_get_enable_secure_boot()
 
@@ -2035,6 +2067,52 @@ def common_set_up_gpu_profile(self):
         )
         self.assertEqual(dec_agentpool_1, ground_truth_agentpool_1)
 
+    def common_set_up_gpu_mig_profiles(self):
+        profiles = ["MIG1g", "MIG2g", "MIG2g", "MIG2g"]
+        dec = AKSPreviewAgentPoolAddDecorator(
+            self.cmd,
+            self.client,
+            {
+                "gpu_mig_strategy": "Mixed",
+                "gpu_mig_profiles": profiles,
+            },
+            self.resource_type,
+            self.agentpool_decorator_mode,
+        )
+        with self.assertRaises(CLIInternalError):
+            dec.set_up_gpu_mig_profiles(None)
+
+        agentpool = self.create_initialized_agentpool_instance(
+            restore_defaults=False,
+            gpu_profile=self.models.GPUProfile(
+                nvidia=self.models.NvidiaGPUProfile(management_mode="Managed")
+            ),
+        )
+        dec.context.attach_agentpool(agentpool)
+        agentpool = dec.set_up_gpu_mig_strategy(agentpool)
+        agentpool = dec.set_up_gpu_mig_profiles(agentpool)
+
+        self.assertEqual(agentpool.gpu_profile.driver, CONST_GPU_DRIVER_INSTALL)
+        self.assertEqual(agentpool.gpu_profile.nvidia.management_mode, "Managed")
+        self.assertEqual(agentpool.gpu_profile.nvidia.mig_strategy, "Mixed")
+        self.assertEqual(agentpool.gpu_profile.nvidia.mig_profiles, profiles)
+        self.assertEqual(
+            agentpool.as_dict()["properties"]["gpuProfile"]["nvidia"]["migProfiles"],
+            profiles,
+        )
+
+        omitted = AKSPreviewAgentPoolAddDecorator(
+            self.cmd,
+            self.client,
+            {"gpu_mig_profiles": None},
+            self.resource_type,
+            self.agentpool_decorator_mode,
+        )
+        agentpool_without_gpu = self.create_initialized_agentpool_instance(restore_defaults=False)
+        omitted.context.attach_agentpool(agentpool_without_gpu)
+        omitted.set_up_gpu_mig_profiles(agentpool_without_gpu)
+        self.assertIsNone(agentpool_without_gpu.gpu_profile)
+
     def common_set_up_secure_boot(self):
         dec_1 = AKSPreviewAgentPoolAddDecorator(
             self.cmd,
@@ -2578,6 +2656,9 @@ def test_set_up_skip_gpu_driver_install(self):
     def test_set_up_gpu_profile(self):
         self.common_set_up_gpu_profile()
 
+    def test_set_up_gpu_mig_profiles(self):
+        self.common_set_up_gpu_mig_profiles()
+
     def test_set_up_secure_boot(self):
         self.common_set_up_secure_boot()
 
@@ -3752,6 +3833,51 @@ def common_update_gpu_profile(self):
         )
         self.assertEqual(dec_agentpool_1, ground_truth_agentpool_1)        
 
+    def common_update_gpu_mig_profiles(self):
+        existing_profiles = ["MIG7g"]
+        agentpool = self.create_initialized_agentpool_instance(
+            gpu_profile=self.models.GPUProfile(
+                driver=CONST_GPU_DRIVER_INSTALL,
+                nvidia=self.models.NvidiaGPUProfile(
+                    management_mode="Managed",
+                    mig_strategy="Single",
+                    mig_profiles=existing_profiles,
+                ),
+            )
+        )
+
+        omitted = AKSPreviewAgentPoolUpdateDecorator(
+            self.cmd,
+            self.client,
+            {"gpu_mig_profiles": None},
+            self.resource_type,
+            self.agentpool_decorator_mode,
+        )
+        with self.assertRaises(CLIInternalError):
+            omitted.update_gpu_mig_profiles(None)
+        omitted.context.attach_agentpool(agentpool)
+        omitted.update_gpu_mig_profiles(agentpool)
+        self.assertEqual(agentpool.gpu_profile.nvidia.mig_profiles, existing_profiles)
+
+        profiles = ["MIG1g", "MIG2g", "MIG2g", "MIG2g"]
+        requested = AKSPreviewAgentPoolUpdateDecorator(
+            self.cmd,
+            self.client,
+            {"gpu_mig_profiles": profiles},
+            self.resource_type,
+            self.agentpool_decorator_mode,
+        )
+        requested.context.attach_agentpool(agentpool)
+        requested.update_gpu_mig_profiles(agentpool)
+        self.assertEqual(agentpool.gpu_profile.nvidia.mig_profiles, profiles)
+        self.assertEqual(agentpool.gpu_profile.nvidia.management_mode, "Managed")
+        self.assertEqual(agentpool.gpu_profile.nvidia.mig_strategy, "Single")
+
+        payload = agentpool.as_dict()
+        if self.agentpool_decorator_mode == AgentPoolDecoratorMode.STANDALONE:
+            payload = payload["properties"]
+        self.assertEqual(payload["gpuProfile"]["nvidia"]["migProfiles"], profiles)
+
     def common_test_process_dns_overrides_helper(self):
         from azext_aks_preview._helpers import process_dns_overrides
 
@@ -3851,6 +3977,9 @@ def test_update_localdns_profile(self):
     def test_update_gpu_profile(self):
         self.common_update_gpu_profile()
 
+    def test_update_gpu_mig_profiles(self):
+        self.common_update_gpu_mig_profiles()
+
     def test_process_dns_overrides_helper(self):
         self.common_test_process_dns_overrides_helper()
 
@@ -3961,6 +4090,9 @@ def test_process_dns_overrides_helper(self):
     def test_update_gpu_profile(self):
         self.common_update_gpu_profile()
 
+    def test_update_gpu_mig_profiles(self):
+        self.common_update_gpu_mig_profiles()
+
     def test_update_agentpool_profile_preview(self):
         import inspect