You are viewing a plain text version of this content. The canonical link for it is here.
Posted to commits@airflow.apache.org by "cruseakshay (via GitHub)" <gi...@apache.org> on 2026/03/03 08:29:20 UTC
[PR] Add deferrable mode to AzureContainerInstancesOperator [airflow]
cruseakshay opened a new pull request, #62772:
URL: https://github.com/apache/airflow/pull/62772
Add `deferrable=True` support to `AzureContainerInstancesOperator` so
tasks release their worker slot while waiting for long-running containers,
offloading polling to the lightweight Triggerer process.
## Changes
- **New `AzureContainerInstanceAsyncHook`** — async counterpart to the
existing sync hook, using `azure.mgmt.containerinstance.aio` for
non-blocking state/log/delete calls.
- **New `AzureContainerInstanceTrigger`** — polls ACI at a configurable
interval and yields a `TriggerEvent` when the container reaches a
terminal state.
- **`AzureContainerInstancesOperator`** — three new parameters:
- `deferrable` (default from `conf`) — enables deferrable mode
- `polling_interval` (default `5.0s`) — trigger poll frequency
- `remove_on_success` (default `True`) — controls cleanup on success
- Also fixes a bug where the `finally` cleanup block would delete a
still-running container group immediately after `self.defer()` raised
`TaskDeferred`.
- **`provider.yaml` + `get_provider_info.py`** updated to register the
new trigger.
closes: #62433
---
##### Was generative AI tooling used to co-author this PR?
- [X] Yes — Cursor, following [the guidelines](https://github.com/apache/airflow/blob/main/contributing-docs/05_pull_requests.rst#gen-ai-assisted-contributions)
--
This is an automated message from the Apache Git Service.
To respond to the message, please log on to GitHub and use the
URL above to go to the specific comment.
To unsubscribe, e-mail: commits-unsubscribe@airflow.apache.org
For queries about this service, please contact Infrastructure at:
users@infra.apache.org
Re: [PR] Add deferrable mode to AzureContainerInstancesOperator [airflow]
Posted by "cruseakshay (via GitHub)" <gi...@apache.org>.
cruseakshay commented on code in PR #62772:
URL: https://github.com/apache/airflow/pull/62772#discussion_r2899290278
##########
providers/microsoft/azure/src/airflow/providers/microsoft/azure/operators/container_instances.py:
##########
@@ -303,12 +315,14 @@ def _ensure_identity(identity: ContainerGroupIdentity | dict | None) -> Containe
)
return identity
+ @cached_property
+ def _ci_hook(self) -> AzureContainerInstanceHook:
+ return AzureContainerInstanceHook(azure_conn_id=self.ci_conn_id)
+
Review Comment:
This is consistent with ADF/Synapse operators. The benefit is lazy initialization for execute() and execute_complete(), since both methods require the hook. With @cached_property, we define it once and avoid duplicating the instantiation logic across both methods.
--
This is an automated message from the Apache Git Service.
To respond to the message, please log on to GitHub and use the
URL above to go to the specific comment.
To unsubscribe, e-mail: commits-unsubscribe@airflow.apache.org
For queries about this service, please contact Infrastructure at:
users@infra.apache.org
Re: [PR] Add deferrable mode to AzureContainerInstancesOperator [airflow]
Posted by "potiuk (via GitHub)" <gi...@apache.org>.
potiuk merged PR #62772:
URL: https://github.com/apache/airflow/pull/62772
--
This is an automated message from the Apache Git Service.
To respond to the message, please log on to GitHub and use the
URL above to go to the specific comment.
To unsubscribe, e-mail: commits-unsubscribe@airflow.apache.org
For queries about this service, please contact Infrastructure at:
users@infra.apache.org
Re: [PR] Add deferrable mode to AzureContainerInstancesOperator [airflow]
Posted by "cruseakshay (via GitHub)" <gi...@apache.org>.
cruseakshay commented on code in PR #62772:
URL: https://github.com/apache/airflow/pull/62772#discussion_r3079165098
##########
providers/microsoft/azure/src/airflow/providers/microsoft/azure/triggers/container_instance.py:
##########
@@ -0,0 +1,123 @@
+# 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:
+ state = cg_state.provisioning_state
+ exit_code = 0
+ detail_status = "Provisioning"
+
+ if state in TERMINAL_STATES:
+ if state in SUCCESS_STATES and exit_code == 0:
Review Comment:
addressed
--
This is an automated message from the Apache Git Service.
To respond to the message, please log on to GitHub and use the
URL above to go to the specific comment.
To unsubscribe, e-mail: commits-unsubscribe@airflow.apache.org
For queries about this service, please contact Infrastructure at:
users@infra.apache.org
Re: [PR] Add deferrable mode to AzureContainerInstancesOperator [airflow]
Posted by "dabla (via GitHub)" <gi...@apache.org>.
dabla commented on code in PR #62772:
URL: https://github.com/apache/airflow/pull/62772#discussion_r2976995590
##########
providers/microsoft/azure/src/airflow/providers/microsoft/azure/hooks/container_instance.py:
##########
@@ -170,3 +177,97 @@ 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:
+ self._async_conn: AsyncContainerInstanceManagementClient | None = None
+ self._async_credential: Any = None
+ super().__init__(azure_conn_id=azure_conn_id)
+
+ async def __aenter__(self) -> AzureContainerInstanceAsyncHook:
+ return self
Review Comment:
Don't really like the contextmanager on hook level, I would prefer to see this on the get_async_conn method instead, as the the connection we want to make sure is closed correctly in case of exception, not the hook, as the hook is just a way to get the connection, but it isn't the resource itself, the connection is.
So more something like this:
```
@asynccontextmanager
async def get_async_conn(self):
credential = ...
client = AsyncContainerInstanceManagementClient(...)
try:
yield client
finally:
await client.close()
if hasattr(credential, "close"):
await credential.close()
```
So later on it can be used like this in the trigger:
```
async with hook.get_async_conn() as client:
await client.container_groups.get(...)
```
But I really like the fact that you implemented a contextmanager, because that's a clean way to cope with connections.
##########
providers/microsoft/azure/src/airflow/providers/microsoft/azure/operators/container_instances.py:
##########
@@ -303,12 +315,14 @@ def _ensure_identity(identity: ContainerGroupIdentity | dict | None) -> Containe
)
return identity
+ @cached_property
+ def _ci_hook(self) -> AzureContainerInstanceHook:
Review Comment:
Why not just call this property hook?
##########
providers/microsoft/azure/src/airflow/providers/microsoft/azure/hooks/container_instance.py:
##########
@@ -170,3 +177,97 @@ 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:
+ self._async_conn: AsyncContainerInstanceManagementClient | None = None
Review Comment:
First call super().__init(), then initialise additional fields
##########
providers/microsoft/azure/src/airflow/providers/microsoft/azure/hooks/container_instance.py:
##########
@@ -170,3 +177,97 @@ 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:
+ self._async_conn: AsyncContainerInstanceManagementClient | None = None
+ self._async_credential: Any = None
+ super().__init__(azure_conn_id=azure_conn_id)
+
+ async def __aenter__(self) -> AzureContainerInstanceAsyncHook:
+ return self
+
+ async def __aexit__(self, exc_type: Any, exc_val: Any, exc_tb: Any) -> None:
+ await self.close()
+
+ async def close(self) -> None:
+ """Close the async connection and credential."""
+ if self._async_conn is not None:
+ await self._async_conn.close()
+ self._async_conn = None
+ if self._async_credential is not None and hasattr(self._async_credential, "close"):
+ await self._async_credential.close()
+ self._async_credential = None
+
+ async def get_async_conn(self) -> AsyncContainerInstanceManagementClient:
+ """Return (or create) the async management client."""
+ if self._async_conn is not None:
+ return self._async_conn
+
+ conn = self.get_connection(self.conn_id)
Review Comment:
I think we need to use the get_async_connection from common compat here instead of self.get_connection(self.conn_id)
--
This is an automated message from the Apache Git Service.
To respond to the message, please log on to GitHub and use the
URL above to go to the specific comment.
To unsubscribe, e-mail: commits-unsubscribe@airflow.apache.org
For queries about this service, please contact Infrastructure at:
users@infra.apache.org
Re: [PR] Add deferrable mode to AzureContainerInstancesOperator [airflow]
Posted by "cruseakshay (via GitHub)" <gi...@apache.org>.
cruseakshay commented on code in PR #62772:
URL: https://github.com/apache/airflow/pull/62772#discussion_r2899310869
##########
providers/microsoft/azure/tests/unit/microsoft/azure/hooks/test_container_instance.py:
##########
@@ -160,3 +167,177 @@ def test_get_conn_fallback_to_default_azure_credential(
credential=mock_credential,
subscription_id="subscription_id",
)
+
+
+@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_cls.return_value = mock_credential
+ mock_client_instance = MagicMock(spec=AsyncContainerInstanceManagementClient)
+ mock_client_cls.return_value = mock_client_instance
+
+ hook = AzureContainerInstanceAsyncHook(azure_conn_id=async_conn_with_credentials.conn_id)
+ conn = await hook.get_async_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 == 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(spec=AsyncClientSecretCredential)
+ mock_default_cred.return_value = mock_credential
+ mock_client_instance = MagicMock(spec=AsyncContainerInstanceManagementClient)
+ mock_client_cls.return_value = mock_client_instance
+
+ hook = AzureContainerInstanceAsyncHook(azure_conn_id=async_conn_without_credentials.conn_id)
+ conn = await hook.get_async_conn()
+
+ mock_default_cred.assert_called_once_with(
+ managed_identity_client_id=None,
+ workload_identity_tenant_id=None,
+ )
+ assert conn == mock_client_instance
+
+ @pytest.mark.asyncio
+ async def test_get_async_conn_returns_cached_connection(self, async_conn_with_credentials):
+ hook = AzureContainerInstanceAsyncHook(azure_conn_id=async_conn_with_credentials.conn_id)
+ mock_conn = MagicMock(spec=AsyncContainerInstanceManagementClient)
+ hook._async_conn = mock_conn
+
+ conn = await hook.get_async_conn()
+ assert conn is mock_conn
+
+ @pytest.mark.asyncio
+ async def test_get_state(self, async_conn_with_credentials):
+ hook = AzureContainerInstanceAsyncHook(azure_conn_id=async_conn_with_credentials.conn_id)
+ mock_client = MagicMock()
+ mock_cg = MagicMock(spec=ContainerGroup)
+ mock_client.container_groups.get = AsyncMock(return_value=mock_cg)
+ hook._async_conn = mock_client
+
+ 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
+
+ @pytest.mark.asyncio
+ async def test_get_logs(self, async_conn_with_credentials):
+ hook = AzureContainerInstanceAsyncHook(azure_conn_id=async_conn_with_credentials.conn_id)
+ mock_client = MagicMock()
+ mock_logs = MagicMock(spec=Logs)
+ mock_logs.content = "line1\nline2\n"
+ mock_client.containers.list_logs = AsyncMock(return_value=mock_logs)
+ hook._async_conn = mock_client
+
+ 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"]
+
+ @pytest.mark.asyncio
+ async def test_get_logs_returns_none_sentinel_when_content_is_none(self, async_conn_with_credentials):
+ hook = AzureContainerInstanceAsyncHook(azure_conn_id=async_conn_with_credentials.conn_id)
+ mock_client = MagicMock()
+ mock_logs = MagicMock(spec=Logs)
+ mock_logs.content = None
+ mock_client.containers.list_logs = AsyncMock(return_value=mock_logs)
+ hook._async_conn = mock_client
+
+ result = await hook.get_logs("my-rg", "my-container")
+ assert result == [None]
+
+ @pytest.mark.asyncio
+ async def test_delete(self, async_conn_with_credentials):
+ hook = AzureContainerInstanceAsyncHook(azure_conn_id=async_conn_with_credentials.conn_id)
+ mock_client = MagicMock()
+ mock_client.container_groups.begin_delete = AsyncMock()
Review Comment:
We can change if we decide against `fire-and-forget`
--
This is an automated message from the Apache Git Service.
To respond to the message, please log on to GitHub and use the
URL above to go to the specific comment.
To unsubscribe, e-mail: commits-unsubscribe@airflow.apache.org
For queries about this service, please contact Infrastructure at:
users@infra.apache.org
Re: [PR] Add deferrable mode to AzureContainerInstancesOperator [airflow]
Posted by "SameerMesiah97 (via GitHub)" <gi...@apache.org>.
SameerMesiah97 commented on code in PR #62772:
URL: https://github.com/apache/airflow/pull/62772#discussion_r2899795798
##########
providers/microsoft/azure/tests/unit/microsoft/azure/hooks/test_container_instance.py:
##########
@@ -160,3 +167,177 @@ def test_get_conn_fallback_to_default_azure_credential(
credential=mock_credential,
subscription_id="subscription_id",
)
+
+
+@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_cls.return_value = mock_credential
+ mock_client_instance = MagicMock(spec=AsyncContainerInstanceManagementClient)
+ mock_client_cls.return_value = mock_client_instance
+
+ hook = AzureContainerInstanceAsyncHook(azure_conn_id=async_conn_with_credentials.conn_id)
+ conn = await hook.get_async_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 == 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(spec=AsyncClientSecretCredential)
+ mock_default_cred.return_value = mock_credential
+ mock_client_instance = MagicMock(spec=AsyncContainerInstanceManagementClient)
+ mock_client_cls.return_value = mock_client_instance
+
+ hook = AzureContainerInstanceAsyncHook(azure_conn_id=async_conn_without_credentials.conn_id)
+ conn = await hook.get_async_conn()
+
+ mock_default_cred.assert_called_once_with(
+ managed_identity_client_id=None,
+ workload_identity_tenant_id=None,
+ )
+ assert conn == mock_client_instance
+
+ @pytest.mark.asyncio
+ async def test_get_async_conn_returns_cached_connection(self, async_conn_with_credentials):
+ hook = AzureContainerInstanceAsyncHook(azure_conn_id=async_conn_with_credentials.conn_id)
+ mock_conn = MagicMock(spec=AsyncContainerInstanceManagementClient)
+ hook._async_conn = mock_conn
+
+ conn = await hook.get_async_conn()
+ assert conn is mock_conn
+
+ @pytest.mark.asyncio
+ async def test_get_state(self, async_conn_with_credentials):
+ hook = AzureContainerInstanceAsyncHook(azure_conn_id=async_conn_with_credentials.conn_id)
+ mock_client = MagicMock()
+ mock_cg = MagicMock(spec=ContainerGroup)
+ mock_client.container_groups.get = AsyncMock(return_value=mock_cg)
+ hook._async_conn = mock_client
+
+ 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
+
+ @pytest.mark.asyncio
+ async def test_get_logs(self, async_conn_with_credentials):
+ hook = AzureContainerInstanceAsyncHook(azure_conn_id=async_conn_with_credentials.conn_id)
+ mock_client = MagicMock()
+ mock_logs = MagicMock(spec=Logs)
+ mock_logs.content = "line1\nline2\n"
+ mock_client.containers.list_logs = AsyncMock(return_value=mock_logs)
+ hook._async_conn = mock_client
+
+ 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"]
+
+ @pytest.mark.asyncio
+ async def test_get_logs_returns_none_sentinel_when_content_is_none(self, async_conn_with_credentials):
+ hook = AzureContainerInstanceAsyncHook(azure_conn_id=async_conn_with_credentials.conn_id)
+ mock_client = MagicMock()
+ mock_logs = MagicMock(spec=Logs)
+ mock_logs.content = None
+ mock_client.containers.list_logs = AsyncMock(return_value=mock_logs)
+ hook._async_conn = mock_client
+
+ result = await hook.get_logs("my-rg", "my-container")
+ assert result == [None]
+
+ @pytest.mark.asyncio
+ async def test_delete(self, async_conn_with_credentials):
+ hook = AzureContainerInstanceAsyncHook(azure_conn_id=async_conn_with_credentials.conn_id)
+ mock_client = MagicMock()
+ mock_client.container_groups.begin_delete = AsyncMock()
Review Comment:
> We can change if we decide against `fire-and-forget`
No need to change now as I think your current approach is ideal.
--
This is an automated message from the Apache Git Service.
To respond to the message, please log on to GitHub and use the
URL above to go to the specific comment.
To unsubscribe, e-mail: commits-unsubscribe@airflow.apache.org
For queries about this service, please contact Infrastructure at:
users@infra.apache.org
Re: [PR] Add deferrable mode to AzureContainerInstancesOperator [airflow]
Posted by "cruseakshay (via GitHub)" <gi...@apache.org>.
cruseakshay commented on code in PR #62772:
URL: https://github.com/apache/airflow/pull/62772#discussion_r3025543169
##########
providers/microsoft/azure/src/airflow/providers/microsoft/azure/hooks/container_instance.py:
##########
@@ -170,3 +177,97 @@ 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:
+ self._async_conn: AsyncContainerInstanceManagementClient | None = None
+ self._async_credential: Any = None
+ super().__init__(azure_conn_id=azure_conn_id)
+
+ async def __aenter__(self) -> AzureContainerInstanceAsyncHook:
+ return self
Review Comment:
@dabla I have added asynccontextmanager, PTAL.
--
This is an automated message from the Apache Git Service.
To respond to the message, please log on to GitHub and use the
URL above to go to the specific comment.
To unsubscribe, e-mail: commits-unsubscribe@airflow.apache.org
For queries about this service, please contact Infrastructure at:
users@infra.apache.org
Re: [PR] Add deferrable mode to AzureContainerInstancesOperator [airflow]
Posted by "cruseakshay (via GitHub)" <gi...@apache.org>.
cruseakshay commented on code in PR #62772:
URL: https://github.com/apache/airflow/pull/62772#discussion_r2899239928
##########
providers/microsoft/azure/src/airflow/providers/microsoft/azure/hooks/container_instance.py:
##########
@@ -170,3 +177,93 @@ 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:
+ self._async_conn: AsyncContainerInstanceManagementClient | None = None
+ super().__init__(azure_conn_id=azure_conn_id)
+
+ async def __aenter__(self) -> AzureContainerInstanceAsyncHook:
+ return self
+
+ async def __aexit__(self, exc_type: Any, exc_val: Any, exc_tb: Any) -> None:
+ await self.close()
+
+ async def close(self) -> None:
+ """Close the async connection."""
+ if self._async_conn is not None:
+ await self._async_conn.close()
+ self._async_conn = None
Review Comment:
Yes, good point.
--
This is an automated message from the Apache Git Service.
To respond to the message, please log on to GitHub and use the
URL above to go to the specific comment.
To unsubscribe, e-mail: commits-unsubscribe@airflow.apache.org
For queries about this service, please contact Infrastructure at:
users@infra.apache.org
Re: [PR] Add deferrable mode to AzureContainerInstancesOperator [airflow]
Posted by "potiuk (via GitHub)" <gi...@apache.org>.
potiuk commented on PR #62772:
URL: https://github.com/apache/airflow/pull/62772#issuecomment-4043131148
@cruseakshay This PR has been converted to **draft** because it does not yet meet our [Pull Request quality criteria](https://github.com/apache/airflow/blob/main/contributing-docs/05_pull_requests.rst#pull-request-quality-criteria).
**Issues found:**
- :x: **Merge conflicts**: This PR has merge conflicts with the `main` branch. Your branch is 340 commits behind `main`. Please rebase your branch (`git fetch origin && git rebase origin/main`), resolve the conflicts, and push again. See [contributing quick start](https://github.com/apache/airflow/blob/main/contributing-docs/03a_contributors_quick_start_beginners.rst).
> **Note:** Your branch is **340 commits behind `main`**. Some check failures may be caused by changes in the base branch rather than by your PR. Please rebase your branch and push again to get up-to-date CI results.
**What to do next:**
- The comment informs you what you need to do.
- Fix each issue, then mark the PR as "Ready for review" in the GitHub UI - but only after making sure that all the issues are fixed.
- Maintainers will then proceed with a normal review.
Converting a PR to draft is **not** a rejection — it is an invitation to bring the PR up to the project's standards so that maintainer review time is spent productively. If you have questions, feel free to ask on the [Airflow Slack](https://s.apache.org/airflow-slack).
--
This is an automated message from the Apache Git Service.
To respond to the message, please log on to GitHub and use the
URL above to go to the specific comment.
To unsubscribe, e-mail: commits-unsubscribe@airflow.apache.org
For queries about this service, please contact Infrastructure at:
users@infra.apache.org
Re: [PR] Add deferrable mode to AzureContainerInstancesOperator [airflow]
Posted by "Subham-KRLX (via GitHub)" <gi...@apache.org>.
Subham-KRLX commented on PR #62772:
URL: https://github.com/apache/airflow/pull/62772#issuecomment-4364398340
LGTM! Clean deferrable implementation with good test coverage backward compatible and follows Airflow patterns well Just document the polling interval defaults and you are good to go.
--
This is an automated message from the Apache Git Service.
To respond to the message, please log on to GitHub and use the
URL above to go to the specific comment.
To unsubscribe, e-mail: commits-unsubscribe@airflow.apache.org
For queries about this service, please contact Infrastructure at:
users@infra.apache.org
Re: [PR] Add deferrable mode to AzureContainerInstancesOperator [airflow]
Posted by "cruseakshay (via GitHub)" <gi...@apache.org>.
cruseakshay commented on code in PR #62772:
URL: https://github.com/apache/airflow/pull/62772#discussion_r2899222616
##########
providers/microsoft/azure/src/airflow/providers/microsoft/azure/hooks/container_instance.py:
##########
@@ -170,3 +177,93 @@ 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:
+ self._async_conn: AsyncContainerInstanceManagementClient | None = None
+ super().__init__(azure_conn_id=azure_conn_id)
+
+ async def __aenter__(self) -> AzureContainerInstanceAsyncHook:
+ return self
+
+ async def __aexit__(self, exc_type: Any, exc_val: Any, exc_tb: Any) -> None:
+ await self.close()
+
+ async def close(self) -> None:
+ """Close the async connection."""
+ if self._async_conn is not None:
+ await self._async_conn.close()
+ self._async_conn = None
+
+ async def get_async_conn(self) -> AsyncContainerInstanceManagementClient:
+ """Return (or create) the async management client."""
+ if self._async_conn is not None:
+ return self._async_conn
+
+ conn = self.get_connection(self.conn_id)
+ tenant = conn.extra_dejson.get("tenantId")
+ subscription_id = cast("str", conn.extra_dejson.get("subscriptionId"))
+
+ if all([conn.login, conn.password, tenant]):
+ credential: Any = 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,
+ )
+
+ self._async_conn = AsyncContainerInstanceManagementClient(
+ credential=credential,
+ subscription_id=subscription_id,
+ )
+ return self._async_conn
+
+ 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
+ """
+ client = await self.get_async_conn()
+ 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
+ """
+ client = await self.get_async_conn()
+ 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
+ """
+ client = await self.get_async_conn()
+ await client.container_groups.begin_delete(resource_group, name)
Review Comment:
Yes, I kept it the same to remain consistent. LMK if you think we should do `await poller.result()` in async and then similarly change the sync hook as well.
--
This is an automated message from the Apache Git Service.
To respond to the message, please log on to GitHub and use the
URL above to go to the specific comment.
To unsubscribe, e-mail: commits-unsubscribe@airflow.apache.org
For queries about this service, please contact Infrastructure at:
users@infra.apache.org
Re: [PR] Add deferrable mode to AzureContainerInstancesOperator [airflow]
Posted by "cruseakshay (via GitHub)" <gi...@apache.org>.
cruseakshay commented on code in PR #62772:
URL: https://github.com/apache/airflow/pull/62772#discussion_r3067525300
##########
providers/microsoft/azure/src/airflow/providers/microsoft/azure/hooks/container_instance.py:
##########
@@ -172,3 +185,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")
Review Comment:
`key_path / key_json` is not a new async limitation here. For the current track-2 azure-mgmt-containerinstance client, the existing sync hook’s get_client_from_auth_file/json_dict path already fails upstream because Azure’s legacy helper rejects autorest-v3 clients.
--
This is an automated message from the Apache Git Service.
To respond to the message, please log on to GitHub and use the
URL above to go to the specific comment.
To unsubscribe, e-mail: commits-unsubscribe@airflow.apache.org
For queries about this service, please contact Infrastructure at:
users@infra.apache.org
Re: [PR] Add deferrable mode to AzureContainerInstancesOperator [airflow]
Posted by "cruseakshay (via GitHub)" <gi...@apache.org>.
cruseakshay commented on code in PR #62772:
URL: https://github.com/apache/airflow/pull/62772#discussion_r3079168945
##########
providers/microsoft/azure/src/airflow/providers/microsoft/azure/triggers/container_instance.py:
##########
@@ -0,0 +1,123 @@
+# 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:
+ state = cg_state.provisioning_state
+ exit_code = 0
+ detail_status = "Provisioning"
+
+ 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)
Review Comment:
added as denfense
--
This is an automated message from the Apache Git Service.
To respond to the message, please log on to GitHub and use the
URL above to go to the specific comment.
To unsubscribe, e-mail: commits-unsubscribe@airflow.apache.org
For queries about this service, please contact Infrastructure at:
users@infra.apache.org
Re: [PR] Add deferrable mode to AzureContainerInstancesOperator [airflow]
Posted by "cruseakshay (via GitHub)" <gi...@apache.org>.
cruseakshay commented on code in PR #62772:
URL: https://github.com/apache/airflow/pull/62772#discussion_r2881816405
##########
providers/microsoft/azure/src/airflow/providers/microsoft/azure/operators/container_instances.py:
##########
@@ -417,6 +459,44 @@ def on_kill(self) -> None:
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 an AirflowException on failure; returns the exit code on success.
+ """
+ if event is None:
+ raise AirflowException("Trigger error: event is None")
Review Comment:
Raising built-in exceptions, thanks!
--
This is an automated message from the Apache Git Service.
To respond to the message, please log on to GitHub and use the
URL above to go to the specific comment.
To unsubscribe, e-mail: commits-unsubscribe@airflow.apache.org
For queries about this service, please contact Infrastructure at:
users@infra.apache.org
Re: [PR] Add deferrable mode to AzureContainerInstancesOperator [airflow]
Posted by "SameerMesiah97 (via GitHub)" <gi...@apache.org>.
SameerMesiah97 commented on code in PR #62772:
URL: https://github.com/apache/airflow/pull/62772#discussion_r2893208264
##########
providers/microsoft/azure/src/airflow/providers/microsoft/azure/hooks/container_instance.py:
##########
@@ -170,3 +177,93 @@ 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:
+ self._async_conn: AsyncContainerInstanceManagementClient | None = None
+ super().__init__(azure_conn_id=azure_conn_id)
+
+ async def __aenter__(self) -> AzureContainerInstanceAsyncHook:
+ return self
+
+ async def __aexit__(self, exc_type: Any, exc_val: Any, exc_tb: Any) -> None:
+ await self.close()
+
+ async def close(self) -> None:
+ """Close the async connection."""
+ if self._async_conn is not None:
+ await self._async_conn.close()
+ self._async_conn = None
Review Comment:
The async credential created in `get_async_conn()` isn't stored or closed in `close()`. Some Azure async credentials support `close()`. Would it make sense to keep a reference and close it here too?
##########
providers/microsoft/azure/src/airflow/providers/microsoft/azure/operators/container_instances.py:
##########
@@ -193,6 +201,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 = 5.0,
Review Comment:
Why 5 seconds? How long do you expect Azure Container Instances to run? It's seconds to minutes, this is fine. But if they can run for hours, I think 5 seconds is a bit too aggressive. 10-30 seconds would be more reasonable.
##########
providers/microsoft/azure/src/airflow/providers/microsoft/azure/hooks/container_instance.py:
##########
@@ -170,3 +177,93 @@ 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:
+ self._async_conn: AsyncContainerInstanceManagementClient | None = None
+ super().__init__(azure_conn_id=azure_conn_id)
+
+ async def __aenter__(self) -> AzureContainerInstanceAsyncHook:
+ return self
+
+ async def __aexit__(self, exc_type: Any, exc_val: Any, exc_tb: Any) -> None:
+ await self.close()
+
+ async def close(self) -> None:
+ """Close the async connection."""
+ if self._async_conn is not None:
+ await self._async_conn.close()
+ self._async_conn = None
+
+ async def get_async_conn(self) -> AsyncContainerInstanceManagementClient:
+ """Return (or create) the async management client."""
+ if self._async_conn is not None:
+ return self._async_conn
+
+ conn = self.get_connection(self.conn_id)
+ tenant = conn.extra_dejson.get("tenantId")
+ subscription_id = cast("str", conn.extra_dejson.get("subscriptionId"))
+
+ if all([conn.login, conn.password, tenant]):
+ credential: Any = 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,
+ )
+
+ self._async_conn = AsyncContainerInstanceManagementClient(
+ credential=credential,
+ subscription_id=subscription_id,
+ )
+ return self._async_conn
+
+ 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
+ """
+ client = await self.get_async_conn()
+ 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
+ """
+ client = await self.get_async_conn()
+ 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
+ """
+ client = await self.get_async_conn()
+ await client.container_groups.begin_delete(resource_group, name)
Review Comment:
Should the async hook mirror the existing sync behavior (fire-and-forget), or should it await the poller to ensure deletion completes and surface errors? Did you consider something like this:
```
poller = await client.container_groups.begin_delete(
resource_group,
name
)
await poller.result()
```
The key thing to understand here is that `begin_delete` does not return the result but an LRO poller object which waits for the result. I don't think it's necessarily wrong to keep the replicate the mechanics of the sync hook here but I was just wondering if you had thought of this approach.
##########
providers/microsoft/azure/tests/unit/microsoft/azure/hooks/test_container_instance.py:
##########
@@ -160,3 +167,177 @@ def test_get_conn_fallback_to_default_azure_credential(
credential=mock_credential,
subscription_id="subscription_id",
)
+
+
+@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_cls.return_value = mock_credential
+ mock_client_instance = MagicMock(spec=AsyncContainerInstanceManagementClient)
+ mock_client_cls.return_value = mock_client_instance
+
+ hook = AzureContainerInstanceAsyncHook(azure_conn_id=async_conn_with_credentials.conn_id)
+ conn = await hook.get_async_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 == 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(spec=AsyncClientSecretCredential)
+ mock_default_cred.return_value = mock_credential
+ mock_client_instance = MagicMock(spec=AsyncContainerInstanceManagementClient)
+ mock_client_cls.return_value = mock_client_instance
+
+ hook = AzureContainerInstanceAsyncHook(azure_conn_id=async_conn_without_credentials.conn_id)
+ conn = await hook.get_async_conn()
+
+ mock_default_cred.assert_called_once_with(
+ managed_identity_client_id=None,
+ workload_identity_tenant_id=None,
+ )
+ assert conn == mock_client_instance
+
+ @pytest.mark.asyncio
+ async def test_get_async_conn_returns_cached_connection(self, async_conn_with_credentials):
+ hook = AzureContainerInstanceAsyncHook(azure_conn_id=async_conn_with_credentials.conn_id)
+ mock_conn = MagicMock(spec=AsyncContainerInstanceManagementClient)
+ hook._async_conn = mock_conn
+
+ conn = await hook.get_async_conn()
+ assert conn is mock_conn
+
+ @pytest.mark.asyncio
+ async def test_get_state(self, async_conn_with_credentials):
+ hook = AzureContainerInstanceAsyncHook(azure_conn_id=async_conn_with_credentials.conn_id)
+ mock_client = MagicMock()
+ mock_cg = MagicMock(spec=ContainerGroup)
+ mock_client.container_groups.get = AsyncMock(return_value=mock_cg)
+ hook._async_conn = mock_client
+
+ 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
+
+ @pytest.mark.asyncio
+ async def test_get_logs(self, async_conn_with_credentials):
+ hook = AzureContainerInstanceAsyncHook(azure_conn_id=async_conn_with_credentials.conn_id)
+ mock_client = MagicMock()
+ mock_logs = MagicMock(spec=Logs)
+ mock_logs.content = "line1\nline2\n"
+ mock_client.containers.list_logs = AsyncMock(return_value=mock_logs)
+ hook._async_conn = mock_client
+
+ 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"]
+
+ @pytest.mark.asyncio
+ async def test_get_logs_returns_none_sentinel_when_content_is_none(self, async_conn_with_credentials):
+ hook = AzureContainerInstanceAsyncHook(azure_conn_id=async_conn_with_credentials.conn_id)
+ mock_client = MagicMock()
+ mock_logs = MagicMock(spec=Logs)
+ mock_logs.content = None
+ mock_client.containers.list_logs = AsyncMock(return_value=mock_logs)
+ hook._async_conn = mock_client
+
+ result = await hook.get_logs("my-rg", "my-container")
+ assert result == [None]
+
+ @pytest.mark.asyncio
+ async def test_delete(self, async_conn_with_credentials):
+ hook = AzureContainerInstanceAsyncHook(azure_conn_id=async_conn_with_credentials.conn_id)
+ mock_client = MagicMock()
+ mock_client.container_groups.begin_delete = AsyncMock()
Review Comment:
If you adjust the implementation for deleting containers, you will have to change this too.
##########
providers/microsoft/azure/src/airflow/providers/microsoft/azure/operators/container_instances.py:
##########
@@ -407,8 +445,12 @@ 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:
+ if self.remove_on_success:
+ self.on_kill()
+ elif self.remove_on_error:
+ self.on_kill()
Review Comment:
This can be made clearer like this:
```
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()
```
##########
providers/microsoft/azure/src/airflow/providers/microsoft/azure/triggers/container_instance.py:
##########
@@ -0,0 +1,124 @@
+# 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 = 5.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."""
+ try:
+ async with AzureContainerInstanceAsyncHook(azure_conn_id=self.ci_conn_id) as hook:
+ while True:
+ cg_state = await hook.get_state(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:
+ state = cg_state.provisioning_state
+ exit_code = 0
+ detail_status = "Provisioning"
+
+ self.log.info("Container group %s/%s state: %s", self.resource_group, self.name, state)
Review Comment:
This does not need to be logged in every iteration. I think you only log container state during transitions. With large numbers of concurrent deferrable tasks, this will result in extreme log pollution in the triggerer.
##########
providers/microsoft/azure/tests/unit/microsoft/azure/hooks/test_container_instance.py:
##########
@@ -160,3 +167,177 @@ def test_get_conn_fallback_to_default_azure_credential(
credential=mock_credential,
subscription_id="subscription_id",
)
+
+
+@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_cls.return_value = mock_credential
+ mock_client_instance = MagicMock(spec=AsyncContainerInstanceManagementClient)
+ mock_client_cls.return_value = mock_client_instance
+
+ hook = AzureContainerInstanceAsyncHook(azure_conn_id=async_conn_with_credentials.conn_id)
+ conn = await hook.get_async_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 == 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(spec=AsyncClientSecretCredential)
+ mock_default_cred.return_value = mock_credential
+ mock_client_instance = MagicMock(spec=AsyncContainerInstanceManagementClient)
+ mock_client_cls.return_value = mock_client_instance
+
+ hook = AzureContainerInstanceAsyncHook(azure_conn_id=async_conn_without_credentials.conn_id)
+ conn = await hook.get_async_conn()
+
+ mock_default_cred.assert_called_once_with(
+ managed_identity_client_id=None,
+ workload_identity_tenant_id=None,
+ )
+ assert conn == mock_client_instance
+
+ @pytest.mark.asyncio
+ async def test_get_async_conn_returns_cached_connection(self, async_conn_with_credentials):
+ hook = AzureContainerInstanceAsyncHook(azure_conn_id=async_conn_with_credentials.conn_id)
+ mock_conn = MagicMock(spec=AsyncContainerInstanceManagementClient)
+ hook._async_conn = mock_conn
+
+ conn = await hook.get_async_conn()
+ assert conn is mock_conn
Review Comment:
It might be worth adding a test that calls `get_async_conn() `twice and verifies that the same client instance is returned, to confirm the caching behavior.
##########
providers/microsoft/azure/src/airflow/providers/microsoft/azure/operators/container_instances.py:
##########
@@ -303,12 +315,14 @@ def _ensure_identity(identity: ContainerGroupIdentity | dict | None) -> Containe
)
return identity
+ @cached_property
+ def _ci_hook(self) -> AzureContainerInstanceHook:
+ return AzureContainerInstanceHook(azure_conn_id=self.ci_conn_id)
+
Review Comment:
Can you explain why you turned this into a cached property?
--
This is an automated message from the Apache Git Service.
To respond to the message, please log on to GitHub and use the
URL above to go to the specific comment.
To unsubscribe, e-mail: commits-unsubscribe@airflow.apache.org
For queries about this service, please contact Infrastructure at:
users@infra.apache.org
Re: [PR] Add deferrable mode to AzureContainerInstancesOperator [airflow]
Posted by "nailo2c (via GitHub)" <gi...@apache.org>.
nailo2c commented on code in PR #62772:
URL: https://github.com/apache/airflow/pull/62772#discussion_r2879977473
##########
providers/microsoft/azure/src/airflow/providers/microsoft/azure/operators/container_instances.py:
##########
@@ -417,6 +459,44 @@ def on_kill(self) -> None:
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 an AirflowException on failure; returns the exit code on success.
+ """
+ if event is None:
+ raise AirflowException("Trigger error: event is None")
Review Comment:
Thanks for implementing this!
Nit: per the contributing guide, we avoid raising `AirflowException` directly. Could you raise a built-in exception here instead? https://github.com/apache/airflow/blob/main/contributing-docs/05_pull_requests.rst#dont-raise-airflowexception-directly
--
This is an automated message from the Apache Git Service.
To respond to the message, please log on to GitHub and use the
URL above to go to the specific comment.
To unsubscribe, e-mail: commits-unsubscribe@airflow.apache.org
For queries about this service, please contact Infrastructure at:
users@infra.apache.org
Re: [PR] Add deferrable mode to AzureContainerInstancesOperator [airflow]
Posted by "cruseakshay (via GitHub)" <gi...@apache.org>.
cruseakshay commented on code in PR #62772:
URL: https://github.com/apache/airflow/pull/62772#discussion_r2899261064
##########
providers/microsoft/azure/src/airflow/providers/microsoft/azure/operators/container_instances.py:
##########
@@ -193,6 +201,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 = 5.0,
Review Comment:
I agree, 5 sec is aggressive. Setting it to 30 sec.
--
This is an automated message from the Apache Git Service.
To respond to the message, please log on to GitHub and use the
URL above to go to the specific comment.
To unsubscribe, e-mail: commits-unsubscribe@airflow.apache.org
For queries about this service, please contact Infrastructure at:
users@infra.apache.org
Re: [PR] Add deferrable mode to AzureContainerInstancesOperator [airflow]
Posted by "SameerMesiah97 (via GitHub)" <gi...@apache.org>.
SameerMesiah97 commented on code in PR #62772:
URL: https://github.com/apache/airflow/pull/62772#discussion_r2899794491
##########
providers/microsoft/azure/src/airflow/providers/microsoft/azure/hooks/container_instance.py:
##########
@@ -170,3 +177,93 @@ 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:
+ self._async_conn: AsyncContainerInstanceManagementClient | None = None
+ super().__init__(azure_conn_id=azure_conn_id)
+
+ async def __aenter__(self) -> AzureContainerInstanceAsyncHook:
+ return self
+
+ async def __aexit__(self, exc_type: Any, exc_val: Any, exc_tb: Any) -> None:
+ await self.close()
+
+ async def close(self) -> None:
+ """Close the async connection."""
+ if self._async_conn is not None:
+ await self._async_conn.close()
+ self._async_conn = None
+
+ async def get_async_conn(self) -> AsyncContainerInstanceManagementClient:
+ """Return (or create) the async management client."""
+ if self._async_conn is not None:
+ return self._async_conn
+
+ conn = self.get_connection(self.conn_id)
+ tenant = conn.extra_dejson.get("tenantId")
+ subscription_id = cast("str", conn.extra_dejson.get("subscriptionId"))
+
+ if all([conn.login, conn.password, tenant]):
+ credential: Any = 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,
+ )
+
+ self._async_conn = AsyncContainerInstanceManagementClient(
+ credential=credential,
+ subscription_id=subscription_id,
+ )
+ return self._async_conn
+
+ 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
+ """
+ client = await self.get_async_conn()
+ 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
+ """
+ client = await self.get_async_conn()
+ 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
+ """
+ client = await self.get_async_conn()
+ await client.container_groups.begin_delete(resource_group, name)
Review Comment:
> Yes, I kept it the same to remain consistent. LMK if you think we should do `await poller.result()` in async and then similarly change the sync hook as well.
On second thought, I think you should keep the current approach. After re-evaluating my suggestion in light of your comment, I cannot see a strong reason to break consistency with sync mode. Both should be fire-and-forget or delete-and-confirm, instead of one being the former and the other being the latter and vice versa.
--
This is an automated message from the Apache Git Service.
To respond to the message, please log on to GitHub and use the
URL above to go to the specific comment.
To unsubscribe, e-mail: commits-unsubscribe@airflow.apache.org
For queries about this service, please contact Infrastructure at:
users@infra.apache.org
Re: [PR] Add deferrable mode to AzureContainerInstancesOperator [airflow]
Posted by "cruseakshay (via GitHub)" <gi...@apache.org>.
cruseakshay commented on code in PR #62772:
URL: https://github.com/apache/airflow/pull/62772#discussion_r2981245515
##########
providers/microsoft/azure/src/airflow/providers/microsoft/azure/hooks/container_instance.py:
##########
@@ -170,3 +177,97 @@ 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:
+ self._async_conn: AsyncContainerInstanceManagementClient | None = None
+ self._async_credential: Any = None
+ super().__init__(azure_conn_id=azure_conn_id)
+
+ async def __aenter__(self) -> AzureContainerInstanceAsyncHook:
+ return self
+
+ async def __aexit__(self, exc_type: Any, exc_val: Any, exc_tb: Any) -> None:
+ await self.close()
+
+ async def close(self) -> None:
+ """Close the async connection and credential."""
+ if self._async_conn is not None:
+ await self._async_conn.close()
+ self._async_conn = None
+ if self._async_credential is not None and hasattr(self._async_credential, "close"):
+ await self._async_credential.close()
+ self._async_credential = None
+
+ async def get_async_conn(self) -> AsyncContainerInstanceManagementClient:
+ """Return (or create) the async management client."""
+ if self._async_conn is not None:
+ return self._async_conn
+
+ conn = self.get_connection(self.conn_id)
Review Comment:
Correct, thanks!
--
This is an automated message from the Apache Git Service.
To respond to the message, please log on to GitHub and use the
URL above to go to the specific comment.
To unsubscribe, e-mail: commits-unsubscribe@airflow.apache.org
For queries about this service, please contact Infrastructure at:
users@infra.apache.org
Re: [PR] Add deferrable mode to AzureContainerInstancesOperator [airflow]
Posted by "cruseakshay (via GitHub)" <gi...@apache.org>.
cruseakshay commented on code in PR #62772:
URL: https://github.com/apache/airflow/pull/62772#discussion_r2989309888
##########
providers/microsoft/azure/src/airflow/providers/microsoft/azure/hooks/container_instance.py:
##########
@@ -170,3 +177,97 @@ 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:
+ self._async_conn: AsyncContainerInstanceManagementClient | None = None
+ self._async_credential: Any = None
+ super().__init__(azure_conn_id=azure_conn_id)
+
+ async def __aenter__(self) -> AzureContainerInstanceAsyncHook:
+ return self
Review Comment:
Thanks @dabla! I will work on `asynccontextmanager` as you mentioned.
--
This is an automated message from the Apache Git Service.
To respond to the message, please log on to GitHub and use the
URL above to go to the specific comment.
To unsubscribe, e-mail: commits-unsubscribe@airflow.apache.org
For queries about this service, please contact Infrastructure at:
users@infra.apache.org
Re: [PR] Add deferrable mode to AzureContainerInstancesOperator [airflow]
Posted by "cruseakshay (via GitHub)" <gi...@apache.org>.
cruseakshay commented on code in PR #62772:
URL: https://github.com/apache/airflow/pull/62772#discussion_r3067525300
##########
providers/microsoft/azure/src/airflow/providers/microsoft/azure/hooks/container_instance.py:
##########
@@ -172,3 +185,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")
Review Comment:
`key_path / key_json` is not a new async limitation here. For the current track-2 azure-mgmt-containerinstance client, the existing sync hook’s `get_client_from_auth_file/json_dict path` already fails upstream because Azure’s legacy helper rejects autorest-v3 clients.
--
This is an automated message from the Apache Git Service.
To respond to the message, please log on to GitHub and use the
URL above to go to the specific comment.
To unsubscribe, e-mail: commits-unsubscribe@airflow.apache.org
For queries about this service, please contact Infrastructure at:
users@infra.apache.org
Re: [PR] Add deferrable mode to AzureContainerInstancesOperator [airflow]
Posted by "Copilot (via GitHub)" <gi...@apache.org>.
Copilot commented on code in PR #62772:
URL: https://github.com/apache/airflow/pull/62772#discussion_r3066478341
##########
providers/microsoft/azure/src/airflow/providers/microsoft/azure/triggers/container_instance.py:
##########
@@ -0,0 +1,123 @@
+# 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:
+ state = cg_state.provisioning_state
+ exit_code = 0
+ detail_status = "Provisioning"
+
+ if state in TERMINAL_STATES:
+ if state in SUCCESS_STATES and exit_code == 0:
Review Comment:
When `instance_view` is `None`, using `cg_state.provisioning_state` to decide terminal success can complete the trigger early (e.g., provisioning state may become `Succeeded` before the container has actually finished). Consider treating `instance_view is None` as non-terminal except for explicit provisioning failure states (e.g., `Failed`/`Unhealthy`), and only using container `current_state` for success/exit-code decisions.
##########
providers/microsoft/azure/src/airflow/providers/microsoft/azure/hooks/container_instance.py:
##########
@@ -172,3 +185,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")
Review Comment:
The sync hook supports `key_path`/`key_json` connection extras, but the async hook currently ignores them and always uses either client secret creds or default creds. This can make `deferrable=True` behave differently (or fail) for existing ACI connections that rely on `key_path`/`key_json`. Either add equivalent support in the async hook or raise a clear `AirflowException` when those extras are present so users understand the limitation.
##########
providers/microsoft/azure/src/airflow/providers/microsoft/azure/operators/container_instances.py:
##########
@@ -381,13 +395,36 @@ 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:
+ cg_state = self.hook.get_state(self.resource_group, self.name)
+ instance_view = cg_state.containers[0].instance_view
+ current_state = (
+ instance_view.current_state.state
+ if instance_view is not None
+ else cg_state.provisioning_state
+ )
+ terminal_states = {"Terminated", "Succeeded", "Failed", "Unhealthy"}
+ if current_state not in terminal_states:
Review Comment:
In deferrable mode, treating `cg_state.provisioning_state` as a terminal state can prevent deferral (e.g., provisioning may become `Succeeded` while the container is still running but `instance_view` isn't populated yet). To keep deferrable behavior correct and consistent with monitoring logic, consider not classifying provisioning `Succeeded` as terminal for deferral decisions; defer until container execution state is terminal.
```suggestion
if instance_view is not None and instance_view.current_state is not None
else None
)
terminal_states = {"Terminated", "Succeeded", "Failed", "Unhealthy"}
if current_state is None or current_state not in terminal_states:
```
##########
providers/microsoft/azure/tests/unit/microsoft/azure/operators/test_container_instances.py:
##########
@@ -642,6 +644,317 @@ 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_when_container_running(self, aci_mock):
+ """When deferrable=True and the container is still running, defer() is called."""
+ running_cg = make_mock_container(state="Running", exit_code=0, detail_status="test")
+ aci_mock.return_value.get_state.return_value = running_cg
+ 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 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"
+ # Container must NOT be deleted when deferring — it is still running
+ assert aci_mock.return_value.delete.call_count == 0
+
+ @mock.patch(
+ "airflow.providers.microsoft.azure.operators.container_instances.AzureContainerInstanceHook",
+ autospec=True,
+ )
+ def test_execute_deferrable_completes_synchronously_if_already_terminated(self, aci_mock):
+ """When deferrable=True but container is already terminal, no deferral — sync completion."""
+ 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="azure_default",
+ registry_conn_id=None,
+ resource_group="resource-group",
+ name="container-name",
+ image="container-image",
+ region="region",
+ task_id="task",
+ deferrable=True,
+ )
+ result = aci.execute(None)
+ assert result == 0
+ assert aci_mock.return_value.create_or_update.call_count == 1
+
+ @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 AirflowException when event has status=error."""
Review Comment:
This test docstring says `execute_complete` raises `AirflowException`, but the assertion expects `RuntimeError`. Please update the docstring (or the implementation) so the documented behavior matches what the test is actually verifying.
```suggestion
"""execute_complete raises RuntimeError when event has status=error."""
```
##########
providers/microsoft/azure/src/airflow/providers/microsoft/azure/triggers/container_instance.py:
##########
@@ -0,0 +1,123 @@
+# 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:
+ state = cg_state.provisioning_state
+ exit_code = 0
+ detail_status = "Provisioning"
+
+ 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)
Review Comment:
Catching `Exception` will also catch `asyncio.CancelledError` on Python versions where it subclasses `Exception`, which can prevent clean trigger cancellation/shutdown and instead emit an error event. Handle `asyncio.CancelledError` explicitly (re-raise) before the generic exception handler.
```suggestion
await asyncio.sleep(self.polling_interval)
except asyncio.CancelledError:
raise
```
##########
providers/microsoft/azure/src/airflow/providers/microsoft/azure/operators/container_instances.py:
##########
@@ -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,
Review Comment:
The PR description states `polling_interval` defaults to `5.0s`, but the operator signature defaults to `30.0`. Please align the code default with the PR description (or update the PR description) to avoid confusing users and reviewers.
##########
providers/microsoft/azure/src/airflow/providers/microsoft/azure/operators/container_instances.py:
##########
@@ -407,24 +444,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}",
+ )
+ )
Review Comment:
The non-deferrable path raises `AirflowException` on container failure, but the deferrable completion path raises a bare `RuntimeError`. For consistency with other operators and easier error handling, prefer raising `AirflowException` (preserving the same message) here as well.
--
This is an automated message from the Apache Git Service.
To respond to the message, please log on to GitHub and use the
URL above to go to the specific comment.
To unsubscribe, e-mail: commits-unsubscribe@airflow.apache.org
For queries about this service, please contact Infrastructure at:
users@infra.apache.org
Re: [PR] Add deferrable mode to AzureContainerInstancesOperator [airflow]
Posted by "SameerMesiah97 (via GitHub)" <gi...@apache.org>.
SameerMesiah97 commented on code in PR #62772:
URL: https://github.com/apache/airflow/pull/62772#discussion_r2899794491
##########
providers/microsoft/azure/src/airflow/providers/microsoft/azure/hooks/container_instance.py:
##########
@@ -170,3 +177,93 @@ 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:
+ self._async_conn: AsyncContainerInstanceManagementClient | None = None
+ super().__init__(azure_conn_id=azure_conn_id)
+
+ async def __aenter__(self) -> AzureContainerInstanceAsyncHook:
+ return self
+
+ async def __aexit__(self, exc_type: Any, exc_val: Any, exc_tb: Any) -> None:
+ await self.close()
+
+ async def close(self) -> None:
+ """Close the async connection."""
+ if self._async_conn is not None:
+ await self._async_conn.close()
+ self._async_conn = None
+
+ async def get_async_conn(self) -> AsyncContainerInstanceManagementClient:
+ """Return (or create) the async management client."""
+ if self._async_conn is not None:
+ return self._async_conn
+
+ conn = self.get_connection(self.conn_id)
+ tenant = conn.extra_dejson.get("tenantId")
+ subscription_id = cast("str", conn.extra_dejson.get("subscriptionId"))
+
+ if all([conn.login, conn.password, tenant]):
+ credential: Any = 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,
+ )
+
+ self._async_conn = AsyncContainerInstanceManagementClient(
+ credential=credential,
+ subscription_id=subscription_id,
+ )
+ return self._async_conn
+
+ 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
+ """
+ client = await self.get_async_conn()
+ 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
+ """
+ client = await self.get_async_conn()
+ 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
+ """
+ client = await self.get_async_conn()
+ await client.container_groups.begin_delete(resource_group, name)
Review Comment:
> Yes, I kept it the same to remain consistent. LMK if you think we should do `await poller.result()` in async and then similarly change the sync hook as well.
On second thought, I think you should keep the current approach. After re-evaluating my suggestion, I cannot see a strong reason to break consistency with sync mode. Both should be fire-and-forget or delete-and-confirm, instead of one being the former and the other being the latter and vice versa.
--
This is an automated message from the Apache Git Service.
To respond to the message, please log on to GitHub and use the
URL above to go to the specific comment.
To unsubscribe, e-mail: commits-unsubscribe@airflow.apache.org
For queries about this service, please contact Infrastructure at:
users@infra.apache.org
Re: [PR] Add deferrable mode to AzureContainerInstancesOperator [airflow]
Posted by "cruseakshay (via GitHub)" <gi...@apache.org>.
cruseakshay commented on code in PR #62772:
URL: https://github.com/apache/airflow/pull/62772#discussion_r3067627555
##########
providers/microsoft/azure/tests/unit/microsoft/azure/operators/test_container_instances.py:
##########
@@ -642,6 +644,317 @@ 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_when_container_running(self, aci_mock):
+ """When deferrable=True and the container is still running, defer() is called."""
+ running_cg = make_mock_container(state="Running", exit_code=0, detail_status="test")
+ aci_mock.return_value.get_state.return_value = running_cg
+ 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 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"
+ # Container must NOT be deleted when deferring — it is still running
+ assert aci_mock.return_value.delete.call_count == 0
+
+ @mock.patch(
+ "airflow.providers.microsoft.azure.operators.container_instances.AzureContainerInstanceHook",
+ autospec=True,
+ )
+ def test_execute_deferrable_completes_synchronously_if_already_terminated(self, aci_mock):
+ """When deferrable=True but container is already terminal, no deferral — sync completion."""
+ 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="azure_default",
+ registry_conn_id=None,
+ resource_group="resource-group",
+ name="container-name",
+ image="container-image",
+ region="region",
+ task_id="task",
+ deferrable=True,
+ )
+ result = aci.execute(None)
+ assert result == 0
+ assert aci_mock.return_value.create_or_update.call_count == 1
+
+ @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 AirflowException when event has status=error."""
Review Comment:
Stale docstring, will correct it.
--
This is an automated message from the Apache Git Service.
To respond to the message, please log on to GitHub and use the
URL above to go to the specific comment.
To unsubscribe, e-mail: commits-unsubscribe@airflow.apache.org
For queries about this service, please contact Infrastructure at:
users@infra.apache.org
Re: [PR] Add deferrable mode to AzureContainerInstancesOperator [airflow]
Posted by "cruseakshay (via GitHub)" <gi...@apache.org>.
cruseakshay commented on code in PR #62772:
URL: https://github.com/apache/airflow/pull/62772#discussion_r3067626645
##########
providers/microsoft/azure/src/airflow/providers/microsoft/azure/operators/container_instances.py:
##########
@@ -407,24 +444,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}",
+ )
+ )
Review Comment:
Added as nailo2c review suggestions: Don't use AirflowException directly — use standard Python exceptions
--
This is an automated message from the Apache Git Service.
To respond to the message, please log on to GitHub and use the
URL above to go to the specific comment.
To unsubscribe, e-mail: commits-unsubscribe@airflow.apache.org
For queries about this service, please contact Infrastructure at:
users@infra.apache.org
Re: [PR] Add deferrable mode to AzureContainerInstancesOperator [airflow]
Posted by "cruseakshay (via GitHub)" <gi...@apache.org>.
cruseakshay commented on code in PR #62772:
URL: https://github.com/apache/airflow/pull/62772#discussion_r3079162204
##########
providers/microsoft/azure/src/airflow/providers/microsoft/azure/operators/container_instances.py:
##########
@@ -381,13 +395,36 @@ 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:
+ cg_state = self.hook.get_state(self.resource_group, self.name)
+ instance_view = cg_state.containers[0].instance_view
+ current_state = (
+ instance_view.current_state.state
+ if instance_view is not None
+ else cg_state.provisioning_state
+ )
+ terminal_states = {"Terminated", "Succeeded", "Failed", "Unhealthy"}
+ if current_state not in terminal_states:
Review Comment:
Removed the pre-defer get_state() check entirely — the operator now defers unconditionally after create_or_update(). The trigger's polling loop now handles all state checking, and treats instance_view is None as non-terminal (keeps polling) except for explicit provisioning failures (Failed/Unhealthy).
--
This is an automated message from the Apache Git Service.
To respond to the message, please log on to GitHub and use the
URL above to go to the specific comment.
To unsubscribe, e-mail: commits-unsubscribe@airflow.apache.org
For queries about this service, please contact Infrastructure at:
users@infra.apache.org