From 09824d0d0913a27ed9deeedf71fa4fffb1bed8ff Mon Sep 17 00:00:00 2001 From: "Daniel (dB.) Doubrovkine" Date: Tue, 13 Aug 2024 10:58:38 -0400 Subject: [PATCH] Add support for ML neural search. (#504) Signed-off-by: dblock --- .cspell | 6 +- CHANGELOG.md | 3 +- spec/namespaces/ml.yaml | 52 +++++ spec/schemas/_common.mapping.yaml | 36 ++++ spec/schemas/_common.query_dsl.yaml | 28 +++ spec/schemas/ml._common.yaml | 14 ++ .../ingest/pipeline/neural_search.yaml | 194 ++++++++++++++++++ .../text_embedding.yaml} | 2 +- 8 files changed, 331 insertions(+), 4 deletions(-) create mode 100644 tests/default/ingest/pipeline/neural_search.yaml rename tests/default/ingest/{pipeline.yaml => pipeline/text_embedding.yaml} (94%) diff --git a/.cspell b/.cspell index 1fcd3199c..65a801cbf 100644 --- a/.cspell +++ b/.cspell @@ -59,6 +59,7 @@ gsub Gsub haasephonetik heteroscedastic +hnsw homoscedastic hotthreads huggingface @@ -74,6 +75,7 @@ kstem kuromoji Kuromoji languageset +localstats Lovins lucene Lucene @@ -176,6 +178,7 @@ tokenfilters translog Translog tubone +Undeploys unigrams Unmanaged unmatch @@ -186,5 +189,4 @@ urldecode vectory whoamiprotected wordnet -Yrtsd -localstats \ No newline at end of file +Yrtsd \ No newline at end of file diff --git a/CHANGELOG.md b/CHANGELOG.md index 25a085e8a..e1f7db7db 100644 --- a/CHANGELOG.md +++ b/CHANGELOG.md @@ -69,7 +69,8 @@ Inspired from [Keep a Changelog](https://keepachangelog.com/en/1.0.0/) - Added support for reusing output variables as keys in payload expectations ([#471](https://github.com/opensearch-project/opensearch-api-specification/pull/471)) - Added support for running tests against Amazon OpenSearch ([#476](https://github.com/opensearch-project/opensearch-api-specification/pull/476)) - Added API spec for security plugin ([#271](https://github.com/opensearch-project/opensearch-api-specification/pull/271)) -- Added `/_plugins/_security/api/certificates/` to API spec ([#439](https://github.com/opensearch-project/opensearch-api-specification/pull/439)) +- Added `/_plugins/_security/api/certificates/` ([#439](https://github.com/opensearch-project/opensearch-api-specification/pull/439)) +- Added `/_plugins/_ml/models/{model_id}/_deploy`, `_undeploy` and `knn_vector` type in `passage_embedding` ([#504](https://github.com/opensearch-project/opensearch-api-specification/pull/504)) ### Changed diff --git a/spec/namespaces/ml.yaml b/spec/namespaces/ml.yaml index 2822ceb89..00054b4cb 100644 --- a/spec/namespaces/ml.yaml +++ b/spec/namespaces/ml.yaml @@ -53,6 +53,26 @@ paths: responses: '200': $ref: '#/components/responses/ml.delete_model@200' + /_plugins/_ml/models/{model_id}/_deploy: + post: + operationId: ml.deploy_model.0 + x-operation-group: ml.deploy_model + description: Deploys a model. + parameters: + - $ref: '#/components/parameters/ml.deploy_model::path.model_id' + responses: + '200': + $ref: '#/components/responses/ml.deploy_model@200' + /_plugins/_ml/models/{model_id}/_undeploy: + post: + operationId: ml.undeploy_model.0 + x-operation-group: ml.undeploy_model + description: Undeploys a model. + parameters: + - $ref: '#/components/parameters/ml.undeploy_model::path.model_id' + responses: + '200': + $ref: '#/components/responses/ml.undeploy_model@200' /_plugins/_ml/tasks/{task_id}: get: operationId: ml.get_task.0 @@ -161,6 +181,26 @@ components: required: - status - task_id + ml.deploy_model@200: + content: + application/json: + schema: + type: object + properties: + status: + type: string + task_id: + type: string + task_type: + type: string + required: + - status + - task_id + ml.undeploy_model@200: + content: + application/json: + schema: + $ref: '../schemas/ml._common.yaml#/components/schemas/UndeployModelResponse' ml.delete_model@200: content: application/json: @@ -195,6 +235,18 @@ components: required: true schema: type: string + ml.deploy_model::path.model_id: + name: model_id + in: path + required: true + schema: + type: string + ml.undeploy_model::path.model_id: + name: model_id + in: path + required: true + schema: + type: string ml.get_task::path.task_id: name: task_id in: path diff --git a/spec/schemas/_common.mapping.yaml b/spec/schemas/_common.mapping.yaml index cc8e8de0a..ef7a31a31 100644 --- a/spec/schemas/_common.mapping.yaml +++ b/spec/schemas/_common.mapping.yaml @@ -208,6 +208,7 @@ components: - $ref: '#/components/schemas/IntegerRangeProperty' - $ref: '#/components/schemas/IpRangeProperty' - $ref: '#/components/schemas/LongRangeProperty' + - $ref: '#/components/schemas/KnnVectorProperty' BinaryProperty: allOf: - $ref: '#/components/schemas/DocValuesPropertyBase' @@ -1129,6 +1130,30 @@ components: type: boolean index: type: boolean + KnnVectorPropertyBase: + type: object + properties: + dimension: + type: number + method: + $ref: '#/components/schemas/KnnVectorMethod' + required: + - dimension + KnnVectorMethod: + type: object + properties: + name: + type: string + space_type: + type: string + engine: + type: string + parameters: + type: object + additionalProperties: + type: object + required: + - name DoubleRangeProperty: allOf: - $ref: '#/components/schemas/RangePropertyBase' @@ -1184,6 +1209,17 @@ components: - long_range required: - type + KnnVectorProperty: + allOf: + - $ref: '#/components/schemas/KnnVectorPropertyBase' + - type: object + properties: + type: + type: string + enum: + - knn_vector + required: + - type MatchType: type: string enum: diff --git a/spec/schemas/_common.query_dsl.yaml b/spec/schemas/_common.query_dsl.yaml index a0d547cd5..a4bb57b06 100644 --- a/spec/schemas/_common.query_dsl.yaml +++ b/spec/schemas/_common.query_dsl.yaml @@ -112,6 +112,8 @@ components: $ref: '#/components/schemas/MultiMatchQuery' nested: $ref: '#/components/schemas/NestedQuery' + neural: + $ref: '#/components/schemas/NeuralQuery' parent_id: $ref: '#/components/schemas/ParentIdQuery' percolate: @@ -1222,6 +1224,32 @@ components: required: - path - query + NeuralQuery: + allOf: + - $ref: '#/components/schemas/QueryBase' + - type: object + additionalProperties: + $ref: '#/components/schemas/NeuralQueryVectorField' + NeuralQueryVectorField: + type: object + properties: + query_text: + type: string + query_image: + type: string + format: binary + model_id: + type: string + k: + type: integer + min_score: + type: number + max_distance: + type: number + filter: + $ref: '#/components/schemas/QueryContainer' + required: + - model_id ParentIdQuery: allOf: - $ref: '#/components/schemas/QueryBase' diff --git a/spec/schemas/ml._common.yaml b/spec/schemas/ml._common.yaml index 8d1b40d3e..400a92fd3 100644 --- a/spec/schemas/ml._common.yaml +++ b/spec/schemas/ml._common.yaml @@ -139,5 +139,19 @@ components: format: int64 is_async: type: boolean + error: + type: string required: - state + UndeployModelResponse: + type: object + additionalProperties: + $ref: '#/components/schemas/UndeployModelResponseModels' + UndeployModelResponseModels: + type: object + properties: + stats: + $ref: '#/components/schemas/UndeployModelResponseStats' + UndeployModelResponseStats: + type: object + additionalProperties: true diff --git a/tests/default/ingest/pipeline/neural_search.yaml b/tests/default/ingest/pipeline/neural_search.yaml new file mode 100644 index 000000000..0d1d89a22 --- /dev/null +++ b/tests/default/ingest/pipeline/neural_search.yaml @@ -0,0 +1,194 @@ +$schema: ../../../../json_schemas/test_story.schema.yaml + +description: Test the creation a neural search ingest pipeline. +prologues: + - path: /_cluster/settings + method: PUT + request: + payload: + persistent: + plugins: + ml_commons: + only_run_on_ml_node: false +epilogues: + - path: /_ingest/pipeline/movies_pipeline + method: DELETE + status: [200, 404] + - path: /movies + method: DELETE + status: [200, 404] + - path: /_plugins/_ml/models/{model_id}/_undeploy + method: POST + parameters: + model_id: ${get_completed_register_model_task.model_id} + status: [200, 404] + - path: /_plugins/_ml/models/{model_id} + parameters: + model_id: ${get_completed_register_model_task.model_id} + method: DELETE + status: [200, 404] + - path: /_plugins/_ml/model_groups/{model_group_id} + method: DELETE + status: [200, 404] + parameters: + model_group_id: ${create_model_group.test_model_group_id} +version: '>= 2.11' +chapters: + - synopsis: Create model group. + id: create_model_group + path: /_plugins/_ml/model_groups/_register + method: POST + request: + payload: + name: NLP_Group + description: Model group for NLP models. + response: + status: 200 + output: + test_model_group_id: payload.model_group_id + - synopsis: Register model. + id: register_model + path: /_plugins/_ml/models/_register + method: POST + request: + payload: + name: huggingface/sentence-transformers/msmarco-distilbert-base-tas-b + version: 1.0.1 + model_format: TORCH_SCRIPT + response: + status: 200 + output: + task_id: payload.task_id + - synopsis: Wait to get completed task. + id: get_completed_register_model_task + path: /_plugins/_ml/tasks/{task_id} + method: GET + parameters: + task_id: ${register_model.task_id} + response: + status: 200 + payload: + state: COMPLETED + output: + model_id: payload.model_id + retry: + count: 3 + wait: 10000 + - synopsis: Deploy a model. + id: deploy_model + path: /_plugins/_ml/models/{model_id}/_deploy + method: POST + parameters: + model_id: ${get_completed_register_model_task.model_id} + output: + task_id: payload.task_id + response: + status: 200 + - synopsis: Wait to get completed task. + id: get_completed_deploy_model_task + path: /_plugins/_ml/tasks/{task_id} + method: GET + parameters: + task_id: ${deploy_model.task_id} + response: + status: 200 + payload: + state: COMPLETED + output: + model_id: payload.model_id + retry: + count: 3 + wait: 10000 + - synopsis: Create ingest pipeline for text embedding. + path: /_ingest/pipeline/{id} + method: PUT + parameters: + id: movies_pipeline + request: + payload: + description: Extracts text from field and embeds it. + processors: + - text_embedding: + model_id: text-embedding-model + field_map: + text: passage_embedding + response: + status: 200 + payload: + acknowledged: true + - synopsis: Create an index using the pipeline. + path: /{index} + method: PUT + parameters: + index: movies + request: + payload: + settings: + index.knn: true + default_pipeline: movies_pipeline + mappings: + properties: + title: + type: text + year: + type: integer + passage_embedding: + type: knn_vector + dimension: 768 + method: + engine: lucene + space_type: l2 + name: hnsw + parameters: {} + response: + status: 200 + payload: + acknowledged: true + - synopsis: Ingest data. + path: /_bulk + method: POST + parameters: + refresh: 'true' + request: + content_type: application/x-ndjson + payload: + - {create: {_index: movies}} + - {director: Bennett Miller, title: Moneyball, year: 2011} + - {create: {_index: movies}} + - {author: Nicolas Winding Refn, title: Drive, year: 1960} + response: + status: 200 + - synopsis: Search. + path: /{index}/_search + method: POST + parameters: + index: movies + request: + payload: + _source: + excludes: [passage_embedding] + query: + bool: + should: + - script_score: + query: + neural: + passage_embedding: + query_text: Money + model_id: ${get_completed_register_model_task.model_id} + k: 100 + script: + source: _score * 1.5 + - script_score: + query: + match: + title: Moneyball + script: + source: _score * 1.7 + - synopsis: Undeploy a model. + path: /_plugins/_ml/models/{model_id}/_undeploy + method: POST + parameters: + model_id: ${get_completed_register_model_task.model_id} + response: + status: 200 diff --git a/tests/default/ingest/pipeline.yaml b/tests/default/ingest/pipeline/text_embedding.yaml similarity index 94% rename from tests/default/ingest/pipeline.yaml rename to tests/default/ingest/pipeline/text_embedding.yaml index 3145dcb62..6ba46e8e9 100644 --- a/tests/default/ingest/pipeline.yaml +++ b/tests/default/ingest/pipeline/text_embedding.yaml @@ -1,4 +1,4 @@ -$schema: ../../../json_schemas/test_story.schema.yaml +$schema: ../../../../json_schemas/test_story.schema.yaml description: Test the creation of an ingest pipeline with a text embedding processor. epilogues: