feat: add disable_artifact_streaming arg

Author: Kanan Mehta

src/aks-preview/HISTORY.rst

@@ -11,6 +11,7 @@ To release a new version, please select a new version number (usually plus 1 to
 
 Pending
 +++++++
+* `az aks nodepool update`: Add `--disable-artifact-streaming` to disable artifact streaming.
 
 19.0.0b27
 +++++++

src/aks-preview/azext_aks_preview/_help.py

@@ -2400,6 +2400,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: --os-sku
           type: string
           short-summary: The os-sku of the agent node pool.

src/aks-preview/azext_aks_preview/_params.py

@@ -2135,6 +2135,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(
             "os_sku",
             arg_type=get_enum_type(node_os_skus_update),

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

@@ -585,8 +585,20 @@ 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 "--enable-artifact-streaming" and "--disable-artifact-streaming" at the same time'
+            )
         return enable_artifact_streaming
 
+    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
@@ -1686,6 +1698,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_os_sku(self, agentpool: AgentPool) -> AgentPool:

src/aks-preview/azext_aks_preview/custom.py

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

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

@@ -257,6 +257,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_pod_ip_allocation_mode(self):
         # default
         ctx_1 = AKSPreviewAgentPoolContext(
@@ -1037,6 +1055,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_pod_ip_allocation_mode(self):
         self.common_get_pod_ip_allocation_mode()
 
@@ -1140,6 +1161,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()
 
@@ -2349,6 +2373,42 @@ def common_update_artifact_streaming(self):
         )
         self.assertEqual(dec_agentpool_2, grond_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_secure_boot(self):
         dec_1 = AKSPreviewAgentPoolUpdateDecorator(
             self.cmd,

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