init: airflow_demo 初始化示例 DAG 集合

xcom_example2.py

@@ -0,0 +1,145 @@
+# 导入 sleep，用于模拟模型训练所需的时间
+from time import sleep
+
+# 导入 DAG，用于定义 Airflow 工作流
+from airflow import DAG
+
+# 导入 PythonOperator，用于执行 Python 函数
+from airflow.operators.python import PythonOperator
+
+# 导入 datetime，用于设置 DAG 的开始日期
+from datetime import datetime
+
+# 导入 random，用于生成随机准确率
+import random
+
+
+# ==================== 1. 模拟模型训练任务 ====================
+def train_model(model_name, **context):
+    """
+    模拟训练一个机器学习模型，并将模型准确率存入 XCom。
+
+    model_name：模型名称，例如 model_A、model_B、model_C
+    **context：Airflow 传入的运行上下文，其中包含 ti 等信息
+    """
+
+    # 模拟模型训练过程，让当前任务等待 10 秒
+    sleep(10)
+
+    # 随机生成一个 0.7 到 1.0 之间的浮点数，模拟模型准确率
+    accuracy = random.uniform(0.7, 1.0)
+
+    # 将模型准确率手动推送到 XCom
+    # key='model_accuracy'：这条数据在 XCom 中的键名
+    # value=accuracy：实际保存的准确率
+    context['ti'].xcom_push(
+        key='model_accuracy',
+        value=accuracy
+    )
+
+    # 将模型名称和准确率输出到任务日志中
+    print(f'{model_name} accuracy: {accuracy}')
+
+    # 注意：这里没有 return accuracy
+    # 因此不会把 accuracy 作为函数返回值保存到默认的 return_value 中
+
+
+# ==================== 2. 选择最佳模型 ====================
+def choose_best_model(**context):
+    """
+    从三个训练任务的 XCom 中读取准确率，并找出准确率最高的模型。
+    """
+
+    # 获取当前任务实例 TaskInstance
+    ti = context['ti']
+
+    # 一次性读取三个上游任务的准确率
+    accuracies = ti.xcom_pull(
+        task_ids=['model_A', 'model_B', 'model_C'],  # 指定要读取的任务
+        key='model_accuracy'                         # 指定要读取的 XCom 键
+    )
+
+    # 输出读取到的准确率列表
+    # 例如：[0.82, 0.95, 0.88]
+    print(f'accuracies: {accuracies}')
+
+    # 找出列表中的最大准确率
+    # max(accuracies) 返回最大值，例如 0.95
+    # accuracies.index(...) 返回该值第一次出现的索引，例如 1
+    best_model_index = accuracies.index(max(accuracies))
+
+    # 输出最高准确率
+    # 注意：这里原代码通过索引拼出了 model_1、model_2 或 model_3，
+    # 并不是实际任务 ID model_A、model_B 或 model_C。
+    print(
+        f'Best model accuracy: {accuracies[best_model_index]} '
+        f'from model_{best_model_index + 1}'
+    )
+
+
+# ==================== 3. 定义 DAG 工作流 ====================
+with DAG(
+    # DAG 的唯一标识符
+    dag_id='xcom_example2',
+
+    # DAG 的开始日期；不代表当天一定会自动执行
+    start_date=datetime(2024, 10, 1),
+
+    # 不设置定时调度，通常需要手动触发 DAG
+    schedule_interval=None,
+
+    # 不补跑历史上未执行的调度任务
+    catchup=False,
+
+    # 为 DAG 添加标签，便于在 Airflow 页面中筛选
+    tags=["demo"]
+) as dag:
+
+    # ==================== 4. 定义模型 A 训练任务 ====================
+    model_A = PythonOperator(
+        # Airflow 中的任务 ID
+        task_id='model_A',
+
+        # 指定要执行的 Python 函数
+        python_callable=train_model,
+
+        # 向 train_model 函数传递额外参数
+        # 相当于调用 train_model(model_name='model_A', ...)
+        op_kwargs={'model_name': 'model_A'}
+    )
+
+    # ==================== 5. 定义模型 B 训练任务 ====================
+    model_B = PythonOperator(
+        task_id='model_B',
+        python_callable=train_model,
+        op_kwargs={'model_name': 'model_B'}
+    )
+
+    # ==================== 6. 定义模型 C 训练任务 ====================
+    model_C = PythonOperator(
+        task_id='model_C',
+        python_callable=train_model,
+        op_kwargs={'model_name': 'model_C'}
+    )
+
+    # ==================== 7. 定义最佳模型选择任务 ====================
+    choose_task = PythonOperator(
+        # 任务 ID
+        task_id='choose_best_model',
+
+        # 指定执行选择最佳模型的函数
+        python_callable=choose_best_model
+    )
+
+    # ==================== 8. 设置任务依赖关系 ====================
+
+    # 三个模型任务互相没有依赖关系，可以并行执行
+    # choose_task 必须等待三个模型任务全部成功完成后才会执行
+    [model_A, model_B, model_C] >> choose_task
+
+
+# 脚本被直接执行时，调用 DAG 的 CLI 入口
+# 这是旧式的命令行调试写法，不是部署 DAG 必需的内容；
+# 正常使用 Airflow 调度器时，通常可以删除这段代码。
+if __name__ == '__main__':
+    dag.cli()