[AKS] az aks nodepool add/update: Add --enable-managed-dranet (#10266)

Author: rbtr

src/aks-preview/HISTORY.rst

@@ -11,6 +11,10 @@ To release a new version, please select a new version number (usually plus 1 to
 
 Pending
 +++++++
+
+22.0.0b6
++++++++++
+* `az aks nodepool add/update`: Add preview `--enable-managed-dranet` to enable Managed DRANET on a node pool.
 * `az aks nodepool add` and `az aks nodepool update`: Add `--managed-gpu-driver-mode` to select `DRA` or `DevicePlugin` when managed GPU is enabled.
 * Add options `Windows2022` and `Windows2025` to `--os-sku` for `az aks nodepool update`, allowing in-place OS SKU upgrades between these Windows Server versions.
 * `az aks create` and `az aks nodepool add`: `--enable-fips-image` is now required and always enabled when `--os-sku` is `Windows2025`; `--disable-fips-image` cannot be used with `Windows2025`.

src/aks-preview/azext_aks_preview/_help.py

@@ -2664,6 +2664,9 @@
             Specify secondary NICs to attach to each node. Accepts inline JSON or `@filename`.
             Example: '[{"type":"Standard","vnetSubnetId":"/subscriptions/.../subnets/mysubnet","enableAcceleratedNetworking":true}]'
             Supported NIC types are "Standard" (requires vnetSubnetId) and "Dynamic".
+        - name: --enable-managed-dranet
+          type: bool
+          short-summary: Enable Managed DRANET on the node pool.
         - name: --upgrade-strategy
           type: string
           short-summary: Upgrade strategy for the node pool. Allowed values are "Rolling" or "BlueGreen". Default is "Rolling".
@@ -2837,6 +2840,9 @@
         - name: --asg-ids
           type: string
           short-summary: The IDs of the application security groups to which the node pool's network interface should belong. When specified, format should be a comma-separated list of IDs. Must use VMSS agent pool type.
+        - name: --enable-managed-dranet
+          type: bool
+          short-summary: Enable Managed DRANET on the node pool.
         - name: --enable-artifact-streaming
           type: bool
           short-summary: Enable artifact streaming for VirtualMachineScaleSets managed by a node pool, to speed up the cold-start of containers on a node through on-demand image loading. To use this feature, container images must also enable artifact streaming on ACR. If not specified, the default is false.

src/aks-preview/azext_aks_preview/_params.py

@@ -2525,6 +2525,11 @@ def load_arguments(self, _):
                  'Example: \'[{"type":"Standard","vnetSubnetId":"/subscriptions/.../subnets/mysubnet"}]\'',
             is_preview=True,
         )
+        c.argument(
+            "enable_managed_dranet",
+            action="store_true",
+            is_preview=True,
+        )
         # prepared image specification
         c.argument(
             'prepared_image_specification_id',
@@ -2579,6 +2584,11 @@ def load_arguments(self, _):
         c.argument(
             "asg_ids", validator=validate_application_security_groups, is_preview=True
         )
+        c.argument(
+            "enable_managed_dranet",
+            action="store_true",
+            is_preview=True,
+        )
         c.argument(
             "enable_artifact_streaming",
             action="store_true",

src/aks-preview/azext_aks_preview/agentpool_decorator.py

@@ -336,6 +336,10 @@ def get_allowed_host_ports(self) -> Union[List[PortRange], None]:
             ))
         return port_ranges
 
+    def get_enable_managed_dranet(self) -> bool:
+        """Obtain the value of enable_managed_dranet."""
+        return self.raw_param.get("enable_managed_dranet")
+
     def get_ip_tags(self) -> Union[List[IPTag], None]:
         ip_tags = self.raw_param.get("node_public_ip_tags")
         res = []
@@ -1454,6 +1458,11 @@ def set_up_agentpool_network_profile(self, agentpool: AgentPool) -> AgentPool:
         if secondary_nics is not None:
             agentpool.network_profile.secondary_network_interfaces = secondary_nics
 
+        if self.context.get_enable_managed_dranet():
+            agentpool.network_profile.dranet = self.models.DRANETProfile(
+                mode="Managed"
+            )
+
         return agentpool
 
     def set_up_taints(self, agentpool: AgentPool) -> AgentPool:
@@ -1959,12 +1968,17 @@ def update_network_profile(self, agentpool: AgentPool) -> AgentPool:
 
         asg_ids = self.context.get_asg_ids()
         allowed_host_ports = self.context.get_allowed_host_ports()
-        if not agentpool.network_profile and (asg_ids or allowed_host_ports):
+        enable_managed_dranet = self.context.get_enable_managed_dranet()
+        if not agentpool.network_profile and (asg_ids is not None or allowed_host_ports is not None or enable_managed_dranet):
             agentpool.network_profile = self.models.AgentPoolNetworkProfile()  # pylint: disable=no-member
         if asg_ids is not None:
             agentpool.network_profile.application_security_groups = asg_ids
         if allowed_host_ports is not None:
             agentpool.network_profile.allowed_host_ports = allowed_host_ports
+        if enable_managed_dranet:
+            agentpool.network_profile.dranet = self.models.DRANETProfile(
+                mode="Managed"
+            )
         return agentpool
 
     def update_gpu_profile(self, agentpool: AgentPool) -> AgentPool:

src/aks-preview/azext_aks_preview/custom.py

@@ -2289,6 +2289,7 @@ def aks_agentpool_add(
     disable_windows_outbound_nat=False,
     allowed_host_ports=None,
     asg_ids=None,
+    enable_managed_dranet=False,
     node_public_ip_tags=None,
     enable_artifact_streaming=False,
     enable_managed_gpu=False,
@@ -2374,6 +2375,7 @@ def aks_agentpool_update(
     # extensions
     allowed_host_ports=None,
     asg_ids=None,
+    enable_managed_dranet=False,
     enable_artifact_streaming=False,
     disable_artifact_streaming=False,
     enable_managed_gpu=False,

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

@@ -254,6 +254,25 @@ def common_get_enable_artifact_streaming(self):
             self.agentpool_decorator_mode,
         )
         self.assertEqual(ctx_2.get_enable_artifact_streaming(), None)
+
+    def common_get_enable_managed_dranet(self):
+        ctx_1 = AKSPreviewAgentPoolContext(
+            self.cmd,
+            AKSAgentPoolParamDict({"enable_managed_dranet": False}),
+            self.models,
+            DecoratorMode.CREATE,
+            self.agentpool_decorator_mode,
+        )
+        self.assertEqual(ctx_1.get_enable_managed_dranet(), False)
+
+        ctx_2 = AKSPreviewAgentPoolContext(
+            self.cmd,
+            AKSAgentPoolParamDict({"enable_managed_dranet": True}),
+            self.models,
+            DecoratorMode.UPDATE,
+            self.agentpool_decorator_mode,
+        )
+        self.assertEqual(ctx_2.get_enable_managed_dranet(), True)
         agentpool_2 = self.create_initialized_agentpool_instance(
             artifact_streaming_profile=self.models.AgentPoolArtifactStreamingProfile(
                 enabled=True
@@ -1416,6 +1435,9 @@ def test_get_workload_runtime(self):
     def test_get_enable_artifact_streaming(self):
         self.common_get_enable_artifact_streaming()
 
+    def test_get_enable_managed_dranet(self):
+        self.common_get_enable_managed_dranet()
+
     def test_get_enable_os_disk_full_caching(self):
         self.common_get_enable_os_disk_full_caching()
 
@@ -1534,6 +1556,9 @@ def test_get_workload_runtime(self):
     def test_get_enable_artifact_streaming(self):
         self.common_get_enable_artifact_streaming()
 
+    def test_get_enable_managed_dranet(self):
+        self.common_get_enable_managed_dranet()
+
     def test_get_enable_os_disk_full_caching(self):
         self.common_get_enable_os_disk_full_caching()
 
@@ -2099,6 +2124,31 @@ def common_set_up_agentpool_gateway_profile(self):
         )
         self.assertEqual(dec_agentpool_1, ground_truth_agentpool_1)
 
+    def common_set_up_managed_dranet(self):
+        dec_1 = AKSPreviewAgentPoolAddDecorator(
+            self.cmd,
+            self.client,
+            {"enable_managed_dranet": False},
+            self.resource_type,
+            self.agentpool_decorator_mode,
+        )
+        agentpool_1 = self.create_initialized_agentpool_instance(restore_defaults=False)
+        dec_1.context.attach_agentpool(agentpool_1)
+        dec_agentpool_1 = dec_1.set_up_agentpool_network_profile(agentpool_1)
+        self.assertIsNone(dec_agentpool_1.network_profile.dranet)
+
+        dec_2 = AKSPreviewAgentPoolAddDecorator(
+            self.cmd,
+            self.client,
+            {"enable_managed_dranet": True},
+            self.resource_type,
+            self.agentpool_decorator_mode,
+        )
+        agentpool_2 = self.create_initialized_agentpool_instance(restore_defaults=False)
+        dec_2.context.attach_agentpool(agentpool_2)
+        dec_agentpool_2 = dec_2.set_up_agentpool_network_profile(agentpool_2)
+        self.assertEqual(dec_agentpool_2.network_profile.dranet.mode, "Managed")
+
     def common_set_up_virtual_machines_profile(self):
         dec_1 = AKSPreviewAgentPoolAddDecorator(
             self.cmd,
@@ -2540,6 +2590,9 @@ def test_set_up_agentpool_windows_profile(self):
     def test_set_up_agentpool_gateway_profile(self):
         self.common_set_up_agentpool_gateway_profile()
 
+    def test_set_up_managed_dranet(self):
+        self.common_set_up_managed_dranet()
+
     def test_set_up_virtual_machines_profile(self):
         self.common_set_up_virtual_machines_profile()
 
@@ -2770,6 +2823,9 @@ def test_set_up_agentpool_windows_profile(self):
     def test_set_up_agentpool_gateway_profile(self):
         self.common_set_up_agentpool_gateway_profile()
 
+    def test_set_up_managed_dranet(self):
+        self.common_set_up_managed_dranet()
+
     def test_set_up_virtual_machines_profile(self):
         self.common_set_up_virtual_machines_profile()
 
@@ -2995,6 +3051,35 @@ def common_update_artifact_streaming(self):
         with self.assertRaises(MutuallyExclusiveArgumentError):
             dec_4.update_artifact_streaming(agentpool_3)
 
+    def common_update_managed_dranet(self):
+        dec_1 = AKSPreviewAgentPoolUpdateDecorator(
+            self.cmd,
+            self.client,
+            {"enable_managed_dranet": False},
+            self.resource_type,
+            self.agentpool_decorator_mode,
+        )
+        agentpool_1 = self.create_initialized_agentpool_instance(
+            network_profile=self.models.AgentPoolNetworkProfile(
+                dranet=self.models.DRANETProfile(mode="Managed")
+            )
+        )
+        dec_1.context.attach_agentpool(agentpool_1)
+        dec_agentpool_1 = dec_1.update_network_profile(agentpool_1)
+        self.assertEqual(dec_agentpool_1.network_profile.dranet.mode, "Managed")
+
+        dec_2 = AKSPreviewAgentPoolUpdateDecorator(
+            self.cmd,
+            self.client,
+            {"enable_managed_dranet": True},
+            self.resource_type,
+            self.agentpool_decorator_mode,
+        )
+        agentpool_2 = self.create_initialized_agentpool_instance()
+        dec_2.context.attach_agentpool(agentpool_2)
+        dec_agentpool_2 = dec_2.update_network_profile(agentpool_2)
+        self.assertEqual(dec_agentpool_2.network_profile.dranet.mode, "Managed")
+
     def common_update_managed_gpu(self):
         dec_1 = AKSPreviewAgentPoolUpdateDecorator(
             self.cmd,
@@ -3679,6 +3764,9 @@ def setUp(self):
     def test_update_artifact_streaming(self):
         self.common_update_artifact_streaming()
 
+    def test_update_managed_dranet(self):
+        self.common_update_managed_dranet()
+
     def test_update_managed_gpu(self):
         self.common_update_managed_gpu()
 
@@ -3782,7 +3870,10 @@ def setUp(self):
 
     def test_update_artifact_streaming(self):
         self.common_update_artifact_streaming()
-    
+
+    def test_update_managed_dranet(self):
+        self.common_update_managed_dranet()
+
     def test_update_managed_gpu(self):
         self.common_update_managed_gpu()
 

src/aks-preview/setup.py

@@ -9,7 +9,7 @@
 
 from setuptools import find_packages, setup
 
-VERSION = "22.0.0b5"
+VERSION = "22.0.0b6"
 
 CLASSIFIERS = [
     "Development Status :: 4 - Beta",