From 3be25198a71d6cc3580e17a48f02e10e2a91d9b4 Mon Sep 17 00:00:00 2001 From: Joshua Palis Date: Tue, 13 Aug 2024 21:46:24 +0000 Subject: [PATCH] Adding reprovision integration tests Signed-off-by: Joshua Palis --- .../FlowFrameworkRestTestCase.java | 19 ++ .../rest/FlowFrameworkRestApiIT.java | 206 ++++++++++++++++++ .../registerlocalmodel-createindex.json | 59 +++++ ...localmodel-ingestpipeline-createindex.json | 86 ++++++++ ...localmodel-ingestpipeline-updateindex.json | 86 ++++++++ .../registerlocalmodel-ingestpipeline.json | 52 +++++ .../template/registerlocalmodel.json | 29 +++ 7 files changed, 537 insertions(+) create mode 100644 src/test/resources/template/registerlocalmodel-createindex.json create mode 100644 src/test/resources/template/registerlocalmodel-ingestpipeline-createindex.json create mode 100644 src/test/resources/template/registerlocalmodel-ingestpipeline-updateindex.json create mode 100644 src/test/resources/template/registerlocalmodel-ingestpipeline.json create mode 100644 src/test/resources/template/registerlocalmodel.json diff --git a/src/test/java/org/opensearch/flowframework/FlowFrameworkRestTestCase.java b/src/test/java/org/opensearch/flowframework/FlowFrameworkRestTestCase.java index dc25f44fa..877b6292a 100644 --- a/src/test/java/org/opensearch/flowframework/FlowFrameworkRestTestCase.java +++ b/src/test/java/org/opensearch/flowframework/FlowFrameworkRestTestCase.java @@ -415,6 +415,25 @@ protected Response createWorkflowValidation(RestClient client, Template template return TestHelpers.makeRequest(client, "POST", WORKFLOW_URI, Collections.emptyMap(), template.toJson(), null); } + /** + * Helper method to invoke the Reprovision Workflow API + * @param client the rest client + * @param workflowId the document id + * @param templateFields the template to reprovision + * @throws Exception if the request fails + * @return a rest response + */ + protected Response reprovisionWorkflow(RestClient client, String workflowId, Template template) throws Exception { + return TestHelpers.makeRequest( + client, + "PUT", + String.format(Locale.ROOT, "%s/%s?reprovision=true", WORKFLOW_URI, workflowId), + Collections.emptyMap(), + template.toJson(), + null + ); + } + /** * Helper method to invoke the Update Workflow API * @param client the rest client diff --git a/src/test/java/org/opensearch/flowframework/rest/FlowFrameworkRestApiIT.java b/src/test/java/org/opensearch/flowframework/rest/FlowFrameworkRestApiIT.java index 6a454dc75..747d9eabf 100644 --- a/src/test/java/org/opensearch/flowframework/rest/FlowFrameworkRestApiIT.java +++ b/src/test/java/org/opensearch/flowframework/rest/FlowFrameworkRestApiIT.java @@ -363,6 +363,212 @@ public void testCreateAndProvisionAgentFrameworkWorkflow() throws Exception { assertBusy(() -> { getAndAssertWorkflowStatusNotFound(client(), workflowId); }, 30, TimeUnit.SECONDS); } + public void testReprovisionWorkflow() throws Exception { + // Begin with a template to register a local pretrained model + Template template = TestHelpers.createTemplateFromFile("registerlocalmodel.json"); + + // Hit Create Workflow API to create agent-framework template, with template validation check and provision parameter + Response response; + if (!indexExistsWithAdminClient(".plugins-ml-config")) { + assertBusy(() -> assertTrue(indexExistsWithAdminClient(".plugins-ml-config")), 40, TimeUnit.SECONDS); + response = createWorkflowWithProvision(client(), template); + } else { + response = createWorkflowWithProvision(client(), template); + } + assertEquals(RestStatus.CREATED, TestHelpers.restStatus(response)); + Map responseMap = entityAsMap(response); + String workflowId = (String) responseMap.get(WORKFLOW_ID); + // wait and ensure state is completed/done + assertBusy( + () -> { getAndAssertWorkflowStatus(client(), workflowId, State.COMPLETED, ProvisioningProgress.DONE); }, + 120, + TimeUnit.SECONDS + ); + + // Wait until provisioning has completed successfully before attempting to retrieve created resources + List resourcesCreated = getResourcesCreated(client(), workflowId, 30); + assertEquals(2, resourcesCreated.size()); + assertEquals("register_local_pretrained_model", resourcesCreated.get(0).workflowStepName()); + + // Reprovision template to add ingest pipeline which uses the model ID + template = TestHelpers.createTemplateFromFile("registerlocalmodel-ingestpipeline.json"); + response = reprovisionWorkflow(client(), workflowId, template); + assertEquals(RestStatus.CREATED, TestHelpers.restStatus(response)); + + resourcesCreated = getResourcesCreated(client(), workflowId, 10); + assertEquals(3, resourcesCreated.size()); + List resourceIds = resourcesCreated.stream().map(x -> x.workflowStepName()).collect(Collectors.toList()); + assertTrue(resourceIds.contains("register_local_pretrained_model")); + assertTrue(resourceIds.contains("create_ingest_pipeline")); + + // Retrieve pipeline by ID to ensure model ID is set correctly + String modelId = resourcesCreated.stream() + .filter(x -> x.workflowStepName().equals("register_local_pretrained_model")) + .findFirst() + .get() + .resourceId(); + String pipelineId = resourcesCreated.stream() + .filter(x -> x.workflowStepName().equals("create_ingest_pipeline")) + .findFirst() + .get() + .resourceId(); + GetPipelineResponse getPipelineResponse = getPipelines(pipelineId); + assertEquals(1, getPipelineResponse.pipelines().size()); + assertTrue(getPipelineResponse.pipelines().get(0).getConfigAsMap().toString().contains(modelId)); + + // Reprovision template to add index which uses default ingest pipeline + template = TestHelpers.createTemplateFromFile("registerlocalmodel-ingestpipeline-createindex.json"); + response = reprovisionWorkflow(client(), workflowId, template); + assertEquals(RestStatus.CREATED, TestHelpers.restStatus(response)); + + resourcesCreated = getResourcesCreated(client(), workflowId, 10); + assertEquals(4, resourcesCreated.size()); + resourceIds = resourcesCreated.stream().map(x -> x.workflowStepName()).collect(Collectors.toList()); + assertTrue(resourceIds.contains("register_local_pretrained_model")); + assertTrue(resourceIds.contains("create_ingest_pipeline")); + assertTrue(resourceIds.contains("create_index")); + + // Retrieve index settings to ensure pipeline ID is set correctly + String indexName = resourcesCreated.stream() + .filter(x -> x.workflowStepName().equals("create_index")) + .findFirst() + .get() + .resourceId(); + Map indexSettings = getIndexSettingsAsMap(indexName); + assertEquals(pipelineId, indexSettings.get("index.default_pipeline")); + + // Reprovision template to remove default ingest pipeline + template = TestHelpers.createTemplateFromFile("registerlocalmodel-ingestpipeline-updateindex.json"); + response = reprovisionWorkflow(client(), workflowId, template); + assertEquals(RestStatus.CREATED, TestHelpers.restStatus(response)); + + resourcesCreated = getResourcesCreated(client(), workflowId, 10); + // resource count should remain unchanged when updating an existing node + assertEquals(4, resourcesCreated.size()); + + // Retrieve index settings to ensure default pipeline has been updated correctly + indexSettings = getIndexSettingsAsMap(indexName); + assertEquals("_none", indexSettings.get("index.default_pipeline")); + } + + public void testReprovisionWorkflowMidNodeAddition() throws Exception { + // Begin with a template to register a local pretrained model and create an index, no edges + Template template = TestHelpers.createTemplateFromFile("registerlocalmodel-createindex.json"); + + // Hit Create Workflow API to create agent-framework template, with template validation check and provision parameter + Response response; + if (!indexExistsWithAdminClient(".plugins-ml-config")) { + assertBusy(() -> assertTrue(indexExistsWithAdminClient(".plugins-ml-config")), 40, TimeUnit.SECONDS); + response = createWorkflowWithProvision(client(), template); + } else { + response = createWorkflowWithProvision(client(), template); + } + assertEquals(RestStatus.CREATED, TestHelpers.restStatus(response)); + Map responseMap = entityAsMap(response); + String workflowId = (String) responseMap.get(WORKFLOW_ID); + // wait and ensure state is completed/done + assertBusy( + () -> { getAndAssertWorkflowStatus(client(), workflowId, State.COMPLETED, ProvisioningProgress.DONE); }, + 120, + TimeUnit.SECONDS + ); + + // Wait until provisioning has completed successfully before attempting to retrieve created resources + List resourcesCreated = getResourcesCreated(client(), workflowId, 30); + assertEquals(3, resourcesCreated.size()); + List resourceIds = resourcesCreated.stream().map(x -> x.workflowStepName()).collect(Collectors.toList()); + assertTrue(resourceIds.contains("register_local_pretrained_model")); + assertTrue(resourceIds.contains("create_index")); + + // Reprovision template to add ingest pipeline which uses the model ID + template = TestHelpers.createTemplateFromFile("registerlocalmodel-ingestpipeline-createindex.json"); + response = reprovisionWorkflow(client(), workflowId, template); + assertEquals(RestStatus.CREATED, TestHelpers.restStatus(response)); + + resourcesCreated = getResourcesCreated(client(), workflowId, 10); + assertEquals(4, resourcesCreated.size()); + resourceIds = resourcesCreated.stream().map(x -> x.workflowStepName()).collect(Collectors.toList()); + assertTrue(resourceIds.contains("register_local_pretrained_model")); + assertTrue(resourceIds.contains("create_ingest_pipeline")); + assertTrue(resourceIds.contains("create_index")); + + // Ensure ingest pipeline configuration contains the model id and index settings have the ingest pipeline as default + String modelId = resourcesCreated.stream() + .filter(x -> x.workflowStepName().equals("register_local_pretrained_model")) + .findFirst() + .get() + .resourceId(); + String pipelineId = resourcesCreated.stream() + .filter(x -> x.workflowStepName().equals("create_ingest_pipeline")) + .findFirst() + .get() + .resourceId(); + GetPipelineResponse getPipelineResponse = getPipelines(pipelineId); + assertEquals(1, getPipelineResponse.pipelines().size()); + assertTrue(getPipelineResponse.pipelines().get(0).getConfigAsMap().toString().contains(modelId)); + + String indexName = resourcesCreated.stream() + .filter(x -> x.workflowStepName().equals("create_index")) + .findFirst() + .get() + .resourceId(); + Map indexSettings = getIndexSettingsAsMap(indexName); + assertEquals(pipelineId, indexSettings.get("index.default_pipeline")); + } + + public void testReprovisionWithNoChange() throws Exception { + Template template = TestHelpers.createTemplateFromFile("registerlocalmodel-ingestpipeline-createindex.json"); + + Response response; + if (!indexExistsWithAdminClient(".plugins-ml-config")) { + assertBusy(() -> assertTrue(indexExistsWithAdminClient(".plugins-ml-config")), 40, TimeUnit.SECONDS); + response = createWorkflowWithProvision(client(), template); + } else { + response = createWorkflowWithProvision(client(), template); + } + assertEquals(RestStatus.CREATED, TestHelpers.restStatus(response)); + Map responseMap = entityAsMap(response); + String workflowId = (String) responseMap.get(WORKFLOW_ID); + // wait and ensure state is completed/done + assertBusy( + () -> { getAndAssertWorkflowStatus(client(), workflowId, State.COMPLETED, ProvisioningProgress.DONE); }, + 120, + TimeUnit.SECONDS + ); + + // Attempt to reprovision the same template with no changes + ResponseException exception = expectThrows(ResponseException.class, () -> reprovisionWorkflow(client(), workflowId, template)); + assertEquals(RestStatus.BAD_REQUEST.getStatus(), exception.getResponse().getStatusLine().getStatusCode()); + assertTrue(exception.getMessage().contains("Template does not contain any modifications")); + } + + public void testReprovisionWithDeletion() throws Exception { + Template template = TestHelpers.createTemplateFromFile("registerlocalmodel-ingestpipeline-createindex.json"); + + Response response; + if (!indexExistsWithAdminClient(".plugins-ml-config")) { + assertBusy(() -> assertTrue(indexExistsWithAdminClient(".plugins-ml-config")), 40, TimeUnit.SECONDS); + response = createWorkflowWithProvision(client(), template); + } else { + response = createWorkflowWithProvision(client(), template); + } + assertEquals(RestStatus.CREATED, TestHelpers.restStatus(response)); + Map responseMap = entityAsMap(response); + String workflowId = (String) responseMap.get(WORKFLOW_ID); + // wait and ensure state is completed/done + assertBusy( + () -> { getAndAssertWorkflowStatus(client(), workflowId, State.COMPLETED, ProvisioningProgress.DONE); }, + 120, + TimeUnit.SECONDS + ); + + // Attempt to reprovision template without ingest pipeline node + Template templateWithoutIngestPipeline = TestHelpers.createTemplateFromFile("registerlocalmodel-createindex.json"); + ResponseException exception = expectThrows(ResponseException.class, () -> reprovisionWorkflow(client(), workflowId, templateWithoutIngestPipeline)); + assertEquals(RestStatus.BAD_REQUEST.getStatus(), exception.getResponse().getStatusLine().getStatusCode()); + assertTrue(exception.getMessage().contains("Workflow Step deletion is not supported when reprovisioning a template.")); + } + public void testTimestamps() throws Exception { Template noopTemplate = TestHelpers.createTemplateFromFile("noop.json"); // Create the template, should have created and updated matching diff --git a/src/test/resources/template/registerlocalmodel-createindex.json b/src/test/resources/template/registerlocalmodel-createindex.json new file mode 100644 index 000000000..8dad29432 --- /dev/null +++ b/src/test/resources/template/registerlocalmodel-createindex.json @@ -0,0 +1,59 @@ +{ + "name": "semantic search with local pretrained model", + "description": "Setting up semantic search, with a local pretrained embedding model", + "use_case": "SEMANTIC_SEARCH", + "version": { + "template": "1.0.0", + "compatibility": [ + "2.12.0", + "3.0.0" + ] + }, + "workflows": { + "provision": { + "nodes": [ + { + "id": "register_local_pretrained_model", + "type": "register_local_pretrained_model", + "user_inputs": { + "name": "huggingface/sentence-transformers/paraphrase-MiniLM-L3-v2", + "version": "1.0.1", + "description": "This is a sentence transformer model", + "model_format": "TORCH_SCRIPT", + "deploy": true + } + }, + { + "id": "create_index", + "type": "create_index", + "user_inputs": { + "index_name": "my-nlp-index", + "configurations": { + "settings": { + "index.knn": true, + "index.number_of_shards": "2" + }, + "mappings": { + "properties": { + "passage_embedding": { + "type": "knn_vector", + "dimension": "768", + "method": { + "engine": "lucene", + "space_type": "l2", + "name": "hnsw", + "parameters": {} + } + }, + "passage_text": { + "type": "text" + } + } + } + } + } + } + ] + } + } + } diff --git a/src/test/resources/template/registerlocalmodel-ingestpipeline-createindex.json b/src/test/resources/template/registerlocalmodel-ingestpipeline-createindex.json new file mode 100644 index 000000000..b60da54a2 --- /dev/null +++ b/src/test/resources/template/registerlocalmodel-ingestpipeline-createindex.json @@ -0,0 +1,86 @@ +{ + "name": "semantic search with local pretrained model", + "description": "Setting up semantic search, with a local pretrained embedding model", + "use_case": "SEMANTIC_SEARCH", + "version": { + "template": "1.0.0", + "compatibility": [ + "2.12.0", + "3.0.0" + ] + }, + "workflows": { + "provision": { + "nodes": [ + { + "id": "register_local_pretrained_model", + "type": "register_local_pretrained_model", + "user_inputs": { + "name": "huggingface/sentence-transformers/paraphrase-MiniLM-L3-v2", + "version": "1.0.1", + "description": "This is a sentence transformer model", + "model_format": "TORCH_SCRIPT", + "deploy": true + } + }, + { + "id": "create_ingest_pipeline", + "type": "create_ingest_pipeline", + "previous_node_inputs": { + "register_local_pretrained_model": "model_id" + }, + "user_inputs": { + "pipeline_id": "nlp-ingest-pipeline", + "configurations": { + "description": "A text embedding pipeline", + "processors": [ + { + "text_embedding": { + "model_id": "${{register_local_pretrained_model.model_id}}", + "field_map": { + "passage_text": "passage_embedding" + } + } + } + ] + } + } + }, + { + "id": "create_index", + "type": "create_index", + "previous_node_inputs": { + "create_ingest_pipeline": "pipeline_id" + }, + "user_inputs": { + "index_name": "my-nlp-index", + "configurations": { + "settings": { + "index.knn": true, + "default_pipeline": "${{create_ingest_pipeline.pipeline_id}}", + "index.number_of_shards": "2" + }, + "mappings": { + "properties": { + "passage_embedding": { + "type": "knn_vector", + "dimension": "768", + "method": { + "engine": "lucene", + "space_type": "l2", + "name": "hnsw", + "parameters": {} + } + }, + "passage_text": { + "type": "text" + } + } + } + } + } + } + ] + } + } + } diff --git a/src/test/resources/template/registerlocalmodel-ingestpipeline-updateindex.json b/src/test/resources/template/registerlocalmodel-ingestpipeline-updateindex.json new file mode 100644 index 000000000..5303af938 --- /dev/null +++ b/src/test/resources/template/registerlocalmodel-ingestpipeline-updateindex.json @@ -0,0 +1,86 @@ +{ + "name": "semantic search with local pretrained model", + "description": "Setting up semantic search, with a local pretrained embedding model", + "use_case": "SEMANTIC_SEARCH", + "version": { + "template": "1.0.0", + "compatibility": [ + "2.12.0", + "3.0.0" + ] + }, + "workflows": { + "provision": { + "nodes": [ + { + "id": "register_local_pretrained_model", + "type": "register_local_pretrained_model", + "user_inputs": { + "name": "huggingface/sentence-transformers/paraphrase-MiniLM-L3-v2", + "version": "1.0.1", + "description": "This is a sentence transformer model", + "model_format": "TORCH_SCRIPT", + "deploy": true + } + }, + { + "id": "create_ingest_pipeline", + "type": "create_ingest_pipeline", + "previous_node_inputs": { + "register_local_pretrained_model": "model_id" + }, + "user_inputs": { + "pipeline_id": "nlp-ingest-pipeline", + "configurations": { + "description": "A text embedding pipeline", + "processors": [ + { + "text_embedding": { + "model_id": "${{register_local_pretrained_model.model_id}}", + "field_map": { + "passage_text": "passage_embedding" + } + } + } + ] + } + } + }, + { + "id": "create_index", + "type": "create_index", + "previous_node_inputs": { + "create_ingest_pipeline": "pipeline_id" + }, + "user_inputs": { + "index_name": "my-nlp-index", + "configurations": { + "settings": { + "index.knn": true, + "default_pipeline": "_none", + "index.number_of_shards": "2" + }, + "mappings": { + "properties": { + "passage_embedding": { + "type": "knn_vector", + "dimension": "768", + "method": { + "engine": "lucene", + "space_type": "l2", + "name": "hnsw", + "parameters": {} + } + }, + "passage_text": { + "type": "text" + } + } + } + } + } + } + ] + } + } + } diff --git a/src/test/resources/template/registerlocalmodel-ingestpipeline.json b/src/test/resources/template/registerlocalmodel-ingestpipeline.json new file mode 100644 index 000000000..a4ceab116 --- /dev/null +++ b/src/test/resources/template/registerlocalmodel-ingestpipeline.json @@ -0,0 +1,52 @@ +{ + "name": "semantic search with local pretrained model", + "description": "Setting up semantic search, with a local pretrained embedding model", + "use_case": "SEMANTIC_SEARCH", + "version": { + "template": "1.0.0", + "compatibility": [ + "2.12.0", + "3.0.0" + ] + }, + "workflows": { + "provision": { + "nodes": [ + { + "id": "register_local_pretrained_model", + "type": "register_local_pretrained_model", + "user_inputs": { + "name": "huggingface/sentence-transformers/paraphrase-MiniLM-L3-v2", + "version": "1.0.1", + "description": "This is a sentence transformer model", + "model_format": "TORCH_SCRIPT", + "deploy": true + } + }, + { + "id": "create_ingest_pipeline", + "type": "create_ingest_pipeline", + "previous_node_inputs": { + "register_local_pretrained_model": "model_id" + }, + "user_inputs": { + "pipeline_id": "nlp-ingest-pipeline", + "configurations": { + "description": "A text embedding pipeline", + "processors": [ + { + "text_embedding": { + "model_id": "${{register_local_pretrained_model.model_id}}", + "field_map": { + "passage_text": "passage_embedding" + } + } + } + ] + } + } + } + ] + } + } + } diff --git a/src/test/resources/template/registerlocalmodel.json b/src/test/resources/template/registerlocalmodel.json new file mode 100644 index 000000000..8394f76ab --- /dev/null +++ b/src/test/resources/template/registerlocalmodel.json @@ -0,0 +1,29 @@ +{ + "name": "semantic search with local pretrained model", + "description": "Setting up semantic search, with a local pretrained embedding model", + "use_case": "SEMANTIC_SEARCH", + "version": { + "template": "1.0.0", + "compatibility": [ + "2.12.0", + "3.0.0" + ] + }, + "workflows": { + "provision": { + "nodes": [ + { + "id": "register_local_pretrained_model", + "type": "register_local_pretrained_model", + "user_inputs": { + "name": "huggingface/sentence-transformers/paraphrase-MiniLM-L3-v2", + "version": "1.0.1", + "description": "This is a sentence transformer model", + "model_format": "TORCH_SCRIPT", + "deploy": true + } + } + ] + } + } + }