4
4
import json
5
5
6
6
import sqlite3
7
+ from typing import Optional , Any , Dict
7
8
8
9
import dask
9
10
from dask .utils import tmpfile
@@ -106,7 +107,6 @@ def __init__(
106
107
self .instance_labels = _instance_labels
107
108
108
109
self .general_zone = "-" .join (self .zone .split ("-" )[:2 ]) # us-east1-c -> us-east1
109
-
110
110
self .service_account = service_account or self .config .get ("service_account" )
111
111
112
112
def create_gcp_config (self ):
@@ -494,6 +494,8 @@ class GCPCluster(VMCluster):
494
494
service_account: str
495
495
Service account that all VMs will run under.
496
496
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
497
499
498
500
Examples
499
501
--------
@@ -587,9 +589,10 @@ def __init__(
587
589
debug = False ,
588
590
instance_labels = None ,
589
591
service_account = None ,
592
+ service_account_credentials : Optional [Dict [str , Any ]] = None ,
590
593
** kwargs ,
591
594
):
592
- self .compute = GCPCompute ()
595
+ self .compute = GCPCompute (service_account_credentials )
593
596
594
597
self .config = dask .config .get ("cloudprovider.gcp" , {})
595
598
self .auto_shutdown = (
@@ -641,9 +644,17 @@ def __init__(
641
644
642
645
643
646
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
+ """
645
655
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 {}
647
658
self ._compute = self .refresh_client ()
648
659
649
660
def refresh_client (self ):
@@ -654,6 +665,13 @@ def refresh_client(self):
654
665
os .environ ["GOOGLE_APPLICATION_CREDENTIALS" ],
655
666
scopes = ["https://www.googleapis.com/auth/cloud-platform" ],
656
667
)
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
+ )
657
675
else :
658
676
import google .auth .credentials # google-auth
659
677
0 commit comments