LinkedIn
Twitter
Expected Behavior
I want to use a GPU on a component of Vertex AI.
Actual Behavior
Unfortunately, `torch.cuda.is_available()` is returning `False`. Also, `nvidia-smi` is not working if ran in the container of Vertex AI.
Note: both commands also don't work locally in the container if I'm not specifying the `--gpus all` flag in the command `docker run --rm -it --gpus all ee97db5bbd98 /bin/bash`. However, I can't find any option to add the `--gpus all` flag for Vertex AI. Would this be required?
Steps to Reproduce the Problem
My YAML file:
name: Processing
description: Process all the found HTML
inputs:
- name: friendly_name
type: String
description: The name of the company
- name: language
type: String
description: The language to process
- name: models
type: String
description: The models that will be used (all, genre, or standard)
implementation:
container:
image: eu.gcr.io/uman-interns/backend:v1.7
command:
[
python,
backend/pages/III_Process_website_data/process_website_data.py,
--friendly_name,
{ inputValue: friendly_name },
--language,
{ inputValue: language },
--models,
{ inputValue: models }
]
My pipeline:
rom kfp.v2 import compiler, dsl
import kfp.components as comp
from config import GCS_ARTIFACT_BUCKET, VARS
processing = comp.load_component_from_file("scraping/components/processing.yaml")
embeddings = comp.load_component_from_file("scraping/components/embeddings.yaml")
def compile_pipeline(file_name: str, tag: str):
@dsl.pipeline(
name="scraping",
description="Scrape a site and extract meaningful topics",
pipeline_root=f"gs://{GCS_ARTIFACT_BUCKET}/scraping/{tag}",
)
def pipeline(
friendly_name: str, url: str, language: str, google: bool, models: str) :
PROJECT_ID = VARS["PROJECT_ID"]
process = (
processing(friendly_name, language, models)
.set_display_name("URL processing")
.set_env_variable("PROJECT_ID", PROJECT_ID)
.set_caching_options(enable_caching=False)
.set_cpu_limit("4")
.set_memory_limit("16G")
.add_node_selector_constraint(
"cloud.google.com/gke-accelerator", "NVIDIA_TESLA_T4"
)
.set_gpu_limit(1)
).after(crawling)
embed = (
embeddings(friendly_name, language, models)
.set_display_name("Create embeddings")
.set_env_variable("PROJECT_ID", PROJECT_ID)
.set_caching_options(enable_caching=False)
.set_cpu_limit("4")
.set_memory_limit("16G")
.add_node_selector_constraint(
"cloud.google.com/gke-accelerator", "NVIDIA_TESLA_T4"
)
.set_gpu_limit(1)
).after(process)
compiler.Compiler().compile(pipeline, file_name)
Visualized in the browser:
