GCP authentication using oauth tokens (#9279)

python/ray/autoscaler/gcp/config.py

@@ -9,6 +9,7 @@ from cryptography.hazmat.primitives.asymmetric import rsa
 from cryptography.hazmat.backends import default_backend
 from googleapiclient import discovery, errors
 from google.oauth2 import service_account
+from google.oauth2.credentials import Credentials as OAuthCredentials
 
 logger = logging.getLogger(__name__)
 
@@ -103,52 +104,66 @@ def generate_rsa_key_pair():
     return public_key, pem
 
 
-def _create_crm(gcp_credentials):
+def _create_crm(gcp_credentials=None):
     return discovery.build(
         "cloudresourcemanager", "v1", credentials=gcp_credentials)
 
 
-def _create_iam(gcp_credentials):
+def _create_iam(gcp_credentials=None):
     return discovery.build("iam", "v1", credentials=gcp_credentials)
 
 
-def _create_compute(gcp_credentials):
+def _create_compute(gcp_credentials=None):
     return discovery.build("compute", "v1", credentials=gcp_credentials)
 
 
-def fetch_gcp_credentials_from_provider_config(provider_config):
+def construct_clients_from_provider_config(provider_config):
     """
     Attempt to fetch and parse the JSON GCP credentials from the provider
     config yaml file.
     """
-    service_account_info_string = provider_config.get("gcp_credentials")
-    if service_account_info_string is None:
-        logger.info("gcp_credentials not found in cluster yaml file. "
-                    "Falling back to GOOGLE_APPLICATION_CREDENTIALS "
-                    "environment variable.")
+    gcp_credentials = provider_config.get("gcp_credentials")
+    if gcp_credentials is None:
+        logger.debug("gcp_credentials not found in cluster yaml file. "
+                     "Falling back to GOOGLE_APPLICATION_CREDENTIALS "
+                     "environment variable.")
         # If gcp_credentials is None, then discovery.build will search for
         # credentials in the local environment.
-        return None
+        return _create_crm(), \
+            _create_iam(), \
+            _create_compute()
 
-    # If parsing the gcp_credentials failed, then the user likely made a
-    # mistake in copying the credentials into the config yaml.
-    try:
-        service_account_info = json.loads(service_account_info_string)
-    except json.decoder.JSONDecodeError:
-        raise RuntimeError("gcp_credentials found in cluster yaml file but "
-                           "formatted improperly.")
-    gcp_credentials = service_account.Credentials.from_service_account_info(
-        service_account_info)
-    return gcp_credentials
+    assert ("type" in gcp_credentials), \
+        "gcp_credentials cluster yaml field missing 'type' field."
+    assert ("credentials" in gcp_credentials), \
+        "gcp_credentials cluster yaml field missing 'credentials' field."
+
+    cred_type = gcp_credentials["type"]
+    credentials_field = gcp_credentials["credentials"]
+
+    if cred_type == "service_account":
+        # If parsing the gcp_credentials failed, then the user likely made a
+        # mistake in copying the credentials into the config yaml.
+        try:
+            service_account_info = json.loads(credentials_field)
+        except json.decoder.JSONDecodeError:
+            raise RuntimeError(
+                "gcp_credentials found in cluster yaml file but "
+                "formatted improperly.")
+        credentials = service_account.Credentials.from_service_account_info(
+            service_account_info)
+    elif cred_type == "credentials_token":
+        # Otherwise the credentials type must be credentials_token.
+        credentials = OAuthCredentials(credentials_field)
+
+    return _create_crm(credentials), \
+        _create_iam(credentials), \
+        _create_compute(credentials)
 
 
 def bootstrap_gcp(config):
-    gcp_credentials = fetch_gcp_credentials_from_provider_config(
-        config["provider"])
-
-    crm = _create_crm(gcp_credentials)
-    iam = _create_iam(gcp_credentials)
-    compute = _create_compute(gcp_credentials)
+    crm, iam, compute = \
+        construct_clients_from_provider_config(config["provider"])
 
     config = _configure_project(config, crm)
     config = _configure_iam_role(config, crm, iam)

python/ray/autoscaler/gcp/node_provider.py

@@ -6,7 +6,7 @@ import logging
 from ray.autoscaler.node_provider import NodeProvider
 from ray.autoscaler.tags import TAG_RAY_CLUSTER_NAME, TAG_RAY_NODE_NAME
 from ray.autoscaler.gcp.config import MAX_POLLS, POLL_INTERVAL, \
-        fetch_gcp_credentials_from_provider_config, _create_compute
+        construct_clients_from_provider_config
 
 logger = logging.getLogger(__name__)
 
@@ -42,11 +42,9 @@ class GCPNodeProvider(NodeProvider):
         NodeProvider.__init__(self, provider_config, cluster_name)
 
         self.lock = RLock()
-        gcp_credentials = fetch_gcp_credentials_from_provider_config(
+        _, _, self.compute = construct_clients_from_provider_config(
             provider_config)
 
-        self.compute = _create_compute(gcp_credentials)
-
         # Cache of node objects from the last nodes() call. This avoids
         # excessive DescribeInstances requests.
         self.cached_nodes = {}

python/ray/autoscaler/ray-schema.json

@@ -141,8 +141,21 @@
                     "description": "GCP globally unique project id"
                 },
                 "gcp_credentials": {
-                    "type": "string",
-                    "description": "JSON string constituting GCP credentials"
+                    "type": "object",
+                    "description": "Credentials for authenticating with the GCP client",
+                    "required": [ "type" ],
+                    "additionalProperties": false,
+                    "properties": {
+                        "type": {
+                            "type": "string",
+                            "enum": ["credentials_token", "service_account"],
+                            "description": "Credentials type: either temporary OAuth 2.0 token or permanent service account credentials blob."
+                        },
+                        "credentials": {
+                            "type": "string",
+                            "description": "Oauth token or JSON string constituting service account credentials"
+                        }
+                    }
                 }
             }
         },