diff --git a/providers/microsoft/azure/provider.yaml b/providers/microsoft/azure/provider.yaml index a2a21cc270c9d..5b9e8f2601260 100644 --- a/providers/microsoft/azure/provider.yaml +++ b/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 diff --git a/providers/microsoft/azure/src/airflow/providers/microsoft/azure/get_provider_info.py b/providers/microsoft/azure/src/airflow/providers/microsoft/azure/get_provider_info.py index 944b296d7085a..0b312ec09a955 100644 --- a/providers/microsoft/azure/src/airflow/providers/microsoft/azure/get_provider_info.py +++ b/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"], diff --git a/providers/microsoft/azure/src/airflow/providers/microsoft/azure/hooks/container_instance.py b/providers/microsoft/azure/src/airflow/providers/microsoft/azure/hooks/container_instance.py index 27941ec53afc8..bf7e7177298ad 100644 --- a/providers/microsoft/azure/src/airflow/providers/microsoft/azure/hooks/container_instance.py +++ b/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` 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) diff --git a/providers/microsoft/azure/src/airflow/providers/microsoft/azure/operators/container_instances.py b/providers/microsoft/azure/src/airflow/providers/microsoft/azure/operators/container_instances.py index 3bcacfd4c0097..3901ff5258e4d 100644 --- a/providers/microsoft/azure/src/airflow/providers/microsoft/azure/operators/container_instances.py +++ b/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,16 +434,57 @@ 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 @@ -424,7 +492,7 @@ def _monitor_logging(self, resource_group: str, name: str) -> int: 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 diff --git a/providers/microsoft/azure/src/airflow/providers/microsoft/azure/triggers/container_instance.py b/providers/microsoft/azure/src/airflow/providers/microsoft/azure/triggers/container_instance.py new file mode 100644 index 0000000000000..db855d1b0e3d2 --- /dev/null +++ b/providers/microsoft/azure/src/airflow/providers/microsoft/azure/triggers/container_instance.py @@ -0,0 +1,130 @@ +# Licensed to the Apache Software Foundation (ASF) under one +# or more contributor license agreements. See the NOTICE file +# distributed with this work for additional information +# regarding copyright ownership. The ASF licenses this file +# to you under the Apache License, Version 2.0 (the +# "License"); you may not use this file except in compliance +# with the License. You may obtain a copy of the License at +# +# http://www.apache.org/licenses/LICENSE-2.0 +# +# Unless required by applicable law or agreed to in writing, +# software distributed under the License is distributed on an +# "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF ANY +# KIND, either express or implied. See the License for the +# specific language governing permissions and limitations +# under the License. +from __future__ import annotations + +import asyncio +from collections.abc import AsyncIterator +from typing import Any + +from airflow.providers.microsoft.azure.hooks.container_instance import AzureContainerInstanceAsyncHook +from airflow.triggers.base import BaseTrigger, TriggerEvent + +TERMINAL_STATES = frozenset({"Terminated", "Succeeded", "Failed", "Unhealthy"}) +SUCCESS_STATES = frozenset({"Terminated", "Succeeded"}) + + +class AzureContainerInstanceTrigger(BaseTrigger): + """ + Poll an Azure Container Instance until it reaches a terminal state. + + :param resource_group: the name of the resource group + :param name: the name of the container group + :param ci_conn_id: connection id of the Azure service principal + :param polling_interval: time in seconds between state polls + """ + + def __init__( + self, + resource_group: str, + name: str, + ci_conn_id: str, + polling_interval: float = 30.0, + ) -> None: + super().__init__() + self.resource_group = resource_group + self.name = name + self.ci_conn_id = ci_conn_id + self.polling_interval = polling_interval + + def serialize(self) -> tuple[str, dict[str, Any]]: + """Serialize trigger arguments and classpath.""" + return ( + "airflow.providers.microsoft.azure.triggers.container_instance.AzureContainerInstanceTrigger", + { + "resource_group": self.resource_group, + "name": self.name, + "ci_conn_id": self.ci_conn_id, + "polling_interval": self.polling_interval, + }, + ) + + async def run(self) -> AsyncIterator[TriggerEvent]: + """Poll ACI until a terminal state is reached, then yield a TriggerEvent.""" + hook = AzureContainerInstanceAsyncHook(azure_conn_id=self.ci_conn_id) + try: + async with hook.get_async_conn() as client: + while True: + cg_state = await client.container_groups.get(self.resource_group, self.name) + instance_view = cg_state.containers[0].instance_view + + if instance_view is not None: + c_state = instance_view.current_state + state = c_state.state + exit_code = c_state.exit_code + detail_status = c_state.detail_status + else: + prov = cg_state.provisioning_state + if prov in ("Failed", "Unhealthy"): + state = prov + exit_code = 1 + detail_status = prov + else: + await asyncio.sleep(self.polling_interval) + continue + + if state in TERMINAL_STATES: + if state in SUCCESS_STATES and exit_code == 0: + yield TriggerEvent( + { + "status": "success", + "exit_code": exit_code, + "detail_status": detail_status, + "resource_group": self.resource_group, + "name": self.name, + } + ) + else: + yield TriggerEvent( + { + "status": "error", + "exit_code": exit_code, + "detail_status": detail_status, + "resource_group": self.resource_group, + "name": self.name, + "message": ( + f"Container group {self.resource_group}/{self.name} " + f"reached state {state!r} with exit code {exit_code} " + f"({detail_status})" + ), + } + ) + return + + await asyncio.sleep(self.polling_interval) + except asyncio.CancelledError: + raise + except Exception as e: + yield TriggerEvent( + { + "status": "error", + "message": str(e), + "resource_group": self.resource_group, + "name": self.name, + "exit_code": 1, + "detail_status": "", + } + ) diff --git a/providers/microsoft/azure/tests/unit/microsoft/azure/hooks/test_container_instance.py b/providers/microsoft/azure/tests/unit/microsoft/azure/hooks/test_container_instance.py index 0bebf4d7afc59..33df9b1d99274 100644 --- a/providers/microsoft/azure/tests/unit/microsoft/azure/hooks/test_container_instance.py +++ b/providers/microsoft/azure/tests/unit/microsoft/azure/hooks/test_container_instance.py @@ -17,19 +17,27 @@ # under the License. from __future__ import annotations -from unittest.mock import MagicMock, patch +from unittest.mock import AsyncMock, MagicMock, patch import pytest from azure.core.exceptions import ResourceNotFoundError from azure.identity import AzureAuthorityHosts +from azure.identity.aio import ClientSecretCredential as AsyncClientSecretCredential +from azure.mgmt.containerinstance.aio import ( + ContainerInstanceManagementClient as AsyncContainerInstanceManagementClient, +) from azure.mgmt.containerinstance.models import ( + ContainerGroup, Logs, ResourceRequests, ResourceRequirements, ) from airflow.models import Connection -from airflow.providers.microsoft.azure.hooks.container_instance import AzureContainerInstanceHook +from airflow.providers.microsoft.azure.hooks.container_instance import ( + AzureContainerInstanceAsyncHook, + AzureContainerInstanceHook, +) @pytest.fixture @@ -252,3 +260,256 @@ def test_get_conn_cloud_environment( base_url=expected_base_url, credential_scopes=expected_scopes, ) + + +@pytest.fixture +def async_conn_with_credentials(create_mock_connection): + return create_mock_connection( + Connection( + conn_id="azure_aci_async_test", + conn_type="azure_container_instance", + login="client-id", + password="client-secret", + extra={ + "tenantId": "tenant-id", + "subscriptionId": "subscription-id", + }, + ) + ) + + +@pytest.fixture +def async_conn_without_credentials(create_mock_connection): + return create_mock_connection( + Connection( + conn_id="azure_aci_async_no_creds", + conn_type="azure_container_instance", + extra={"subscriptionId": "subscription-id"}, + ) + ) + + +class TestAzureContainerInstanceAsyncHook: + @patch( + "airflow.providers.microsoft.azure.hooks.container_instance.AsyncContainerInstanceManagementClient" + ) + @patch("airflow.providers.microsoft.azure.hooks.container_instance.AsyncClientSecretCredential") + @pytest.mark.asyncio + async def test_get_async_conn_with_client_secret( + self, + mock_credential_cls, + mock_client_cls, + async_conn_with_credentials, + ): + mock_credential = MagicMock(spec=AsyncClientSecretCredential) + mock_credential.close = AsyncMock() + mock_credential_cls.return_value = mock_credential + mock_client_instance = MagicMock(spec=AsyncContainerInstanceManagementClient) + mock_client_instance.close = AsyncMock() + mock_client_cls.return_value = mock_client_instance + + hook = AzureContainerInstanceAsyncHook(azure_conn_id=async_conn_with_credentials.conn_id) + async with hook.get_async_conn() as conn: + mock_credential_cls.assert_called_once_with( + client_id="client-id", + client_secret="client-secret", + tenant_id="tenant-id", + ) + mock_client_cls.assert_called_once_with( + credential=mock_credential, + subscription_id="subscription-id", + ) + assert conn is mock_client_instance + + @patch( + "airflow.providers.microsoft.azure.hooks.container_instance.AsyncContainerInstanceManagementClient" + ) + @patch("airflow.providers.microsoft.azure.hooks.container_instance.get_async_default_azure_credential") + @pytest.mark.asyncio + async def test_get_async_conn_with_default_credential( + self, + mock_default_cred, + mock_client_cls, + async_conn_without_credentials, + ): + mock_credential = MagicMock() + mock_credential.close = AsyncMock() + mock_default_cred.return_value = mock_credential + mock_client_instance = MagicMock(spec=AsyncContainerInstanceManagementClient) + mock_client_instance.close = AsyncMock() + mock_client_cls.return_value = mock_client_instance + + hook = AzureContainerInstanceAsyncHook(azure_conn_id=async_conn_without_credentials.conn_id) + async with hook.get_async_conn() as conn: + mock_default_cred.assert_called_once_with( + managed_identity_client_id=None, + workload_identity_tenant_id=None, + ) + assert conn is mock_client_instance + + @patch( + "airflow.providers.microsoft.azure.hooks.container_instance.AsyncContainerInstanceManagementClient" + ) + @patch("airflow.providers.microsoft.azure.hooks.container_instance.AsyncClientSecretCredential") + @pytest.mark.asyncio + async def test_get_async_conn_closes_client_and_credential( + self, + mock_credential_cls, + mock_client_cls, + async_conn_with_credentials, + ): + mock_credential = AsyncMock(spec=AsyncClientSecretCredential) + mock_credential_cls.return_value = mock_credential + mock_client_instance = AsyncMock(spec=AsyncContainerInstanceManagementClient) + mock_client_cls.return_value = mock_client_instance + + hook = AzureContainerInstanceAsyncHook(azure_conn_id=async_conn_with_credentials.conn_id) + async with hook.get_async_conn() as conn: + assert conn is mock_client_instance + + mock_client_instance.close.assert_called_once() + mock_credential.close.assert_called_once() + + @patch( + "airflow.providers.microsoft.azure.hooks.container_instance.AsyncContainerInstanceManagementClient" + ) + @patch("airflow.providers.microsoft.azure.hooks.container_instance.get_async_default_azure_credential") + @pytest.mark.asyncio + async def test_get_async_conn_skips_credential_close_when_not_closeable( + self, + mock_default_cred, + mock_client_cls, + async_conn_without_credentials, + ): + # Credential without a close method (spec=[] removes all attributes) + mock_credential = MagicMock(spec=[]) + mock_default_cred.return_value = mock_credential + mock_client_instance = AsyncMock(spec=AsyncContainerInstanceManagementClient) + mock_client_cls.return_value = mock_client_instance + + hook = AzureContainerInstanceAsyncHook(azure_conn_id=async_conn_without_credentials.conn_id) + async with hook.get_async_conn(): + pass + + mock_client_instance.close.assert_called_once() + assert not hasattr(mock_credential, "close") + + @patch( + "airflow.providers.microsoft.azure.hooks.container_instance.AsyncContainerInstanceManagementClient" + ) + @patch("airflow.providers.microsoft.azure.hooks.container_instance.AsyncClientSecretCredential") + @pytest.mark.asyncio + async def test_get_async_conn_closes_on_exception( + self, + mock_credential_cls, + mock_client_cls, + async_conn_with_credentials, + ): + mock_credential = AsyncMock(spec=AsyncClientSecretCredential) + mock_credential_cls.return_value = mock_credential + mock_client_instance = AsyncMock(spec=AsyncContainerInstanceManagementClient) + mock_client_cls.return_value = mock_client_instance + + hook = AzureContainerInstanceAsyncHook(azure_conn_id=async_conn_with_credentials.conn_id) + with pytest.raises(RuntimeError, match="boom"): + async with hook.get_async_conn(): + raise RuntimeError("boom") + + mock_client_instance.close.assert_called_once() + mock_credential.close.assert_called_once() + + @patch( + "airflow.providers.microsoft.azure.hooks.container_instance.AsyncContainerInstanceManagementClient" + ) + @patch("airflow.providers.microsoft.azure.hooks.container_instance.AsyncClientSecretCredential") + @pytest.mark.asyncio + async def test_get_state( + self, + mock_credential_cls, + mock_client_cls, + async_conn_with_credentials, + ): + mock_credential_cls.return_value = AsyncMock(spec=AsyncClientSecretCredential) + mock_cg = MagicMock(spec=ContainerGroup) + mock_client = MagicMock() + mock_client.close = AsyncMock() + mock_client.container_groups.get = AsyncMock(return_value=mock_cg) + mock_client_cls.return_value = mock_client + + hook = AzureContainerInstanceAsyncHook(azure_conn_id=async_conn_with_credentials.conn_id) + result = await hook.get_state("my-rg", "my-container") + + mock_client.container_groups.get.assert_called_once_with("my-rg", "my-container") + assert result is mock_cg + + @patch( + "airflow.providers.microsoft.azure.hooks.container_instance.AsyncContainerInstanceManagementClient" + ) + @patch("airflow.providers.microsoft.azure.hooks.container_instance.AsyncClientSecretCredential") + @pytest.mark.asyncio + async def test_get_logs( + self, + mock_credential_cls, + mock_client_cls, + async_conn_with_credentials, + ): + mock_credential_cls.return_value = AsyncMock(spec=AsyncClientSecretCredential) + mock_logs = MagicMock(spec=Logs) + mock_logs.content = "line1\nline2\n" + mock_client = MagicMock() + mock_client.close = AsyncMock() + mock_client.containers.list_logs = AsyncMock(return_value=mock_logs) + mock_client_cls.return_value = mock_client + + hook = AzureContainerInstanceAsyncHook(azure_conn_id=async_conn_with_credentials.conn_id) + result = await hook.get_logs("my-rg", "my-container") + + mock_client.containers.list_logs.assert_called_once_with( + "my-rg", "my-container", "my-container", tail=1000 + ) + assert result == ["line1\n", "line2\n"] + + @patch( + "airflow.providers.microsoft.azure.hooks.container_instance.AsyncContainerInstanceManagementClient" + ) + @patch("airflow.providers.microsoft.azure.hooks.container_instance.AsyncClientSecretCredential") + @pytest.mark.asyncio + async def test_get_logs_returns_none_sentinel_when_content_is_none( + self, + mock_credential_cls, + mock_client_cls, + async_conn_with_credentials, + ): + mock_credential_cls.return_value = AsyncMock(spec=AsyncClientSecretCredential) + mock_logs = MagicMock(spec=Logs) + mock_logs.content = None + mock_client = MagicMock() + mock_client.close = AsyncMock() + mock_client.containers.list_logs = AsyncMock(return_value=mock_logs) + mock_client_cls.return_value = mock_client + + hook = AzureContainerInstanceAsyncHook(azure_conn_id=async_conn_with_credentials.conn_id) + result = await hook.get_logs("my-rg", "my-container") + assert result == [None] + + @patch( + "airflow.providers.microsoft.azure.hooks.container_instance.AsyncContainerInstanceManagementClient" + ) + @patch("airflow.providers.microsoft.azure.hooks.container_instance.AsyncClientSecretCredential") + @pytest.mark.asyncio + async def test_delete( + self, + mock_credential_cls, + mock_client_cls, + async_conn_with_credentials, + ): + mock_credential_cls.return_value = AsyncMock(spec=AsyncClientSecretCredential) + mock_client = MagicMock() + mock_client.close = AsyncMock() + mock_client.container_groups.begin_delete = AsyncMock() + mock_client_cls.return_value = mock_client + + hook = AzureContainerInstanceAsyncHook(azure_conn_id=async_conn_with_credentials.conn_id) + await hook.delete("my-rg", "my-container") + + mock_client.container_groups.begin_delete.assert_called_once_with("my-rg", "my-container") diff --git a/providers/microsoft/azure/tests/unit/microsoft/azure/operators/test_container_instances.py b/providers/microsoft/azure/tests/unit/microsoft/azure/operators/test_container_instances.py index 9325b318bce5d..d21070de0c3c8 100644 --- a/providers/microsoft/azure/tests/unit/microsoft/azure/operators/test_container_instances.py +++ b/providers/microsoft/azure/tests/unit/microsoft/azure/operators/test_container_instances.py @@ -32,8 +32,10 @@ Event, ) +from airflow.exceptions import TaskDeferred from airflow.providers.common.compat.sdk import AirflowException from airflow.providers.microsoft.azure.operators.container_instances import AzureContainerInstancesOperator +from airflow.providers.microsoft.azure.triggers.container_instance import AzureContainerInstanceTrigger from tests_common.test_utils.compat import Context @@ -642,6 +644,292 @@ def test_execute_with_identity_dict(self, aci_mock): # user_assigned_identities should contain the resource id as a key assert resource_id in (called_cg.identity.user_assigned_identities or {}) + @mock.patch( + "airflow.providers.microsoft.azure.operators.container_instances.AzureContainerInstanceHook", + autospec=True, + ) + def test_execute_deferrable_defers_unconditionally(self, aci_mock): + """When deferrable=True, defer() is called immediately after create_or_update.""" + aci_mock.return_value.exists.return_value = False + + aci = AzureContainerInstancesOperator( + ci_conn_id="azure_default", + registry_conn_id=None, + resource_group="resource-group", + name="container-name", + image="container-image", + region="region", + task_id="task", + deferrable=True, + polling_interval=10.0, + ) + with pytest.raises(TaskDeferred) as exc_info: + aci.execute(None) + + assert aci_mock.return_value.create_or_update.call_count == 1 + assert isinstance(exc_info.value.trigger, AzureContainerInstanceTrigger) + assert exc_info.value.trigger.resource_group == "resource-group" + assert exc_info.value.trigger.name == "container-name" + assert exc_info.value.trigger.polling_interval == 10.0 + assert exc_info.value.method_name == "execute_complete" + assert aci_mock.return_value.delete.call_count == 0 + assert aci_mock.return_value.get_state.call_count == 0 + + @mock.patch( + "airflow.providers.microsoft.azure.operators.container_instances.AzureContainerInstanceHook", + autospec=True, + ) + def test_execute_complete_success(self, aci_mock): + """execute_complete succeeds, does not raise, and deletes the container group.""" + aci_mock.return_value.get_logs.return_value = None + + aci = AzureContainerInstancesOperator( + ci_conn_id="azure_default", + registry_conn_id=None, + resource_group="resource-group", + name="container-name", + image="container-image", + region="region", + task_id="task", + deferrable=True, + remove_on_success=True, + ) + result = aci.execute_complete( + context=None, + event={ + "status": "success", + "exit_code": 0, + "detail_status": "Completed", + "resource_group": "resource-group", + "name": "container-name", + }, + ) + assert result == 0 + assert aci_mock.return_value.delete.call_count == 1 + + @mock.patch( + "airflow.providers.microsoft.azure.operators.container_instances.AzureContainerInstanceHook", + autospec=True, + ) + def test_execute_complete_success_with_remove_on_success_false(self, aci_mock): + """execute_complete with remove_on_success=False should NOT delete the container.""" + aci_mock.return_value.get_logs.return_value = None + + aci = AzureContainerInstancesOperator( + ci_conn_id="azure_default", + registry_conn_id=None, + resource_group="resource-group", + name="container-name", + image="container-image", + region="region", + task_id="task", + deferrable=True, + remove_on_success=False, + ) + aci.execute_complete( + context=None, + event={ + "status": "success", + "exit_code": 0, + "detail_status": "Completed", + "resource_group": "resource-group", + "name": "container-name", + }, + ) + assert aci_mock.return_value.delete.call_count == 0 + + @mock.patch( + "airflow.providers.microsoft.azure.operators.container_instances.AzureContainerInstanceHook", + autospec=True, + ) + def test_execute_complete_error_event_raises(self, aci_mock): + """execute_complete raises RuntimeError when event has status=error.""" + aci = AzureContainerInstancesOperator( + ci_conn_id="azure_default", + registry_conn_id=None, + resource_group="resource-group", + name="container-name", + image="container-image", + region="region", + task_id="task", + deferrable=True, + remove_on_error=False, + ) + with pytest.raises(RuntimeError, match="Container group failed"): + aci.execute_complete( + context=None, + event={ + "status": "error", + "exit_code": 1, + "detail_status": "OOMKilled", + "message": "Container group failed with exit code 1", + "resource_group": "resource-group", + "name": "container-name", + }, + ) + + @mock.patch( + "airflow.providers.microsoft.azure.operators.container_instances.AzureContainerInstanceHook", + autospec=True, + ) + def test_execute_complete_error_event_removes_on_error(self, aci_mock): + """execute_complete with remove_on_error=True deletes the container on error.""" + aci = AzureContainerInstancesOperator( + ci_conn_id="azure_default", + registry_conn_id=None, + resource_group="resource-group", + name="container-name", + image="container-image", + region="region", + task_id="task", + deferrable=True, + remove_on_error=True, + ) + with pytest.raises(RuntimeError): + aci.execute_complete( + context=None, + event={ + "status": "error", + "exit_code": 1, + "detail_status": "Failed", + "message": "Container group failed with exit code 1", + "resource_group": "resource-group", + "name": "container-name", + }, + ) + assert aci_mock.return_value.delete.call_count == 1 + + @mock.patch( + "airflow.providers.microsoft.azure.operators.container_instances.AzureContainerInstanceHook", + autospec=True, + ) + def test_execute_complete_none_event_raises(self, aci_mock): + """execute_complete raises ValueError when event is None.""" + aci = AzureContainerInstancesOperator( + ci_conn_id="azure_default", + registry_conn_id=None, + resource_group="resource-group", + name="container-name", + image="container-image", + region="region", + task_id="task", + deferrable=True, + ) + with pytest.raises(ValueError, match="event is None"): + aci.execute_complete(context=None, event=None) + + @mock.patch( + "airflow.providers.microsoft.azure.operators.container_instances.AzureContainerInstanceHook", + autospec=True, + ) + def test_execute_complete_pushes_all_logs_to_xcom(self, aci_mock): + """execute_complete pushes all logs to XCom when xcom_all=True.""" + aci_mock.return_value.get_logs.return_value = ["line1\n", "line2\n"] + + aci = AzureContainerInstancesOperator( + ci_conn_id="azure_default", + registry_conn_id=None, + resource_group="resource-group", + name="container-name", + image="container-image", + region="region", + task_id="task", + deferrable=True, + xcom_all=True, + remove_on_success=False, + ) + context = Context(ti=XcomMock()) + aci.execute_complete( + context=context, + event={ + "status": "success", + "exit_code": 0, + "detail_status": "Completed", + "resource_group": "resource-group", + "name": "container-name", + }, + ) + assert context["ti"].xcom_pull(key="logs") == ["line1\n", "line2\n"] + + @mock.patch( + "airflow.providers.microsoft.azure.operators.container_instances.AzureContainerInstanceHook", + autospec=True, + ) + def test_execute_complete_pushes_last_log_to_xcom(self, aci_mock): + """execute_complete pushes only last log line to XCom when xcom_all=False.""" + aci_mock.return_value.get_logs.return_value = ["line1\n", "line2\n"] + + aci = AzureContainerInstancesOperator( + ci_conn_id="azure_default", + registry_conn_id=None, + resource_group="resource-group", + name="container-name", + image="container-image", + region="region", + task_id="task", + deferrable=True, + xcom_all=False, + remove_on_success=False, + ) + context = Context(ti=XcomMock()) + aci.execute_complete( + context=context, + event={ + "status": "success", + "exit_code": 0, + "detail_status": "Completed", + "resource_group": "resource-group", + "name": "container-name", + }, + ) + assert context["ti"].xcom_pull(key="logs") == ["line2\n"] + + @mock.patch( + "airflow.providers.microsoft.azure.operators.container_instances.AzureContainerInstanceHook", + autospec=True, + ) + def test_execute_remove_on_success_false_sync_path(self, aci_mock): + """Non-deferrable: remove_on_success=False should NOT delete the container on success.""" + terminated_cg = make_mock_container(state="Terminated", exit_code=0, detail_status="test") + aci_mock.return_value.get_state.return_value = terminated_cg + aci_mock.return_value.exists.return_value = False + + aci = AzureContainerInstancesOperator( + ci_conn_id=None, + registry_conn_id=None, + resource_group="resource-group", + name="container-name", + image="container-image", + region="region", + task_id="task", + remove_on_success=False, + ) + aci.execute(None) + assert aci_mock.return_value.delete.call_count == 0 + + @mock.patch( + "airflow.providers.microsoft.azure.operators.container_instances.AzureContainerInstanceHook", + autospec=True, + ) + def test_execute_remove_on_success_true_sync_path(self, aci_mock): + """Non-deferrable: remove_on_success=True (default) deletes the container on success.""" + terminated_cg = make_mock_container(state="Terminated", exit_code=0, detail_status="test") + aci_mock.return_value.get_state.return_value = terminated_cg + aci_mock.return_value.exists.return_value = False + + aci = AzureContainerInstancesOperator( + ci_conn_id=None, + registry_conn_id=None, + resource_group="resource-group", + name="container-name", + image="container-image", + region="region", + task_id="task", + remove_on_success=True, + ) + aci.execute(None) + assert aci_mock.return_value.delete.call_count == 1 + class XcomMock: def __init__(self) -> None: diff --git a/providers/microsoft/azure/tests/unit/microsoft/azure/triggers/test_container_instance.py b/providers/microsoft/azure/tests/unit/microsoft/azure/triggers/test_container_instance.py new file mode 100644 index 0000000000000..e4d26060ed2bd --- /dev/null +++ b/providers/microsoft/azure/tests/unit/microsoft/azure/triggers/test_container_instance.py @@ -0,0 +1,248 @@ +# +# Licensed to the Apache Software Foundation (ASF) under one +# or more contributor license agreements. See the NOTICE file +# distributed with this work for additional information +# regarding copyright ownership. The ASF licenses this file +# to you under the Apache License, Version 2.0 (the +# "License"); you may not use this file except in compliance +# with the License. You may obtain a copy of the License at +# +# http://www.apache.org/licenses/LICENSE-2.0 +# +# Unless required by applicable law or agreed to in writing, +# software distributed under the License is distributed on an +# "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF ANY +# KIND, either express or implied. See the License for the +# specific language governing permissions and limitations +# under the License. +from __future__ import annotations + +from contextlib import asynccontextmanager +from unittest import mock +from unittest.mock import AsyncMock, MagicMock + +import pytest +from azure.mgmt.containerinstance.models import ( + Container, + ContainerGroup, + ContainerPropertiesInstanceView, + ContainerState, +) + +from airflow.providers.microsoft.azure.triggers.container_instance import AzureContainerInstanceTrigger +from airflow.triggers.base import TriggerEvent + +RESOURCE_GROUP = "test-resource-group" +CONTAINER_NAME = "test-container" +CI_CONN_ID = "azure_container_instance_test" +POLLING_INTERVAL = 5.0 + +MODULE = "airflow.providers.microsoft.azure.triggers.container_instance" + + +def _make_cg_state(state: str, exit_code: int = 0, detail_status: str = "test") -> ContainerGroup: + """Build a minimal fake ContainerGroup using real SDK models (read-only attrs use type: ignore).""" + container_state = ContainerState() + container_state.state = state # type: ignore[assignment] + container_state.exit_code = exit_code # type: ignore[assignment] + container_state.detail_status = detail_status # type: ignore[assignment] + instance_view = ContainerPropertiesInstanceView() + instance_view.current_state = container_state # type: ignore[assignment] + container = Container(name="test", image="test", resources="test") # type: ignore[arg-type] + container.instance_view = instance_view # type: ignore[assignment] + cg = ContainerGroup(containers=[container], os_type="Linux") + cg.provisioning_state = state # type: ignore[assignment] + return cg + + +def _make_provisioning_state(state: str) -> ContainerGroup: + """Build a ContainerGroup where instance_view is None (provisioning phase).""" + container = Container(name="test", image="test", resources="test") # type: ignore[arg-type] + container.instance_view = None # type: ignore[assignment] + cg = ContainerGroup(containers=[container], os_type="Linux") + cg.provisioning_state = state # type: ignore[assignment] + return cg + + +def _make_mock_hook(mock_client): + """Create a mock hook whose get_async_conn yields the given mock client.""" + mock_hook = MagicMock() + + @asynccontextmanager + async def _get_async_conn(): + yield mock_client + + mock_hook.get_async_conn = _get_async_conn + return mock_hook + + +class TestAzureContainerInstanceTrigger: + def test_serialize(self): + trigger = AzureContainerInstanceTrigger( + resource_group=RESOURCE_GROUP, + name=CONTAINER_NAME, + ci_conn_id=CI_CONN_ID, + polling_interval=POLLING_INTERVAL, + ) + classpath, kwargs = trigger.serialize() + + assert classpath == ( + "airflow.providers.microsoft.azure.triggers.container_instance.AzureContainerInstanceTrigger" + ) + assert kwargs == { + "resource_group": RESOURCE_GROUP, + "name": CONTAINER_NAME, + "ci_conn_id": CI_CONN_ID, + "polling_interval": POLLING_INTERVAL, + } + + @pytest.mark.asyncio + @mock.patch(f"{MODULE}.AzureContainerInstanceAsyncHook") + async def test_run_yields_success_on_terminated_with_exit_code_zero(self, mock_hook_cls): + mock_client = MagicMock() + mock_client.container_groups.get = AsyncMock(return_value=_make_cg_state("Terminated", exit_code=0)) + mock_hook_cls.return_value = _make_mock_hook(mock_client) + + trigger = AzureContainerInstanceTrigger( + resource_group=RESOURCE_GROUP, + name=CONTAINER_NAME, + ci_conn_id=CI_CONN_ID, + polling_interval=POLLING_INTERVAL, + ) + events = [event async for event in trigger.run()] + + assert len(events) == 1 + assert isinstance(events[0], TriggerEvent) + assert events[0].payload["status"] == "success" + assert events[0].payload["exit_code"] == 0 + + @pytest.mark.asyncio + @mock.patch(f"{MODULE}.AzureContainerInstanceAsyncHook") + async def test_run_yields_success_on_succeeded_state(self, mock_hook_cls): + mock_client = MagicMock() + mock_client.container_groups.get = AsyncMock(return_value=_make_cg_state("Succeeded", exit_code=0)) + mock_hook_cls.return_value = _make_mock_hook(mock_client) + + trigger = AzureContainerInstanceTrigger( + resource_group=RESOURCE_GROUP, + name=CONTAINER_NAME, + ci_conn_id=CI_CONN_ID, + ) + events = [event async for event in trigger.run()] + + assert events[0].payload["status"] == "success" + + @pytest.mark.asyncio + @mock.patch(f"{MODULE}.AzureContainerInstanceAsyncHook") + async def test_run_yields_error_on_terminated_with_nonzero_exit_code(self, mock_hook_cls): + mock_client = MagicMock() + mock_client.container_groups.get = AsyncMock(return_value=_make_cg_state("Terminated", exit_code=1)) + mock_hook_cls.return_value = _make_mock_hook(mock_client) + + trigger = AzureContainerInstanceTrigger( + resource_group=RESOURCE_GROUP, + name=CONTAINER_NAME, + ci_conn_id=CI_CONN_ID, + ) + events = [event async for event in trigger.run()] + + assert events[0].payload["status"] == "error" + assert events[0].payload["exit_code"] == 1 + + @pytest.mark.asyncio + @mock.patch(f"{MODULE}.AzureContainerInstanceAsyncHook") + @pytest.mark.parametrize("terminal_state", ["Failed", "Unhealthy"]) + async def test_run_yields_error_on_failed_states(self, mock_hook_cls, terminal_state): + mock_client = MagicMock() + mock_client.container_groups.get = AsyncMock(return_value=_make_cg_state(terminal_state, exit_code=1)) + mock_hook_cls.return_value = _make_mock_hook(mock_client) + + trigger = AzureContainerInstanceTrigger( + resource_group=RESOURCE_GROUP, + name=CONTAINER_NAME, + ci_conn_id=CI_CONN_ID, + ) + events = [event async for event in trigger.run()] + + assert events[0].payload["status"] == "error" + + @pytest.mark.asyncio + @mock.patch(f"{MODULE}.asyncio.sleep", new_callable=AsyncMock) + @mock.patch(f"{MODULE}.AzureContainerInstanceAsyncHook") + async def test_run_polls_through_non_terminal_states(self, mock_hook_cls, mock_sleep): + mock_client = MagicMock() + mock_client.container_groups.get = AsyncMock( + side_effect=[ + _make_provisioning_state("Creating"), + _make_cg_state("Running", exit_code=0), + _make_cg_state("Terminated", exit_code=0), + ] + ) + mock_hook_cls.return_value = _make_mock_hook(mock_client) + + trigger = AzureContainerInstanceTrigger( + resource_group=RESOURCE_GROUP, + name=CONTAINER_NAME, + ci_conn_id=CI_CONN_ID, + polling_interval=POLLING_INTERVAL, + ) + events = [event async for event in trigger.run()] + + assert mock_client.container_groups.get.call_count == 3 + assert mock_sleep.call_count == 2 + mock_sleep.assert_called_with(POLLING_INTERVAL) + assert events[0].payload["status"] == "success" + + @pytest.mark.asyncio + @mock.patch(f"{MODULE}.AzureContainerInstanceAsyncHook") + async def test_run_yields_error_event_on_exception(self, mock_hook_cls): + mock_client = MagicMock() + mock_client.container_groups.get = AsyncMock(side_effect=Exception("Azure API error")) + mock_hook_cls.return_value = _make_mock_hook(mock_client) + + trigger = AzureContainerInstanceTrigger( + resource_group=RESOURCE_GROUP, + name=CONTAINER_NAME, + ci_conn_id=CI_CONN_ID, + ) + events = [event async for event in trigger.run()] + + assert len(events) == 1 + assert events[0].payload["status"] == "error" + assert "Azure API error" in events[0].payload["message"] + + @pytest.mark.asyncio + @mock.patch(f"{MODULE}.asyncio.sleep", new_callable=AsyncMock) + @mock.patch(f"{MODULE}.AzureContainerInstanceAsyncHook") + async def test_run_provisioning_state_does_not_yield_until_terminal(self, mock_hook_cls, mock_sleep): + """Container in 'Creating' (no instance_view) should keep polling.""" + mock_client = MagicMock() + mock_client.container_groups.get = AsyncMock( + side_effect=[ + _make_provisioning_state("Creating"), + _make_cg_state("Terminated", exit_code=0), + ] + ) + mock_hook_cls.return_value = _make_mock_hook(mock_client) + + trigger = AzureContainerInstanceTrigger( + resource_group=RESOURCE_GROUP, + name=CONTAINER_NAME, + ci_conn_id=CI_CONN_ID, + ) + events = [event async for event in trigger.run()] + + assert mock_client.container_groups.get.call_count == 2 + assert events[0].payload["status"] == "success" + + def test_serialize_deserialize_roundtrip(self): + original = AzureContainerInstanceTrigger( + resource_group=RESOURCE_GROUP, + name=CONTAINER_NAME, + ci_conn_id=CI_CONN_ID, + polling_interval=10.0, + ) + classpath, kwargs = original.serialize() + reconstructed = AzureContainerInstanceTrigger(**kwargs) + _, kwargs2 = reconstructed.serialize() + assert kwargs == kwargs2