From 17ef159fffe52e1db0d960547bb7650c66c38bf0 Mon Sep 17 00:00:00 2001 From: "Sanjay Kumar.C" <55217181+jay90099@users.noreply.github.com> Date: Tue, 9 Nov 2021 23:10:07 -0800 Subject: [PATCH] Uses v1 constants for Vertex pipeline client. (#4448) PiperOrigin-RevId: 408775112 (cherry picked from commit 5d53523c04237f7ac273cf6fcc0fe6df0a9c3616) Co-authored-by: jiyongjung --- .../kubeflow/v2/vertex_client_utils.py | 13 +++++++------ 1 file changed, 7 insertions(+), 6 deletions(-) diff --git a/tfx/orchestration/kubeflow/v2/vertex_client_utils.py b/tfx/orchestration/kubeflow/v2/vertex_client_utils.py index e59c35ad39..e39d1285db 100644 --- a/tfx/orchestration/kubeflow/v2/vertex_client_utils.py +++ b/tfx/orchestration/kubeflow/v2/vertex_client_utils.py @@ -18,13 +18,14 @@ from absl import logging from google.cloud.aiplatform import pipeline_jobs -from google.cloud.aiplatform.compat.types import pipeline_state_v1beta1 +from google.cloud.aiplatform_v1.types import pipeline_state + _PIPELINE_COMPLETE_STATES = frozenset([ - pipeline_state_v1beta1.PipelineState.PIPELINE_STATE_SUCCEEDED, - pipeline_state_v1beta1.PipelineState.PIPELINE_STATE_FAILED, - pipeline_state_v1beta1.PipelineState.PIPELINE_STATE_CANCELLED, - pipeline_state_v1beta1.PipelineState.PIPELINE_STATE_PAUSED, + pipeline_state.PipelineState.PIPELINE_STATE_SUCCEEDED, + pipeline_state.PipelineState.PIPELINE_STATE_FAILED, + pipeline_state.PipelineState.PIPELINE_STATE_CANCELLED, + pipeline_state.PipelineState.PIPELINE_STATE_PAUSED, ]) @@ -49,7 +50,7 @@ def poll_job_status(job_id: str, timeout: datetime.timedelta, job = pipeline_jobs.PipelineJob.get(resource_name=job_id) if (job.state == - pipeline_state_v1beta1.PipelineState.PIPELINE_STATE_SUCCEEDED): + pipeline_state.PipelineState.PIPELINE_STATE_SUCCEEDED): logging.info('Job succeeded: %s', job) return elif job.state in _PIPELINE_COMPLETE_STATES: