如何在Airflow中将AzureContainerInstancesOperator状态存入XCOM

一、直接获取Azure容器实例的运行状态（最简便方案） #

Airflow的AzureContainerInstancesOperator搭配原生钩子就能直接查询容器状态，无需额外复杂依赖：

• 利用AzureContainerInstanceHook，在ACI任务执行后，通过容器组名称+资源组查询运行状态（如Succeeded/Failed/Running）、退出码，甚至容器日志。
• 示例代码：

from airflow.providers.microsoft.azure.hooks.azure_container_instance import AzureContainerInstanceHook
from airflow.providers.microsoft.azure.operators.azure_container_instances import AzureContainerInstancesOperator
from airflow.operators.python import PythonOperator

def extract_aci_status(**context):
    hook = AzureContainerInstanceHook(azure_conn_id="azure_default")
    # 从上游任务获取创建的容器组名称
    container_group = context["ti"].xcom_pull(task_ids="execute_aci")["container_group_name"]
    # 查询容器组详情
    group_details = hook.get_container_group(resource_group_name="your_resource_group", name=container_group)
    # 提取容器状态与退出码并存入XCOM
    container_state = group_details.containers[0].instance_view.current_state
    context["ti"].xcom_push(key="aci_status", value=container_state.state)
    context["ti"].xcom_push(key="aci_exit_code", value=container_state.exit_code)
    # 若需要日志，可调用get_container_logs方法
    # logs = hook.get_container_logs(resource_group_name="your_resource_group", container_group_name=container_group, container_name=group_details.containers[0].name)
    # context["ti"].xcom_push(key="aci_logs", value=logs)

execute_aci = AzureContainerInstancesOperator(
    task_id="execute_aci",
    azure_conn_id="azure_default",
    resource_group_name="your_resource_group",
    name="test-aci-group",
    image="your-custom-image",
    # 其他必要参数如cpu/memory等
)

fetch_aci_status = PythonOperator(
    task_id="fetch_aci_status",
    python_callable=extract_aci_status,
    provide_context=True,
)

execute_aci >> fetch_aci_status

• 只要Airflow节点能访问Azure API，就能直接完成状态拉取，是替代DockerOperator XCOM功能的最直接方案。

二、Airflow中便捷获取Azure Service Bus消息内容 #

针对按关联ID过滤消息的需求，可通过原生钩子+SDK扩展实现，无需复杂自定义订阅：

• 编写Python函数，利用AzureServiceBusHook建立连接后，接收消息并按关联ID筛选：

from airflow.providers.microsoft.azure.hooks.azure_service_bus import AzureServiceBusHook
from airflow.operators.python import PythonOperator

def get_sb_message_by_correlation_id(**context):
    hook = AzureServiceBusHook(azure_conn_id="azure_sb_default")
    sb_client = hook.get_conn()
    # 订阅指定主题
    receiver = sb_client.get_subscription_receiver(
        topic_name="your-topic",
        subscription_name="your-subscription"
    )
    # 用任务ID作为关联ID示例，可替换为自定义标识
    target_correlation_id = context["ti"].task_id
    with receiver:
        # 批量拉取消息并筛选
        messages = receiver.receive_messages(max_message_count=10, max_wait_time=5)
        for msg in messages:
            if msg.correlation_id == target_correlation_id:
                context["ti"].xcom_push(key="sb_message", value=msg.body.decode())
                receiver.complete_message(msg)
                break

• 若需长期复用，可将此逻辑封装为自定义Operator；如果只是基础消息接收，也可用原生AzureServiceBusReceiveMessageOperator拉取后再过滤。

三、你的三个方案优化参考 #

• 方案1（直接获取容器输出）：用get_container_logs()即可实现，无需额外网络配置，适合输出内容不大的场景。
• 方案2（共享存储回传）：适合大体积输出，容器将结果写入Azure Blob后，用AzureBlobStorageHook读取，比直接拉日志更稳定。
• 方案3（Service Bus消息）：适合异步触发场景，容器任务完成后主动发消息到Service Bus，Airflow端用上述自定义函数接收筛选即可。

内容的提问来源于stack exchange，提问作者Tim Perry