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

Author: lilypan26

src/azure-cli/azure/cli/command_modules/acs/_help.py

@@ -2126,6 +2126,9 @@
   - name: --localdns-config
     type: string
     short-summary: Set the localDNS Profile for a nodepool with a JSON config file.
+  - 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/azure-cli/azure/cli/command_modules/acs/_params.py

@@ -1089,6 +1089,7 @@ def load_arguments(self, _):
         c.argument("if_match")
         c.argument("if_none_match")
         c.argument('localdns_config', help='Path to a JSON file to configure the local DNS profile for a new nodepool.')
+        c.argument('gpu_driver', arg_type=get_enum_type(gpu_driver_install_modes))
 
     with self.argument_context('aks nodepool upgrade') as c:
         c.argument('max_surge', validator=validate_max_surge)

src/azure-cli/azure/cli/command_modules/acs/agentpool_decorator.py

@@ -2622,6 +2622,18 @@ def update_vtpm(self, agentpool: AgentPool) -> AgentPool:
 
         return agentpool
 
+    def update_gpu_profile(self, agentpool: AgentPool) -> AgentPool:
+        self._ensure_agentpool(agentpool)
+
+        gpu_driver = self.context.get_gpu_driver()
+
+        # Construct AgentPoolGPUProfile if one of the fields has been set
+        if gpu_driver:
+            agentpool.gpu_profile = self.models.GPUProfile()
+            agentpool.gpu_profile.driver = gpu_driver
+
+        return agentpool
+
     def update_agentpool_profile_default(self, agentpools: List[AgentPool] = None) -> AgentPool:
         """The overall controller used to update AgentPool profile by default.
 
@@ -2652,6 +2664,8 @@ def update_agentpool_profile_default(self, agentpools: List[AgentPool] = None) -
         agentpool = self.update_secure_boot(agentpool)
         # update local DNS profile
         agentpool = self.update_localdns_profile(agentpool)
+        # update gpu profile
+        agentpool = self.update_gpu_profile(agentpool)
         return agentpool
 
     def update_agentpool(self, agentpool: AgentPool) -> AgentPool:

src/azure-cli/azure/cli/command_modules/acs/custom.py

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

src/azure-cli/azure/cli/command_modules/acs/tests/latest/test_agentpool_decorator.py

@@ -3241,6 +3241,31 @@ def common_update_vm_properties(self):
         )
         self.assertEqual(dec_agentpool_1, grond_truth_agentpool_1)
 
+    def common_update_gpu_profile(self):
+        dec_1 = AKSAgentPoolUpdateDecorator(
+            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_update_fips_image(self):
         dec_1 = AKSAgentPoolUpdateDecorator(
             self.cmd,
@@ -3481,6 +3506,9 @@ def test_update_vm_properties(self):
     def test_update_fips_image(self):
         self.common_update_fips_image()
 
+    def test_update_gpu_profile(self):
+        self.common_update_gpu_profile()
+
     def test_update_agentpool_profile_default(self):
         import inspect