chore: fix Ray client logging and internal clean up

yinghsienwu · copybara-github · commit 13442f896824 · 2023-09-22T10:33:09.000-07:00
FUTURE_COPYBARA_INTEGRATE_REVIEW=#2604 from googleapis:release-please--branches--main a84695c PiperOrigin-RevId: 567439389
diff --git a/google/cloud/aiplatform/constants/base.py b/google/cloud/aiplatform/constants/base.py
@@ -14,7 +14,7 @@
 # See the License for the specific language governing permissions and
 # limitations under the License.
 #
-import os
+
 from google.cloud.aiplatform import version as aiplatform_version
 
 
@@ -56,8 +56,8 @@
         "us-west4",
     }
 )
-# This env variable injection is for testing, but not considered to be a public API.
-API_BASE_PATH = os.environ.get("_VERTEX_API_BASE_PATH") or "aiplatform.googleapis.com"
+
+API_BASE_PATH = "aiplatform.googleapis.com"
 PREDICTION_API_BASE_PATH = API_BASE_PATH
 
 # Batch Prediction
diff --git a/google/cloud/aiplatform/initializer.py b/google/cloud/aiplatform/initializer.py
@@ -421,13 +421,13 @@ def create_client(
             "google-cloud-aiplatform",
         ).version
 
+        if appended_gapic_version:
+            gapic_version = f"{gapic_version}+{appended_gapic_version}"
+
         user_agent = f"{constants.USER_AGENT_PRODUCT}/{gapic_version}"
         if appended_user_agent:
             user_agent = f"{user_agent} {' '.join(appended_user_agent)}"
 
-        if appended_gapic_version:
-            gapic_version = f"{gapic_version}+{appended_gapic_version}"
-
         client_info = gapic_v1.client_info.ClientInfo(
             gapic_version=gapic_version,
             user_agent=user_agent,
diff --git a/google/cloud/aiplatform/preview/vertex_ray/client_builder.py b/google/cloud/aiplatform/preview/vertex_ray/client_builder.py
@@ -14,19 +14,17 @@
 # See the License for the specific language governing permissions and
 # limitations under the License.
 #
-
+import logging
 from typing import Dict
 from typing import Optional
 from google.cloud import aiplatform
-from google.cloud.aiplatform import base
 from ray import client_builder
 from .render import VertexRayTemplate
 from .util import _validation_utils
 from .util import _gapic_utils
 
 
 VERTEX_SDK_VERSION = aiplatform.__version__
-_LOGGER = base.Logger(__name__)
 
 
 class _VertexRayClientContext(client_builder.ClientContext):
@@ -84,7 +82,7 @@ def __init__(self, address: Optional[str]) -> None:
         _validation_utils.valid_resource_name(address)
 
         self.vertex_address = address
-        _LOGGER.info(
+        print(
             "[Ray on Vertex AI]: Using cluster resource name to access head address with GAPIC API"
         )
 
@@ -112,12 +110,12 @@ def __init__(self, address: Optional[str]) -> None:
                 address,
                 " failed to start Head node properly because custom service account isn't supported.",
             )
-        _LOGGER.debug("[Ray on Vertex AI]: Resolved head node ip: %s", address)
+        logging.debug("[Ray on Vertex AI]: Resolved head node ip: %s", address)
         super().__init__(address)
 
     def connect(self) -> _VertexRayClientContext:
         # Can send any other params to ray cluster here
-        _LOGGER.info("[Ray on Vertex AI]: Connecting...")
+        print("[Ray on Vertex AI]: Connecting...")
         ray_client_context = super().connect()
         ray_head_uris = self.response.resource_runtime.access_uris
 
diff --git a/setup.py b/setup.py
@@ -74,7 +74,7 @@
 ]
 
 vizier_extra_require = [
-    "google-vizier==0.0.4; python_version<'3.11'",
+    "google-vizier==0.0.11; python_version>='3.9', '<'3.11'",
     "google-vizier>=0.1.6; python_version>='3.11'",
 ]
 
diff --git a/tests/unit/aiplatform/test_explain_lit.py b/tests/unit/aiplatform/test_explain_lit.py
@@ -29,21 +29,30 @@
     prediction_service as gca_prediction_service,
     explanation as gca_explanation,
 )
-from google.cloud.aiplatform.explain.lit import (
+
+# TODO (b/301592787): update testing_extra_require deps to numpy >= 1.22.0 when it doesn't cause conflicts
+try:
+    from lit_nlp import notebook  # noqa: F401
+except ImportError:
+    pytest.skip(
+        "Skipping test_explain_list due to dependency conflict with numpy",
+        allow_module_level=True,
+    )
+
+from google.cloud.aiplatform.explain.lit import (  # noqa: E402
     create_lit_dataset,
     create_lit_model,
     create_lit_model_from_endpoint,
     open_lit,
     set_up_and_open_lit,
 )
-from google.cloud.aiplatform.compat.services import (
+from google.cloud.aiplatform.compat.services import (  # noqa: E402
     endpoint_service_client,
     prediction_service_client,
 )
-from importlib import reload
-from lit_nlp.api import types as lit_types
-from lit_nlp import notebook
-from unittest import mock
+from importlib import reload  # noqa: E402
+from lit_nlp.api import types as lit_types  # noqa: E402
+from unittest import mock  # noqa: E402
 
 _TEST_PROJECT = "test-project"
 _TEST_LOCATION = "us-central1"

Original file line number	Diff line number	Diff line change
`@@ -14,7 +14,7 @@`
`14`	`14`	`# See the License for the specific language governing permissions and`
`15`	`15`	`# limitations under the License.`
`16`	`16`	`#`
`17`		`-import os`
	`17`	`+`
`18`	`18`	`from google.cloud.aiplatform import version as aiplatform_version`
`19`	`19`
`20`	`20`
`@@ -56,8 +56,8 @@`
`56`	`56`	`"us-west4",`
`57`	`57`	`}`
`58`	`58`	`)`
`59`		`-# This env variable injection is for testing, but not considered to be a public API.`
`60`		`-API_BASE_PATH = os.environ.get("_VERTEX_API_BASE_PATH") or "aiplatform.googleapis.com"`
	`59`	`+`
	`60`	`+API_BASE_PATH = "aiplatform.googleapis.com"`
`61`	`61`	`PREDICTION_API_BASE_PATH = API_BASE_PATH`
`62`	`62`
`63`	`63`	`# Batch Prediction`
Original file line number	Diff line number	Diff line change
`@@ -74,7 +74,7 @@`
`74`	`74`	`]`
`75`	`75`
`76`	`76`	`vizier_extra_require = [`
`77`		`- "google-vizier==0.0.4; python_version<'3.11'",`
	`77`	`+ "google-vizier==0.0.11; python_version>='3.9', '<'3.11'",`
`78`	`78`	`"google-vizier>=0.1.6; python_version>='3.11'",`
`79`	`79`	`]`
`80`	`80`