From 1e1cbfef76b7c5c8db705d5c4c17c3691de7b032 Mon Sep 17 00:00:00 2001 From: Alexey Volkov Date: Thu, 6 Jan 2022 15:27:42 -0800 Subject: [PATCH 1/9] 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, From 5b63ac32932abb07078098854dbf2e29497e50ef Mon Sep 17 00:00:00 2001 From: Alexey Volkov Date: Wed, 2 Feb 2022 18:19:48 -0800 Subject: [PATCH 2/9] Stopped depending on the google-api-python-client package --- google/cloud/aiplatform/initializer.py | 4 ++++ setup.py | 1 - 2 files changed, 4 insertions(+), 1 deletion(-) diff --git a/google/cloud/aiplatform/initializer.py b/google/cloud/aiplatform/initializer.py index eada34bd99..0c6933bb9f 100644 --- a/google/cloud/aiplatform/initializer.py +++ b/google/cloud/aiplatform/initializer.py @@ -170,6 +170,10 @@ def project(self) -> str: .execute()["projectId"] ) self._project = project_id + except ImportError: + logging.debug( + f"Need google-api-python-client to convert project number to project ID." + ) except Exception as e: logging.warning(f"Error converting project number to project ID: {e}") return self._project diff --git a/setup.py b/setup.py index e11347b5af..5ceb81b60e 100644 --- a/setup.py +++ b/setup.py @@ -89,7 +89,6 @@ "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, From 6fe3f2aab472ad6912c9efdde69d2d141d5d523c Mon Sep 17 00:00:00 2001 From: Alexey Volkov Date: Fri, 4 Feb 2022 17:58:47 -0800 Subject: [PATCH 3/9] Fixed a lint warning --- google/cloud/aiplatform/initializer.py | 2 +- 1 file changed, 1 insertion(+), 1 deletion(-) diff --git a/google/cloud/aiplatform/initializer.py b/google/cloud/aiplatform/initializer.py index 0c6933bb9f..cf67463e02 100644 --- a/google/cloud/aiplatform/initializer.py +++ b/google/cloud/aiplatform/initializer.py @@ -172,7 +172,7 @@ def project(self) -> str: self._project = project_id except ImportError: logging.debug( - f"Need google-api-python-client to convert project number to project ID." + "Need google-api-python-client to convert project number to project ID." ) except Exception as e: logging.warning(f"Error converting project number to project ID: {e}") From 46f6324b8f79eb395e12f2c416df02cc788166f4 Mon Sep 17 00:00:00 2001 From: Alexey Volkov Date: Thu, 10 Mar 2022 17:13:44 -0800 Subject: [PATCH 4/9] Extracted getting the project ID to another module --- google/cloud/aiplatform/initializer.py | 25 ++++++++----------------- 1 file changed, 8 insertions(+), 17 deletions(-) diff --git a/google/cloud/aiplatform/initializer.py b/google/cloud/aiplatform/initializer.py index cf67463e02..937bc5fb6e 100644 --- a/google/cloud/aiplatform/initializer.py +++ b/google/cloud/aiplatform/initializer.py @@ -32,6 +32,7 @@ from google.cloud.aiplatform.constants import base as constants from google.cloud.aiplatform import utils from google.cloud.aiplatform.metadata import metadata +from google.cloud.aiplatform.utils import resource_manager_utils from google.cloud.aiplatform.compat.types import ( encryption_spec as gca_encryption_spec_compat, @@ -156,27 +157,17 @@ def project(self) -> str: # 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"] + project_id = resource_manager_utils.get_project_id( + project_number=project_number, credentials=self.credentials, ) - self._project = project_id - except ImportError: - logging.debug( - "Need google-api-python-client to convert project number to project ID." + return project_id + except: + logging.getLogger(__name__).warning( + "Failed to convert project number to project ID.", exc_info=True ) - except Exception as e: - logging.warning(f"Error converting project number to project ID: {e}") - return self._project + return project_number project_not_found_exception_str = ( "Unable to find your project. Please provide a project ID by:" From 036362e7cf9210c827930ce1da55e3fdab21f265 Mon Sep 17 00:00:00 2001 From: Alexey Volkov Date: Sat, 12 Mar 2022 17:27:27 -0800 Subject: [PATCH 5/9] Specified the exception type --- google/cloud/aiplatform/initializer.py | 2 +- 1 file changed, 1 insertion(+), 1 deletion(-) diff --git a/google/cloud/aiplatform/initializer.py b/google/cloud/aiplatform/initializer.py index 937bc5fb6e..bb4b165d31 100644 --- a/google/cloud/aiplatform/initializer.py +++ b/google/cloud/aiplatform/initializer.py @@ -163,7 +163,7 @@ def project(self) -> str: project_number=project_number, credentials=self.credentials, ) return project_id - except: + except Exception: logging.getLogger(__name__).warning( "Failed to convert project number to project ID.", exc_info=True ) From 1f79bcefdbf4329e7fdb2f69b179f635423a0625 Mon Sep 17 00:00:00 2001 From: Alexey Volkov Date: Wed, 16 Mar 2022 15:07:45 -0700 Subject: [PATCH 6/9] tests: Added project ID inference test --- .../aiplatform/test_project_id_inference.py | 91 +++++++++++++++++++ 1 file changed, 91 insertions(+) create mode 100644 tests/system/aiplatform/test_project_id_inference.py diff --git a/tests/system/aiplatform/test_project_id_inference.py b/tests/system/aiplatform/test_project_id_inference.py new file mode 100644 index 0000000000..dc047eb350 --- /dev/null +++ b/tests/system/aiplatform/test_project_id_inference.py @@ -0,0 +1,91 @@ +# -*- coding: utf-8 -*- + +# Copyright 2021 Google LLC +# +# Licensed under the Apache License, Version 2.0 (the "License"); +# you may not use this file except in compliance with the License. +# You may obtain a copy of the License at +# +# http://www.apache.org/licenses/LICENSE-2.0 +# +# Unless required by applicable law or agreed to in writing, software +# distributed under the License is distributed on an "AS IS" BASIS, +# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +# See the License for the specific language governing permissions and +# limitations under the License. +# + +import pytest + +from google.cloud import aiplatform +from google.cloud.aiplatform.compat.types import pipeline_state as gca_pipeline_state +from tests.system.aiplatform import e2e_base + + +@pytest.mark.usefixtures("prepare_staging_bucket", "delete_staging_bucket") +class TestProjectIDInference(e2e_base.TestEndToEnd): + + _temp_prefix = "temp-vertex-sdk-project-id-inference" + + def test_project_id_inference(self, shared_state): + # Collection of resources generated by this test, to be deleted during teardown + shared_state["resources"] = [] + + aiplatform.init( + location=e2e_base._LOCATION, + staging_bucket=shared_state["staging_bucket_name"], + ) + + worker_pool_specs = [ + { + "machine_spec": {"machine_type": "n1-standard-2"}, + "replica_count": 1, + "container_spec": { + "image_uri": "python:3.9", + "command": [ + "sh", + "-exc", + """python3 -m pip install git+https://github.com/Ark-kun/python-aiplatform@fix--Fixed-getitng-project-ID-when-running-on-Vertex-AI#egg=google-cloud-aiplatform&subdirectory=. + "$0" "$@" + """, + "python3", + "-c", + """ + from google.cloud import aiplatform + # Not initializing the Vertex SDK explicitly + # Checking teh project ID + print(aiplatform.initializer.global_config.project) + assert not aiplatform.initializer.global_config.project.endswith("-tp") + # Testing ability to list resources + endpoints = aiplatform.Endpoint.list() + print(endpoints) + """, + ], + "args": [], + }, + } + ] + + custom_job = aiplatform.CustomJob( + display_name=self._make_display_name("custom"), + worker_pool_specs=worker_pool_specs, + ) + custom_job.run( + enable_web_access=True, sync=False, + ) + + shared_state["resources"].append(custom_job) + + in_progress_done_check = custom_job.done() + custom_job.wait_for_resource_creation() + + completion_done_check = custom_job.done() + + assert ( + custom_job.state + == gca_pipeline_state.PipelineState.PIPELINE_STATE_SUCCEEDED + ) + + # Check done() method works correctly + assert in_progress_done_check is False + assert completion_done_check is True From de1bb246b7df7f223cc65331c4fb8f0d5f35af50 Mon Sep 17 00:00:00 2001 From: Alexey Volkov Date: Wed, 16 Mar 2022 21:59:45 -0700 Subject: [PATCH 7/9] Added unit test --- tests/unit/aiplatform/test_initializer.py | 10 ++++++++++ 1 file changed, 10 insertions(+) diff --git a/tests/unit/aiplatform/test_initializer.py b/tests/unit/aiplatform/test_initializer.py index e52dfef3aa..b721648eb0 100644 --- a/tests/unit/aiplatform/test_initializer.py +++ b/tests/unit/aiplatform/test_initializer.py @@ -28,6 +28,7 @@ from google.cloud.aiplatform.metadata.metadata import metadata_service from google.cloud.aiplatform.constants import base as constants from google.cloud.aiplatform import utils +from google.cloud.aiplatform.utils import resource_manager_utils from google.cloud.aiplatform_v1.services.model_service import ( client as model_service_client, @@ -61,6 +62,15 @@ def mock_auth_default(): monkeypatch.setattr(google.auth, "default", mock_auth_default) assert initializer.global_config.project == _TEST_PROJECT + def test_infer_project_id(self): + def mock_get_project_id(): + return _TEST_PROJECT + + with mock.patch.object( + target=os, attribute="cpu_count", new=mock_get_project_id + ): + assert initializer.global_config.project == _TEST_PROJECT + def test_init_location_sets_location(self): initializer.global_config.init(location=_TEST_LOCATION) assert initializer.global_config.location == _TEST_LOCATION From bfe1e980e8d118cbeada2741bff487917bfa3b01 Mon Sep 17 00:00:00 2001 From: Alexey Volkov Date: Thu, 17 Mar 2022 12:23:13 -0700 Subject: [PATCH 8/9] Update test_initializer.py --- tests/unit/aiplatform/test_initializer.py | 4 +++- 1 file changed, 3 insertions(+), 1 deletion(-) diff --git a/tests/unit/aiplatform/test_initializer.py b/tests/unit/aiplatform/test_initializer.py index b721648eb0..0c695f3fad 100644 --- a/tests/unit/aiplatform/test_initializer.py +++ b/tests/unit/aiplatform/test_initializer.py @@ -67,7 +67,9 @@ def mock_get_project_id(): return _TEST_PROJECT with mock.patch.object( - target=os, attribute="cpu_count", new=mock_get_project_id + target=resource_manager_utils, + attribute="get_project_id", + new=mock_get_project_id, ): assert initializer.global_config.project == _TEST_PROJECT From 4f20ae0716e7b9b997dbac63522025241585216f Mon Sep 17 00:00:00 2001 From: Alexey Volkov Date: Thu, 17 Mar 2022 14:25:56 -0700 Subject: [PATCH 9/9] Patched os.environ in the unit test --- tests/unit/aiplatform/test_initializer.py | 7 ++++++- 1 file changed, 6 insertions(+), 1 deletion(-) diff --git a/tests/unit/aiplatform/test_initializer.py b/tests/unit/aiplatform/test_initializer.py index 0c695f3fad..a85e8257bc 100644 --- a/tests/unit/aiplatform/test_initializer.py +++ b/tests/unit/aiplatform/test_initializer.py @@ -63,13 +63,18 @@ def mock_auth_default(): assert initializer.global_config.project == _TEST_PROJECT def test_infer_project_id(self): - def mock_get_project_id(): + cloud_project_number = "123" + + def mock_get_project_id(project_number: str, **_): + assert project_number == cloud_project_number return _TEST_PROJECT with mock.patch.object( target=resource_manager_utils, attribute="get_project_id", new=mock_get_project_id, + ), mock.patch.dict( + os.environ, {"CLOUD_ML_PROJECT_ID": cloud_project_number}, clear=True ): assert initializer.global_config.project == _TEST_PROJECT