How to use GPU on AI Platform Pipelines

Question

How do I use a GPU on AI Platform Pipelines? My pipeline uses set_gpu_limit(1) in one of the ops but I end up getting a This step is in Pending state with this message: Unschedulable: 0/3 nodes are available: 3 Insufficient nvidia.com/gpu. error.

Answer 1

Got it a few mins later... I followed the normal Kubeflow on GPU instructions

export GPU_POOL_NAME=gpu-pool
export CLUSTER_NAME=cluster-1

gcloud container node-pools create ${GPU_POOL_NAME} \
--accelerator type=nvidia-tesla-k80,count=1 \
--zone us-central1-a --cluster ${CLUSTER_NAME} \
--num-nodes=0 --machine-type=n1-standard-4 --min-nodes=0 --max-nodes=1 --enable-autoscaling

kubectl apply -f https://raw.githubusercontent.com/GoogleCloudPlatform/container-engine-accelerators/master/nvidia-driver-installer/cos/daemonset-preloaded.yaml