feat: Support global network parameter.

COPYBARA_INTEGRATE_REVIEW=#1702 from nayaknishant:nn-network 32d0972
PiperOrigin-RevId: 485422918

Author: nayaknishant

google/cloud/aiplatform/initializer.py

@@ -1,6 +1,6 @@
 # -*- coding: utf-8 -*-
 
-# Copyright 2020 Google LLC
+# Copyright 2022 Google LLC
 #
 # Licensed under the Apache License, Version 2.0 (the "License");
 # you may not use this file except in compliance with the License.
@@ -51,6 +51,7 @@ def __init__(self):
         self._staging_bucket = None
         self._credentials = None
         self._encryption_spec_key_name = None
+        self._network = None
 
     def init(
         self,
@@ -65,6 +66,7 @@ def init(
         staging_bucket: Optional[str] = None,
         credentials: Optional[auth_credentials.Credentials] = None,
         encryption_spec_key_name: Optional[str] = None,
+        network: Optional[str] = None,
     ):
         """Updates common initialization parameters with provided options.
 
@@ -95,6 +97,12 @@ def init(
                 resource is created.
 
                 If set, this resource and all sub-resources will be secured by this key.
+            network (str):
+                Optional. The full name of the Compute Engine network to which jobs
+                and resources should be peered. E.g. "projects/12345/global/networks/myVPC".
+                Private services access must already be configured for the network.
+                If specified, all eligible jobs and resources created will be peered
+                with this VPC.
         Raises:
             ValueError:
                 If experiment_description is provided but experiment is not.
@@ -130,6 +138,8 @@ def init(
             self._credentials = credentials
         if encryption_spec_key_name:
             self._encryption_spec_key_name = encryption_spec_key_name
+        if network is not None:
+            self._network = network
 
         if experiment:
             metadata._experiment_tracker.set_experiment(
@@ -237,6 +247,11 @@ def encryption_spec_key_name(self) -> Optional[str]:
         """Default encryption spec key name, if provided."""
         return self._encryption_spec_key_name
 
+    @property
+    def network(self) -> Optional[str]:
+        """Default Compute Engine network to peer to, if provided."""
+        return self._network
+
     @property
     def experiment_name(self) -> Optional[str]:
         """Default experiment name, if provided."""

google/cloud/aiplatform/jobs.py

@@ -1515,7 +1515,6 @@ def from_local_script(
             staging_bucket=staging_bucket,
         )
 
-    @base.optional_sync()
     def run(
         self,
         service_account: Optional[str] = None,
@@ -1537,7 +1536,8 @@ def run(
                 Optional. The full name of the Compute Engine network to which the job
                 should be peered. For example, projects/12345/global/networks/myVPC.
                 Private services access must already be configured for the network.
-                If left unspecified, the job is not peered with any network.
+                If left unspecified, the network set in aiplatform.init will be used.
+                Otherwise, the job is not peered with any network.
             timeout (int):
                 The maximum job running time in seconds. The default is 7 days.
             restart_job_on_worker_restart (bool):
@@ -1570,7 +1570,73 @@ def run(
             create_request_timeout (float):
                 Optional. The timeout for the create request in seconds.
         """
+        network = network or initializer.global_config.network
+
+        self._run(
+            service_account=service_account,
+            network=network,
+            timeout=timeout,
+            restart_job_on_worker_restart=restart_job_on_worker_restart,
+            enable_web_access=enable_web_access,
+            tensorboard=tensorboard,
+            sync=sync,
+            create_request_timeout=create_request_timeout,
+        )
+
+    @base.optional_sync()
+    def _run(
+        self,
+        service_account: Optional[str] = None,
+        network: Optional[str] = None,
+        timeout: Optional[int] = None,
+        restart_job_on_worker_restart: bool = False,
+        enable_web_access: bool = False,
+        tensorboard: Optional[str] = None,
+        sync: bool = True,
+        create_request_timeout: Optional[float] = None,
+    ) -> None:
+        """Helper method to ensure network synchronization and to run the configured CustomJob.
+
+        Args:
+            service_account (str):
+                Optional. Specifies the service account for workload run-as account.
+                Users submitting jobs must have act-as permission on this run-as account.
+            network (str):
+                Optional. The full name of the Compute Engine network to which the job
+                should be peered. For example, projects/12345/global/networks/myVPC.
+                Private services access must already be configured for the network.
+            timeout (int):
+                The maximum job running time in seconds. The default is 7 days.
+            restart_job_on_worker_restart (bool):
+                Restarts the entire CustomJob if a worker
+                gets restarted. This feature can be used by
+                distributed training jobs that are not resilient
+                to workers leaving and joining a job.
+            enable_web_access (bool):
+                Whether you want Vertex AI to enable interactive shell access
+                to training containers.
+                https://cloud.google.com/vertex-ai/docs/training/monitor-debug-interactive-shell
+            tensorboard (str):
+                Optional. The name of a Vertex AI
+                [Tensorboard][google.cloud.aiplatform.v1beta1.Tensorboard]
+                resource to which this CustomJob will upload Tensorboard
+                logs. Format:
+                ``projects/{project}/locations/{location}/tensorboards/{tensorboard}``
 
+                The training script should write Tensorboard to following Vertex AI environment
+                variable:
+
+                AIP_TENSORBOARD_LOG_DIR
+
+                `service_account` is required with provided `tensorboard`.
+                For more information on configuring your service account please visit:
+                https://cloud.google.com/vertex-ai/docs/experiments/tensorboard-training
+            sync (bool):
+                Whether to execute this method synchronously. If False, this method
+                will unblock and it will be executed in a concurrent Future.
+            create_request_timeout (float):
+                Optional. The timeout for the create request in seconds.
+        """
         if service_account:
             self._gca_resource.job_spec.service_account = service_account
 
@@ -1907,7 +1973,6 @@ def _log_web_access_uris(self):
                     )
                     self._logged_web_access_uris.add(uri)
 
