如何从KF容器化Python组件将模型及指标正确导入Vertex AI并在UI可见

一、将训练生成的模型文件上传至Vertex AI Model Registry #

1. 准备符合要求的模型结构 #

确保训练输出的模型文件遵循Vertex AI规范：

• TensorFlow模型需导出为SavedModel格式，存放在独立目录中
• PyTorch模型需转换为TorchScript格式，或搭配Vertex AI官方PyTorch推理容器使用
• 自定义模型需附带model.yaml配置文件（若用自定义推理容器），明确模型输入输出格式

2. 在KF组件中调用Vertex AI SDK上传 #

在容器化Python组件代码中，直接使用google-cloud-aiplatform SDK完成上传（容器会自动继承KF pipeline的服务账号权限，无需手动配置密钥）：

import google.cloud.aiplatform as aiplatform

def train_and_upload_model(project_id: str, region: str, model_gcs_path: str):
    # 自定义训练逻辑，生成模型文件到指定GCS路径
    # ...
    
    # 初始化Vertex AI客户端
    aiplatform.init(project=project_id, location=region)
    
    # 上传模型至Model Registry
    uploaded_model = aiplatform.Model.upload(
        display_name="customer-churn-model",
        artifact_uri=model_gcs_path,
        # 指定适配模型的推理容器，可使用官方预构建容器或自定义容器
        serving_container_image_uri="us-docker.pkg.dev/vertex-ai/prediction/tf2-cpu.2-11:latest",
    )
    
    # 可选：为模型添加标签，方便后续检索
    uploaded_model.add_labels({"pipeline-run-id": "kf-v2-run-001", "model-type": "classification"})

3. 配置组件输出与权限 #

• 在KF V2组件定义中，将model_gcs_path声明为输出参数（类型为str或kfp.v2.dsl.Model）
• 确保KF pipeline使用的服务账号拥有storage.objectAdmin（读写GCS）和aiplatform.models.uploader（上传模型）权限

二、将训练指标上传至Vertex AI Experiments #

1. 关联Vertex AI实验运行 #

在组件中初始化ExperimentRun，关联到指定的Vertex AI实验：

from google.cloud.aiplatform import ExperimentRun

def train_and_log_metrics(project_id: str, region: str, experiment_name: str, run_name: str):
    # 初始化实验运行
    run = ExperimentRun(
        experiment=experiment_name,
        run_name=run_name,
        project=project_id,
        location=region
    )
    
    # 训练循环中记录时序指标
    for epoch in range(20):
        train_loss = 0.5 * (20 - epoch) / 20
        val_accuracy = 0.7 + 0.015 * epoch
        
        # 记录带步骤的时序指标
        run.log_metric("training_loss", train_loss, step=epoch)
        run.log_metric("validation_accuracy", val_accuracy, step=epoch)
    
    # 记录最终指标
    run.log_metric("final_validation_accuracy", val_accuracy)

2. KF V2集成优化 #

KF V2与Vertex AI Experiments原生集成，若组件是KF pipeline的一部分，也可通过kfp.v2.dsl.Metric类型的输出参数传递指标，系统会自动同步至对应实验运行中。

三、适配Kubeflow V2的关键注意事项 #

• 组件定义语法：使用kfp.v2.dsl.component装饰器替代V1的create_component_from_func，明确指定输入输出类型（如str, float, kfp.v2.dsl.Model）
• 存储路径规范：所有组件的输入输出必须指向GCS路径，本地路径会被自动挂载到容器的/gcs目录下
• SDK版本兼容性：容器中需安装google-cloud-aiplatform>=1.20.0版本，确保适配Vertex AI V2的API
• 服务账号权限：需为KF pipeline服务账号添加aiplatform.experiments.writer权限，确保能写入指标数据
• 自定义容器依赖：自定义训练容器需包含kfp>=1.8.14（适配V2的稳定版本），避免依赖冲突

内容的提问来源于stack exchange，提问作者Maxim Volgin