Use discovery for ML API client instead of static file

src/python/tensorflow_cloud/core/deploy.py

@@ -34,6 +34,7 @@ def deploy_job(
     enable_stream_logs,
     job_labels=None,
     service_account=None,
+    region=None,
 ):
     """Deploys job with the given parameters to Google Cloud.
 
@@ -57,12 +58,15 @@ def deploy_job(
             to be used for training instead of the service account that AI
             Platform Training uses by default. See [custom service account](
             https://cloud.google.com/ai-platform/training/docs/custom-service-account)
+        region: Target region for running the AI Platform job.
     Returns:
         ID of the invoked remote Cloud AI Platform job.
 
     Raises:
         RuntimeError, if there was an error submitting the job.
     """
+    region = region or gcp.get_region()
+
     job_id = _generate_job_id()
     project_id = gcp.get_project_name()
     ml_apis = discovery.build(
@@ -74,7 +78,7 @@ def deploy_job(
 
     request_dict = _create_request_dict(
         job_id,
-        gcp.get_region(),
+        region,
         image_uri,
         chief_config,
         worker_count,

src/python/tensorflow_cloud/core/run.py

@@ -21,6 +21,7 @@
 from . import containerize
 from . import deploy
 from . import docker_config as docker_config_module
+from . import gcp
 from . import machine_config
 from . import preprocess
 from . import validate
@@ -73,6 +74,7 @@ def run_cloudtuner(num_jobs=1, **kwargs):
     stream_logs = kwargs.pop("stream_logs", False)
     job_labels = kwargs.pop("job_labels", None)
     service_account = kwargs.pop("service_account", None)
+    region = kwargs.pop("region", gcp.get_region())
 
     job_ids = [run_results["job_id"]]
     for _ in range(1, num_jobs):
@@ -96,6 +98,7 @@ def run_cloudtuner(num_jobs=1, **kwargs):
                 stream_logs,
                 job_labels=job_labels,
                 service_account=service_account,
+                region=region,
             )
         ])
 
@@ -117,6 +120,7 @@ def run(
     stream_logs=False,
     job_labels=None,
     service_account=None,
+    region=None,
     **kwargs
 ):
     """Runs your Tensorflow code in Google Cloud Platform.
@@ -195,6 +199,7 @@ def run(
             to be used for training instead of the service account that AI
             Platform Training uses by default. see [custom-service-account](
             https://cloud.google.com/ai-platform/training/docs/custom-service-account)
+        region: Target region for running the AI Platform Training job.
         **kwargs: Additional keyword arguments.
 
     Returns:
@@ -329,6 +334,7 @@ def run(
         stream_logs,
         job_labels=job_labels,
         service_account=service_account,
+        region=region,
     )
 
     # Call `exit` to prevent training the Keras model in the local env.

src/python/tensorflow_cloud/tuner/constants.py

@@ -14,13 +14,6 @@
 # limitations under the License.
 """Constants definitions for tuner sub module."""
 
-import os
-
-# API definition of Cloud AI Platform Vizier service
-OPTIMIZER_API_DOCUMENT_FILE = os.path.join(
-    os.path.dirname(os.path.abspath(__file__)),
-    "api/ml_public_google_rest_v1.json")
-
 # By default, the Tuner worker(s) always requests one trial at a time because
 # we would parallelize the tuning loop themselves as opposed to getting multiple
 # trial suggestions in one tuning loop.

src/python/tensorflow_cloud/tuner/vizier_client.py

@@ -15,10 +15,10 @@
 """A thin client for the Cloud AI Platform Vizier Service."""
 import datetime
 import http
-import json
 import time
 from typing import Any, Dict, List, Mapping, Optional, Text, Union
 
+from google.api_core.client_options import ClientOptions
 from googleapiclient import discovery
 from googleapiclient import errors
 import tensorflow as tf
@@ -427,11 +427,20 @@ def create_or_load_study(
     # Build the API client
     # Note that Vizier service is exposed as a regional endpoint. As such,
     # an API client needs to be created separately from the default.
-    with open(constants.OPTIMIZER_API_DOCUMENT_FILE) as f:
-        service_client = discovery.build_from_document(
-            service=json.load(f),
-            requestBuilder=google_api_client.TFCloudHttpRequest,
-        )
+    ml_endpoint = f"https://{region}-ml.googleapis.com"
+    client_options = ClientOptions(
+        api_endpoint=ml_endpoint,
+    )
+
+    # Disabling cache discovery to suppress noisy warning. More details at:
+    # https://github.com/googleapis/google-api-python-client/issues/299
+    service_client = discovery.build(
+        serviceName="ml",
+        version="v1",
+        requestBuilder=google_api_client.TFCloudHttpRequest,
+        cache_discovery=False,
+        client_options=client_options,
+    )
 
     # Creates or loads a study.
     study_parent = "projects/{}/locations/{}".format(project_id, region)