Skip to content

Commit

Permalink
feat: Add endpoind_id arg to Endpoint#create (#1168)
Browse files Browse the repository at this point in the history
Thank you for opening a Pull Request! Before submitting your PR, there are a few things you can do to make sure it goes smoothly:
- [x] Make sure to open an issue as a [bug/issue](https://github.com/googleapis/python-aiplatform/issues/new/choose) before writing your code!  That way we can discuss the change, evaluate designs, and agree on the general idea
- [x] Ensure the tests and linter pass
- [x] Code coverage does not decrease (if any source code was changed)
- [x] Appropriate docs were updated (if necessary)

Fixes #1169 🦕
  • Loading branch information
mai-nakagawa authored Apr 20, 2022
1 parent 10f95cd commit 4c21993
Show file tree
Hide file tree
Showing 2 changed files with 56 additions and 0 deletions.
26 changes: 26 additions & 0 deletions google/cloud/aiplatform/models.py
Original file line number Diff line number Diff line change
Expand Up @@ -208,6 +208,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 @@ -257,6 +258,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 @@ -288,6 +300,7 @@ def create(
),
sync=sync,
create_request_timeout=create_request_timeout,
endpoint_id=endpoint_id,
)

@classmethod
Expand All @@ -305,6 +318,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 @@ -350,6 +364,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 @@ -369,6 +394,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

0 comments on commit 4c21993

Please sign in to comment.