如何让Vertex AI自定义作业自动获取配置参数（无需硬编码）

1. 利用Vertex AI自动注入的环境变量 #

Vertex AI会自动将核心作业配置注入到运行环境的环境变量中，无需额外配置即可直接读取：

• 对应baseOutputDirectory.outputUriPrefix的环境变量为AIP_OUTPUT_DIR，代码读取示例：

  import os
  output_directory = os.environ.get("AIP_OUTPUT_DIR")
  

• 其他常用自动注入的环境变量：
  • AIP_MODEL_DIR：模型保存的目标目录
  • AIP_TRAINING_DATA_URI：训练数据的存储路径（若作业配置了训练数据）
  • AIP_VALIDATION_DATA_URI：验证数据的存储路径（若作业配置了验证数据）

2. 通过Vertex AI Metadata服务获取完整作业配置 #

如果需要获取workerPoolSpecs中的机器规格、磁盘配置等详细参数，可通过官方SDK访问作业元数据：

1. 确保训练镜像中已安装google-cloud-aiplatform库（Vertex AI官方训练镜像通常已预装）
2. 在Python代码中获取当前作业的完整配置：

   from google.cloud import aiplatform
   import os
   
   # 从环境变量获取当前作业的资源名称
   job_resource_name = os.environ.get("AIP_TRAINING_JOB_NAME")
   aiplatform.init()
   
   # 获取作业对象并读取配置
   training_job = aiplatform.CustomJob.get(job_resource_name)
   worker_pool_spec = training_job.worker_pool_specs[0]
   
   # 提取具体参数
   machine_type = worker_pool_spec["machine_spec"]["machine_type"]
   disk_size_gb = worker_pool_spec["disk_spec"]["boot_disk_size_gb"]
   executor_image = worker_pool_spec["python_package_spec"]["executor_image_uri"]
   

3. 自定义参数传递（扩展场景） #

如果需要添加自定义参数，可通过pythonPackageSpec的args字段传递，再在代码中解析：

• 修改作业配置JSON，添加args字段：

  "pythonPackageSpec": {
    "executorImageUri": "europe-docker.pkg.dev/vertex-ai/training/sklearn-cpu.1-0:latest",
    "packageUris": ["gs://my_bucket_for_vertexai/my_package_directory/trainer-0.1.tar.gz"],
    "pythonModule": "trainer.task",
    "args": ["--replica-count", "1", "--disk-type", "pd-standard"]
  }
  

• 在Python代码中使用argparse解析参数：

  import argparse
  
  def main():
      parser = argparse.ArgumentParser()
      parser.add_argument("--replica-count", type=int)
      parser.add_argument("--disk-type", type=str)
      args = parser.parse_args()
  
      # 使用解析后的参数
      print(f"Replica count: {args.replica_count}")
      print(f"Disk type: {args.disk_type}")
  
  if __name__ == "__main__":
      main()
  

内容的提问来源于stack exchange，提问作者Maxime Oriol