feat: add disable_artifact_streaming arg for nodepool update in aks preview (#9679)

Author: KananMehta

src/aks-preview/HISTORY.rst

@@ -22,6 +22,7 @@ Pending
 * `az aks list-vm-skus`: New command to list available VM SKUs for AKS clusters in a given region.
 * Add managed GPU enablement option to node pool property in `az aks nodepool add` and `az aks nodepool update`.
 * `az aks namespace update`: Fix location should use existing namespace location.
+* `az aks nodepool update`: Add `--disable-artifact-streaming` to disable artifact streaming.
 
 19.0.0b27
 +++++++

src/aks-preview/azext_aks_preview/_help.py

@@ -2437,6 +2437,9 @@
         - 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.
+        - name: --disable-artifact-streaming
+          type: bool
+          short-summary: Disable artifact streaming for VirtualMachineScaleSets managed by a node pool.
         - name: --enable-managed-gpu
           type: bool
           short-summary: Enable the Managed GPU experience, which installs additional components like DCGM metrics for monitoring on top of the GPU driver. For more details, visit aka.ms/aks/managed-gpu.

src/aks-preview/azext_aks_preview/_params.py

@@ -2227,6 +2227,12 @@ def load_arguments(self, _):
             validator=validate_artifact_streaming,
             is_preview=True,
         )
+        c.argument(
+            "disable_artifact_streaming",
+            action="store_true",
+            validator=validate_artifact_streaming,
+            is_preview=True,
+        )
         c.argument(
             "enable_managed_gpu",
             action="store_true",

src/aks-preview/azext_aks_preview/_validators.py

@@ -954,10 +954,20 @@ def validate_asm_egress_name(namespace):
 
 
 def validate_artifact_streaming(namespace):
-    """Validates that artifact streaming enablement can only be used on Linux."""
-    if namespace.enable_artifact_streaming:
-        if hasattr(namespace, 'os_type') and str(namespace.os_type).lower() == "windows":
+    """Validates artifact streaming flags for mutual exclusivity and OS support."""
+    enable_artifact_streaming = getattr(namespace, "enable_artifact_streaming", False)
+    disable_artifact_streaming = getattr(namespace, "disable_artifact_streaming", False)
+
+    if enable_artifact_streaming and disable_artifact_streaming:
+        raise MutuallyExclusiveArgumentError(
+            "Cannot specify both --enable-artifact-streaming and --disable-artifact-streaming at the same time."
+        )
+
+    if hasattr(namespace, "os_type") and str(namespace.os_type).lower() == "windows":
+        if enable_artifact_streaming:
             raise ArgumentUsageError('--enable-artifact-streaming can only be set for Linux nodepools')
+        if disable_artifact_streaming:
+            raise ArgumentUsageError('--disable-artifact-streaming can only be set for Linux nodepools')
 
 
 def validate_custom_endpoints(namespace):

src/aks-preview/azext_aks_preview/agentpool_decorator.py

@@ -588,6 +588,11 @@ def get_enable_artifact_streaming(self) -> bool:
                 self.agentpool.artifact_streaming_profile.enabled is not None
             ):
                 enable_artifact_streaming = self.agentpool.artifact_streaming_profile.enabled
+
+        if enable_artifact_streaming and self.get_disable_artifact_streaming():
+            raise MutuallyExclusiveArgumentError(
+                'Cannot specify both --enable-artifact-streaming and --disable-artifact-streaming.'
+            )
         return enable_artifact_streaming
 
     def get_enable_managed_gpu(self) -> Union[bool, None]:
@@ -611,6 +616,13 @@ def get_enable_managed_gpu(self) -> Union[bool, None]:
                 )
         return enable_managed_gpu
 
