Skip to content

Commit 4753d3c

Browse files
RubenBBlazquezrubenbblazquez
and
rubenbblazquez
authored
[feat]: Adapt GCPCluster to be able to initialize the cluster by passing credentials without obtaining from the environment (#432)
Co-authored-by: rubenbblazquez <[email protected]>
1 parent 5babf64 commit 4753d3c

File tree

1 file changed

+22
-4
lines changed

1 file changed

+22
-4
lines changed

dask_cloudprovider/gcp/instances.py

Lines changed: 22 additions & 4 deletions
Original file line numberDiff line numberDiff line change
@@ -4,6 +4,7 @@
44
import json
55

66
import sqlite3
7+
from typing import Optional, Any, Dict
78

89
import dask
910
from dask.utils import tmpfile
@@ -106,7 +107,6 @@ def __init__(
106107
self.instance_labels = _instance_labels
107108

108109
self.general_zone = "-".join(self.zone.split("-")[:2]) # us-east1-c -> us-east1
109-
110110
self.service_account = service_account or self.config.get("service_account")
111111

112112
def create_gcp_config(self):
@@ -494,6 +494,8 @@ class GCPCluster(VMCluster):
494494
service_account: str
495495
Service account that all VMs will run under.
496496
Defaults to the default Compute Engine service account for your GCP project.
497+
service_account_credentials: Optional[Dict[str, Any]]
498+
Service account credentials to create the compute engine Vms
497499
498500
Examples
499501
--------
@@ -587,9 +589,10 @@ def __init__(
587589
debug=False,
588590
instance_labels=None,
589591
service_account=None,
592+
service_account_credentials: Optional[Dict[str, Any]] = None,
590593
**kwargs,
591594
):
592-
self.compute = GCPCompute()
595+
self.compute = GCPCompute(service_account_credentials)
593596

594597
self.config = dask.config.get("cloudprovider.gcp", {})
595598
self.auto_shutdown = (
@@ -641,9 +644,17 @@ def __init__(
641644

642645

643646
class GCPCompute:
644-
"""Wrapper for the ``googleapiclient`` compute object."""
647+
"""
648+
Wrapper for the ``googleapiclient`` compute object.
649+
650+
Attributes
651+
----------
652+
service_account_credentials: Optional[dict]
653+
Service account credentials to create the compute engine Vms
654+
"""
645655

646-
def __init__(self):
656+
def __init__(self, service_account_credentials: Optional[dict[str, Any]] = None):
657+
self.service_account_credentials = service_account_credentials or {}
647658
self._compute = self.refresh_client()
648659

649660
def refresh_client(self):
@@ -654,6 +665,13 @@ def refresh_client(self):
654665
os.environ["GOOGLE_APPLICATION_CREDENTIALS"],
655666
scopes=["https://www.googleapis.com/auth/cloud-platform"],
656667
)
668+
elif self.service_account_credentials:
669+
import google.oauth2.service_account # google-auth
670+
671+
creds = google.oauth2.service_account.Credentials.from_service_account_info(
672+
self.service_account_credentials,
673+
scopes=["https://www.googleapis.com/auth/cloud-platform"],
674+
)
657675
else:
658676
import google.auth.credentials # google-auth
659677

0 commit comments

Comments
 (0)