基于GCP/Airflow实现Sklearn机器学习模型自动化

一、核心思路 #

将Jupyter Notebook转换为可自动执行的Python脚本，依托Cloud Composer（GCP托管Airflow）实现定时调度，结合GCP原生服务完成模型训练、预测及BigQuery数据写入，全程无需手动干预。

二、分步实现 #

1. 转换Notebook为可执行脚本 #

• 清理训练/预测Notebook中的交互式代码（如input()、手动调试单元格），确保逻辑可从头至尾自动运行
• 使用jupyter nbconvert命令完成格式转换：

  jupyter nbconvert --to script training_notebook.ipynb
  jupyter nbconvert --to script prediction_notebook.ipynb
  

• 将转换后的.py脚本上传至GCS存储桶，便于Airflow调用及版本管理

2. 配置Cloud Composer权限与环境 #

• 确保Composer环境与AI Notebook实例处于同一VPC（若需访问私有资源）
• 为Composer服务账号添加以下角色：
  • roles/aiplatform.user（AI平台操作权限）
  • roles/bigquery.dataEditor（BigQuery数据写入权限）
  • roles/storage.objectViewer（GCS脚本读取权限）
  • roles/compute.instanceAdmin.v1（Notebook实例启停权限，可选）

3. 编写Airflow调度DAG #

（1）每月模型训练DAG

from airflow import DAG
from airflow.operators.bash import BashOperator
from airflow.providers.google.cloud.operators.compute import ComputeEngineStartInstanceOperator, ComputeEngineStopInstanceOperator
from airflow.utils.dates import days_ago
from datetime import timedelta

default_args = {
    'owner': 'airflow',
    'depends_on_past': False,
    'retries': 1,
    'retry_delay': timedelta(minutes=5),
}

with DAG(
    'monthly_model_training',
    default_args=default_args,
    description='每月自动执行模型训练',
    schedule_interval='0 0 1 * *',  # 每月1日凌晨触发
    start_date=days_ago(1),
    catchup=False,
    tags=['model-training', 'monthly'],
) as dag:

    # 可选：启动闲置的Notebook实例以节省成本
    start_notebook = ComputeEngineStartInstanceOperator(
        task_id='start_ai_notebook',
        project_id='你的GCP项目ID',
        zone='Notebook所在可用区',
        resource_id='Notebook实例名称'
    )

    # 远程执行训练脚本
    run_training = BashOperator(
        task_id='execute_training_script',
        bash_command='gcloud compute ssh Notebook实例名称 --zone 可用区 --command "cd /脚本存储路径 && python training_script.py"'
    )

    # 可选：训练完成后关闭Notebook实例
    stop_notebook = ComputeEngineStopInstanceOperator(
        task_id='stop_ai_notebook',
        project_id='你的GCP项目ID',
        zone='Notebook所在可用区',
        resource_id='Notebook实例名称'
    )

    start_notebook >> run_training >> stop_notebook

（2）每周预测并写入BigQuery DAG

from airflow import DAG
from airflow.operators.bash import BashOperator
from airflow.operators.dummy import DummyOperator
from airflow.providers.google.cloud.operators.compute import ComputeEngineStartInstanceOperator, ComputeEngineStopInstanceOperator
from airflow.providers.google.cloud.operators.bigquery import BigQueryInsertJobOperator
from airflow.utils.dates import days_ago
from datetime import timedelta

default_args = {
    'owner': 'airflow',
    'depends_on_past': False,
    'retries': 1,
    'retry_delay': timedelta(minutes=5),
}

with DAG(
    'weekly_prediction_to_bigquery',
    default_args=default_args,
    description='每周执行预测并写入BigQuery',
    schedule_interval='0 0 * * 1',  # 每周一凌晨触发
    start_date=days_ago(1),
    catchup=False,
    tags=['prediction', 'weekly', 'bigquery'],
) as dag:

    # 关联已有的数据准备任务（替换为你实际的任务ID）
    data_prep = DummyOperator(task_id='data_prep_completed')

    start_notebook = ComputeEngineStartInstanceOperator(
        task_id='start_ai_notebook',
        project_id='你的GCP项目ID',
        zone='Notebook所在可用区',
        resource_id='Notebook实例名称'
    )

    # 执行预测脚本（脚本内置BigQuery写入逻辑）
    run_prediction = BashOperator(
        task_id='execute_prediction_script',
        bash_command='gcloud compute ssh Notebook实例名称 --zone 可用区 --command "cd /脚本存储路径 && python prediction_script.py"'
    )

    # 验证BigQuery写入结果
    validate_bq = BigQueryInsertJobOperator(
        task_id='check_bq_records',
        configuration={
            "query": {
                "query": "SELECT COUNT(*) FROM `你的项目.数据集.预测结果表`",
                "useLegacySql": False
            }
        },
        project_id='你的GCP项目ID'
    )

    stop_notebook = ComputeEngineStopInstanceOperator(
        task_id='stop_ai_notebook',
        project_id='你的GCP项目ID',
        zone='Notebook所在可用区',
        resource_id='Notebook实例名称'
    )

    data_prep >> start_notebook >> run_prediction >> validate_bq >> stop_notebook

4. 优化预测脚本的BigQuery写入逻辑 #

在prediction_script.py中使用官方库完成数据写入：

import pandas as pd
from google.cloud import bigquery

# 假设predict_results为预测结果DataFrame
predict_results = pd.DataFrame(...)

# 初始化BigQuery客户端
client = bigquery.Client(project='你的GCP项目ID')

# 写入数据（追加模式，避免覆盖历史数据）
job = client.load_table_from_dataframe(
    predict_results,
    '你的项目.数据集.预测结果表',
    write_disposition='WRITE_APPEND'
)
job.result()  # 等待写入完成

5. 测试与调试 #

• 手动触发单个任务（如start_notebook、run_training）验证脚本执行逻辑
• 通过Cloud Composer日志排查权限、路径等执行错误
• 直接查看BigQuery表数据，验证写入结果是否符合预期

三、生产环境替代方案 #

若无需依赖Notebook实例，可将脚本包装为Vertex AI自定义Job，通过Airflow的VertexAIStartJobOperator调度，这种方式无需管理虚拟机，更适合大规模生产场景。

内容的提问来源于stack exchange，提问作者Łukasz Minkisiewicz