Skip to content

Commit c7f57ad

Browse files
nayaknishantcopybara-github
nayaknishant
authored and
copybara-github
committed
feat: Support global network parameter.
COPYBARA_INTEGRATE_REVIEW=#1702 from nayaknishant:nn-network 32d0972 PiperOrigin-RevId: 485422918
1 parent 98dbe5c commit c7f57ad

8 files changed

Lines changed: 338 additions & 40 deletions

File tree

‎google/cloud/aiplatform/initializer.py‎

Lines changed: 16 additions & 1 deletion
Original file line numberDiff line numberDiff line change
@@ -1,6 +1,6 @@
11
# -*- coding: utf-8 -*-
22

3-
# Copyright 2020 Google LLC
3+
# Copyright 2022 Google LLC
44
#
55
# Licensed under the Apache License, Version 2.0 (the "License");
66
# you may not use this file except in compliance with the License.
@@ -51,6 +51,7 @@ def __init__(self):
5151
self._staging_bucket = None
5252
self._credentials = None
5353
self._encryption_spec_key_name = None
54+
self._network = None
5455

5556
def init(
5657
self,
@@ -65,6 +66,7 @@ def init(
6566
staging_bucket: Optional[str] = None,
6667
credentials: Optional[auth_credentials.Credentials] = None,
6768
encryption_spec_key_name: Optional[str] = None,
69+
network: Optional[str] = None,
6870
):
6971
"""Updates common initialization parameters with provided options.
7072
@@ -95,6 +97,12 @@ def init(
9597
resource is created.
9698
9799
If set, this resource and all sub-resources will be secured by this key.
100+
network (str):
101+
Optional. The full name of the Compute Engine network to which jobs
102+
and resources should be peered. E.g. "projects/12345/global/networks/myVPC".
103+
Private services access must already be configured for the network.
104+
If specified, all eligible jobs and resources created will be peered
105+
with this VPC.
98106
Raises:
99107
ValueError:
100108
If experiment_description is provided but experiment is not.
@@ -130,6 +138,8 @@ def init(
130138
self._credentials = credentials
131139
if encryption_spec_key_name:
132140
self._encryption_spec_key_name = encryption_spec_key_name
141+
if network is not None:
142+
self._network = network
133143

134144
if experiment:
135145
metadata._experiment_tracker.set_experiment(
@@ -237,6 +247,11 @@ def encryption_spec_key_name(self) -> Optional[str]:
237247
"""Default encryption spec key name, if provided."""
238248
return self._encryption_spec_key_name
239249

250+
@property
251+
def network(self) -> Optional[str]:
252+
"""Default Compute Engine network to peer to, if provided."""
253+
return self._network
254+
240255
@property
241256
def experiment_name(self) -> Optional[str]:
242257
"""Default experiment name, if provided."""

‎google/cloud/aiplatform/jobs.py‎

Lines changed: 136 additions & 4 deletions
Original file line numberDiff line numberDiff line change
@@ -1515,7 +1515,6 @@ def from_local_script(
15151515
staging_bucket=staging_bucket,
15161516
)
15171517

1518-
@base.optional_sync()
15191518
def run(
15201519
self,
15211520
service_account: Optional[str] = None,
@@ -1537,7 +1536,8 @@ def run(
15371536
Optional. The full name of the Compute Engine network to which the job
15381537
should be peered. For example, projects/12345/global/networks/myVPC.
15391538
Private services access must already be configured for the network.
1540-
If left unspecified, the job is not peered with any network.
1539+
If left unspecified, the network set in aiplatform.init will be used.
1540+
Otherwise, the job is not peered with any network.
15411541
timeout (int):
15421542
The maximum job running time in seconds. The default is 7 days.
15431543
restart_job_on_worker_restart (bool):
@@ -1570,7 +1570,73 @@ def run(
15701570
create_request_timeout (float):
15711571
Optional. The timeout for the create request in seconds.
15721572
"""
1573+
network = network or initializer.global_config.network
1574+
1575+
self._run(
1576+
service_account=service_account,
1577+
network=network,
1578+
timeout=timeout,
1579+
restart_job_on_worker_restart=restart_job_on_worker_restart,
1580+
enable_web_access=enable_web_access,
1581+
tensorboard=tensorboard,
1582+
sync=sync,
1583+
create_request_timeout=create_request_timeout,
1584+
)
1585+
1586+
@base.optional_sync()
1587+
def _run(
1588+
self,
1589+
service_account: Optional[str] = None,
1590+
network: Optional[str] = None,
1591+
timeout: Optional[int] = None,
1592+
restart_job_on_worker_restart: bool = False,
1593+
enable_web_access: bool = False,
1594+
tensorboard: Optional[str] = None,
1595+
sync: bool = True,
1596+
create_request_timeout: Optional[float] = None,
1597+
) -> None:
1598+
"""Helper method to ensure network synchronization and to run the configured CustomJob.
1599+
1600+
Args:
1601+
service_account (str):
1602+
Optional. Specifies the service account for workload run-as account.
1603+
Users submitting jobs must have act-as permission on this run-as account.
1604+
network (str):
1605+
Optional. The full name of the Compute Engine network to which the job
1606+
should be peered. For example, projects/12345/global/networks/myVPC.
1607+
Private services access must already be configured for the network.
1608+
timeout (int):
1609+
The maximum job running time in seconds. The default is 7 days.
1610+
restart_job_on_worker_restart (bool):
1611+
Restarts the entire CustomJob if a worker
1612+
gets restarted. This feature can be used by
1613+
distributed training jobs that are not resilient
1614+
to workers leaving and joining a job.
1615+
enable_web_access (bool):
1616+
Whether you want Vertex AI to enable interactive shell access
1617+
to training containers.
1618+
https://cloud.google.com/vertex-ai/docs/training/monitor-debug-interactive-shell
1619+
tensorboard (str):
1620+
Optional. The name of a Vertex AI
1621+
[Tensorboard][google.cloud.aiplatform.v1beta1.Tensorboard]
1622+
resource to which this CustomJob will upload Tensorboard
1623+
logs. Format:
1624+
``projects/{project}/locations/{location}/tensorboards/{tensorboard}``
15731625
1626+
The training script should write Tensorboard to following Vertex AI environment
1627+
variable:
1628+
1629+
AIP_TENSORBOARD_LOG_DIR
1630+
1631+
`service_account` is required with provided `tensorboard`.
1632+
For more information on configuring your service account please visit:
1633+
https://cloud.google.com/vertex-ai/docs/experiments/tensorboard-training
1634+
sync (bool):
1635+
Whether to execute this method synchronously. If False, this method
1636+
will unblock and it will be executed in a concurrent Future.
1637+
create_request_timeout (float):
1638+
Optional. The timeout for the create request in seconds.
1639+
"""
15741640
if service_account:
15751641
self._gca_resource.job_spec.service_account = service_account
15761642

@@ -1907,7 +1973,6 @@ def _log_web_access_uris(self):
19071973
)
19081974
self._logged_web_access_uris.add(uri)
19091975

