[Containerapp] az containerapp sessionpool create/update: Add health probe support (#9139)

Author: ShichaoQiu

src/containerapp/HISTORY.rst

@@ -9,6 +9,7 @@ upcoming
 * 'az containerapp sessionpool create/update': Support `--lifecycle-type` and `--max-alive-period`
 * 'az containerapp up': support for `--kind` parameter
 * 'az containerapp env premium-ingress': Deprecate `--min-replicas` and `--max-replicas` parameters, use workload profile scale instead.
+* 'az containerapp sessionpool create/update': Support `--probe-yaml`
 
 1.2.0b3
 ++++++

src/containerapp/azext_containerapp/_help.py

@@ -2010,6 +2010,11 @@
           az containerapp sessionpool create -n mysessionpool -g MyResourceGroup \\
               --environment MyEnvironment --cpu 0.5 --memory 1Gi --target-port 80 --container-type CustomContainer \\
               --cooldown-period 360 --location eastasia
+    - name: Create or update a Session Pool with container type CustomContainer with container probes
+      text: |
+          az containerapp sessionpool create -n mysessionpool -g MyResourceGroup \\
+              --environment MyEnvironment --cpu 0.5 --memory 1Gi --target-port 80 --container-type CustomContainer \\
+              --probe-yaml config.yaml --location eastasia
 """
 
 helps['containerapp sessionpool update'] = """
@@ -2019,6 +2024,9 @@
     - name: Update a session pool's max concurrent sessions configuration and image.
       text: |
           az containerapp sessionpool update -n mysessionpool -g MyResourceGroup --max-sessions 20 --image MyNewImage
+    - name: Update the container probes of a CustomContainer type session pool.
+      text: |
+          az containerapp sessionpool update -n mysessionpool -g MyResourceGroup --probe-yaml config.yaml
 """
 
 helps['containerapp sessionpool delete'] = """

src/containerapp/azext_containerapp/_params.py

@@ -471,6 +471,7 @@ def load_arguments(self, _):
         c.argument('startup_command', nargs='*', options_list=['--command'], help="A list of supported commands on the container that will executed during startup. Space-separated values e.g. \"/bin/queue\" \"mycommand\". Empty string to clear existing values")
         c.argument('args', nargs='*', help="A list of container startup command argument(s). Space-separated values e.g. \"-c\" \"mycommand\". Empty string to clear existing values")
         c.argument('target_port', type=int, validator=validate_target_port_range, help="The session port used for ingress traffic.")
+        c.argument('probe_yaml', type=file_type, help='Path to a .yaml file with the configuration of the container probes. Only applicable for Custom Container Session Pool', is_preview=True)
 
     with self.argument_context('containerapp sessionpool', arg_group='Registry') as c:
         c.argument('registry_server', validator=validate_registry_server, help="The container registry server hostname, e.g. myregistry.azurecr.io.")

src/containerapp/azext_containerapp/containerapp_sessionpool_decorator.py

@@ -28,6 +28,7 @@
                                                            safe_set, safe_get, _ensure_identity_resource_id)
 from azure.cli.command_modules.containerapp._clients import ManagedEnvironmentClient
 from azure.cli.command_modules.containerapp._client_factory import handle_non_404_status_code_exception
+from azure.cli.command_modules.containerapp._decorator_utils import load_yaml_file
 from azure.cli.command_modules.containerapp._utils import is_registry_msi_system
 from azure.cli.core.commands.client_factory import get_subscription_id
 
@@ -156,10 +157,43 @@ def get_argument_system_assigned(self):
     def get_argument_user_assigned(self):
         return self.get_param("mi_user_assigned")
 
+    def get_argument_probe_yaml(self):
+        return self.get_param("probe_yaml")
+
     # pylint: disable=no-self-use
     def get_environment_client(self):
         return ManagedEnvironmentClient
 
