From 4d2a1cf173c0131d4e11f8df927da29881252280 Mon Sep 17 00:00:00 2001 From: Tommy Li Date: Fri, 12 May 2023 14:49:28 -0700 Subject: [PATCH] fix(sdk): Optimize wait_for_run_completion sdk function to skip an extra wait (#9407) --- sdk/python/kfp/client/client.py | 7 ++++--- 1 file changed, 4 insertions(+), 3 deletions(-) diff --git a/sdk/python/kfp/client/client.py b/sdk/python/kfp/client/client.py index 5cfede865c..e9b125fc9b 100644 --- a/sdk/python/kfp/client/client.py +++ b/sdk/python/kfp/client/client.py @@ -1370,8 +1370,8 @@ class Client: if isinstance(timeout, datetime.timedelta): timeout = timeout.total_seconds() is_valid_token = False - while (state is None or state.lower() - not in ['succeeded', 'failed', 'skipped', 'error']): + finish_states = ['succeeded', 'failed', 'skipped', 'error'] + while True: try: get_run_response = self._run_api.get_run(run_id=run_id) is_valid_token = True @@ -1390,8 +1390,9 @@ class Client: logging.info('Waiting for the job to complete...') if elapsed_time > timeout: raise TimeoutError('Run timeout') + if state is not None and state.lower() in finish_states: + return get_run_response time.sleep(sleep_duration) - return get_run_response def upload_pipeline( self,