145 lines
4.7 KiB
Python
145 lines
4.7 KiB
Python
# 导入 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() |