Add optional GPU accelerator config to Vertex pipelines

dags/iris_batch_inference_prod_dag.py

@@ -33,6 +33,11 @@
         "bq_feature_table": Param("iris_features", type="string"),
         "bq_table_predictions": Param("iris_predictions", type="string",
                                       description="BQ table for predictions"),
+        "accelerator_type": Param("", type="string",
+                                  description="GPU type to attach to the inference step, "
+                                              "e.g. NVIDIA_TESLA_T4 (blank for CPU-only)"),
+        "accelerator_count": Param("0", type="string",
+                                   description="Number of GPUs (0, 1, 2, 4, 8, 16)"),
     },
 ) as dag:
 
@@ -48,6 +53,8 @@
             "--bq-dataset", "{{ params.bq_dataset }}",
             "--bq-feature-table", "{{ params.bq_feature_table }}",
             "--bq-table-predictions", "{{ params.bq_table_predictions }}",
+            "--accelerator-type", "{{ params.accelerator_type }}",
+            "--accelerator-count", "{{ params.accelerator_count }}",
         ],
         startup_timeout_seconds=300,
         container_resources=k8s.V1ResourceRequirements(

dags/iris_batch_inference_staging_dag.py

@@ -33,6 +33,11 @@
         "bq_feature_table": Param("iris_features", type="string"),
         "bq_table_predictions": Param("iris_predictions_staging", type="string",
                                       description="BQ table for predictions"),
+        "accelerator_type": Param("", type="string",
+                                  description="GPU type to attach to the inference step, "
+                                              "e.g. NVIDIA_TESLA_T4 (blank for CPU-only)"),
+        "accelerator_count": Param("0", type="string",
+                                   description="Number of GPUs (0, 1, 2, 4, 8, 16)"),
     },
 ) as dag:
 
@@ -48,6 +53,8 @@
             "--bq-dataset", "{{ params.bq_dataset }}",
             "--bq-feature-table", "{{ params.bq_feature_table }}",
             "--bq-table-predictions", "{{ params.bq_table_predictions }}",
+            "--accelerator-type", "{{ params.accelerator_type }}",
+            "--accelerator-count", "{{ params.accelerator_count }}",
         ],
         startup_timeout_seconds=300,
         container_resources=k8s.V1ResourceRequirements(

dags/iris_training_prod_dag.py

@@ -35,6 +35,11 @@
         "bq_table": Param("iris", type="string"),
         "bq_feature_table": Param("iris_features", type="string"),
         "service_account": Param(SERVICE_ACCOUNT, type="string"),
+        "accelerator_type": Param("", type="string",
+                                  description="GPU type to attach to training steps, "
+                                              "e.g. NVIDIA_TESLA_T4 (blank for CPU-only)"),
+        "accelerator_count": Param("0", type="string",
+                                   description="Number of GPUs (0, 1, 2, 4, 8, 16)"),
     },
 ) as dag:
 
@@ -56,6 +61,8 @@
             "--bq-table", "{{ params.bq_table }}",
             "--bq-feature-table", "{{ params.bq_feature_table }}",
             "--service-account", "{{ params.service_account }}",
+            "--accelerator-type", "{{ params.accelerator_type }}",
+            "--accelerator-count", "{{ params.accelerator_count }}",
         ],
         startup_timeout_seconds=300,
         container_resources=k8s.V1ResourceRequirements(

dags/iris_training_staging_dag.py

@@ -35,6 +35,11 @@
         "bq_table": Param("iris", type="string"),
         "bq_feature_table": Param("iris_features", type="string"),
         "service_account": Param(SERVICE_ACCOUNT, type="string"),
+        "accelerator_type": Param("", type="string",
+                                  description="GPU type to attach to training steps, "
+                                              "e.g. NVIDIA_TESLA_T4 (blank for CPU-only)"),
+        "accelerator_count": Param("0", type="string",
+                                   description="Number of GPUs (0, 1, 2, 4, 8, 16)"),
     },
 ) as dag:
 
@@ -56,6 +61,8 @@
             "--bq-table", "{{ params.bq_table }}",
             "--bq-feature-table", "{{ params.bq_feature_table }}",
             "--service-account", "{{ params.service_account }}",
