Skip to content
New issue

Have a question about this project? Sign up for a free GitHub account to open an issue and contact its maintainers and the community.

By clicking “Sign up for GitHub”, you agree to our terms of service and privacy statement. We’ll occasionally send you account related emails.

Already on GitHub? Sign in to your account

feat: Add endpoind_id arg to Endpoint#create #1168

Merged
merged 3 commits into from
Apr 20, 2022
Merged
Show file tree
Hide file tree
Changes from 2 commits
Commits
File filter

Filter by extension

Filter by extension

Conversations
Failed to load comments.
Loading
Jump to
Jump to file
Failed to load files.
Loading
Diff view
Diff view
26 changes: 26 additions & 0 deletions google/cloud/aiplatform/models.py
Original file line number Diff line number Diff line change
Expand Up @@ -207,6 +207,7 @@ def create(
encryption_spec_key_name: Optional[str] = None,
sync=True,
create_request_timeout: Optional[float] = None,
endpoint_id: Optional[str] = None,
) -> "Endpoint":
"""Creates a new endpoint.

Expand Down Expand Up @@ -256,6 +257,17 @@ def create(
be immediately returned and synced when the Future has completed.
create_request_timeout (float):
Optional. The timeout for the create request in seconds.
endpoint_id (str):
Optional. The ID to use for endpoint, which will become
the final component of the endpoint resource name. If
not provided, Vertex AI will generate a value for this
ID.

This value should be 1-10 characters, and valid
characters are /[0-9]/. When using HTTP/JSON, this field
is populated based on a query string argument, such as
``?endpoint_id=12345``. This is the fallback for fields
that are not included in either the URI or the body.
Returns:
endpoint (endpoint.Endpoint):
Created endpoint.
Expand Down Expand Up @@ -287,6 +299,7 @@ def create(
),
sync=sync,
create_request_timeout=create_request_timeout,
endpoint_id=endpoint_id,
)

@classmethod
Expand All @@ -304,6 +317,7 @@ def _create(
encryption_spec: Optional[gca_encryption_spec.EncryptionSpec] = None,
sync=True,
create_request_timeout: Optional[float] = None,
endpoint_id: Optional[str] = None,
) -> "Endpoint":
"""Creates a new endpoint by calling the API client.

Expand Down Expand Up @@ -349,6 +363,17 @@ def _create(
Whether to create this endpoint synchronously.
create_request_timeout (float):
Optional. The timeout for the create request in seconds.
endpoint_id (str):
Optional. The ID to use for endpoint, which will become
the final component of the endpoint resource name. If
not provided, Vertex AI will generate a value for this
ID.

This value should be 1-10 characters, and valid
characters are /[0-9]/. When using HTTP/JSON, this field
is populated based on a query string argument, such as
``?endpoint_id=12345``. This is the fallback for fields
that are not included in either the URI or the body.
Returns:
endpoint (endpoint.Endpoint):
Created endpoint.
Expand All @@ -368,6 +393,7 @@ def _create(
operation_future = api_client.create_endpoint(
parent=parent,
endpoint=gapic_endpoint,
endpoint_id=endpoint_id,
metadata=metadata,
timeout=create_request_timeout,
)
Expand Down
30 changes: 30 additions & 0 deletions tests/unit/aiplatform/test_endpoints.py
Original file line number Diff line number Diff line change
Expand Up @@ -547,6 +547,7 @@ def test_init_aiplatform_with_encryption_key_name_and_create_endpoint(
create_endpoint_mock.assert_called_once_with(
parent=_TEST_PARENT,
endpoint=expected_endpoint,
endpoint_id=None,
metadata=(),
timeout=None,
)
Expand All @@ -573,13 +574,39 @@ def test_create(self, create_endpoint_mock, sync):
create_endpoint_mock.assert_called_once_with(
parent=_TEST_PARENT,
endpoint=expected_endpoint,
endpoint_id=None,
metadata=(),
timeout=None,
)

expected_endpoint.name = _TEST_ENDPOINT_NAME
assert my_endpoint._gca_resource == expected_endpoint

@pytest.mark.usefixtures("get_endpoint_mock")
@pytest.mark.parametrize("sync", [True, False])
def test_create_with_endpoint_id(self, create_endpoint_mock, sync):
my_endpoint = models.Endpoint.create(
display_name=_TEST_DISPLAY_NAME,
endpoint_id=_TEST_ID,
description=_TEST_DESCRIPTION,
sync=sync,
create_request_timeout=None,
)
if not sync:
my_endpoint.wait()

expected_endpoint = gca_endpoint.Endpoint(
display_name=_TEST_DISPLAY_NAME,
description=_TEST_DESCRIPTION,
)
create_endpoint_mock.assert_called_once_with(
parent=_TEST_PARENT,
endpoint=expected_endpoint,
endpoint_id=_TEST_ID,
metadata=(),
timeout=None,
)

@pytest.mark.usefixtures("get_endpoint_mock")
@pytest.mark.parametrize("sync", [True, False])
def test_create_with_timeout(self, create_endpoint_mock, sync):
Expand All @@ -599,6 +626,7 @@ def test_create_with_timeout(self, create_endpoint_mock, sync):
create_endpoint_mock.assert_called_once_with(
parent=_TEST_PARENT,
endpoint=expected_endpoint,
endpoint_id=None,
metadata=(),
timeout=180.0,
)
Expand Down Expand Up @@ -642,6 +670,7 @@ def test_create_with_description(self, create_endpoint_mock, sync):
create_endpoint_mock.assert_called_once_with(
parent=_TEST_PARENT,
endpoint=expected_endpoint,
endpoint_id=None,
metadata=(),
timeout=None,
)
Expand All @@ -665,6 +694,7 @@ def test_create_with_labels(self, create_endpoint_mock, sync):
create_endpoint_mock.assert_called_once_with(
parent=_TEST_PARENT,
endpoint=expected_endpoint,
endpoint_id=None,
metadata=(),
timeout=None,
)
Expand Down