Training spaCy model as a Vertex AI Pipeline "Component"

After some rehearsals, I think I have figured out what my code was missing. Actually, the train component definition was correct (with some minor tweaks relative to what was originally posted); however the pipeline was missing the GPU definition. I will first include a dummy example code, which trains a NER model using spaCy, and orchestrates everything via Vertex AI Pipeline:

from kfp.v2 import compiler
from kfp.v2.dsl import pipeline, component, Dataset, Input, Output, OutputPath, InputPath
from datetime import datetime
from google.cloud import aiplatform
from typing import NamedTuple


# Component definition

@component(
    packages_to_install=[
        "setuptools",
        "wheel", 
        "spacy[cuda113,transformers,lookups]",
    ],
    base_image="gcr.io/deeplearning-platform-release/base-cu113",
    output_component_file="generate.yaml"
)
def generate_spacy_file(train_path: OutputPath(), dev_path: OutputPath()):
    """
    Generates a small, dummy 'train.spacy' & 'dev.spacy' file
    
    Returns:
    -------
    train_path : Relative location in GCS, for the "train.spacy" file.
    dev_path: Relative location in GCS, for the "dev.spacy" file.
    """
    import spacy
    from spacy.training import Example
    from spacy.tokens import DocBin

    td = [    # Train (dummy) dataset, in 'spacy V2 presentation'
              ("Walmart is a leading e-commerce company", {"entities": [(0, 7, "ORG")]}),
              ("I reached Chennai yesterday.", {"entities": [(19, 28, "GPE")]}),
              ("I recently ordered a book from Amazon", {"entities": [(24,32, "ORG")]}),
              ("I was driving a BMW", {"entities": [(16,19, "PRODUCT")]}),
              ("I ordered this from ShopClues", {"entities": [(20,29, "ORG")]}),
              ("Fridge can be ordered in Amazon ", {"entities": [(0,6, "PRODUCT")]}),
              ("I bought a new Washer", {"entities": [(16,22, "PRODUCT")]}),
              ("I bought a old table", {"entities": [(16,21, "PRODUCT")]}),
              ("I bought a fancy dress", {"entities": [(18,23, "PRODUCT")]}),
              ("I rented a camera", {"entities": [(12,18, "PRODUCT")]}),
              ("I rented a tent for our trip", {"entities": [(12,16, "PRODUCT")]}),
              ("I rented a screwdriver from our neighbour", {"entities": [(12,22, "PRODUCT")]}),
              ("I repaired my computer", {"entities": [(15,23, "PRODUCT")]}),
              ("I got my clock fixed", {"entities": [(16,21, "PRODUCT")]}),
              ("I got my truck fixed", {"entities": [(16,21, "PRODUCT")]}),
    ]
    
    dd = [    # Development (dummy) dataset (CV), in 'spacy V2 presentation'
              ("Flipkart started it's journey from zero", {"entities": [(0,8, "ORG")]}),
              ("I recently ordered from Max", {"entities": [(24,27, "ORG")]}),
              ("Flipkart is recognized as leader in market",{"entities": [(0,8, "ORG")]}),
              ("I recently ordered from Swiggy", {"entities": [(24,29, "ORG")]})
    ]

    
    # Converting Train & Development datasets, from 'spaCy V2' to 'spaCy V3'
    nlp = spacy.blank("en")
    db_train = DocBin()
    db_dev = DocBin()

    for text, annotations in td:
        example = Example.from_dict(nlp.make_doc(text), annotations)
        db_train.add(example.reference)
        
    for text, annotations in dd:
        example = Example.from_dict(nlp.make_doc(text), annotations)
        db_dev.add(example.reference)
    
    db_train.to_disk(train_path + ".spacy")  # <== Obtaining and storing "train.spacy"
    db_dev.to_disk(dev_path + ".spacy")      # <== Obtaining and storing "dev.spacy"
    

# ----------------------- ORIGINALLY POSTED CODE -----------------------

@component(
    packages_to_install=[
        "setuptools",
        "wheel", 
        "spacy[cuda113,transformers,lookups]",
    ],
    base_image="gcr.io/deeplearning-platform-release/base-cu113",
    output_component_file="train.yaml"
)
def train(train_path: InputPath(), dev_path: InputPath(), output_path: OutputPath()):
    """
    Trains a spacy model
    
    Parameters:
    ----------
    train_path : Relative location in GCS, for the "train.spacy" file.
    dev_path: Relative location in GCS, for the "dev.spacy" file.
    
    Returns:
    -------
    output : Destination path of the saved model.
    """
    import spacy
    import subprocess
    
    spacy.require_gpu()  # <=== IMAGE NOW MANAGES TO GET BUILT!

    # Presets for training
    subprocess.run(["python", "-m", "spacy", "init", "fill-config", "gcs/secret_path_to_config/base_config.cfg", "config.cfg"])

    # Training model
    subprocess.run(["python", "-m", "spacy", "train", "config.cfg",
                    "--output", output_path,
                    "--paths.train", "{}.spacy".format(train_path),
                    "--paths.dev", "{}.spacy".format(dev_path),
                    "--gpu-id", "0"])

# ----------------------------------------------------------------------
    

# Pipeline definition

@pipeline(
    pipeline_root=PIPELINE_ROOT,
    name="spacy-dummy-pipeline",
)
def spacy_pipeline():
    """
    Builds a custom pipeline
    """
    # Generating dummy "train.spacy" + "dev.spacy"
    train_dev_sets = generate_spacy_file()
    # With the output of the previous component, train a spaCy modeL    
    model = train(
        train_dev_sets.outputs["train_path"],
        train_dev_sets.outputs["dev_path"]
    
    # ------ !!! THIS SECTION DOES THE TRICK !!! ------
    ).add_node_selector_constraint(
        label_name="cloud.google.com/gke-accelerator",
        value="NVIDIA_TESLA_T4"
    ).set_gpu_limit(1).set_memory_limit('32G')
    # -------------------------------------------------

# Pipeline compilation   

compiler.Compiler().compile(
    pipeline_func=spacy_pipeline, package_path="pipeline_spacy_job.json"
)


# Pipeline run

TIMESTAMP = datetime.now().strftime("%Y%m%d%H%M%S")

run = aiplatform.PipelineJob(  # Include your own naming here
    display_name="spacy-dummy-pipeline",
    template_path="pipeline_spacy_job.json",
    job_id="ml-pipeline-spacydummy-small-{0}".format(TIMESTAMP),
    parameter_values={},
    enable_caching=True,
)


# Pipeline gets submitted

run.submit()

Now, the explanation; according to Google:

By default, the component will run on as a Vertex AI CustomJob using an e2-standard-4 machine, with 4 core CPUs and 16GB memory.

Therefore, when the train component gets compiled, it fails as "it was not seeing any GPU available as resource"; in the same link however, all the available settings for both CPU and GPU are mentioned. In my case as you can see, I set train component to run under ONE (1) NVIDIA_TESLA_T4 GPU card, and I also increased my CPU memory, to 32GB. With these modifications, the resulting pipeline looks as follows:

And as you can see, it gets compiled successfully, as well as trains (and eventually obtains) a functional spaCy model. From here, you can tweak this code, to fit your own needs.

I hope this helps to anyone who might be interested.

Thank you.