+            "--accelerator-type", "{{ params.accelerator_type }}",
+            "--accelerator-count", "{{ params.accelerator_count }}",
         ],
         startup_timeout_seconds=300,
         container_resources=k8s.V1ResourceRequirements(

plan/gpu_vertex_pipelines.md

@@ -1,5 +1,21 @@
 # Plan: GPU Support for Vertex AI Pipelines (Training & Inference)
 
+> **Implementation notes (2026-08-04, Option A shipped).** Two deviations from the
+> original plan were required for KFP 2.13 / GCPC 2.20:
+>
+> 1. **API.** `PipelineTask` in KFP 2.13 has no `set_machine_type()` or
+>    `set_accelerator_count()`. The real methods are `set_accelerator_type()` +
+>    `set_accelerator_limit(n)` (n ∈ {0,1,2,4,8,16}). There is **no** way to pin an
+>    exact machine-type string on a lightweight component — Vertex auto-selects a
+>    compatible `n1-*` machine for the accelerator. So `--machine-type` / the
+>    `machine_type` DAG param were **dropped**; only `--accelerator-type` /
+>    `--accelerator-count` remain. (Exact machine-type control, incl. `g2`/`NVIDIA_L4`,
+>    is the Option B upgrade path via `create_custom_training_job_from_component`.)
+> 2. **Compile timing.** `@kfp.dsl.pipeline` builds the graph **eagerly at decoration
+>    (import) time**, so reassigning module constants in `__main__` before `.compile()`
+>    has no effect. The decorator is therefore applied **programmatically inside
+>    `__main__`** (after CLI args are parsed) so the GPU config is baked into the graph.
+
 ## Context
 
 The current KFP pipeline components (`models.py`, `inference.py`, `evaluation.py`) run on Vertex AI's default CPU-only `e2-standard-4` machine type. This plan adds the ability to run specific pipeline steps on GPU-accelerated machines, configured at the KFP component level so Vertex AI provisions the right hardware automatically.

src/ml_pipelines_kfp/iris_xgboost/pipelines/iris_pipeline_inference.py

@@ -16,7 +16,20 @@
     BQ_FEATURE_TABLE,
     BQ_TABLE_PREDICTIONS,
 )
