diff --git a/dags/iris_batch_inference_prod_dag.py b/dags/iris_batch_inference_prod_dag.py index f95f3e0..f7111a0 100644 --- a/dags/iris_batch_inference_prod_dag.py +++ b/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( diff --git a/dags/iris_batch_inference_staging_dag.py b/dags/iris_batch_inference_staging_dag.py index 5209e16..0d0bb74 100644 --- a/dags/iris_batch_inference_staging_dag.py +++ b/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( diff --git a/dags/iris_training_prod_dag.py b/dags/iris_training_prod_dag.py index 3d467ac..e6e769f 100644 --- a/dags/iris_training_prod_dag.py +++ b/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( diff --git a/dags/iris_training_staging_dag.py b/dags/iris_training_staging_dag.py index e4b363c..1129a5e 100644 --- a/dags/iris_training_staging_dag.py +++ b/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( diff --git a/plan/gpu_vertex_pipelines.md b/plan/gpu_vertex_pipelines.md index 68949a6..10839bd 100644 --- a/plan/gpu_vertex_pipelines.md +++ b/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. diff --git a/src/ml_pipelines_kfp/iris_xgboost/pipelines/iris_pipeline_inference.py b/src/ml_pipelines_kfp/iris_xgboost/pipelines/iris_pipeline_inference.py index 35054db..a3e87d5 100644 --- a/src/ml_pipelines_kfp/iris_xgboost/pipelines/iris_pipeline_inference.py +++ b/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", ) diff --git a/src/ml_pipelines_kfp/iris_xgboost/pipelines/iris_pipeline_training.py b/src/ml_pipelines_kfp/iris_xgboost/pipelines/iris_pipeline_training.py index b70d2ce..6e04041 100644 --- a/src/ml_pipelines_kfp/iris_xgboost/pipelines/iris_pipeline_training.py +++ b/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, )