在顶点AI上使用带有Kubeflow管道的Tesla A100 GPU

Question

我使用以下代码行来指定Kubeflow管道(KFP)上所需的机器类型和加速器/GPU，这将通过顶点AI/管道以无服务器的方式运行。

op().
set_cpu_limit(8).
set_memory_limit(50G).
add_node_selector_constraint('cloud.google.com/gke-accelerator', 'nvidia-tesla-k80').
set_gpu_limit(1)

它也适用于其他GPU，如Pascal，Tesla，Volta卡。

但是，对于最新的加速器类型，即Tesla A100，我不能这样做，因为它需要一种特殊的机器类型，它至少是一个a2-highgpu-1g。

当我在顶点上运行这个特定组件时，如何确保它运行在a2-highgpu-1g之上？

如果我只是简单地遵循旧GPU的方法：

op().
set_cpu_limit(12). # max for A2-highgpu-1g
set_memory_limit(85G). # max for A2-highgpu-1g
add_node_selector_constraint('cloud.google.com/gke-accelerator', 'nvidia-tesla-a100').
set_gpu_limit(1)

它在运行/部署时引发错误，因为生成的计算机类型是通用类型，即N1-Highmem-*。

当我没有指定cpu和内存限制时，也发生了同样的事情，希望它能够根据加速器约束自动选择正确的machnie类型。

    op().
    add_node_selector_constraint('cloud.google.com/gke-accelerator', 'nvidia-tesla-a100').
    set_gpu_limit(1)

错误："NVIDIA_TESLA_A100" is not supported for machine type "n1-highmem-2",

Answer 1

目前，GCP不支持普通KF组件的A2机器类型。现在一个可能的解决办法是使用，您可以显式地指定机器类型。