feat: Added matching engine integration

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",
+)