如何在使用Airflow的DataprocInstantiateInlineWorkflowTemplateOperator时启用GPU

你需要在Workflow模板的cluster_config里添加gce_cluster_config.accelerators配置，对应命令行的--master-accelerator和--worker-accelerator参数。以下是修改后的完整模板示例：

{
    "placement": {
      "managed_cluster": {
        "cluster_name": "gpu-cluster-example",
        "config": {
          "gce_cluster_config": {
            "accelerators": [
              {
                "accelerator_type_uri": "nvidia-tesla-t4",
                "accelerator_count": 2,
                "scope": "MASTER"
              },
              {
                "accelerator_type_uri": "nvidia-tesla-t4",
                "accelerator_count": 1,
                "scope": "WORKER"
              }
            ]
          },
          "master_config": {
            "machine_type_uri": "n1-standard-8"
          },
          "worker_config": {
            "machine_type_uri": "n1-standard-8",
            "num_instances": 2
          },
          "software_config": {
            "image_version": "2.1-debian11"
          }
        }
      }
    },
    "jobs": [
      {
        "step_id": "pyspark-gpu-task",
        "pyspark_job": {
          "args": ["input-path", "output-path"],
          "main_python_file_uri": "gs://your-bucket/scripts/gpu-processing.py",
          "jar_file_uris": ["gs://your-bucket/jars/udf.jar"]
        }
      }
    ]
}

关键配置说明 #

• accelerators数组：每个元素对应一组节点的加速器配置：
  • accelerator_type_uri：直接填写GPU型号（如nvidia-tesla-t4），对应命令行的type参数
  • accelerator_count：指定单节点的GPU数量，对应命令行的count参数，默认值为1
  • scope：指定加速器作用的节点组，可选值为MASTER（主节点）、WORKER（工作节点）或ALL（所有节点）

额外注意事项 #

• 确保选择的Dataproc镜像版本支持GPU（推荐使用2.0及以上的Debian11镜像）
• Dataproc默认会自动安装GPU驱动，无需额外配置；如需自定义驱动，可在software_config.properties中添加dataproc:dataproc.gpu.driver.enable=true（默认已开启）

内容的提问来源于stack exchange，提问作者Anirban Saha