diff --git a/docker/autoscaler/requirements.txt b/docker/autoscaler/requirements.txt index ff76312ce..fa8308f6b 100644 --- a/docker/autoscaler/requirements.txt +++ b/docker/autoscaler/requirements.txt @@ -22,4 +22,7 @@ tabulate mlflow kubernetes boto3 -ipython \ No newline at end of file +ipython +google-oauth +google-api-python-client +oauth2client \ No newline at end of file diff --git a/python/ray/autoscaler/gcp/config.py b/python/ray/autoscaler/gcp/config.py index 0584f8382..94fcb5692 100644 --- a/python/ray/autoscaler/gcp/config.py +++ b/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) diff --git a/python/ray/autoscaler/gcp/node_provider.py b/python/ray/autoscaler/gcp/node_provider.py index 4b46c186f..58f8fff38 100644 --- a/python/ray/autoscaler/gcp/node_provider.py +++ b/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 = {} diff --git a/python/ray/autoscaler/ray-schema.json b/python/ray/autoscaler/ray-schema.json index 9fc256bca..ad299ebf5 100644 --- a/python/ray/autoscaler/ray-schema.json +++ b/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" + } + } } } },