Airflow跨任务传递类变量的合规实现方案问询

我在遵循Airflow最佳实践的前提下，对跨任务复用类实例存在困惑。场景如下：

• 在Task_One_Example任务中实例化了exampleClass，该类通过__init__从DAG上下文（dag_id、run_id等）初始化变量
• 目前在Task_Two_Example中只能重新创建类实例，但如果类本身包含大量状态或逻辑，重复初始化会带来额外开销
• 考虑过用XCom传递，但XCom更适合小数据，类较大时扩展性不足
• 若在DAG任务外初始化类，会违反Airflow最佳实践（导致变量在调度器持续重复初始化）

需要一个符合最佳实践的合理解决方案。

示例DAG代码 #

from airflow.decorators import dag, task
from airflow.operators.dummy_operator import DummyOperator
from airflow2_utils import environment_scheduler
from os import environ
from datetime import datetime, timedelta
import pendulum

# Default settings applied to all tasks
default_args = {
    'catchup': False,
    'depends_on_past': False,
    'email_on_failure': False,
    'email_on_retry': False,
    'retries': 1,
    'retry_delay': timedelta(minutes=1)
}

@dag(
    start_date=datetime(2023, 2, 21, tzinfo=pendulum.timezone("Europe/London")),
    schedule_interval=None,
    default_args=default_args,
    tags=["Client: Test", "Version: 2.2.5"]
)
def test_dag():
    start = DummyOperator(task_id='start')

    @task(task_id='Task_One_Example')
    def test_one():
        import PythonClassTest as pct
        class_instance = pct.exampleClass()
        class_instance.print_dag_info()
        
    @task(task_id='Task_Two_Example')
    def test_two():
        import PythonClassTest as pct
        class_instance = pct.exampleClass()
        class_instance.print_dag_info()
        
    end = DummyOperator(task_id='end')

    start >> test_one() >> test_two() >> end

dag = test_dag()

外部类代码 #

from airflow.operators.python import get_current_context

class exampleClass():
    def __init__(self):
        context = get_current_context()
        ti = context['ti']
        self.dag_id = ti.dag_id
        self.run_id = ti.run_id

    def print_dag_info(self):
        context = get_current_context()
        ti = context['ti']
        task_id = ti.task_id
        print(f"""
        dag_id: {self.dag_id}
        run_id: {self.run_id}
        task_id: {task_id}
        """)

解决方案 #

方案1：重构类，分离状态与逻辑（推荐） #

将类拆分为数据载体和逻辑方法两部分，用XCom传递轻量状态数据，避免传递整个类实例：

• 把需要跨任务共享的状态（如dag_id、run_id及其他可序列化字段）提取为简单的数据类/字典
• 第一个任务执行逻辑后，将状态数据通过XCom传递
• 第二个任务接收XCom数据，直接传入类初始化方法，跳过重复的上下文读取或状态计算

重构后的外部类

from airflow.operators.python import get_current_context
from dataclasses import dataclass

@dataclass
class ExampleState:
    dag_id: str
    run_id: str
    # 新增其他需要共享的状态字段

class exampleClass():
    def __init__(self, state: ExampleState = None):
        if state:
            self.dag_id = state.dag_id
            self.run_id = state.run_id
        else:
            context = get_current_context()
            ti = context['ti']
            self.dag_id = ti.dag_id
            self.run_id = ti.run_id

    def get_state(self) -> ExampleState:
        return ExampleState(dag_id=self.dag_id, run_id=self.run_id)

    def print_dag_info(self):
        context = get_current_context()
        ti = context['ti']
        task_id = ti.task_id
        print(f"""
        dag_id: {self.dag_id}
        run_id: {self.run_id}
        task_id: {task_id}
        """)

修改后的DAG任务

@task(task_id='Task_One_Example')
def test_one():
    import PythonClassTest as pct
    class_instance = pct.exampleClass()
    class_instance.print_dag_info()
    # 返回状态数据到XCom
    return class_instance.get_state()

@task(task_id='Task_Two_Example')
def test_two(shared_state):
    import PythonClassTest as pct
    # 用XCom传递的状态初始化类
    class_instance = pct.exampleClass(state=shared_state)
    class_instance.print_dag_info()

# 传递XCom数据
start >> test_one() >> test_two(test_one.output) >> end

方案2：使用共享存储序列化类实例 #

如果类的状态无法拆分为轻量数据，可将实例序列化后存储到共享存储（如S3、共享磁盘、内部数据库）：

• 第一个任务将类实例序列化（用pickle或其他工具）后写入共享存储，把存储路径通过XCom传递
• 第二个任务读取路径并反序列化实例，完成逻辑后清理存储文件避免冗余

示例思路：

import pickle
import os
from airflow.operators.python import get_current_context

# Task One
@task(task_id='Task_One_Example')
def test_one():
    import PythonClassTest as pct
    ti = get_current_context()['ti']
    class_instance = pct.exampleClass()
    class_instance.print_dag_info()
    
    # 序列化到共享路径
    storage_path = f"/shared/airflow/{ti.run_id}_example_class.pkl"
    with open(storage_path, 'wb') as f:
        pickle.dump(class_instance, f)
    return storage_path

# Task Two
@task(task_id='Task_Two_Example')
def test_two(storage_path):
    import PythonClassTest as pct
    with open(storage_path, 'rb') as f:
        class_instance = pickle.load(f)
    class_instance.print_dag_info()
    # 清理存储文件
    os.remove(storage_path)

方案3：直接重新初始化类（适合轻量类） #

如果类的初始化开销极低（仅读取几个上下文变量），直接在每个任务中重新初始化是完全合规的选择：

• Airflow任务是独立执行单元，每个任务运行在独立进程/容器中，重新初始化不会带来显著性能问题
• 这种方式避免了跨任务数据传递的复杂度，代码更简洁

关键原则提醒 #

• 绝对不要在DAG定义层（任务外）初始化类：调度器会周期性解析DAG文件，导致类被重复初始化，可能引发资源泄漏或状态混乱
• XCom仅用于小体积数据：超过48KB的数据会被存储到数据库Blob字段，性能下降，不适合传递大对象
• 保持任务独立性：Airflow任务设计为无状态、可重试的单元，跨任务共享状态需确保状态的可靠性和可复用性

内容的提问来源于stack exchange，提问作者Murray L