Repository: kubeflow/pipelines

Issue:
[sdk] Access token used in wait_for_run_completion expires before it is refreshed
### Environment

*  KFP version:
<!-- For more information, see an overview of KFP installation options: https://www.kubeflow.org/docs/pipelines/installation/overview/. -->
1.8.9
*  KFP SDK version: 
<!-- Specify the version of Kubeflow Pipelines that you are using. The version number appears in the left side navigation of user interface.
To find the version number, See version number shows on bottom of KFP UI left sidenav. -->
1.7.0


### Steps to reproduce
There is no concrete way to reproduce this bug. In our case, we 'd like to check a long running pipeline's status with `wait_for_run_completion`.  This method tries to refresh the access token at every `55` minutes, but my access token sometimes has a shorter expiration time, e.g. 40 or 45 minutes. Then the status checking step fails with `401 Unauthorized`

The pipelines are run in AI Platform in GCP, the access tokens are generated from a service account. We have many different pipelines running under this same service account.

<!--
Specify how to reproduce the problem. 
This may include information such as: a description of the process, code snippets, log output, or screenshots.
-->

### Expected result
`wait_for_run_completion` should work no matter how much the access token's expiration time is.
<!-- What should the correct behavior be? -->


<!-- Don't delete message below to encourage users to support your issue! -->
Impacted by this bug? Give it a 👍. We prioritise the issues with the most 👍.

I had a fix on my local to refresh the access token whenever we receive a `401 Unauthorized` error, I can open a PR with that if it makes sense.



Discussion:

Resolving patch:
diff --git a/sdk/python/kfp/_client.py b/sdk/python/kfp/_client.py
--- a/sdk/python/kfp/_client.py
+++ b/sdk/python/kfp/_client.py
@@ -33,13 +33,8 @@
 from kfp.compiler._k8s_helper import sanitize_k8s_name
 
 from kfp._auth import get_auth_token, get_gcp_access_token
+from kfp_server_api import ApiException
 
-# TTL of the access token associated with the client. This is needed because
-# `gcloud auth print-access-token` generates a token with TTL=1 hour, after
-# which the authentication expires. This TTL is needed for kfp.Client()
-# initialized with host=<inverse proxy endpoint>.
-# Set to 55 mins to provide some safe margin.
-_GCP_ACCESS_TOKEN_TIMEOUT = datetime.timedelta(minutes=55)
 # Operators on scalar values. Only applies to one of |int_value|,
 # |long_value|, |string_value| or |timestamp_value|.
 _FILTER_OPERATIONS = {
@@ -1222,18 +1217,23 @@ def wait_for_run_completion(self, run_id, timeout):
         """
         status = 'Running:'
         start_time = datetime.datetime.now()
-        last_token_refresh_time = datetime.datetime.now()
         if isinstance(timeout, datetime.timedelta):
             timeout = timeout.total_seconds()
+        is_valid_token = False
         while (status is None or status.lower()
                not in ['succeeded', 'failed', 'skipped', 'error']):
-            # Refreshes the access token before it hits the TTL.
-            if (datetime.datetime.now() - last_token_refresh_time >
-                    _GCP_ACCESS_TOKEN_TIMEOUT):
-                self._refresh_api_client_token()
-                last_token_refresh_time = datetime.datetime.now()
-
-            get_run_response = self._run_api.get_run(run_id=run_id)
+            try:
+                get_run_response = self._run_api.get_run(run_id=run_id)
+                is_valid_token = True
+            except ApiException as api_ex:
+                # if the token is valid but receiving 401 Unauthorized error
+                # then refresh the token
+                if is_valid_token and api_ex.status == 401:
+                    logging.info('Access token has expired !!! Refreshing ...')
+                    self._refresh_api_client_token()
+                    continue
+                else:
+                    raise api_ex
             status = get_run_response.run.status
             elapsed_time = (datetime.datetime.now() -
                             start_time).total_seconds()


Test patch: