feat: Added matching engine integration (#995)

* Added matching engine index

* Added matching engine index endpoint

* Added deploy_index and undeploy_index to MatchingEngineIndexEndpoint

* Added mutate_deployed_index to MatchingEngineIndexEndpoint

* Fixed logging

* Added unit tests

* Removed unneeded parameter

* Fixed tests

* Fixed compat types for matching engine

* Added test_matching_engine_index_endpoint.py

* Initial integration test for matching engine added

* Added metadata parameters to MatchingEngineIndex

* Added unit tests for config

* Fixes to integration test

* Added missing config file

* Cleaned up docstrings and removed immutable update parameters

* Removed metadata_schema_uri

* Various matching engine fixes

* Finished main logic for MatchingEngineIndex

* Fixed deployment test

* Added matching engine deployment and tests

* Added undeploy_all tests

* Fixed e2e matching engine test

* Improved tests and docstrings

* Test cleanup and removal of unused mutate_deployed_index parameters

* fix: Updated docstrings

* Updated copyright

* Improved comments

* Removed index_id paramater

* Fixed update bug

* Removed unneeded id parameters

* Added prediction service

* Update tests

* Ran linter

* Removed enable_private_service_connect parameter as it will be deprecated soon

* Cleaned up tests

* Fixed test

* Added teardown and network code

* Fixed undeploy

* Match tweaks

* Removed batch match test for now

* Added MatchingEngineIndexConfig to init

* Removed unneeded index_id parameter on mutate_deployed_index

* Fixed imports

* Switched to returning MatchNeighbor's

* Import tweaks

* Ran linter

* Cleaned up tests

* Added comments

* Fixed lint issues

* Lint

* Updated copyright

* Added remaining docstrings

* Reverted unwanted change

* Fixed comment

* Added undeploy step for index_endpoint system test

* Added more system tests

* Fixed lint issues

* Added unique id to prevent conflicts

* Addressed PR comments

* Reordered imports

* Fixed init file

* Fix lint and brute_force issue

* Addressed PR comments

* Remove unused parameter

* Fixed system test

* Removed vpc reliant system tests

* Fixed unit tests

* Fixed lint error

* Fixed test

* Fixed lint issues

Author: ivanmkc

google/cloud/aiplatform/__init__.py

@@ -38,6 +38,10 @@
     Feature,
     Featurestore,
 )
+from google.cloud.aiplatform.matching_engine import (
+    MatchingEngineIndex,
+    MatchingEngineIndexEndpoint,
+)
 from google.cloud.aiplatform.metadata import metadata
 from google.cloud.aiplatform.models import Endpoint
 from google.cloud.aiplatform.models import Model
@@ -104,6 +108,8 @@
     "EntityType",
     "Feature",
     "Featurestore",
+    "MatchingEngineIndex",
+    "MatchingEngineIndexEndpoint",
     "ImageDataset",
     "HyperparameterTuningJob",
     "Model",

google/cloud/aiplatform/compat/__init__.py

@@ -76,6 +76,11 @@
     types.job_state = types.job_state_v1beta1
     types.machine_resources = types.machine_resources_v1beta1
     types.manual_batch_tuning_parameters = types.manual_batch_tuning_parameters_v1beta1
+    types.matching_engine_deployed_index_ref = (
+        types.matching_engine_deployed_index_ref_v1beta1
+    )
+    types.matching_engine_index = types.matching_engine_index_v1beta1
+    types.matching_engine_index_endpoint = types.matching_engine_index_endpoint_v1beta1
     types.metadata_service = types.metadata_service_v1beta1
     types.metadata_store = types.metadata_store_v1beta1
     types.model = types.model_v1beta1
@@ -147,6 +152,11 @@
     types.job_state = types.job_state_v1
     types.machine_resources = types.machine_resources_v1
     types.manual_batch_tuning_parameters = types.manual_batch_tuning_parameters_v1
+    types.matching_engine_deployed_index_ref = (
+        types.matching_engine_deployed_index_ref_v1
+    )
+    types.matching_engine_index = types.matching_engine_index_v1
+    types.matching_engine_index_endpoint = types.matching_engine_index_endpoint_v1
     types.metadata_service = types.metadata_service_v1
     types.metadata_store = types.metadata_store_v1
     types.model = types.model_v1

google/cloud/aiplatform/compat/services/__init__.py

@@ -27,6 +27,12 @@
 from google.cloud.aiplatform_v1beta1.services.featurestore_service import (
     client as featurestore_service_client_v1beta1,
 )
+from google.cloud.aiplatform_v1beta1.services.index_service import (
+    client as index_service_client_v1beta1,
+)
+from google.cloud.aiplatform_v1beta1.services.index_endpoint_service import (
+    client as index_endpoint_service_client_v1beta1,
+)
 from google.cloud.aiplatform_v1beta1.services.job_service import (
     client as job_service_client_v1beta1,
 )
@@ -61,6 +67,12 @@
 from google.cloud.aiplatform_v1.services.featurestore_service import (
     client as featurestore_service_client_v1,
 )
+from google.cloud.aiplatform_v1.services.index_service import (
+    client as index_service_client_v1,
+)
+from google.cloud.aiplatform_v1.services.index_endpoint_service import (
+    client as index_endpoint_service_client_v1,
+)
 from google.cloud.aiplatform_v1.services.job_service import (
     client as job_service_client_v1,
 )
@@ -89,6 +101,8 @@
     endpoint_service_client_v1,
     featurestore_online_serving_service_client_v1,
     featurestore_service_client_v1,
+    index_service_client_v1,
+    index_endpoint_service_client_v1,
     job_service_client_v1,
     metadata_service_client_v1,
     model_service_client_v1,
@@ -101,6 +115,8 @@
     endpoint_service_client_v1beta1,
     featurestore_online_serving_service_client_v1beta1,
     featurestore_service_client_v1beta1,
+    index_service_client_v1beta1,
+    index_endpoint_service_client_v1beta1,
     job_service_client_v1beta1,
     model_service_client_v1beta1,
     pipeline_service_client_v1beta1,

google/cloud/aiplatform/compat/types/__init__.py

@@ -28,6 +28,7 @@
     data_labeling_job as data_labeling_job_v1beta1,
     dataset as dataset_v1beta1,
     dataset_service as dataset_service_v1beta1,
+    deployed_index_ref as matching_engine_deployed_index_ref_v1beta1,
     deployed_model_ref as deployed_model_ref_v1beta1,
     encryption_spec as encryption_spec_v1beta1,
     endpoint as endpoint_v1beta1,
@@ -45,6 +46,8 @@
     featurestore_monitoring as featurestore_monitoring_v1beta1,
     featurestore_online_service as featurestore_online_service_v1beta1,
     featurestore_service as featurestore_service_v1beta1,
+    index as matching_engine_index_v1beta1,
+    index_endpoint as matching_engine_index_endpoint_v1beta1,
     hyperparameter_tuning_job as hyperparameter_tuning_job_v1beta1,
     io as io_v1beta1,
     job_service as job_service_v1beta1,
@@ -86,6 +89,7 @@
     data_labeling_job as data_labeling_job_v1,
     dataset as dataset_v1,
     dataset_service as dataset_service_v1,
+    deployed_index_ref as matching_engine_deployed_index_ref_v1,
     deployed_model_ref as deployed_model_ref_v1,
     encryption_spec as encryption_spec_v1,
     endpoint as endpoint_v1,
@@ -103,6 +107,8 @@
     featurestore_online_service as featurestore_online_service_v1,
     featurestore_service as featurestore_service_v1,
     hyperparameter_tuning_job as hyperparameter_tuning_job_v1,
+    index as matching_engine_index_v1,
+    index_endpoint as matching_engine_index_endpoint_v1,
     io as io_v1,
     job_service as job_service_v1,
     job_state as job_state_v1,
@@ -167,6 +173,9 @@
     job_state_v1,
     machine_resources_v1,
     manual_batch_tuning_parameters_v1,
+    matching_engine_deployed_index_ref_v1,
+    matching_engine_index_v1,
+    matching_engine_index_endpoint_v1,
     metadata_service_v1,
     metadata_store_v1,
     model_v1,
@@ -223,6 +232,9 @@
     job_state_v1beta1,
     machine_resources_v1beta1,
     manual_batch_tuning_parameters_v1beta1,
+    matching_engine_deployed_index_ref_v1beta1,
+    matching_engine_index_v1beta1,
+    matching_engine_index_endpoint_v1beta1,
     metadata_service_v1beta1,
     metadata_store_v1beta1,
     model_v1beta1,

google/cloud/aiplatform/matching_engine/__init__.py

@@ -0,0 +1,36 @@
+# -*- coding: utf-8 -*-
+
+# Copyright 2022 Google LLC
+#
+# Licensed under the Apache License, Version 2.0 (the "License");
+# you may not use this file except in compliance with the License.
+# You may obtain a copy of the License at
+#
+#     http://www.apache.org/licenses/LICENSE-2.0
+#
+# Unless required by applicable law or agreed to in writing, software
+# distributed under the License is distributed on an "AS IS" BASIS,
+# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+# See the License for the specific language governing permissions and
+# limitations under the License.
+#
+
+from google.cloud.aiplatform.matching_engine.matching_engine_index import (
+    MatchingEngineIndex,
+)
+from google.cloud.aiplatform.matching_engine.matching_engine_index_config import (
+    BruteForceConfig as MatchingEngineBruteForceAlgorithmConfig,
+    MatchingEngineIndexConfig as MatchingEngineIndexConfig,
+    TreeAhConfig as MatchingEngineTreeAhAlgorithmConfig,
+)
+from google.cloud.aiplatform.matching_engine.matching_engine_index_endpoint import (
+    MatchingEngineIndexEndpoint,
+)
+
+__all__ = (
+    "MatchingEngineIndex",
+    "MatchingEngineIndexEndpoint",
+    "MatchingEngineIndexConfig",
+    "MatchingEngineBruteForceAlgorithmConfig",
+    "MatchingEngineTreeAhAlgorithmConfig",
+)