Add deferrable mode to AzureContainerInstancesOperator

providers/microsoft/azure/provider.yaml

@@ -311,6 +311,9 @@ triggers:
   - integration-name: Microsoft Azure Compute
     python-modules:
       - airflow.providers.microsoft.azure.triggers.compute
+  - integration-name: Microsoft Azure Container Instances
+    python-modules:
+      - airflow.providers.microsoft.azure.triggers.container_instance
   - integration-name: Microsoft Azure Data Factory
     python-modules:
       - airflow.providers.microsoft.azure.triggers.data_factory

providers/microsoft/azure/src/airflow/providers/microsoft/azure/get_provider_info.py

@@ -288,6 +288,10 @@ def get_provider_info():
                 "integration-name": "Microsoft Azure Compute",
                 "python-modules": ["airflow.providers.microsoft.azure.triggers.compute"],
             },
+            {
+                "integration-name": "Microsoft Azure Container Instances",
+                "python-modules": ["airflow.providers.microsoft.azure.triggers.container_instance"],
+            },
             {
                 "integration-name": "Microsoft Azure Data Factory",
                 "python-modules": ["airflow.providers.microsoft.azure.triggers.data_factory"],

providers/microsoft/azure/src/airflow/providers/microsoft/azure/hooks/container_instance.py

@@ -17,17 +17,30 @@
 # under the License.
 from __future__ import annotations
 
+from collections.abc import AsyncGenerator
+from contextlib import asynccontextmanager
 from functools import cached_property
 from typing import TYPE_CHECKING, Any, cast
 
 from azure.common.client_factory import get_client_from_auth_file, get_client_from_json_dict
 from azure.core.exceptions import ResourceNotFoundError
 from azure.identity import ClientSecretCredential, DefaultAzureCredential
+from azure.identity.aio import (
+    ClientSecretCredential as AsyncClientSecretCredential,
+    DefaultAzureCredential as AsyncDefaultAzureCredential,
+)
 from azure.mgmt.containerinstance import ContainerInstanceManagementClient
+from azure.mgmt.containerinstance.aio import (
+    ContainerInstanceManagementClient as AsyncContainerInstanceManagementClient,
+)
 
+from airflow.providers.common.compat.connection import get_async_connection
 from airflow.providers.common.compat.sdk import AirflowException
 from airflow.providers.microsoft.azure.hooks.base_azure import _AZURE_CLOUD_ENVIRONMENTS, AzureBaseHook
-from airflow.providers.microsoft.azure.utils import get_sync_default_azure_credential
+from airflow.providers.microsoft.azure.utils import (
+    get_async_default_azure_credential,
+    get_sync_default_azure_credential,
+)
 
 if TYPE_CHECKING:
     from azure.mgmt.containerinstance.models import (
@@ -180,3 +193,84 @@ def test_connection(self):
             return False, str(e)
 
         return True, "Successfully connected to Azure Container Instance."
+
+
+class AzureContainerInstanceAsyncHook(AzureContainerInstanceHook):
+    """
+    An async hook for communicating with Azure Container Instances.
+
+    :param azure_conn_id: :ref:`Azure connection id<howto/connection:azure>` of
+        a service principal which will be used to start the container instance.
+    """
+
+    def __init__(self, azure_conn_id: str = AzureContainerInstanceHook.default_conn_name) -> None:
+        super().__init__(azure_conn_id=azure_conn_id)
+
+    @asynccontextmanager
+    async def get_async_conn(self) -> AsyncGenerator[AsyncContainerInstanceManagementClient, None]:
+        """Create an async management client bound to a single credential."""
+        conn = await get_async_connection(self.conn_id)
+        tenant = conn.extra_dejson.get("tenantId")
+        subscription_id = cast("str", conn.extra_dejson.get("subscriptionId"))
+
+        credential: AsyncClientSecretCredential | AsyncDefaultAzureCredential
+        if all([conn.login, conn.password, tenant]):
+            credential = AsyncClientSecretCredential(
+                client_id=cast("str", conn.login),
+                client_secret=cast("str", conn.password),
+                tenant_id=cast("str", tenant),
+            )
+        else:
+            managed_identity_client_id = conn.extra_dejson.get("managed_identity_client_id")
+            workload_identity_tenant_id = conn.extra_dejson.get("workload_identity_tenant_id")
+            credential = get_async_default_azure_credential(
+                managed_identity_client_id=managed_identity_client_id,
+                workload_identity_tenant_id=workload_identity_tenant_id,
+            )
+
+        client = AsyncContainerInstanceManagementClient(
+            credential=credential,
+            subscription_id=subscription_id,
+        )
+        try:
+            yield client
+        finally:
+            await client.close()
+            if hasattr(credential, "close"):
+                await credential.close()
+
+    async def get_state(self, resource_group: str, name: str) -> ContainerGroup:  # type: ignore[override]
+        """
+        Get the state of a container group asynchronously.
+
+        :param resource_group: the name of the resource group
+        :param name: the name of the container group
+        :return: ContainerGroup
+        """
+        async with self.get_async_conn() as client:
+            return await client.container_groups.get(resource_group, name)
+
+    async def get_logs(self, resource_group: str, name: str, tail: int = 1000) -> list:  # type: ignore[override]
+        """
+        Get the tail from logs of a container group asynchronously.
+
+        :param resource_group: the name of the resource group
+        :param name: the name of the container group
+        :param tail: the size of the tail
+        :return: A list of log messages
+        """
+        async with self.get_async_conn() as client:
+            logs = await client.containers.list_logs(resource_group, name, name, tail=tail)
+            if logs.content is None:
+                return [None]
+            return logs.content.splitlines(True)
+
+    async def delete(self, resource_group: str, name: str) -> None:  # type: ignore[override]
+        """
+        Delete a container group asynchronously.
+
+        :param resource_group: the name of the resource group
+        :param name: the name of the container group
+        """
+        async with self.get_async_conn() as client:
+            await client.container_groups.begin_delete(resource_group, name)

providers/microsoft/azure/src/airflow/providers/microsoft/azure/operators/container_instances.py

@@ -21,6 +21,7 @@
 import time
 from collections import namedtuple
 from collections.abc import Sequence
+from functools import cached_property
 from typing import TYPE_CHECKING, Any, cast
 
 from azure.mgmt.containerinstance.models import (
@@ -42,10 +43,11 @@
 )
 from msrestazure.azure_exceptions import CloudError
 
-from airflow.providers.common.compat.sdk import AirflowException, AirflowTaskTimeout, BaseOperator
+from airflow.providers.common.compat.sdk import AirflowException, AirflowTaskTimeout, BaseOperator, conf
 from airflow.providers.microsoft.azure.hooks.container_instance import AzureContainerInstanceHook
 from airflow.providers.microsoft.azure.hooks.container_registry import AzureContainerRegistryHook
 from airflow.providers.microsoft.azure.hooks.container_volume import AzureContainerVolumeHook
+from airflow.providers.microsoft.azure.triggers.container_instance import AzureContainerInstanceTrigger
 
 if TYPE_CHECKING:
     from airflow.sdk import Context
@@ -105,6 +107,10 @@ class AzureContainerInstancesOperator(BaseOperator):
     :param diagnostics: Container group diagnostic information (Log Analytics).
     :param priority: Container group priority, Possible values include: 'Regular', 'Spot'
     :param identity: List of User/System assigned identities for the container group.
+    :param deferrable: Run in deferrable mode, releasing the worker slot while the container
+        runs. Defaults to ``[operators] default_deferrable`` in ``airflow.cfg``.
+    :param remove_on_success: Delete the container group after a successful run. Default ``True``.
+    :param polling_interval: Seconds between status polls in deferrable mode. Default ``30.0``.
 
     **Example**::
 
@@ -181,6 +187,7 @@ def __init__(
         gpu: Any | None = None,
         command: list[str] | None = None,
         remove_on_error: bool = True,
+        remove_on_success: bool = True,
         fail_if_exists: bool = True,
         tags: dict[str, str] | None = None,
         xcom_all: bool | None = None,
@@ -193,6 +200,8 @@ def __init__(
         diagnostics: ContainerGroupDiagnostics | None = None,
         priority: str | None = "Regular",
         identity: ContainerGroupIdentity | dict | None = None,
+        deferrable: bool = conf.getboolean("operators", "default_deferrable", fallback=False),
+        polling_interval: float = 30.0,
         **kwargs,
     ) -> None:
         super().__init__(**kwargs)
@@ -211,8 +220,8 @@ def __init__(
         self.gpu = gpu
         self.command = command
         self.remove_on_error = remove_on_error
+        self.remove_on_success = remove_on_success
         self.fail_if_exists = fail_if_exists
-        self._ci_hook: Any = None
         self.tags = tags
         self.xcom_all = xcom_all
         self.os_type = os_type
@@ -242,6 +251,8 @@ def __init__(
                 "Please set 'Regular' or 'Spot' as the priority. "
                 f"Found `{self.priority}`."
             )
+        self.deferrable = deferrable
+        self.polling_interval = polling_interval
 
     # helper to accept dict (user-friendly) or ContainerGroupIdentity (SDK object)
     @staticmethod
@@ -303,15 +314,17 @@ def _ensure_identity(identity: ContainerGroupIdentity | dict | None) -> Containe
             )
         return identity
 
+    @cached_property
+    def hook(self) -> AzureContainerInstanceHook:
+        return AzureContainerInstanceHook(azure_conn_id=self.ci_conn_id)
+
     def execute(self, context: Context) -> int:
         # Check name again in case it was templated.
         self._check_name(self.name)
 
-        self._ci_hook = AzureContainerInstanceHook(azure_conn_id=self.ci_conn_id)
-
         if self.fail_if_exists:
             self.log.info("Testing if container group already exists")
-            if self._ci_hook.exists(self.resource_group, self.name):
+            if self.hook.exists(self.resource_group, self.name):
                 raise AirflowException("Container group exists")
 
         if self.registry_conn_id:
@@ -340,6 +353,7 @@ def execute(self, context: Context) -> int:
             volume_mounts.append(VolumeMount(name=mount_name, mount_path=mount_path, read_only=read_only))
 
         exit_code = 1
+        _cleanup = True
         try:
             self.log.info("Starting container group with %.1f cpu %.1f mem", self.cpu, self.memory_in_gb)
             if self.gpu:
@@ -381,13 +395,26 @@ def execute(self, context: Context) -> int:
                 identity=self.identity,
             )
 
-            self._ci_hook.create_or_update(self.resource_group, self.name, container_group)
+            self.hook.create_or_update(self.resource_group, self.name, container_group)
 
             self.log.info("Container group started %s/%s", self.resource_group, self.name)
 
+            if self.deferrable:
+                _cleanup = False
+                self.defer(
+                    trigger=AzureContainerInstanceTrigger(
+                        resource_group=self.resource_group,
+                        name=self.name,
+                        ci_conn_id=self.ci_conn_id,
+                        polling_interval=self.polling_interval,
+                    ),
+                    method_name=self.execute_complete.__name__,
+                    timeout=self.execution_timeout,
+                )
+
             exit_code = self._monitor_logging(self.resource_group, self.name)
             if self.xcom_all is not None:
-                logs = self._ci_hook.get_logs(self.resource_group, self.name)
+                logs = self.hook.get_logs(self.resource_group, self.name)
                 if logs is None:
                     context["ti"].xcom_push(key="logs", value=[])
                 else:
@@ -407,24 +434,65 @@ def execute(self, context: Context) -> int:
             raise AirflowException("Could not start container group")
 
         finally:
-            if exit_code == 0 or self.remove_on_error:
-                self.on_kill()
+            if _cleanup:
+                if exit_code == 0 and self.remove_on_success:
+                    self.on_kill()
+                elif exit_code != 0 and self.remove_on_error:
+                    self.on_kill()
 
     def on_kill(self) -> None:
         self.log.info("Deleting container group")
         try:
-            self._ci_hook.delete(self.resource_group, self.name)
+            self.hook.delete(self.resource_group, self.name)
         except Exception:
             self.log.exception("Could not delete container group")
 
+    def execute_complete(self, context: Context, event: dict[str, Any] | None) -> int:
+        """
+        Handle the trigger event after deferral.
+
+        Called by the Triggerer when the container reaches a terminal state.
+        Raises on failure; returns the exit code on success.
+        """
+        if event is None:
+            raise ValueError("Trigger error: event is None")
+
+        exit_code: int = event.get("exit_code", 1)
+
+        if event["status"] == "error":
+            if self.remove_on_error:
+                self.on_kill()
+            raise RuntimeError(
+                event.get(
+                    "message",
+                    f"Container group {self.resource_group}/{self.name} failed with exit code {exit_code}",
+                )
+            )
+
+        try:
+            if self.xcom_all is not None:
+                logs = self.hook.get_logs(self.resource_group, self.name)
+                if logs is None:
+                    context["ti"].xcom_push(key="logs", value=[])
+                elif self.xcom_all:
+                    context["ti"].xcom_push(key="logs", value=logs)
+                else:
+                    context["ti"].xcom_push(key="logs", value=logs[-1:])
+
+            self.log.info("Container had exit code: %s", exit_code)
+            return exit_code
+        finally:
+            if self.remove_on_success:
+                self.on_kill()
+
     def _monitor_logging(self, resource_group: str, name: str) -> int:
         last_state = None
         last_message_logged = None
         last_line_logged = None
 
         while True:
             try:
-                cg_state = self._ci_hook.get_state(resource_group, name)
+                cg_state = self.hook.get_state(resource_group, name)
                 instance_view = cg_state.containers[0].instance_view
                 # If there is no instance view, we show the provisioning state
                 if instance_view is not None:
@@ -449,7 +517,7 @@ def _monitor_logging(self, resource_group: str, name: str) -> int:
 
                 if state in ["Running", "Terminated", "Succeeded"]:
                     try:
-                        logs = self._ci_hook.get_logs(resource_group, name)
+                        logs = self.hook.get_logs(resource_group, name)
                         if logs and logs[0] is None:
                             self.log.error("Container log is broken, marking as failed.")
                             return 1