[AKS] az aks nodepool update: Add GPU driver install options install and none for --gpu-driver parameter (#9472)

Author: lilypan26

src/aks-preview/HISTORY.rst

@@ -15,6 +15,7 @@ Pending
 19.0.0b18
 +++++++
 * Vendor new SDK and bump API version to 2025-10-02-preview.
+* Add option `--gpu-driver` to `az aks nodepool update` to select skipping GPU driver installation.
 
 19.0.0b17
 +++++++

src/aks-preview/azext_aks_preview/_help.py

@@ -2427,6 +2427,9 @@
         - name: --final-soak-duration
           type: int
           short-summary: Wait time (in minutes) after all old nodes are drained before removing them. Default is 60 minutes. Only for blue-green upgrades.
+        - name: --gpu-driver
+          type: string
+          short-summary: Whether to install driver for GPU node pool. Possible values are "Install" or "None".
     examples:
       - name: Reconcile the nodepool back to its current state.
         text: az aks nodepool update -g MyResourceGroup -n nodepool1 --cluster-name MyManagedCluster

src/aks-preview/azext_aks_preview/_params.py

@@ -2129,6 +2129,10 @@ def load_arguments(self, _):
             options_list=["--node-vm-size", "-s"],
             completer=get_vm_size_completion_list,
         )
+        c.argument(
+            "gpu_driver",
+            arg_type=get_enum_type(gpu_driver_install_modes)
+        )
 
     with self.argument_context("aks nodepool upgrade") as c:
         # upgrade strategy

src/aks-preview/azext_aks_preview/agentpool_decorator.py

@@ -1644,6 +1644,17 @@ def update_network_profile(self, agentpool: AgentPool) -> AgentPool:
             agentpool.network_profile.allowed_host_ports = allowed_host_ports
         return agentpool
 