+    def set_up_probes(self):
+        probes_def = load_yaml_file(self.get_argument_probe_yaml())
+        if not isinstance(probes_def, dict) or 'probes' not in probes_def:
+            raise ValidationError("The probe YAML file must be a dictionary containing a 'probes' key.")
+
+        probes_list = probes_def.get('probes')
+        if probes_list is None:
+            return []
+        if not isinstance(probes_list, list):
+            raise ValidationError("The 'probes' key in the probe YAML file must be a list of probes.")
+
+        return probes_list
+
+    def check_container_related_arguments(self):
+        container_related_args = {
+            '--args': self.get_argument_args(),
+            '--command': self.get_argument_startup_command(),
+            '--container-name': self.get_argument_container_name(),
+            '--cpu': self.get_argument_cpu(),
+            '--env-vars': self.get_argument_env_vars(),
+            '--image or -i': self.get_argument_image(),
+            '--memory': self.get_argument_memory(),
+            '--probe-yaml': self.get_argument_probe_yaml(),
+            '--target-port': self.get_argument_target_port(),
+        }
+
+        for arg_name, arg_value in container_related_args.items():
+            if arg_value is not None:
+                raise ValidationError(f"'{arg_name}' can not be set when container type is not 'CustomContainer'.")
+
 
 class SessionPoolCreateDecorator(SessionPoolPreviewDecorator):
     def validate_arguments(self):
@@ -339,6 +373,8 @@ def set_up_container(self):
         if self.get_argument_args() is not None:
             container_def["args"] = self.get_argument_args()
         container_def["resources"] = self.set_up_resource()
+        if self.get_argument_probe_yaml() is not None:
+            container_def["probes"] = self.set_up_probes()
         return container_def
 
     def set_up_secrets(self):
@@ -475,6 +511,10 @@ def validate_arguments(self):
                 current_lifecycle_type.lower() != LifecycleType.Timed.name.lower():
             raise ValidationError(f"--cooldown-period is not supported for the current --lifecycle-type '{current_lifecycle_type}'.")
 
+        # Validate container related arguments
+        if safe_get(self.existing_pool_def, "properties", "containerType").lower() != ContainerType.CustomContainer.name.lower():
+            self.check_container_related_arguments()
+
     def update(self):
         try:
             return self.client.update(
@@ -623,6 +663,8 @@ def set_up_container(self, customer_container_template):
                 container_def["resources"]["cpu"] = self.get_argument_cpu()
             if self.get_argument_memory() is not None:
                 container_def["resources"]["memory"] = self.get_argument_memory()
+        if self.get_argument_probe_yaml() is not None:
+            container_def["probes"] = self.set_up_probes()
         return container_def
 
     def set_up_registry_auth_configuration(self, secrets_def, customer_container_template):
@@ -677,7 +719,8 @@ def has_container_change(self):
                 self.get_argument_memory() is not None or
                 self.get_argument_env_vars() is not None or
                 self.get_argument_args() is not None or
-                self.get_argument_startup_command() is not None)
+                self.get_argument_startup_command() is not None or
+                self.get_argument_probe_yaml() is not None)
 
     def has_registry_change(self):
         return (self.get_argument_registry_server() is not None or

src/containerapp/azext_containerapp/custom.py

@@ -3117,7 +3117,8 @@ def create_session_pool(cmd,
                         registry_user=None,
                         mi_user_assigned=None,
                         registry_identity=None,
-                        mi_system_assigned=False):
+                        mi_system_assigned=False,
+                        probe_yaml=None):
     raw_parameters = locals()
     session_pool_decorator = SessionPoolCreateDecorator(
         cmd=cmd,
@@ -3158,7 +3159,8 @@ def update_session_pool(cmd,
                         registry_user=None,
                         mi_user_assigned=None,
                         registry_identity=None,
-                        mi_system_assigned=False):
+                        mi_system_assigned=False,
+                        probe_yaml=None):
     raw_parameters = locals()
     session_pool_decorator = SessionPoolUpdateDecorator(
         cmd=cmd,