1910-
@base.optional_sync()
19111976
def run(
19121977
self,
19131978
service_account: Optional[str] = None,
@@ -1929,7 +1994,8 @@ def run(
19291994
Optional. The full name of the Compute Engine network to which the job
19301995
should be peered. For example, projects/12345/global/networks/myVPC.
19311996
Private services access must already be configured for the network.
1932-
If left unspecified, the job is not peered with any network.
1997+
If left unspecified, the network set in aiplatform.init will be used.
1998+
Otherwise, the job is not peered with any network.
19331999
timeout (int):
19342000
Optional. The maximum job running time in seconds. The default is 7 days.
19352001
restart_job_on_worker_restart (bool):
@@ -1962,7 +2028,73 @@ def run(
19622028
create_request_timeout (float):
19632029
Optional. The timeout for the create request in seconds.
19642030
"""
2031+
network = network or initializer.global_config.network
2032+
2033+
self._run(
2034+
service_account=service_account,
2035+
network=network,
2036+
timeout=timeout,
2037+
restart_job_on_worker_restart=restart_job_on_worker_restart,
2038+
enable_web_access=enable_web_access,
2039+
tensorboard=tensorboard,
2040+
sync=sync,
2041+
create_request_timeout=create_request_timeout,
2042+
)
2043+
2044+
@base.optional_sync()
2045+
def _run(
2046+
self,
2047+
service_account: Optional[str] = None,
2048+
network: Optional[str] = None,
2049+
timeout: Optional[int] = None, # seconds
2050+
restart_job_on_worker_restart: bool = False,
2051+
enable_web_access: bool = False,
2052+
tensorboard: Optional[str] = None,
2053+
sync: bool = True,
2054+
create_request_timeout: Optional[float] = None,
2055+
) -> None:
2056+
"""Helper method to ensure network synchronization and to run the configured CustomJob.
2057+
2058+
Args:
2059+
service_account (str):
2060+
Optional. Specifies the service account for workload run-as account.
2061+
Users submitting jobs must have act-as permission on this run-as account.
2062+
network (str):
2063+
Optional. The full name of the Compute Engine network to which the job
2064+
should be peered. For example, projects/12345/global/networks/myVPC.
2065+
Private services access must already be configured for the network.
2066+
timeout (int):
2067+
Optional. The maximum job running time in seconds. The default is 7 days.
2068+
restart_job_on_worker_restart (bool):
2069+
Restarts the entire CustomJob if a worker
2070+
gets restarted. This feature can be used by
2071+
distributed training jobs that are not resilient
2072+
to workers leaving and joining a job.
2073+
enable_web_access (bool):
2074+
Whether you want Vertex AI to enable interactive shell access
2075+
to training containers.
2076+
https://cloud.google.com/vertex-ai/docs/training/monitor-debug-interactive-shell
2077+
tensorboard (str):
2078+
Optional. The name of a Vertex AI
2079+
[Tensorboard][google.cloud.aiplatform.v1beta1.Tensorboard]
2080+
resource to which this CustomJob will upload Tensorboard
2081+
logs. Format:
2082+
``projects/{project}/locations/{location}/tensorboards/{tensorboard}``
19652083
2084+
The training script should write Tensorboard to following Vertex AI environment
2085+
variable:
2086+
2087+
AIP_TENSORBOARD_LOG_DIR
2088+
2089+
`service_account` is required with provided `tensorboard`.
2090+
For more information on configuring your service account please visit:
2091+
https://cloud.google.com/vertex-ai/docs/experiments/tensorboard-training
2092+
sync (bool):
2093+
Whether to execute this method synchronously. If False, this method
2094+
will unblock and it will be executed in a concurrent Future.
2095+
create_request_timeout (float):
2096+
Optional. The timeout for the create request in seconds.
2097+
"""
19662098
if service_account:
19672099
self._gca_resource.trial_job_spec.service_account = service_account
19682100

‎google/cloud/aiplatform/matching_engine/matching_engine_index_endpoint.py‎

Lines changed: 96 additions & 11 deletions
Original file line numberDiff line numberDiff line change
@@ -126,11 +126,10 @@ def __init__(
126126
self._gca_resource = self._get_gca_resource(resource_name=index_endpoint_name)
127127

128128
@classmethod
129-
@base.optional_sync()
130129
def create(
131130
cls,
132131
display_name: str,
133-
network: str,
132+
network: Optional[str] = None,
134133
description: Optional[str] = None,
135134
labels: Optional[Dict[str, str]] = None,
136135
project: Optional[str] = None,
@@ -153,13 +152,12 @@ def create(
153152
The name can be up to 128 characters long and
154153
can be consist of any UTF-8 characters.
155154
network (str):
156-
Required. The full name of the Google Compute Engine
155+
Optional. The full name of the Google Compute Engine
157156
`network <https://cloud.google.com/compute/docs/networks-and-firewalls#networks>`__
158157
to which the IndexEndpoint should be peered.
159158
160-
Private services access must already be configured for the
161-
network. If left unspecified, the Endpoint is not peered
162-
with any network.
159+
Private services access must already be configured for the network.
160+
If left unspecified, the network set with aiplatform.init will be used.
163161
164162
`Format <https://cloud.google.com/compute/docs/reference/rest/v1/networks/insert>`__:
165163
projects/{project}/global/networks/{network}. Where
@@ -182,13 +180,13 @@ def create(
182180
System reserved label keys are prefixed with
183181
"aiplatform.googleapis.com/" and are immutable.
184182
project (str):
185-
Optional. Project to create EntityType in. If not set, project
183+
Optional. Project to create IndexEndpoint in. If not set, project
186184
set in aiplatform.init will be used.
187185
location (str):
188-
Optional. Location to create EntityType in. If not set, location
186+
Optional. Location to create IndexEndpoint in. If not set, location
189187
set in aiplatform.init will be used.
190188
credentials (auth_credentials.Credentials):
191-
Optional. Custom credentials to use to create EntityTypes. Overrides
189+
Optional. Custom credentials to use to create IndexEndpoints. Overrides
192190
credentials set in aiplatform.init.
193191
request_metadata (Sequence[Tuple[str, str]]):
194192
Optional. Strings which should be sent along with the request as metadata.
@@ -200,11 +198,98 @@ def create(
200198
Returns:
201199
MatchingEngineIndexEndpoint - IndexEndpoint resource object
202200
201+
Raises:
202+
ValueError: A network must be instantiated when creating a IndexEndpoint.
203203
"""
204-
gapic_index_endpoint = gca_matching_engine_index_endpoint.IndexEndpoint(
204+
network = network or initializer.global_config.network
205+
206+
if not network:
207+
raise ValueError(
208+
"Please provide `network` argument or set network"
209+
"using aiplatform.init(network=...)"
210+
)
211+
212+
return cls._create(
205213
display_name=display_name,
206-
description=description,
207214
network=network,
215+
description=description,
216+
labels=labels,
217+
project=project,
218+
location=location,
219+
credentials=credentials,
220+
request_metadata=request_metadata,
221+
sync=sync,
222+
)
223+
224+
@classmethod
225+
@base.optional_sync()
226+
def _create(
227+
cls,
228+
display_name: str,
229+
network: Optional[str] = None,
230+
description: Optional[str] = None,
231+
labels: Optional[Dict[str, str]] = None,
232+
project: Optional[str] = None,
233+
location: Optional[str] = None,
234+
credentials: Optional[auth_credentials.Credentials] = None,
235+
request_metadata: Optional[Sequence[Tuple[str, str]]] = (),
236+
sync: bool = True,
237+
) -> "MatchingEngineIndexEndpoint":
238+
"""Helper method to ensure network synchronization and to
239+
create a MatchingEngineIndexEndpoint resource.
240+
241+
Args:
242+
display_name (str):
243+
Required. The display name of the IndexEndpoint.
244+
The name can be up to 128 characters long and
245+
can be consist of any UTF-8 characters.
246+
network (str):
247+
Optional. The full name of the Google Compute Engine
248+
`network <https://cloud.google.com/compute/docs/networks-and-firewalls#networks>`__
249+
to which the IndexEndpoint should be peered.
250+
Private services access must already be configured for the network.
251+
252+
`Format <https://cloud.google.com/compute/docs/reference/rest/v1/networks/insert>`__:
253+
projects/{project}/global/networks/{network}. Where
254+
{project} is a project number, as in '12345', and {network}
255+
is network name.
256+
description (str):
257+
Optional. The description of the IndexEndpoint.
258+
labels (Dict[str, str]):
259+
Optional. The labels with user-defined
260+
metadata to organize your IndexEndpoint.
261+
Label keys and values can be no longer than 64
262+
characters (Unicode codepoints), can only
263+
contain lowercase letters, numeric characters,
264+
underscores and dashes. International characters
265+
are allowed.
266+
See https://goo.gl/xmQnxf for more information
267+
on and examples of labels. No more than 64 user
268+
labels can be associated with one
269+
IndexEndpoint (System labels are excluded)."
270+
System reserved label keys are prefixed with
271+
"aiplatform.googleapis.com/" and are immutable.
272+
project (str):
273+
Optional. Project to create IndexEndpoint in. If not set, project
274+
set in aiplatform.init will be used.
275+
location (str):
276+
Optional. Location to create IndexEndpoint in. If not set, location
277+
set in aiplatform.init will be used.
278+
credentials (auth_credentials.Credentials):
279+
Optional. Custom credentials to use to create IndexEndpoints. Overrides
280+
credentials set in aiplatform.init.
281+
request_metadata (Sequence[Tuple[str, str]]):
282+
Optional. Strings which should be sent along with the request as metadata.
283+
sync (bool):
284+
Optional. Whether to execute this creation synchronously. If False, this method
285+
will be executed in concurrent Future and any downstream object will
286+
be immediately returned and synced when the Future has completed.
287+
288+
Returns:
289+
MatchingEngineIndexEndpoint - IndexEndpoint resource object
290+
"""
291+
gapic_index_endpoint = gca_matching_engine_index_endpoint.IndexEndpoint(
292+
display_name=display_name, description=description, network=network
208293
)
209294

210295
if labels:

0 commit comments

Comments
 (0)