+    def update_gpu_profile(self, agentpool: AgentPool) -> AgentPool:
+        self._ensure_agentpool(agentpool)
+
+        gpu_driver = self.context.get_gpu_driver()
+        driver_type = self.context.get_driver_type()
+        if not agentpool.gpu_profile and (gpu_driver or driver_type):
+            agentpool.gpu_profile = self.models.GPUProfile()
+        if gpu_driver is not None:
+            agentpool.gpu_profile.driver = gpu_driver
+        return agentpool
+
     def update_artifact_streaming(self, agentpool: AgentPool) -> AgentPool:
         """Update artifact streaming property for the AgentPool object.
         :return: the AgentPool object
@@ -1807,6 +1818,9 @@ def update_agentpool_profile_preview(self, agentpools: List[AgentPool] = None) -
         # update blue-green upgrade settings
         agentpool = self.update_blue_green_upgrade_settings(agentpool)
 
+        # update gpu profile
+        agentpool = self.update_gpu_profile(agentpool)
+
         return agentpool
 
     def update_auto_scaler_properties(self, agentpool: AgentPool) -> AgentPool:

src/aks-preview/azext_aks_preview/custom.py

@@ -1977,6 +1977,7 @@ def aks_agentpool_update(
     # local DNS
     localdns_config=None,
     node_vm_size=None,
+    gpu_driver=None,
 ):
     # DO NOT MOVE: get all the original parameters and save them as a dictionary
     raw_parameters = locals()

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

@@ -2660,6 +2660,31 @@ def common_update_localdns_profile(self):
         finally:
             os.unlink(config_file_path)
 
+    def common_update_gpu_profile(self):
+        dec_1 = AKSPreviewAgentPoolUpdateDecorator(
+            self.cmd,
+            self.client,
+            {"gpu_driver": "None"},
+            self.resource_type,
+            self.agentpool_decorator_mode,
+        )
+        # fail on passing the wrong agentpool object
+        with self.assertRaises(CLIInternalError):
+            dec_1.update_gpu_profile(None)
+        agentpool_1 = self.create_initialized_agentpool_instance(
+            gpu_profile=self.models.GPUProfile(
+                driver="Install",
+            )
+        )
+        dec_1.context.attach_agentpool(agentpool_1)
+        dec_agentpool_1 = dec_1.update_gpu_profile(agentpool_1)
+        ground_truth_agentpool_1 = self.create_initialized_agentpool_instance(
+            gpu_profile=self.models.GPUProfile(
+                driver="None",
+            )
+        )
+        self.assertEqual(dec_agentpool_1, ground_truth_agentpool_1)        
+
     def common_test_process_dns_overrides_helper(self):
         from azext_aks_preview._helpers import process_dns_overrides
 
@@ -2744,6 +2769,9 @@ def test_update_blue_green_upgrade_settings(self):
     def test_update_localdns_profile(self):
         self.common_update_localdns_profile()
 
+    def test_update_gpu_profile(self):
+        self.common_update_gpu_profile()
+
     def test_process_dns_overrides_helper(self):
         self.common_test_process_dns_overrides_helper()
 
@@ -2837,6 +2865,9 @@ def test_update_localdns_profile(self):
     def test_process_dns_overrides_helper(self):
         self.common_test_process_dns_overrides_helper()
 
+    def test_update_gpu_profile(self):
+        self.common_update_gpu_profile()
+
     def test_update_agentpool_profile_preview(self):
         import inspect
 

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

@@ -132,6 +132,7 @@ def test_update_agentpool_profile_preview_default_behavior(self):
         decorator.update_auto_scaler_properties_vms = Mock(return_value=agentpool)
         decorator.update_upgrade_strategy = Mock(return_value=agentpool)
         decorator.update_blue_green_upgrade_settings = Mock(return_value=agentpool)
+        decorator.update_gpu_profile = Mock(return_value=agentpool)
 
         # Act
         result = decorator.update_agentpool_profile_preview()
@@ -154,6 +155,7 @@ def test_update_agentpool_profile_preview_default_behavior(self):
         decorator.update_auto_scaler_properties_vms.assert_called_once_with(agentpool)
         decorator.update_upgrade_strategy.assert_called_once_with(agentpool)
         decorator.update_blue_green_upgrade_settings.assert_called_once_with(agentpool)
+        decorator.update_gpu_profile.assert_called_once_with(agentpool)
 
     def test_update_agentpool_profile_preview_with_agentpools_parameter(self):
         """Test update_agentpool_profile_preview with agentpools parameter."""
@@ -194,6 +196,7 @@ def test_update_agentpool_profile_preview_with_agentpools_parameter(self):
         decorator.update_auto_scaler_properties_vms = Mock(return_value=agentpool)
         decorator.update_upgrade_strategy = Mock(return_value=agentpool)
         decorator.update_blue_green_upgrade_settings = Mock(return_value=agentpool)
+        decorator.update_gpu_profile = Mock(return_value=agentpool)
 
         # Act
         result = decorator.update_agentpool_profile_preview(agentpools)
@@ -244,6 +247,7 @@ def test_update_agentpool_profile_preview_managed_system_mode(self):
         decorator.update_auto_scaler_properties_vms = Mock()
         decorator.update_upgrade_strategy = Mock()
         decorator.update_blue_green_upgrade_settings = Mock()
+        decorator.update_gpu_profile = Mock()
 
         # Act
         result = decorator.update_agentpool_profile_preview()
@@ -272,6 +276,7 @@ def test_update_agentpool_profile_preview_managed_system_mode(self):
         decorator.update_auto_scaler_properties_vms.assert_not_called()
         decorator.update_upgrade_strategy.assert_not_called()
         decorator.update_blue_green_upgrade_settings.assert_not_called()
+        decorator.update_gpu_profile.assert_not_called()
 
     def test_update_agentpool_profile_preview_managed_system_mode_with_agentpools(self):
         """Test update_agentpool_profile_preview with ManagedSystem mode and agentpools parameter."""
@@ -349,6 +354,7 @@ def test_update_agentpool_profile_preview_system_mode_regular_flow(self):
         decorator.update_auto_scaler_properties_vms = Mock(return_value=agentpool)
         decorator.update_upgrade_strategy = Mock(return_value=agentpool)
         decorator.update_blue_green_upgrade_settings = Mock(return_value=agentpool)
+        decorator.update_gpu_profile = Mock(return_value=agentpool)
 
         # Act
         result = decorator.update_agentpool_profile_preview()
@@ -369,6 +375,7 @@ def test_update_agentpool_profile_preview_system_mode_regular_flow(self):
         decorator.update_auto_scaler_properties_vms.assert_called_once_with(agentpool)
         decorator.update_upgrade_strategy.assert_called_once_with(agentpool)
         decorator.update_blue_green_upgrade_settings.assert_called_once_with(agentpool)
+        decorator.update_gpu_profile.assert_called_once_with(agentpool)
 
     def test_update_agentpool_profile_preview_execution_order(self):
         """Test that update methods are called in the correct order."""
@@ -414,6 +421,7 @@ def mock_method(pool):
         decorator.update_auto_scaler_properties_vms = create_mock_update_method("update_auto_scaler_properties_vms")
         decorator.update_upgrade_strategy = create_mock_update_method("update_upgrade_strategy")
         decorator.update_blue_green_upgrade_settings = create_mock_update_method("update_blue_green_upgrade_settings")
+        decorator.update_gpu_profile = create_mock_update_method("update_gpu_profile")
 
         # Act
         decorator.update_agentpool_profile_preview()
@@ -430,7 +438,8 @@ def mock_method(pool):
             "update_localdns_profile",
             "update_auto_scaler_properties_vms",
             "update_upgrade_strategy",
-            "update_blue_green_upgrade_settings"
+            "update_blue_green_upgrade_settings",
+            "update_gpu_profile",
         ]
         self.assertEqual(call_order, expected_order)
 
@@ -478,6 +487,7 @@ def track_and_return(pool):
         decorator.update_auto_scaler_properties_vms = create_tracking_mock("update_auto_scaler_properties_vms")
         decorator.update_upgrade_strategy = create_tracking_mock("update_upgrade_strategy")
         decorator.update_blue_green_upgrade_settings = create_tracking_mock("update_blue_green_upgrade_settings")        
+        decorator.update_gpu_profile = create_tracking_mock("update_gpu_profile")
 
         # Act
         result = decorator.update_agentpool_profile_preview()
@@ -540,7 +550,7 @@ def test_update_agentpool_profile_preview_mixed_modes_scenario(self):
                     'update_network_profile', 'update_artifact_streaming',
                     'update_secure_boot', 'update_vtpm', 'update_os_sku', 'update_fips_image',
                     'update_ssh_access', 'update_localdns_profile', 'update_auto_scaler_properties_vms', 
-                    'update_upgrade_strategy', 'update_blue_green_upgrade_settings'
+                    'update_upgrade_strategy', 'update_blue_green_upgrade_settings', 'update_gpu_profile'
                 ]
 
                 for method_name in update_methods:
@@ -612,6 +622,7 @@ def test_update_agentpool_profile_preview_managed_cluster_mode(self):
         decorator.update_auto_scaler_properties_vms = Mock(return_value=agentpool)
         decorator.update_upgrade_strategy = Mock(return_value=agentpool)
         decorator.update_blue_green_upgrade_settings = Mock(return_value=agentpool)
+        decorator.update_gpu_profile = Mock(return_value=agentpool)
 
         # Act
         result = decorator.update_agentpool_profile_preview(agentpools)

src/aks-preview/setup.py

@@ -9,7 +9,7 @@
 
 from setuptools import find_packages, setup
 
-VERSION = "19.0.0b17"
+VERSION = "19.0.0b18"
 
 CLASSIFIERS = [
     "Development Status :: 4 - Beta",