+    def get_disable_artifact_streaming(self) -> bool:
+        """Obtain the value of disable_artifact_streaming.
+        :return: bool
+        """
+
+        return self.raw_param.get("disable_artifact_streaming")
+
     def get_pod_ip_allocation_mode(self: bool = False) -> Union[str, None]:
         """Get the value of pod_ip_allocation_mode.
         :return: str or None
@@ -1743,6 +1755,11 @@ def update_artifact_streaming(self, agentpool: AgentPool) -> AgentPool:
             if agentpool.artifact_streaming_profile is None:
                 agentpool.artifact_streaming_profile = self.models.AgentPoolArtifactStreamingProfile()  # pylint: disable=no-member
             agentpool.artifact_streaming_profile.enabled = True
+
+        if self.context.get_disable_artifact_streaming():
+            if agentpool.artifact_streaming_profile is None:
+                agentpool.artifact_streaming_profile = self.models.AgentPoolArtifactStreamingProfile()  # pylint: disable=no-member
+            agentpool.artifact_streaming_profile.enabled = False
         return agentpool
 
     def update_managed_gpu(self, agentpool: AgentPool) -> AgentPool:

src/aks-preview/azext_aks_preview/custom.py

@@ -2005,6 +2005,7 @@ def aks_agentpool_update(
     allowed_host_ports=None,
     asg_ids=None,
     enable_artifact_streaming=False,
+    disable_artifact_streaming=False,
     enable_managed_gpu=False,
     os_sku=None,
     ssh_access=None,

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

@@ -258,6 +258,24 @@ def common_get_enable_artifact_streaming(self):
         ctx_2.attach_agentpool(agentpool_2)
         self.assertEqual(ctx_2.get_enable_artifact_streaming(), None)
 
+    def common_get_disable_artifact_streaming(self):
+        # default
+        ctx_1 = AKSPreviewAgentPoolContext(
+            self.cmd,
+            AKSAgentPoolParamDict({"disable_artifact_streaming": True}),
+            self.models,
+            DecoratorMode.UPDATE,
+            self.agentpool_decorator_mode,
+        )
+        self.assertEqual(ctx_1.get_disable_artifact_streaming(), True)
+        agentpool_1 = self.create_initialized_agentpool_instance(
+            artifact_streaming_profile=self.models.AgentPoolArtifactStreamingProfile(
+                enabled=True
+            )
+        )
+        ctx_1.attach_agentpool(agentpool_1)
+        self.assertEqual(ctx_1.get_disable_artifact_streaming(), True)
+
     def common_get_enable_managed_gpu(self):
         # default
         ctx_1 = AKSPreviewAgentPoolContext(
@@ -1077,6 +1095,9 @@ def test_get_workload_runtime(self):
     def test_get_enable_artifact_streaming(self):
         self.common_get_enable_artifact_streaming()
 
+    def test_get_disable_artifact_streaming(self):
+        self.common_get_disable_artifact_streaming()
+
     def test_get_enable_managed_gpu(self):
         self.common_get_enable_managed_gpu()
 
@@ -1186,6 +1207,9 @@ def test_get_os_sku(self):
     def test_get_enable_artifact_streaming(self):
         self.common_get_enable_artifact_streaming()
 
+    def test_get_disable_artifact_streaming(self):
+        self.common_get_disable_artifact_streaming()
+
     def test_get_enable_secure_boot(self):
         self.common_get_enable_secure_boot()
 
@@ -2425,6 +2449,42 @@ def common_update_artifact_streaming(self):
             )
         )
         self.assertEqual(dec_agentpool_2, ground_truth_agentpool_2)
+        
+        dec_3 = AKSPreviewAgentPoolUpdateDecorator(
+            self.cmd,
+            self.client,
+            {"disable_artifact_streaming": True},
+            self.resource_type,
+            self.agentpool_decorator_mode,
+        )
+        # fail on passing the wrong agentpool object
+        with self.assertRaises(CLIInternalError):
+            dec_3.update_artifact_streaming(None)
+        agentpool_3 = self.create_initialized_agentpool_instance(
+            artifact_streaming_profile=self.models.AgentPoolArtifactStreamingProfile(
+                enabled=True
+            )
+        )
+        dec_3.context.attach_agentpool(agentpool_3)
+        dec_agentpool_3 = dec_3.update_artifact_streaming(agentpool_3)
+        grond_truth_agentpool_3 = self.create_initialized_agentpool_instance(
+            artifact_streaming_profile=self.models.AgentPoolArtifactStreamingProfile(
+                enabled=False
+            )
+        )
+        self.assertEqual(dec_agentpool_3, grond_truth_agentpool_3)
+
+        # Should error if both set
+        dec_4 = AKSPreviewAgentPoolUpdateDecorator(
+            self.cmd,
+            self.client,
+            {"enable_artifact_streaming": True, "disable_artifact_streaming": True},
+            self.resource_type,
+            self.agentpool_decorator_mode,
+        )
+        dec_4.context.attach_agentpool(agentpool_3)
+        with self.assertRaises(MutuallyExclusiveArgumentError):
+            dec_4.update_artifact_streaming(agentpool_3)
 
     def common_update_managed_gpu(self):
         dec_1 = AKSPreviewAgentPoolUpdateDecorator(

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

@@ -16674,6 +16674,102 @@ def test_aks_nodepool_update_with_artifact_streaming(
             ],
         )
 
+    @live_only()
+    @AllowLargeResponse()
+    @AKSCustomResourceGroupPreparer(
+        random_name_length=17, name_prefix="clitest", location="eastus"
+    )
+    def test_aks_nodepool_update_with_disable_artifact_streaming(
+        self, resource_group, resource_group_location
+    ):
+        aks_name = self.create_random_name("cliakstest", 16)
+        nodepool_name = self.create_random_name("n", 6)
+
+        self.kwargs.update(
+            {
+                "resource_group": resource_group,
+                "name": aks_name,
+                "location": resource_group_location,
+                "ssh_key_value": self.generate_ssh_keys(),
+                "node_pool_name": nodepool_name,
+                "node_vm_size": "standard_d2s_v3",
+            }
+        )
+
+        self.cmd(
+            "aks create "
+            "--resource-group={resource_group} "
+            "--name={name} "
+            "--location={location} "
+            "--ssh-key-value={ssh_key_value} "
+            "--nodepool-name={node_pool_name} "
+            "--node-count=1 "
+            "--node-vm-size={node_vm_size} "
+            "--aks-custom-headers AKSHTTPCustomFeatures=Microsoft.ContainerService/ArtifactStreamingPreview",
+            checks=[
+                self.check("provisioningState", "Succeeded"),
+            ],
+        )
+
+        # disable artifact streaming on a nodepool that never had it enabled
+        self.cmd(
+            "aks nodepool update "
+            "--resource-group={resource_group} "
+            "--cluster-name={name} "
+            "--name={node_pool_name} "
+            "--disable-artifact-streaming "
+            "--aks-custom-headers AKSHTTPCustomFeatures=Microsoft.ContainerService/ArtifactStreamingPreview",
+            checks=[
+                self.check("provisioningState", "Succeeded"),
+                self.check(
+                    "artifactStreamingProfile.enabled", False
+                ),
+            ],
+        )
+
+        # enable artifact streaming
+        self.cmd(
+            "aks nodepool update "
+            "--resource-group={resource_group} "
+            "--cluster-name={name} "
+            "--name={node_pool_name} "
+            "--enable-artifact-streaming "
+            "--aks-custom-headers AKSHTTPCustomFeatures=Microsoft.ContainerService/ArtifactStreamingPreview",
+            checks=[
+                self.check("provisioningState", "Succeeded"),
+                self.check(
+                    "artifactStreamingProfile.enabled", True
+                ),
+            ],
+        )
+
+        # disable artifact streaming
+        self.cmd(
+            "aks nodepool update "
+            "--resource-group={resource_group} "
+            "--cluster-name={name} "
+            "--name={node_pool_name} "
+            "--disable-artifact-streaming "
+            "--aks-custom-headers AKSHTTPCustomFeatures=Microsoft.ContainerService/ArtifactStreamingPreview",
+            checks=[
+                self.check("provisioningState", "Succeeded"),
+                self.check(
+                    "artifactStreamingProfile.enabled", False
+                ),
+            ],
+        )
+
+        # delete
+        cmd = (
+            "aks delete --resource-group={resource_group} --name={name} --yes --no-wait"
+        )
+        self.cmd(
+            cmd,
+            checks=[
+                self.is_empty(),
+            ],
+        )
+
     @live_only()
     @AllowLargeResponse()
     @AKSCustomResourceGroupPreparer(

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

@@ -149,6 +149,13 @@ def __init__(self, os_type, disable_windows_outbound_nat):
         self.disable_windows_outbound_nat = disable_windows_outbound_nat
 
 
+class ArtifactStreamingNamespace:
+    def __init__(self, os_type, enable_artifact_streaming=False, disable_artifact_streaming=False):
+        self.os_type = os_type
+        self.enable_artifact_streaming = enable_artifact_streaming
+        self.disable_artifact_streaming = disable_artifact_streaming
+
+
 class TestMaxSurge(unittest.TestCase):
     def test_valid_cases(self):
         valid = ["5", "33%", "1", "100%"]
@@ -385,6 +392,52 @@ def test_fail_if_os_type_invalid(self):
         )
 
 
+class TestArtifactStreaming(unittest.TestCase):
+    def test_valid_linux_enable(self):
+        validators.validate_artifact_streaming(
+            ArtifactStreamingNamespace("Linux", enable_artifact_streaming=True)
+        )
+
+    def test_valid_linux_disable(self):
+        validators.validate_artifact_streaming(
+            ArtifactStreamingNamespace("Linux", disable_artifact_streaming=True)
+        )
+
+    def test_fail_if_enable_and_disable_are_set(self):
+        with self.assertRaises(MutuallyExclusiveArgumentError) as cm:
+            validators.validate_artifact_streaming(
+                ArtifactStreamingNamespace(
+                    "Linux",
+                    enable_artifact_streaming=True,
+                    disable_artifact_streaming=True,
+                )
+            )
+        self.assertEqual(
+            str(cm.exception),
+            "Cannot specify both --enable-artifact-streaming and --disable-artifact-streaming at the same time.",
+        )
+
+    def test_fail_if_enable_for_windows(self):
+        with self.assertRaises(ArgumentUsageError) as cm:
+            validators.validate_artifact_streaming(
+                ArtifactStreamingNamespace("Windows", enable_artifact_streaming=True)
+            )
+        self.assertEqual(
+            str(cm.exception),
+            "--enable-artifact-streaming can only be set for Linux nodepools",
+        )
+
+    def test_fail_if_disable_for_windows(self):
+        with self.assertRaises(ArgumentUsageError) as cm:
+            validators.validate_artifact_streaming(
+                ArtifactStreamingNamespace("Windows", disable_artifact_streaming=True)
+            )
+        self.assertEqual(
+            str(cm.exception),
+            "--disable-artifact-streaming can only be set for Linux nodepools",
+        )
+
+
 class ValidateAddonsNamespace:
     def __init__(self, addons):
         self.addons = addons

src/aks-preview/linter_exclusions.yml

@@ -446,6 +446,9 @@ aks nodepool update:
     enable_artifact_streaming:
       rule_exclusions:
         - option_length_too_long
+    disable_artifact_streaming:
+      rule_exclusions:
+        - option_length_too_long
     enable_secure_boot:
       rule_exclusions:
         - option_length_too_long