From 1e1cbfef76b7c5c8db705d5c4c17c3691de7b032 Mon Sep 17 00:00:00 2001 From: Alexey Volkov Date: Thu, 6 Jan 2022 15:27:42 -0800 Subject: [PATCH] fix: Fixed getitng project ID when running on Vertex AI When project ID is not explicitly specified in `aiplatform.init()` call, the SDK uses `google.auth.default()` to infer the project ID. However when running under Vertex AI (CustomJob, PipelineJob), the project returned by `google.auth.default()` is not the correct user project. See https://github.com/googleapis/python-aiplatform/issues/852 See https://github.com/googleapis/google-auth-library-python/issues/924 This PR fixes the fallback to get the project ID from the `CLOUD_ML_PROJECT_ID` environment variable. --- google/cloud/aiplatform/initializer.py | 25 +++++++++++++++++++++++++ setup.py | 1 + 2 files changed, 26 insertions(+) diff --git a/google/cloud/aiplatform/initializer.py b/google/cloud/aiplatform/initializer.py index 2aa98b1600..eada34bd99 100644 --- a/google/cloud/aiplatform/initializer.py +++ b/google/cloud/aiplatform/initializer.py @@ -149,6 +149,31 @@ def project(self) -> str: if self._project: return self._project + # Project is not set. Trying to get it from the environment. + # See https://github.com/googleapis/python-aiplatform/issues/852 + # See https://github.com/googleapis/google-auth-library-python/issues/924 + # TODO: Remove when google.auth.default() learns the + # CLOUD_ML_PROJECT_ID env variable or Vertex AI starts setting GOOGLE_CLOUD_PROJECT env variable. + project_number = os.environ.get("CLOUD_ML_PROJECT_ID") + if project_number: + self._project = project_number + # Try to convert project number to project ID which is more readable. + try: + from googleapiclient import discovery + + cloud_resource_manager_service = discovery.build( + "cloudresourcemanager", "v3" + ) + project_id = ( + cloud_resource_manager_service.projects() + .get(name=f"projects/{project_number}") + .execute()["projectId"] + ) + self._project = project_id + except Exception as e: + logging.warning(f"Error converting project number to project ID: {e}") + return self._project + project_not_found_exception_str = ( "Unable to find your project. Please provide a project ID by:" "\n- Passing a constructor argument" diff --git a/setup.py b/setup.py index 5ceb81b60e..e11347b5af 100644 --- a/setup.py +++ b/setup.py @@ -89,6 +89,7 @@ "packaging >= 14.3", "google-cloud-storage >= 1.32.0, < 2.0.0dev", "google-cloud-bigquery >= 1.15.0, < 3.0.0dev", + "google-api-python-client >= 2.29", # for API discovery ), extras_require={ "full": full_extra_require,