-@kfp.dsl.pipeline(name=f"{PIPELINE_NAME}-inference", pipeline_root=PIPELINE_ROOT)
+
+
+def coalesce(*args):
+    return next((a for a in args if a is not None), None)
+
+
+# GPU resource config for the inference step. Module-level so the pipeline body reads
+# them when the graph is built. KFP's @dsl.pipeline decorator builds the graph eagerly
+# at decoration time, so these are reassigned from CLI args in __main__ *before* the
+# decorator is applied there. ACCELERATOR_TYPE == "" means CPU-only.
+ACCELERATOR_TYPE = ""
+ACCELERATOR_COUNT = 0
+
+
 def pipeline(
     project_id: str,
     location: str,
@@ -49,6 +62,12 @@ def pipeline(
         .after(get_model_op)
     )
 
+    # Attach a GPU to the inference step when configured. Vertex AI auto-selects a
+    # compatible (n1-*) machine for the accelerator.
+    if ACCELERATOR_TYPE:
+        inference_op.set_accelerator_type(ACCELERATOR_TYPE)
+        inference_op.set_accelerator_limit(ACCELERATOR_COUNT)
+
 
 if __name__ == "__main__":
     parser = argparse.ArgumentParser(
@@ -59,16 +78,35 @@ def pipeline(
     parser.add_argument("--bq-dataset", default=BQ_DATASET)
     parser.add_argument("--bq-feature-table", default=BQ_FEATURE_TABLE)
     parser.add_argument("--bq-table-predictions", default=BQ_TABLE_PREDICTIONS)
+    parser.add_argument("--accelerator-type", default="",
+                        help="GPU type to attach to the inference step "
+                             "(e.g., NVIDIA_TESLA_T4, NVIDIA_L4). "
+                             "Blank = CPU-only. Vertex auto-selects a compatible machine.")
+    parser.add_argument("--accelerator-count", default="0",
+                        help="Number of GPUs per step (one of 0, 1, 2, 4, 8, 16)")
     cli = parser.parse_args()
 
+    # Resolve GPU config into the module-level globals read by pipeline() at build time.
+    ACCELERATOR_TYPE = coalesce(cli.accelerator_type, "")
+    ACCELERATOR_COUNT = int(coalesce(cli.accelerator_count, "0"))
+    if ACCELERATOR_TYPE and ACCELERATOR_COUNT not in (1, 2, 4, 8, 16):
+        parser.error("--accelerator-count must be one of 1, 2, 4, 8, 16 when a GPU is set")
+
     credentials, _ = google.auth.default(
         scopes=["https://www.googleapis.com/auth/cloud-platform"],
     )
 
     aip.init(project=cli.project_id, credentials=credentials)
 
+    # Apply the KFP pipeline decorator here (not at module top) so the graph is built
+    # AFTER the GPU config above is resolved — KFP builds the graph eagerly at
+    # decoration time, so decorating at import would freeze it with CPU-only config.
+    pipeline_func = kfp.dsl.pipeline(
+        name=f"{PIPELINE_NAME}-inference", pipeline_root=PIPELINE_ROOT
+    )(pipeline)
+
     kfp.compiler.Compiler().compile(
-        pipeline_func=pipeline,
+        pipeline_func=pipeline_func,
         package_path="pipeline.yaml",
         pipeline_name=f"{PIPELINE_NAME}-inference",
     )

src/ml_pipelines_kfp/iris_xgboost/pipelines/iris_pipeline_training.py

@@ -25,7 +25,15 @@ def coalesce(*args):
     return next((a for a in args if a is not None), None)
 
 
-@kfp.dsl.pipeline(name=f"{PIPELINE_NAME}-training", pipeline_root=PIPELINE_ROOT)
+# GPU resource config for the model-training steps. Module-level so the pipeline
+# body reads them when the graph is built. NOTE: KFP's @dsl.pipeline decorator builds
+# the graph eagerly at decoration time, so these are reassigned from CLI args in the
+# __main__ block *before* the decorator is applied (see the programmatic decoration
+# there). ACCELERATOR_TYPE == "" means CPU-only (no GPU attached).
+ACCELERATOR_TYPE = ""
+ACCELERATOR_COUNT = 0
+
+
 def pipeline(project_id: str, location: str, bq_dataset: str, bq_feature_table: str):
 
     # Import components
@@ -64,6 +72,14 @@ def pipeline(project_id: str, location: str, bq_dataset: str, bq_feature_table:
         train_dataset=data_op.outputs["train_dataset"]
     ).set_display_name("Random Forest")
 
+    # Attach a GPU to the training steps when configured. Plain-Python conditional
+    # (not a KFP dsl.If) — ACCELERATOR_TYPE is a concrete value at graph-build time.
+    # Vertex AI auto-selects a compatible (n1-*) machine for the accelerator.
+    if ACCELERATOR_TYPE:
+        for op in (dt_op, rf_op):
+            op.set_accelerator_type(ACCELERATOR_TYPE)
+            op.set_accelerator_limit(ACCELERATOR_COUNT)
+
     choose_model_op = (
         choose_best_model(
             test_dataset=data_op.outputs["test_dataset"],
@@ -111,6 +127,12 @@ def pipeline(project_id: str, location: str, bq_dataset: str, bq_feature_table:
     parser.add_argument("--bq-table")
     parser.add_argument("--bq-feature-table")
     parser.add_argument("--service-account")
+    parser.add_argument("--accelerator-type", default="",
+                        help="GPU type to attach to training steps "
+                             "(e.g., NVIDIA_TESLA_T4, NVIDIA_TESLA_V100). "
+                             "Blank = CPU-only. Vertex auto-selects a compatible machine.")
+    parser.add_argument("--accelerator-count", default="0",
+                        help="Number of GPUs per step (one of 0, 1, 2, 4, 8, 16)")
     cli = parser.parse_args()
 
     # Resolve each param: CLI > env var > constants.py
@@ -126,14 +148,29 @@ def pipeline(project_id: str, location: str, bq_dataset: str, bq_feature_table:
     project_id = coalesce(cli.project_id, os.getenv("PROJECT_ID"), PROJECT_ID)
     region = coalesce(cli.region, os.getenv("REGION"), REGION)
     sa_email = coalesce(cli.service_account, os.getenv("SERVICE_ACCOUNT"), SERVICE_ACCOUNT)
+
+    # Resolve GPU config into the module-level globals read by pipeline() at build time.
+    ACCELERATOR_TYPE = coalesce(cli.accelerator_type, "")
+    ACCELERATOR_COUNT = int(coalesce(cli.accelerator_count, "0"))
+    if ACCELERATOR_TYPE and ACCELERATOR_COUNT not in (1, 2, 4, 8, 16):
+        parser.error("--accelerator-count must be one of 1, 2, 4, 8, 16 when a GPU is set")
+
     credentials, _ = google.auth.default(
         scopes=["https://www.googleapis.com/auth/cloud-platform"],
     )
 
     aip.init(project=project_id, location=region, credentials=credentials)
 
+    # Apply the KFP pipeline decorator here (not at module top) so the graph is built
+    # AFTER the CLI args above have been resolved into module globals. KFP builds the
+    # graph eagerly at decoration time, so decorating at import would freeze it with the
+    # default (CPU) config and ignore --accelerator-type/--model-name/--image-name.
+    pipeline_func = kfp.dsl.pipeline(
+        name=f"{PIPELINE_NAME}-training", pipeline_root=pipeline_root
+    )(pipeline)
+
     kfp.compiler.Compiler().compile(
-        pipeline_func=pipeline,
+        pipeline_func=pipeline_func,
         package_path="pipeline.yaml",
         pipeline_name=pipeline_name,
     )