From 998a76b85c5485996f1f9b6572df54d436bb904d Mon Sep 17 00:00:00 2001 From: Dmitry Balabka Date: Wed, 10 Jul 2024 15:44:00 +0300 Subject: [PATCH 1/9] Use gcp sdk provided flow for obtaininng application default credentials (#429) --- dask_cloudprovider/gcp/instances.py | 32 +++++++---------------------- 1 file changed, 7 insertions(+), 25 deletions(-) diff --git a/dask_cloudprovider/gcp/instances.py b/dask_cloudprovider/gcp/instances.py index d322b984..c8458f1e 100644 --- a/dask_cloudprovider/gcp/instances.py +++ b/dask_cloudprovider/gcp/instances.py @@ -19,6 +19,7 @@ try: import googleapiclient.discovery + import google.auth from googleapiclient.errors import HttpError except ImportError as e: msg = ( @@ -658,38 +659,19 @@ def __init__(self, service_account_credentials: Optional[dict[str, Any]] = None) self._compute = self.refresh_client() def refresh_client(self): - if os.environ.get("GOOGLE_APPLICATION_CREDENTIALS", False): + if self.service_account_credentials: import google.oauth2.service_account # google-auth - creds = google.oauth2.service_account.Credentials.from_service_account_file( - os.environ["GOOGLE_APPLICATION_CREDENTIALS"], - scopes=["https://www.googleapis.com/auth/cloud-platform"], - ) - elif self.service_account_credentials: - import google.oauth2.service_account # google-auth - - creds = google.oauth2.service_account.Credentials.from_service_account_info( + credentials = google.oauth2.service_account.Credentials.from_service_account_info( self.service_account_credentials, scopes=["https://www.googleapis.com/auth/cloud-platform"], ) else: - import google.auth.credentials # google-auth + # Obtain Application Default Credentials + credentials, _ = google.auth.default() - path = os.path.join( - os.path.expanduser("~"), ".config/gcloud/credentials.db" - ) - if not os.path.exists(path): - raise GCPCredentialsError() - conn = sqlite3.connect(path) - creds_rows = conn.execute("select * from credentials").fetchall() - with tmpfile() as f: - with open(f, "w") as f_: - # take first row - f_.write(creds_rows[0][1]) - creds, _ = google.auth.load_credentials_from_file(filename=f) - return googleapiclient.discovery.build( - "compute", "v1", credentials=creds, requestBuilder=build_request(creds) - ) + # Use the credentials to build a service client + return googleapiclient.discovery.build('compute', 'v1', credentials=credentials) def instances(self): try: From 3c5a3f4e3378af0e1580efe3569ad3010dc3bbb2 Mon Sep 17 00:00:00 2001 From: Dmitry Balabka Date: Wed, 10 Jul 2024 16:54:32 +0300 Subject: [PATCH 2/9] Update instances.py --- dask_cloudprovider/gcp/instances.py | 2 +- 1 file changed, 1 insertion(+), 1 deletion(-) diff --git a/dask_cloudprovider/gcp/instances.py b/dask_cloudprovider/gcp/instances.py index c8458f1e..5ad1b425 100644 --- a/dask_cloudprovider/gcp/instances.py +++ b/dask_cloudprovider/gcp/instances.py @@ -671,7 +671,7 @@ def refresh_client(self): credentials, _ = google.auth.default() # Use the credentials to build a service client - return googleapiclient.discovery.build('compute', 'v1', credentials=credentials) + return googleapiclient.discovery.build("compute", "v1", credentials=credentials) def instances(self): try: From 46bf4a2f959c97907930bc37ae128c29dcf70385 Mon Sep 17 00:00:00 2001 From: Dmitry Balabka Date: Thu, 19 Sep 2024 18:05:07 +0300 Subject: [PATCH 3/9] Add missing request builder --- dask_cloudprovider/gcp/instances.py | 23 ++++++++++++++++------- 1 file changed, 16 insertions(+), 7 deletions(-) diff --git a/dask_cloudprovider/gcp/instances.py b/dask_cloudprovider/gcp/instances.py index 5ad1b425..ef8e467d 100644 --- a/dask_cloudprovider/gcp/instances.py +++ b/dask_cloudprovider/gcp/instances.py @@ -629,9 +629,11 @@ def __init__( "gpu_instance": self.gpu_instance, "bootstrap": self.bootstrap, "auto_shutdown": self.auto_shutdown, - "preemptible": preemptible - if preemptible is not None - else self.config.get("preemptible"), + "preemptible": ( + preemptible + if preemptible is not None + else self.config.get("preemptible") + ), "instance_labels": instance_labels or self.config.get("instance_labels"), "service_account": service_account or self.config.get("service_account"), } @@ -662,16 +664,23 @@ def refresh_client(self): if self.service_account_credentials: import google.oauth2.service_account # google-auth - credentials = google.oauth2.service_account.Credentials.from_service_account_info( - self.service_account_credentials, - scopes=["https://www.googleapis.com/auth/cloud-platform"], + credentials = ( + google.oauth2.service_account.Credentials.from_service_account_info( + self.service_account_credentials, + scopes=["https://www.googleapis.com/auth/cloud-platform"], + ) ) else: # Obtain Application Default Credentials credentials, _ = google.auth.default() # Use the credentials to build a service client - return googleapiclient.discovery.build("compute", "v1", credentials=credentials) + return googleapiclient.discovery.build( + "compute", + "v1", + credentials=credentials, + requestBuilder=build_request(credentials), + ) def instances(self): try: From 042be9453953c7c9d8be377b3e37f6604232eb02 Mon Sep 17 00:00:00 2001 From: Dmitry Balabka Date: Thu, 19 Sep 2024 18:39:47 +0300 Subject: [PATCH 4/9] Fix tests --- dask_cloudprovider/gcp/instances.py | 2 +- 1 file changed, 1 insertion(+), 1 deletion(-) diff --git a/dask_cloudprovider/gcp/instances.py b/dask_cloudprovider/gcp/instances.py index ef8e467d..a92b2a35 100644 --- a/dask_cloudprovider/gcp/instances.py +++ b/dask_cloudprovider/gcp/instances.py @@ -19,7 +19,7 @@ try: import googleapiclient.discovery - import google.auth + import google.auth.default from googleapiclient.errors import HttpError except ImportError as e: msg = ( From 1d4546bac5a682889618dc593b6e23e1db77df95 Mon Sep 17 00:00:00 2001 From: Dmitry Balabka Date: Thu, 19 Sep 2024 18:45:20 +0300 Subject: [PATCH 5/9] Fix tests --- dask_cloudprovider/gcp/instances.py | 4 ++-- 1 file changed, 2 insertions(+), 2 deletions(-) diff --git a/dask_cloudprovider/gcp/instances.py b/dask_cloudprovider/gcp/instances.py index a92b2a35..843deffb 100644 --- a/dask_cloudprovider/gcp/instances.py +++ b/dask_cloudprovider/gcp/instances.py @@ -19,7 +19,7 @@ try: import googleapiclient.discovery - import google.auth.default + from google.auth import default as google_auth_default from googleapiclient.errors import HttpError except ImportError as e: msg = ( @@ -672,7 +672,7 @@ def refresh_client(self): ) else: # Obtain Application Default Credentials - credentials, _ = google.auth.default() + credentials, _ = google_auth_default() # Use the credentials to build a service client return googleapiclient.discovery.build( From ab317d2c6bfb4e6e6ec965cc352688e4d4f8ce86 Mon Sep 17 00:00:00 2001 From: Dmitry Balabka Date: Thu, 19 Sep 2024 18:48:49 +0300 Subject: [PATCH 6/9] Fix tests --- dask_cloudprovider/gcp/instances.py | 4 ++-- 1 file changed, 2 insertions(+), 2 deletions(-) diff --git a/dask_cloudprovider/gcp/instances.py b/dask_cloudprovider/gcp/instances.py index 843deffb..625b13e4 100644 --- a/dask_cloudprovider/gcp/instances.py +++ b/dask_cloudprovider/gcp/instances.py @@ -19,7 +19,6 @@ try: import googleapiclient.discovery - from google.auth import default as google_auth_default from googleapiclient.errors import HttpError except ImportError as e: msg = ( @@ -671,8 +670,9 @@ def refresh_client(self): ) ) else: + import google.auth # Obtain Application Default Credentials - credentials, _ = google_auth_default() + credentials, _ = google.auth.default() # Use the credentials to build a service client return googleapiclient.discovery.build( From dd5ef54f46da27c009db4dd67fdd0c6099846560 Mon Sep 17 00:00:00 2001 From: Dmitry Balabka Date: Thu, 19 Sep 2024 18:56:08 +0300 Subject: [PATCH 7/9] Fix tests --- dask_cloudprovider/gcp/instances.py | 1 + 1 file changed, 1 insertion(+) diff --git a/dask_cloudprovider/gcp/instances.py b/dask_cloudprovider/gcp/instances.py index 625b13e4..09fee60d 100644 --- a/dask_cloudprovider/gcp/instances.py +++ b/dask_cloudprovider/gcp/instances.py @@ -671,6 +671,7 @@ def refresh_client(self): ) else: import google.auth + # Obtain Application Default Credentials credentials, _ = google.auth.default() From 7609a6e83a6488262a0fe02e0850dc294a729758 Mon Sep 17 00:00:00 2001 From: Dmitry Balabka Date: Thu, 19 Sep 2024 19:02:24 +0300 Subject: [PATCH 8/9] Fix tests --- dask_cloudprovider/gcp/instances.py | 5 ++++- 1 file changed, 4 insertions(+), 1 deletion(-) diff --git a/dask_cloudprovider/gcp/instances.py b/dask_cloudprovider/gcp/instances.py index 09fee60d..0f9ff17c 100644 --- a/dask_cloudprovider/gcp/instances.py +++ b/dask_cloudprovider/gcp/instances.py @@ -673,7 +673,10 @@ def refresh_client(self): import google.auth # Obtain Application Default Credentials - credentials, _ = google.auth.default() + try: + credentials, _ = google.auth.default() + except google.auth.exceptions.DefaultCredentialsError: + credentials = None # Use the credentials to build a service client return googleapiclient.discovery.build( From 8be2193a7be091d75910cd9e89b8443559007162 Mon Sep 17 00:00:00 2001 From: Dmitry Balabka Date: Thu, 19 Sep 2024 19:25:59 +0300 Subject: [PATCH 9/9] Add right error handling --- dask_cloudprovider/gcp/instances.py | 9 +++------ 1 file changed, 3 insertions(+), 6 deletions(-) diff --git a/dask_cloudprovider/gcp/instances.py b/dask_cloudprovider/gcp/instances.py index 0f9ff17c..957e7985 100644 --- a/dask_cloudprovider/gcp/instances.py +++ b/dask_cloudprovider/gcp/instances.py @@ -1,13 +1,10 @@ import asyncio -import os import uuid import json -import sqlite3 from typing import Optional, Any, Dict import dask -from dask.utils import tmpfile from dask_cloudprovider.generic.vmcluster import ( VMCluster, VMInterface, @@ -672,11 +669,11 @@ def refresh_client(self): else: import google.auth - # Obtain Application Default Credentials + # Obtain Application Default Credentials (ADC) try: credentials, _ = google.auth.default() - except google.auth.exceptions.DefaultCredentialsError: - credentials = None + except google.auth.exceptions.DefaultCredentialsError as e: + raise GCPCredentialsError() from e # Use the credentials to build a service client return googleapiclient.discovery.build(