Skip to content

Commit ecc4f09

Browse files
vertex-sdk-botcopybara-github
authored andcommitted
fix: Include DeploymentResourcePool class in aiplatform top-level sdk module
PiperOrigin-RevId: 652541626
1 parent cfe0cc6 commit ecc4f09

File tree

3 files changed

+12
-9
lines changed

3 files changed

+12
-9
lines changed

google/cloud/aiplatform/__init__.py

+2
Original file line numberDiff line numberDiff line change
@@ -44,6 +44,7 @@
4444
)
4545
from google.cloud.aiplatform import metadata
4646
from google.cloud.aiplatform.tensorboard import uploader_tracker
47+
from google.cloud.aiplatform.models import DeploymentResourcePool
4748
from google.cloud.aiplatform.models import Endpoint
4849
from google.cloud.aiplatform.models import PrivateEndpoint
4950
from google.cloud.aiplatform.models import Model
@@ -153,6 +154,7 @@
153154
"CustomTrainingJob",
154155
"CustomContainerTrainingJob",
155156
"CustomPythonPackageTrainingJob",
157+
"DeploymentResourcePool",
156158
"Endpoint",
157159
"EntityType",
158160
"Execution",

tests/unit/aiplatform/test_deployment_resource_pools.py

+8-9
Original file line numberDiff line numberDiff line change
@@ -26,7 +26,6 @@
2626
from google.cloud import aiplatform
2727
from google.cloud.aiplatform import base
2828
from google.cloud.aiplatform import initializer
29-
from google.cloud.aiplatform import models
3029

3130
from google.cloud.aiplatform.compat.services import (
3231
deployment_resource_pool_service_client,
@@ -232,7 +231,7 @@ def test_constructor_gets_drp(self, get_drp_mock):
232231
location=_TEST_LOCATION,
233232
credentials=_TEST_CREDENTIALS,
234233
)
235-
models.DeploymentResourcePool(_TEST_DRP_NAME)
234+
aiplatform.DeploymentResourcePool(_TEST_DRP_NAME)
236235
get_drp_mock.assert_called_once_with(
237236
name=_TEST_DRP_NAME, retry=base._DEFAULT_RETRY
238237
)
@@ -242,7 +241,7 @@ def test_constructor_with_conflicting_location_fails(self):
242241
"""Passing a full resource name with `_TEST_LOCATION` and providing `_TEST_LOCATION_2` as location"""
243242

244243
with pytest.raises(RuntimeError) as err:
245-
models.DeploymentResourcePool(_TEST_DRP_NAME, location=_TEST_LOCATION_2)
244+
aiplatform.DeploymentResourcePool(_TEST_DRP_NAME, location=_TEST_LOCATION_2)
246245

247246
assert err.match(
248247
regexp=r"is provided, but different from the resource location"
@@ -251,7 +250,7 @@ def test_constructor_with_conflicting_location_fails(self):
251250
@pytest.mark.usefixtures("create_drp_mock")
252251
@pytest.mark.parametrize("sync", [True, False])
253252
def test_create(self, create_drp_mock, sync):
254-
test_drp = models.DeploymentResourcePool.create(
253+
test_drp = aiplatform.DeploymentResourcePool.create(
255254
deployment_resource_pool_id=_TEST_ID,
256255
machine_type=_TEST_MACHINE_TYPE,
257256
min_replica_count=10,
@@ -285,7 +284,7 @@ def test_create(self, create_drp_mock, sync):
285284
@pytest.mark.usefixtures("create_drp_mock")
286285
@pytest.mark.parametrize("sync", [True, False])
287286
def test_create_with_timeout(self, create_drp_mock, sync):
288-
test_drp = models.DeploymentResourcePool.create(
287+
test_drp = aiplatform.DeploymentResourcePool.create(
289288
deployment_resource_pool_id=_TEST_ID,
290289
machine_type=_TEST_MACHINE_TYPE,
291290
min_replica_count=10,
@@ -319,17 +318,17 @@ def test_create_with_timeout(self, create_drp_mock, sync):
319318

320319
@pytest.mark.usefixtures("list_drp_mock")
321320
def test_list(self, list_drp_mock):
322-
drp_list = models.DeploymentResourcePool.list()
321+
drp_list = aiplatform.DeploymentResourcePool.list()
323322

324323
list_drp_mock.assert_called_once()
325324

326325
for drp in drp_list:
327-
assert isinstance(drp, models.DeploymentResourcePool)
326+
assert isinstance(drp, aiplatform.DeploymentResourcePool)
328327

329328
@pytest.mark.usefixtures("delete_drp_mock", "get_drp_mock")
330329
@pytest.mark.parametrize("sync", [True, False])
331330
def test_delete(self, delete_drp_mock, get_drp_mock, sync):
332-
test_drp = models.DeploymentResourcePool(
331+
test_drp = aiplatform.DeploymentResourcePool(
333332
deployment_resource_pool_name=_TEST_DRP_NAME
334333
)
335334
test_drp.delete(sync=sync)
@@ -341,7 +340,7 @@ def test_delete(self, delete_drp_mock, get_drp_mock, sync):
341340

342341
@pytest.mark.usefixtures("query_deployed_models_mock", "get_drp_mock")
343342
def test_query_deployed_models(self, query_deployed_models_mock, get_drp_mock):
344-
test_drp = models.DeploymentResourcePool(
343+
test_drp = aiplatform.DeploymentResourcePool(
345344
deployment_resource_pool_name=_TEST_DRP_NAME
346345
)
347346
dm_refs = test_drp.query_deployed_models()

vertexai/resources/__init__.py

+2
Original file line numberDiff line numberDiff line change
@@ -39,6 +39,7 @@
3939
)
4040
from google.cloud.aiplatform import metadata
4141
from google.cloud.aiplatform.tensorboard import uploader_tracker
42+
from google.cloud.aiplatform.models import DeploymentResourcePool
4243
from google.cloud.aiplatform.models import Endpoint
4344
from google.cloud.aiplatform.models import PrivateEndpoint
4445
from google.cloud.aiplatform.models import Model
@@ -148,6 +149,7 @@
148149
"CustomTrainingJob",
149150
"CustomContainerTrainingJob",
150151
"CustomPythonPackageTrainingJob",
152+
"DeploymentResourcePool",
151153
"Endpoint",
152154
"EntityType",
153155
"Execution",

0 commit comments

Comments
 (0)