Skip to content

Commit

Permalink
fix: Fixed getitng project ID when running on Vertex AI
Browse files Browse the repository at this point in the history
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 googleapis#852
See googleapis/google-auth-library-python#924
This PR fixes the fallback to get the project ID from the `CLOUD_ML_PROJECT_ID` environment variable.
  • Loading branch information
Ark-kun committed Jan 6, 2022
1 parent b72067b commit 1e1cbfe
Show file tree
Hide file tree
Showing 2 changed files with 26 additions and 0 deletions.
25 changes: 25 additions & 0 deletions google/cloud/aiplatform/initializer.py
Original file line number Diff line number Diff line change
Expand Up @@ -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"
Expand Down
1 change: 1 addition & 0 deletions setup.py
Original file line number Diff line number Diff line change
Expand Up @@ -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,
Expand Down

0 comments on commit 1e1cbfe

Please sign in to comment.