-    @base.optional_sync()
     def run(
         self,
         service_account: Optional[str] = None,
@@ -1929,7 +1994,8 @@ def run(
                 Optional. The full name of the Compute Engine network to which the job
                 should be peered. For example, projects/12345/global/networks/myVPC.
                 Private services access must already be configured for the network.
-                If left unspecified, the job is not peered with any network.
+                If left unspecified, the network set in aiplatform.init will be used.
+                Otherwise, the job is not peered with any network.
             timeout (int):
                 Optional. The maximum job running time in seconds. The default is 7 days.
             restart_job_on_worker_restart (bool):
@@ -1962,7 +2028,73 @@ def run(
             create_request_timeout (float):
                 Optional. The timeout for the create request in seconds.
         """
+        network = network or initializer.global_config.network
+
+        self._run(
+            service_account=service_account,
+            network=network,
+            timeout=timeout,
+            restart_job_on_worker_restart=restart_job_on_worker_restart,
+            enable_web_access=enable_web_access,
+            tensorboard=tensorboard,
+            sync=sync,
+            create_request_timeout=create_request_timeout,
+        )
+
+    @base.optional_sync()
+    def _run(
+        self,
+        service_account: Optional[str] = None,
+        network: Optional[str] = None,
+        timeout: Optional[int] = None,  # seconds
+        restart_job_on_worker_restart: bool = False,
+        enable_web_access: bool = False,
+        tensorboard: Optional[str] = None,
+        sync: bool = True,
+        create_request_timeout: Optional[float] = None,
+    ) -> None:
+        """Helper method to ensure network synchronization and to run the configured CustomJob.
+
+        Args:
+            service_account (str):
+                Optional. Specifies the service account for workload run-as account.
+                Users submitting jobs must have act-as permission on this run-as account.
+            network (str):
+                Optional. The full name of the Compute Engine network to which the job
+                should be peered. For example, projects/12345/global/networks/myVPC.
+                Private services access must already be configured for the network.
+            timeout (int):
+                Optional. The maximum job running time in seconds. The default is 7 days.
+            restart_job_on_worker_restart (bool):
+                Restarts the entire CustomJob if a worker
+                gets restarted. This feature can be used by
+                distributed training jobs that are not resilient
+                to workers leaving and joining a job.
+            enable_web_access (bool):
+                Whether you want Vertex AI to enable interactive shell access
+                to training containers.
+                https://cloud.google.com/vertex-ai/docs/training/monitor-debug-interactive-shell
+            tensorboard (str):
+                Optional. The name of a Vertex AI
+                [Tensorboard][google.cloud.aiplatform.v1beta1.Tensorboard]
+                resource to which this CustomJob will upload Tensorboard
+                logs. Format:
+                ``projects/{project}/locations/{location}/tensorboards/{tensorboard}``
 
+                The training script should write Tensorboard to following Vertex AI environment
+                variable:
+
+                AIP_TENSORBOARD_LOG_DIR
+
+                `service_account` is required with provided `tensorboard`.
+                For more information on configuring your service account please visit:
+                https://cloud.google.com/vertex-ai/docs/experiments/tensorboard-training
+            sync (bool):
+                Whether to execute this method synchronously. If False, this method
+                will unblock and it will be executed in a concurrent Future.
+            create_request_timeout (float):
+                Optional. The timeout for the create request in seconds.
+        """
         if service_account:
             self._gca_resource.trial_job_spec.service_account = service_account
 

google/cloud/aiplatform/matching_engine/matching_engine_index_endpoint.py

@@ -126,11 +126,10 @@ def __init__(
         self._gca_resource = self._get_gca_resource(resource_name=index_endpoint_name)
 
     @classmethod
-    @base.optional_sync()
     def create(
         cls,
         display_name: str,
-        network: str,
+        network: Optional[str] = None,
         description: Optional[str] = None,
         labels: Optional[Dict[str, str]] = None,
         project: Optional[str] = None,
@@ -153,13 +152,12 @@ def create(
                 The name can be up to 128 characters long and
                 can be consist of any UTF-8 characters.
             network (str):
-                Required. The full name of the Google Compute Engine
+                Optional. The full name of the Google Compute Engine
                 `network <https://cloud.google.com/compute/docs/networks-and-firewalls#networks>`__
                 to which the IndexEndpoint should be peered.
 
-                Private services access must already be configured for the
-                network. If left unspecified, the Endpoint is not peered
-                with any network.
+                Private services access must already be configured for the network.
+                If left unspecified, the network set with aiplatform.init will be used.
 
                 `Format <https://cloud.google.com/compute/docs/reference/rest/v1/networks/insert>`__:
                 projects/{project}/global/networks/{network}. Where
@@ -182,13 +180,13 @@ def create(
                 System reserved label keys are prefixed with
                 "aiplatform.googleapis.com/" and are immutable.
             project (str):
-                Optional. Project to create EntityType in. If not set, project
+                Optional. Project to create IndexEndpoint in. If not set, project
                 set in aiplatform.init will be used.
             location (str):
-                Optional. Location to create EntityType in. If not set, location
+                Optional. Location to create IndexEndpoint in. If not set, location
                 set in aiplatform.init will be used.
             credentials (auth_credentials.Credentials):
-                Optional. Custom credentials to use to create EntityTypes. Overrides
+                Optional. Custom credentials to use to create IndexEndpoints. Overrides
                 credentials set in aiplatform.init.
             request_metadata (Sequence[Tuple[str, str]]):
                 Optional. Strings which should be sent along with the request as metadata.
@@ -200,11 +198,98 @@ def create(
         Returns:
             MatchingEngineIndexEndpoint - IndexEndpoint resource object
 
+        Raises:
+            ValueError: A network must be instantiated when creating a IndexEndpoint.
         """
-        gapic_index_endpoint = gca_matching_engine_index_endpoint.IndexEndpoint(
+        network = network or initializer.global_config.network
+
+        if not network:
+            raise ValueError(
+                "Please provide `network` argument or set network"
+                "using aiplatform.init(network=...)"
+            )
+
+        return cls._create(
             display_name=display_name,
-            description=description,
             network=network,
+            description=description,
+            labels=labels,
+            project=project,
+            location=location,
+            credentials=credentials,
+            request_metadata=request_metadata,
+            sync=sync,
+        )
+
+    @classmethod
+    @base.optional_sync()
+    def _create(
+        cls,
+        display_name: str,
+        network: Optional[str] = None,
+        description: Optional[str] = None,
+        labels: Optional[Dict[str, str]] = None,
+        project: Optional[str] = None,
+        location: Optional[str] = None,
+        credentials: Optional[auth_credentials.Credentials] = None,
+        request_metadata: Optional[Sequence[Tuple[str, str]]] = (),
+        sync: bool = True,
+    ) -> "MatchingEngineIndexEndpoint":
+        """Helper method to ensure network synchronization and to
+        create a MatchingEngineIndexEndpoint resource.
+
+        Args:
+            display_name (str):
+                Required. The display name of the IndexEndpoint.
+                The name can be up to 128 characters long and
+                can be consist of any UTF-8 characters.
+            network (str):
+                Optional. The full name of the Google Compute Engine
+                `network <https://cloud.google.com/compute/docs/networks-and-firewalls#networks>`__
+                to which the IndexEndpoint should be peered.
+                Private services access must already be configured for the network.
+
+                `Format <https://cloud.google.com/compute/docs/reference/rest/v1/networks/insert>`__:
+                projects/{project}/global/networks/{network}. Where
+                {project} is a project number, as in '12345', and {network}
+                is network name.
+            description (str):
+                Optional. The description of the IndexEndpoint.
+            labels (Dict[str, str]):
+                Optional. The labels with user-defined
+                metadata to organize your IndexEndpoint.
+                Label keys and values can be no longer than 64
+                characters (Unicode codepoints), can only
+                contain lowercase letters, numeric characters,
+                underscores and dashes. International characters
+                are allowed.
+                See https://goo.gl/xmQnxf for more information
+                on and examples of labels. No more than 64 user
+                labels can be associated with one
+                IndexEndpoint (System labels are excluded)."
+                System reserved label keys are prefixed with
+                "aiplatform.googleapis.com/" and are immutable.
+            project (str):
+                Optional. Project to create IndexEndpoint in. If not set, project
+                set in aiplatform.init will be used.
+            location (str):
+                Optional. Location to create IndexEndpoint in. If not set, location
+                set in aiplatform.init will be used.
+            credentials (auth_credentials.Credentials):
+                Optional. Custom credentials to use to create IndexEndpoints. Overrides
+                credentials set in aiplatform.init.
+            request_metadata (Sequence[Tuple[str, str]]):
+                Optional. Strings which should be sent along with the request as metadata.
+            sync (bool):
+                Optional. Whether to execute this creation synchronously. If False, this method
+                will be executed in concurrent Future and any downstream object will
+                be immediately returned and synced when the Future has completed.
+
+        Returns:
+            MatchingEngineIndexEndpoint - IndexEndpoint resource object
+        """
+        gapic_index_endpoint = gca_matching_engine_index_endpoint.IndexEndpoint(
+            display_name=display_name, description=description, network=network
         )
 
         if labels: