Add optional GPU accelerator config to Vertex pipelines - #65
Merged
Conversation
…lines Implements Option A of plan/gpu_vertex_pipelines.md: attach a GPU to the model-training steps and the inference step via KFP's native resource API. - Add --accelerator-type / --accelerator-count CLI args to both pipelines and wire them through all four Airflow DAGs (staging + prod) as params. Blank accelerator = CPU-only (unchanged default behavior). - Use set_accelerator_type() + set_accelerator_limit() (real KFP 2.13 API); the plan's set_machine_type()/set_accelerator_count() do not exist. Vertex auto-selects a compatible n1-* machine, so machine_type was dropped. - Apply @kfp.dsl.pipeline programmatically in __main__ after args are parsed: KFP builds the graph eagerly at decoration time, so decorating at import would freeze it with CPU-only config and ignore the CLI overrides. - Validate accelerator count is one of 1/2/4/8/16 when a GPU is requested. - Document both deviations in the plan doc. Verified: both pipelines compile CPU-only (no accelerator block) and with a T4 GPU (correct accelerator spec on each targeted step). Co-Authored-By: Claude Opus 4.8 (1M context) <noreply@anthropic.com>
This file contains hidden or bidirectional Unicode text that may be interpreted or compiled differently than what appears below. To review, open the file in an editor that reveals hidden Unicode characters.
Learn more about bidirectional Unicode characters
Sign up for free
to join this conversation on GitHub.
Already have an account?
Sign in to comment
1 participant
Add this suggestion to a batch that can be applied as a single commit.This suggestion is invalid because no changes were made to the code.Suggestions cannot be applied while the pull request is closed.Suggestions cannot be applied while viewing a subset of changes.Only one suggestion per line can be applied in a batch.Add this suggestion to a batch that can be applied as a single commit.Applying suggestions on deleted lines is not supported.You must change the existing code in this line in order to create a valid suggestion.Outdated suggestions cannot be applied.This suggestion has been applied or marked resolved.Suggestions cannot be applied from pending reviews.Suggestions cannot be applied on multi-line comments.Suggestions cannot be applied while the pull request is queued to merge.Suggestion cannot be applied right now. Please check back later.
Summary
Implements Option A of
plan/gpu_vertex_pipelines.md: lets specific Vertex AI pipeline steps run on GPU-accelerated hardware, configurable per-run without code changes. Default behavior is unchanged (CPU-only).--accelerator-type/--accelerator-countCLI args to the training and inference pipelines, wired through all four Airflow DAGs (staging + prod) as params.decision_tree,random_forest) and theinference_modelstep when configured; blank accelerator = CPU-only.Deviations from the plan (KFP 2.13 / GCPC 2.20)
PipelineTaskhas noset_machine_type()/set_accelerator_count(). Used the real methodsset_accelerator_type()+set_accelerator_limit(n)(n ∈ {0,1,2,4,8,16}). There's no way to pin an exact machine-type string on a lightweight component — Vertex auto-selects a compatiblen1-*machine — somachine_typewas dropped. Exact machine-type control (incl.g2/NVIDIA_L4) is the Option B upgrade path viacreate_custom_training_job_from_component.@kfp.dsl.pipelinebuilds the graph eagerly at decoration/import time, so reassigning module constants in__main__before.compile()is ignored. The decorator is now applied programmatically in__main__after CLI args are parsed, so the GPU config is baked into the graph. (This also makes the existing--model-name/--image-nameoverrides authoritative.)Testing
Compiled both pipelines offline:
acceleratorblock in the compiled spec (no behavior change).accelerator: {type: NVIDIA_TESLA_T4, count: '1'}on each targeted step.python -m py_compilepasses on all 6 changed files.Note: current Iris models are scikit-learn (no GPU benefit per plan §7); this is plumbing for when models move to XGBoost
gpu_hist/ deep learning.🤖 Generated with Claude Code