diff --git a/sdk/python/kfp/compiler/pipeline_spec_builder.py b/sdk/python/kfp/compiler/pipeline_spec_builder.py index 0395146c80b..afc014530fa 100644 --- a/sdk/python/kfp/compiler/pipeline_spec_builder.py +++ b/sdk/python/kfp/compiler/pipeline_spec_builder.py @@ -667,6 +667,10 @@ def convert_to_placeholder(input_value: str) -> str: task.container_spec.resources.accelerator_type), resource_count=convert_to_placeholder( task.container_spec.resources.accelerator_count), + type=convert_to_placeholder( + task.container_spec.resources.accelerator_type), + count=convert_to_placeholder( + int(task.container_spec.resources.accelerator_count)), )) return container_spec diff --git a/sdk/python/test_data/pipelines/pipeline_with_resource_spec.yaml b/sdk/python/test_data/pipelines/pipeline_with_resource_spec.yaml index 07a4ca82975..380151586e9 100644 --- a/sdk/python/test_data/pipelines/pipeline_with_resource_spec.yaml +++ b/sdk/python/test_data/pipelines/pipeline_with_resource_spec.yaml @@ -61,6 +61,8 @@ deploymentSpec: image: gcr.io/my-project/my-fancy-trainer resources: accelerator: + count: '1' + type: tpu-v3 resourceCount: '1' resourceType: tpu-v3 resourceCpuLimit: '4'