feat: Add matching engine deployment tier parameter and new shard size

PiperOrigin-RevId: 817012629

Author: vertex-sdk-bot

google/cloud/aiplatform/matching_engine/matching_engine_index.py

@@ -588,6 +588,7 @@ def create_tree_ah_index(
                     SHARD_SIZE_SMALL
                     SHARD_SIZE_MEDIUM
                     SHARD_SIZE_LARGE
+                    SHARD_SIZE_SO_DYNAMIC
 
 
         Returns:
@@ -740,6 +741,7 @@ def create_brute_force_index(
                     SHARD_SIZE_SMALL
                     SHARD_SIZE_MEDIUM
                     SHARD_SIZE_LARGE
+                    SHARD_SIZE_SO_DYNAMIC
 
         Returns:
             MatchingEngineIndex - Index resource object

google/cloud/aiplatform/matching_engine/matching_engine_index_endpoint.py

@@ -944,6 +944,7 @@ def _build_deployed_index(
         auth_config_audiences: Optional[Sequence[str]] = None,
         auth_config_allowed_issuers: Optional[Sequence[str]] = None,
         psc_automation_configs: Optional[Sequence[Tuple[str, str]]] = None,
+        deployment_tier: Optional[str] = None,
     ) -> gca_matching_engine_index_endpoint.DeployedIndex:
         """Builds a DeployedIndex.
 
@@ -1046,6 +1047,8 @@ def _build_deployed_index(
                 projects/{project}/global/networks/{network}, where
                 {project} is a project number, as in '12345', and {network}
                 is network name.
+            deployment_tier (str):
+                Optional. The deployment tier that the index is deployed to.
 
         """
 
@@ -1056,6 +1059,7 @@ def _build_deployed_index(
             enable_access_logging=enable_access_logging,
             reserved_ip_ranges=reserved_ip_ranges,
             deployment_group=deployment_group,
+            deployment_tier=deployment_tier,
         )
 
         if auth_config_audiences and auth_config_allowed_issuers:
@@ -1115,6 +1119,7 @@ def deploy_index(
         sync: bool = True,
         deploy_request_timeout: Optional[float] = None,
         psc_automation_configs: Optional[Sequence[Tuple[str, str]]] = None,
+        deployment_tier: Optional[str] = None,
     ) -> "MatchingEngineIndexEndpoint":
         """Deploys an existing index resource to this endpoint resource.
 
@@ -1231,6 +1236,8 @@ def deploy_index(
                 [(project_id_1, network_1), (project_id_1, network_2))] will enable
                 PSC automation for the index to be deployed to project_id_1's network_1
                 and network_2 and can be queried within these networks.
+            deployment_tier (str):
+                Optional. The deployment tier that the index is deployed to.
         Returns:
             MatchingEngineIndexEndpoint - IndexEndpoint resource object
         """
@@ -1250,6 +1257,7 @@ def deploy_index(
             sync=sync,
             deploy_request_timeout=deploy_request_timeout,
             psc_automation_configs=psc_automation_configs,
+            deployment_tier=deployment_tier,
         )
 
     @base.optional_sync(return_input_arg="self")
@@ -1270,6 +1278,7 @@ def _deploy_index(
         sync: bool = True,
         deploy_request_timeout: Optional[float] = None,
         psc_automation_configs: Optional[Sequence[Tuple[str, str]]] = None,
+        deployment_tier: Optional[str] = None,
     ) -> "MatchingEngineIndexEndpoint":
         """Helper method to deploy an existing index resource to this endpoint resource.
 
@@ -1386,6 +1395,8 @@ def _deploy_index(
                 [(project_id_1, network_1), (project_id_1, network_2))] will enable
                 PSC automation for the index to be deployed to project_id_1's network_1
                 and network_2 and can be queried within these networks.
+            deployment_tier (str):
+                Optional. The deployment tier that the index is deployed to.
         Returns:
             MatchingEngineIndexEndpoint - IndexEndpoint resource object
         """
@@ -1411,6 +1422,7 @@ def _deploy_index(
             auth_config_audiences=auth_config_audiences,
             auth_config_allowed_issuers=auth_config_allowed_issuers,
             psc_automation_configs=psc_automation_configs,
+            deployment_tier=deployment_tier,
         )
 
         deploy_lro = self.api_client.deploy_index(

tests/unit/aiplatform/test_matching_engine_index.py

@@ -77,7 +77,12 @@
 _TEST_INDEX_APPROXIMATE_NEIGHBORS_COUNT = 150
 _TEST_LEAF_NODE_EMBEDDING_COUNT = 123
 _TEST_LEAF_NODES_TO_SEARCH_PERCENT = 50
-_TEST_SHARD_SIZES = ["SHARD_SIZE_SMALL", "SHARD_SIZE_LARGE", "SHARD_SIZE_MEDIUM"]
+_TEST_SHARD_SIZES = [
+    "SHARD_SIZE_SMALL",
+    "SHARD_SIZE_LARGE",
+    "SHARD_SIZE_MEDIUM",
+    "SHARD_SIZE_SO_DYNAMIC",
+]
 
 _TEST_INDEX_DESCRIPTION = test_constants.MatchingEngineConstants._TEST_INDEX_DESCRIPTION
 

tests/unit/aiplatform/test_matching_engine_index_endpoint.py

@@ -103,6 +103,7 @@
     "service-account-name-1@project-id.iam.gserviceaccount.com",
     "service-account-name-2@project-id.iam.gserviceaccount.com",
 ]
+_TEST_DEPLOYMENT_TIER = "STORAGE"
 _TEST_SIGNED_JWT = "signed_jwt"
 _TEST_AUTHORIZATION_METADATA = (("authorization", f"Bearer: {_TEST_SIGNED_JWT}"),)
 
@@ -1324,6 +1325,47 @@ def test_deploy_index_psc_automation_configs(self, deploy_index_mock, sync):
             timeout=_TEST_TIMEOUT,
         )
 
+    @pytest.mark.usefixtures("get_index_endpoint_mock", "get_index_mock")
+    @pytest.mark.parametrize("sync", [True, False])
+    def test_deploy_index_deployment_tier(self, deploy_index_mock, sync):
+        aiplatform.init(project=_TEST_PROJECT)
+
+        my_index_endpoint = aiplatform.MatchingEngineIndexEndpoint(
+            index_endpoint_name=_TEST_INDEX_ENDPOINT_ID
+        )
+
+        # Get index
+        my_index = aiplatform.MatchingEngineIndex(index_name=_TEST_INDEX_NAME)
+
+        my_index_endpoint = my_index_endpoint.deploy_index(
+            index=my_index,
+            deployed_index_id=_TEST_DEPLOYED_INDEX_ID,
+            display_name=_TEST_DEPLOYED_INDEX_DISPLAY_NAME,
+            deployment_tier=_TEST_DEPLOYMENT_TIER,
+            request_metadata=_TEST_REQUEST_METADATA,
+            sync=sync,
+            deploy_request_timeout=_TEST_TIMEOUT,
+        )
+
+        if not sync:
+            my_index_endpoint.wait()
+
+        deploy_index_mock.assert_called_once_with(
+            index_endpoint=my_index_endpoint.resource_name,
+            deployed_index=gca_index_endpoint.DeployedIndex(
+                id=_TEST_DEPLOYED_INDEX_ID,
+                index=my_index.resource_name,
+                display_name=_TEST_DEPLOYED_INDEX_DISPLAY_NAME,
+                automatic_resources={
+                    "min_replica_count": None,
+                    "max_replica_count": None,
+                },
+                deployment_tier=_TEST_DEPLOYMENT_TIER,
+            ),
+            metadata=_TEST_REQUEST_METADATA,
+            timeout=_TEST_TIMEOUT,
+        )
+
     @pytest.mark.usefixtures("get_index_endpoint_mock", "get_index_mock")
     def test_mutate_deployed_index(self, mutate_deployed_index_mock):
         aiplatform.init(project=_TEST_PROJECT)