I need to run a custom Docker container with GPU support using Vertex AI Pipelines, and I'm not seeing a clear way to do that. This requires several components:
- Applications (pytorch in my case)
- CUDA toolkit
- CUDA driver
- NVIDIA GPUs
I can use a NVIDIA Docker base image for #1 and #2, and a GCP accelerator for #4, but how do I install the CUDA driver in a Vertex AI pipeline? There's documentation on how to install NVIDIA drivers on GCE instances and GKE nodes, but nothing for Vertex AI.
One option could be to create a derivative container based on a GCP Deep Learning Container, but then I have to use a GCP container and don't have as much control over the environment.