From f0e7630118391403fea43179688f30b210f1997f Mon Sep 17 00:00:00 2001 From: Owais Kazi Date: Fri, 5 Jan 2024 14:43:04 -0800 Subject: [PATCH 1/2] Added a new workflow step to delete model group Signed-off-by: Owais Kazi --- .../common/WorkflowResources.java | 3 +- .../workflow/DeleteModelGroupStep.java | 40 +++++++++++++++++++ .../workflow/ModelGroupStep.java | 2 +- .../workflow/WorkflowStepFactory.java | 1 + .../resources/mappings/workflow-steps.json | 7 +++- .../workflow/DeleteModelGroupTests.java | 31 ++++++++++++++ 6 files changed, 81 insertions(+), 3 deletions(-) create mode 100644 src/main/java/org/opensearch/flowframework/workflow/DeleteModelGroupStep.java create mode 100644 src/test/java/org/opensearch/flowframework/workflow/DeleteModelGroupTests.java diff --git a/src/main/java/org/opensearch/flowframework/common/WorkflowResources.java b/src/main/java/org/opensearch/flowframework/common/WorkflowResources.java index 329cb6ce4..870489351 100644 --- a/src/main/java/org/opensearch/flowframework/common/WorkflowResources.java +++ b/src/main/java/org/opensearch/flowframework/common/WorkflowResources.java @@ -17,6 +17,7 @@ import org.opensearch.flowframework.workflow.CreateIngestPipelineStep; import org.opensearch.flowframework.workflow.DeleteAgentStep; import org.opensearch.flowframework.workflow.DeleteConnectorStep; +import org.opensearch.flowframework.workflow.DeleteModelGroupStep; import org.opensearch.flowframework.workflow.DeleteModelStep; import org.opensearch.flowframework.workflow.DeployModelStep; import org.opensearch.flowframework.workflow.ModelGroupStep; @@ -41,7 +42,7 @@ public enum WorkflowResources { /** Workflow steps for registering/deleting a local model and associated created resource */ REGISTER_LOCAL_MODEL(RegisterLocalModelStep.NAME, WorkflowResources.MODEL_ID, DeleteModelStep.NAME), /** Workflow steps for registering a model group and associated created resource */ - REGISTER_MODEL_GROUP(ModelGroupStep.NAME, WorkflowResources.MODEL_GROUP_ID, null), // TODO delete step + REGISTER_MODEL_GROUP(ModelGroupStep.NAME, WorkflowResources.MODEL_GROUP_ID, DeleteModelGroupStep.NAME), /** Workflow steps for deploying/undeploying a model and associated created resource */ DEPLOY_MODEL(DeployModelStep.NAME, WorkflowResources.MODEL_ID, UndeployModelStep.NAME), /** Workflow steps for creating an ingest-pipeline and associated created resource */ diff --git a/src/main/java/org/opensearch/flowframework/workflow/DeleteModelGroupStep.java b/src/main/java/org/opensearch/flowframework/workflow/DeleteModelGroupStep.java new file mode 100644 index 000000000..5edea653f --- /dev/null +++ b/src/main/java/org/opensearch/flowframework/workflow/DeleteModelGroupStep.java @@ -0,0 +1,40 @@ +/* + * Copyright OpenSearch Contributors + * SPDX-License-Identifier: Apache-2.0 + * + * The OpenSearch Contributors require contributions made to + * this file be licensed under the Apache-2.0 license or a + * compatible open source license. + */ +package org.opensearch.flowframework.workflow; + +import java.util.Map; +import java.util.concurrent.CompletableFuture; + +/** + * Step to delete a model group + */ +public class DeleteModelGroupStep implements WorkflowStep { + + /** Instantiate this class */ + public DeleteModelGroupStep() {} + + /** The name of this step, used as a key in the template and the {@link WorkflowStepFactory} */ + public static final String NAME = "delete_model_group"; + + @Override + public CompletableFuture execute( + String currentNodeId, + WorkflowData currentNodeInputs, + Map outputs, + Map previousNodeInputs + ) { + return CompletableFuture.completedFuture(WorkflowData.EMPTY); + } + + @Override + public String getName() { + return NAME; + } + +} diff --git a/src/main/java/org/opensearch/flowframework/workflow/ModelGroupStep.java b/src/main/java/org/opensearch/flowframework/workflow/ModelGroupStep.java index 71408fd77..4192bb8f7 100644 --- a/src/main/java/org/opensearch/flowframework/workflow/ModelGroupStep.java +++ b/src/main/java/org/opensearch/flowframework/workflow/ModelGroupStep.java @@ -48,7 +48,7 @@ public class ModelGroupStep implements WorkflowStep { private final FlowFrameworkIndicesHandler flowFrameworkIndicesHandler; /** The name of this step, used as a key in the template and the {@link WorkflowStepFactory} */ - public static final String NAME = "register_model_group"; + public static final String NAME = "model_group"; /** * Instantiate this class diff --git a/src/main/java/org/opensearch/flowframework/workflow/WorkflowStepFactory.java b/src/main/java/org/opensearch/flowframework/workflow/WorkflowStepFactory.java index 839b8e9b8..08ccfc2a8 100644 --- a/src/main/java/org/opensearch/flowframework/workflow/WorkflowStepFactory.java +++ b/src/main/java/org/opensearch/flowframework/workflow/WorkflowStepFactory.java @@ -63,6 +63,7 @@ public WorkflowStepFactory( stepMap.put(CreateConnectorStep.NAME, () -> new CreateConnectorStep(mlClient, flowFrameworkIndicesHandler)); stepMap.put(DeleteConnectorStep.NAME, () -> new DeleteConnectorStep(mlClient)); stepMap.put(ModelGroupStep.NAME, () -> new ModelGroupStep(mlClient, flowFrameworkIndicesHandler)); + stepMap.put(DeleteModelGroupStep.NAME, DeleteModelGroupStep::new); stepMap.put(ToolStep.NAME, ToolStep::new); stepMap.put(RegisterAgentStep.NAME, () -> new RegisterAgentStep(mlClient, flowFrameworkIndicesHandler)); stepMap.put(DeleteAgentStep.NAME, () -> new DeleteAgentStep(mlClient)); diff --git a/src/main/resources/mappings/workflow-steps.json b/src/main/resources/mappings/workflow-steps.json index 6431b2aa6..0fca9055f 100644 --- a/src/main/resources/mappings/workflow-steps.json +++ b/src/main/resources/mappings/workflow-steps.json @@ -123,7 +123,7 @@ "opensearch-ml" ] }, - "register_model_group": { + "model_group": { "inputs":[ "name" ], @@ -135,6 +135,11 @@ "opensearch-ml" ] }, + "delete_model_group": { + "inputs":[], + "outputs":[], + "required_plugins":[] + }, "register_agent": { "inputs":[ "name", diff --git a/src/test/java/org/opensearch/flowframework/workflow/DeleteModelGroupTests.java b/src/test/java/org/opensearch/flowframework/workflow/DeleteModelGroupTests.java new file mode 100644 index 000000000..e55eaea70 --- /dev/null +++ b/src/test/java/org/opensearch/flowframework/workflow/DeleteModelGroupTests.java @@ -0,0 +1,31 @@ +/* + * Copyright OpenSearch Contributors + * SPDX-License-Identifier: Apache-2.0 + * + * The OpenSearch Contributors require contributions made to + * this file be licensed under the Apache-2.0 license or a + * compatible open source license. + */ +package org.opensearch.flowframework.workflow; + +import org.opensearch.test.OpenSearchTestCase; + +import java.io.IOException; +import java.util.Collections; +import java.util.concurrent.CompletableFuture; + +public class DeleteModelGroupTests extends OpenSearchTestCase { + + public void testDeleteModelGroup() throws IOException { + DeleteModelGroupStep deleteModelGroupStep = new DeleteModelGroupStep(); + assertEquals(DeleteModelGroupStep.NAME, deleteModelGroupStep.getName()); + CompletableFuture future = deleteModelGroupStep.execute( + "nodeId", + WorkflowData.EMPTY, + Collections.emptyMap(), + Collections.emptyMap() + ); + assertTrue(future.isDone()); + assertFalse(future.isCompletedExceptionally()); + } +} From 7f76221bbb065dffd4357ab8dddea8f3fd371874 Mon Sep 17 00:00:00 2001 From: Owais Kazi Date: Fri, 5 Jan 2024 15:08:52 -0800 Subject: [PATCH 2/2] Renamed file Signed-off-by: Owais Kazi --- .../flowframework/common/WorkflowResources.java | 6 +++--- .../{ModelGroupStep.java => RegisterModelGroupStep.java} | 8 ++++---- .../flowframework/workflow/WorkflowStepFactory.java | 2 +- src/main/resources/mappings/workflow-steps.json | 2 +- .../flowframework/workflow/ModelGroupStepTests.java | 6 +++--- 5 files changed, 12 insertions(+), 12 deletions(-) rename src/main/java/org/opensearch/flowframework/workflow/{ModelGroupStep.java => RegisterModelGroupStep.java} (95%) diff --git a/src/main/java/org/opensearch/flowframework/common/WorkflowResources.java b/src/main/java/org/opensearch/flowframework/common/WorkflowResources.java index 870489351..4245bc3a9 100644 --- a/src/main/java/org/opensearch/flowframework/common/WorkflowResources.java +++ b/src/main/java/org/opensearch/flowframework/common/WorkflowResources.java @@ -20,9 +20,9 @@ import org.opensearch.flowframework.workflow.DeleteModelGroupStep; import org.opensearch.flowframework.workflow.DeleteModelStep; import org.opensearch.flowframework.workflow.DeployModelStep; -import org.opensearch.flowframework.workflow.ModelGroupStep; import org.opensearch.flowframework.workflow.RegisterAgentStep; import org.opensearch.flowframework.workflow.RegisterLocalModelStep; +import org.opensearch.flowframework.workflow.RegisterModelGroupStep; import org.opensearch.flowframework.workflow.RegisterRemoteModelStep; import org.opensearch.flowframework.workflow.UndeployModelStep; @@ -41,8 +41,8 @@ public enum WorkflowResources { REGISTER_REMOTE_MODEL(RegisterRemoteModelStep.NAME, WorkflowResources.MODEL_ID, DeleteModelStep.NAME), /** Workflow steps for registering/deleting a local model and associated created resource */ REGISTER_LOCAL_MODEL(RegisterLocalModelStep.NAME, WorkflowResources.MODEL_ID, DeleteModelStep.NAME), - /** Workflow steps for registering a model group and associated created resource */ - REGISTER_MODEL_GROUP(ModelGroupStep.NAME, WorkflowResources.MODEL_GROUP_ID, DeleteModelGroupStep.NAME), + /** Workflow steps for registering/deleting a model group and associated created resource */ + REGISTER_MODEL_GROUP(RegisterModelGroupStep.NAME, WorkflowResources.MODEL_GROUP_ID, DeleteModelGroupStep.NAME), /** Workflow steps for deploying/undeploying a model and associated created resource */ DEPLOY_MODEL(DeployModelStep.NAME, WorkflowResources.MODEL_ID, UndeployModelStep.NAME), /** Workflow steps for creating an ingest-pipeline and associated created resource */ diff --git a/src/main/java/org/opensearch/flowframework/workflow/ModelGroupStep.java b/src/main/java/org/opensearch/flowframework/workflow/RegisterModelGroupStep.java similarity index 95% rename from src/main/java/org/opensearch/flowframework/workflow/ModelGroupStep.java rename to src/main/java/org/opensearch/flowframework/workflow/RegisterModelGroupStep.java index 4192bb8f7..518e2879e 100644 --- a/src/main/java/org/opensearch/flowframework/workflow/ModelGroupStep.java +++ b/src/main/java/org/opensearch/flowframework/workflow/RegisterModelGroupStep.java @@ -39,23 +39,23 @@ /** * Step to register a model group */ -public class ModelGroupStep implements WorkflowStep { +public class RegisterModelGroupStep implements WorkflowStep { - private static final Logger logger = LogManager.getLogger(ModelGroupStep.class); + private static final Logger logger = LogManager.getLogger(RegisterModelGroupStep.class); private final MachineLearningNodeClient mlClient; private final FlowFrameworkIndicesHandler flowFrameworkIndicesHandler; /** The name of this step, used as a key in the template and the {@link WorkflowStepFactory} */ - public static final String NAME = "model_group"; + public static final String NAME = "register_model_group"; /** * Instantiate this class * @param mlClient client to instantiate MLClient * @param flowFrameworkIndicesHandler FlowFrameworkIndicesHandler class to update system indices */ - public ModelGroupStep(MachineLearningNodeClient mlClient, FlowFrameworkIndicesHandler flowFrameworkIndicesHandler) { + public RegisterModelGroupStep(MachineLearningNodeClient mlClient, FlowFrameworkIndicesHandler flowFrameworkIndicesHandler) { this.mlClient = mlClient; this.flowFrameworkIndicesHandler = flowFrameworkIndicesHandler; } diff --git a/src/main/java/org/opensearch/flowframework/workflow/WorkflowStepFactory.java b/src/main/java/org/opensearch/flowframework/workflow/WorkflowStepFactory.java index 08ccfc2a8..d19f51128 100644 --- a/src/main/java/org/opensearch/flowframework/workflow/WorkflowStepFactory.java +++ b/src/main/java/org/opensearch/flowframework/workflow/WorkflowStepFactory.java @@ -62,7 +62,7 @@ public WorkflowStepFactory( stepMap.put(UndeployModelStep.NAME, () -> new UndeployModelStep(mlClient)); stepMap.put(CreateConnectorStep.NAME, () -> new CreateConnectorStep(mlClient, flowFrameworkIndicesHandler)); stepMap.put(DeleteConnectorStep.NAME, () -> new DeleteConnectorStep(mlClient)); - stepMap.put(ModelGroupStep.NAME, () -> new ModelGroupStep(mlClient, flowFrameworkIndicesHandler)); + stepMap.put(RegisterModelGroupStep.NAME, () -> new RegisterModelGroupStep(mlClient, flowFrameworkIndicesHandler)); stepMap.put(DeleteModelGroupStep.NAME, DeleteModelGroupStep::new); stepMap.put(ToolStep.NAME, ToolStep::new); stepMap.put(RegisterAgentStep.NAME, () -> new RegisterAgentStep(mlClient, flowFrameworkIndicesHandler)); diff --git a/src/main/resources/mappings/workflow-steps.json b/src/main/resources/mappings/workflow-steps.json index 0fca9055f..6690a9039 100644 --- a/src/main/resources/mappings/workflow-steps.json +++ b/src/main/resources/mappings/workflow-steps.json @@ -123,7 +123,7 @@ "opensearch-ml" ] }, - "model_group": { + "register_model_group": { "inputs":[ "name" ], diff --git a/src/test/java/org/opensearch/flowframework/workflow/ModelGroupStepTests.java b/src/test/java/org/opensearch/flowframework/workflow/ModelGroupStepTests.java index c930d42d4..749a7b2d5 100644 --- a/src/test/java/org/opensearch/flowframework/workflow/ModelGroupStepTests.java +++ b/src/test/java/org/opensearch/flowframework/workflow/ModelGroupStepTests.java @@ -73,7 +73,7 @@ public void testRegisterModelGroup() throws ExecutionException, InterruptedExcep String modelGroupId = MODEL_GROUP_ID; String status = MLTaskState.CREATED.name(); - ModelGroupStep modelGroupStep = new ModelGroupStep(machineLearningNodeClient, flowFrameworkIndicesHandler); + RegisterModelGroupStep modelGroupStep = new RegisterModelGroupStep(machineLearningNodeClient, flowFrameworkIndicesHandler); @SuppressWarnings("unchecked") ArgumentCaptor> actionListenerCaptor = ArgumentCaptor.forClass(ActionListener.class); @@ -107,7 +107,7 @@ public void testRegisterModelGroup() throws ExecutionException, InterruptedExcep } public void testRegisterModelGroupFailure() throws IOException { - ModelGroupStep modelGroupStep = new ModelGroupStep(machineLearningNodeClient, flowFrameworkIndicesHandler); + RegisterModelGroupStep modelGroupStep = new RegisterModelGroupStep(machineLearningNodeClient, flowFrameworkIndicesHandler); @SuppressWarnings("unchecked") ArgumentCaptor> actionListenerCaptor = ArgumentCaptor.forClass(ActionListener.class); @@ -135,7 +135,7 @@ public void testRegisterModelGroupFailure() throws IOException { } public void testRegisterModelGroupWithNoName() throws IOException { - ModelGroupStep modelGroupStep = new ModelGroupStep(machineLearningNodeClient, flowFrameworkIndicesHandler); + RegisterModelGroupStep modelGroupStep = new RegisterModelGroupStep(machineLearningNodeClient, flowFrameworkIndicesHandler); CompletableFuture future = modelGroupStep.execute( inputDataWithNoName.getNodeId(),