Skip to content

Commit

Permalink
feat(components): add strategy to v1 GCPC custom job components/utils
Browse files Browse the repository at this point in the history
Signed-off-by: Ze Mao <[email protected]>
PiperOrigin-RevId: 684613251
  • Loading branch information
Ze Mao authored and Google Cloud Pipeline Components maintainers committed Nov 15, 2024
1 parent 634aadf commit bae5e94
Show file tree
Hide file tree
Showing 3 changed files with 7 additions and 0 deletions.
1 change: 1 addition & 0 deletions components/google-cloud/RELEASE.md
Original file line number Diff line number Diff line change
Expand Up @@ -6,6 +6,7 @@
* Update Docker image.
* Bump supported KFP versions to `kfp>=2.6.0,<=2.10.0`.
* Support Python versions 3.12 and 3.13.
* add strategy to v1 GCPC custom job components/utils

## Release 2.17.0
* Fix Gemini batch prediction support to `v1.model_evaluation.autosxs_pipeline` after output schema change.
Expand Down
Original file line number Diff line number Diff line change
Expand Up @@ -38,6 +38,7 @@ def custom_training_job(
encryption_spec_key_name: str = '',
persistent_resource_id: str = _placeholders.PERSISTENT_RESOURCE_ID_PLACEHOLDER,
project: str = _placeholders.PROJECT_ID_PLACEHOLDER,
strategy: str = 'STANDARD',
):
# fmt: off
"""Launch a Vertex AI [custom training job](https://cloud.google.com/vertex-ai/docs/training/create-custom-job) using the [CustomJob](https://cloud.google.com/vertex-ai/docs/reference/rest/v1/projects.locations.customJobs) API. See [Create custom training jobs ](https://cloud.google.com/vertex-ai/docs/training/create-custom-job) for more information.
Expand All @@ -58,6 +59,7 @@ def custom_training_job(
encryption_spec_key_name: Customer-managed encryption key options for the CustomJob. If this is set, then all resources created by the CustomJob will be encrypted with the provided encryption key.
persistent_resource_id: The ID of the PersistentResource in the same Project and Location which to run. The default value is a placeholder that will be resolved to the PipelineJob [RuntimeConfig](https://cloud.google.com/vertex-ai/docs/reference/rest/v1/projects.locations.pipelineJobs#PipelineJob.RuntimeConfig)'s persistent resource id at runtime. However, if the PipelineJob doesn't set Persistent Resource as the job level runtime, the placedholder will be resolved to an empty string and the custom job will be run on demand. If the value is set explicitly, the custom job will runs in the specified persistent resource, in this case, please note the network and CMEK configs on the job should be consistent with those on the PersistentResource, otherwise, the job will be rejected.
project: Project to create the custom training job in. Defaults to the project in which the PipelineJob is run.
strategy: The strategy to use for the custom training job. The default is 'STANDARD'. See [more information](https://cloud.google.com/vertex-ai/docs/reference/rest/v1/CustomJobSpec#Strategy).
Returns:
gcp_resources: Serialized JSON of `gcp_resources` [proto](https://github.com/kubeflow/pipelines/tree/master/components/google-cloud/google_cloud_pipeline_components/proto) which tracks the CustomJob.
Expand All @@ -75,6 +77,7 @@ def custom_training_job(
'restart_job_on_worker_restart': (
restart_job_on_worker_restart
),
'strategy': strategy,
},
'service_account': service_account,
'tensorboard': tensorboard,
Expand Down
Original file line number Diff line number Diff line change
Expand Up @@ -71,6 +71,7 @@ def create_custom_training_job_from_component(
labels: Optional[Dict[str, str]] = None,
persistent_resource_id: str = _placeholders.PERSISTENT_RESOURCE_ID_PLACEHOLDER,
env: Optional[List[Dict[str, str]]] = None,
strategy: str = 'STANDARD',
reservation_affinity_type: Optional[str] = None,
reservation_affinity_key: Optional[str] = None,
reservation_affinity_values: Optional[List[str]] = None,
Expand Down Expand Up @@ -102,6 +103,7 @@ def create_custom_training_job_from_component(
labels: The labels with user-defined metadata to organize the CustomJob. See [more information](https://goo.gl/xmQnxf).
persistent_resource_id: The ID of the PersistentResource in the same Project and Location which to run. The default value is a placeholder that will be resolved to the PipelineJob [RuntimeConfig](https://cloud.google.com/vertex-ai/docs/reference/rest/v1/projects.locations.pipelineJobs#PipelineJob.RuntimeConfig)'s persistent resource id at runtime. However, if the PipelineJob doesn't set Persistent Resource as the job level runtime, the placedholder will be resolved to an empty string and the custom job will be run on demand. If the value is set explicitly, the custom job will runs in the specified persistent resource, in this case, please note the network and CMEK configs on the job should be consistent with those on the PersistentResource, otherwise, the job will be rejected.
env: Environment variables to be passed to the container. Takes the form `[{'name': '...', 'value': '...'}]`. Maximum limit is 100.
startegy: The strategy to use for the custom training job. The default is 'STANDARD'. See [more information](https://cloud.google.com/vertex-ai/docs/reference/rest/v1/CustomJobSpec#Strategy).
reservation_affinity_type: The type of [reservation affinity](https://cloud.google.com/vertex-ai/docs/reference/rest/v1/MachineSpec#reservationaffinity). Valid values are "NO_RESERVATION", "ANY_RESERVATION", "SPECIFIC_RESERVATION".
reservation_affinity_key: Corresponds to the label key of a reservation resource. To target a SPECIFIC_RESERVATION by name, use compute.googleapis.com/reservation-name as the key and specify the name of your reservation as its value.
reservation_affinity_values: Corresponds to the label values of a reservation resource. This must be the full resource name of the reservation.
Expand Down Expand Up @@ -214,6 +216,7 @@ def create_custom_training_job_from_component(
'worker_pool_specs': worker_pool_specs,
'timeout': timeout,
'restart_job_on_worker_restart': restart_job_on_worker_restart,
'strategy': strategy,
'service_account': service_account,
'tensorboard': tensorboard,
'enable_web_access': enable_web_access,
Expand Down

0 comments on commit bae5e94

Please sign in to comment.