From 3b94d22a65f8a4590180ac74ae7127838d847144 Mon Sep 17 00:00:00 2001 From: Amit Galitzky Date: Tue, 30 Apr 2024 18:26:26 -0700 Subject: [PATCH 1/3] adding context.restore when missed and fixing model group parsing Signed-off-by: Amit Galitzky --- .../transport/CreateWorkflowTransportAction.java | 1 + .../org/opensearch/flowframework/util/EncryptorUtils.java | 4 +++- .../flowframework/workflow/RegisterModelGroupStep.java | 3 ++- 3 files changed, 6 insertions(+), 2 deletions(-) diff --git a/src/main/java/org/opensearch/flowframework/transport/CreateWorkflowTransportAction.java b/src/main/java/org/opensearch/flowframework/transport/CreateWorkflowTransportAction.java index 9dabd0c4a..720fa1ea6 100644 --- a/src/main/java/org/opensearch/flowframework/transport/CreateWorkflowTransportAction.java +++ b/src/main/java/org/opensearch/flowframework/transport/CreateWorkflowTransportAction.java @@ -311,6 +311,7 @@ void checkMaxWorkflows(TimeValue requestTimeOut, Integer maxWorkflow, ActionList try (ThreadContext.StoredContext context = client.threadPool().getThreadContext().stashContext()) { logger.info("Querying existing workflows to count the max"); client.search(searchRequest, ActionListener.wrap(searchResponse -> { + context.restore(); internalListener.onResponse(searchResponse.getHits().getTotalHits().value < maxWorkflow); }, exception -> { String errorMessage = "Unable to fetch the workflows"; diff --git a/src/main/java/org/opensearch/flowframework/util/EncryptorUtils.java b/src/main/java/org/opensearch/flowframework/util/EncryptorUtils.java index 946c68908..f17f26717 100644 --- a/src/main/java/org/opensearch/flowframework/util/EncryptorUtils.java +++ b/src/main/java/org/opensearch/flowframework/util/EncryptorUtils.java @@ -250,7 +250,6 @@ public void initializeMasterKey(ActionListener listener) { GetRequest getRequest = new GetRequest(CONFIG_INDEX).id(MASTER_KEY); client.get(getRequest, ActionListener.wrap(getResponse -> { - if (!getResponse.isExists()) { // Generate new key and index @@ -260,6 +259,7 @@ public void initializeMasterKey(ActionListener listener) { .setRefreshPolicy(WriteRequest.RefreshPolicy.IMMEDIATE); client.index(masterKeyIndexRequest, ActionListener.wrap(indexResponse -> { + context.restore(); // Set generated key to master logger.info("Config has been initialized successfully"); this.masterKey = generatedKey; @@ -270,6 +270,7 @@ public void initializeMasterKey(ActionListener listener) { })); } else { + context.restore(); // Set existing key to master logger.info("Config has already been initialized"); this.masterKey = (String) getResponse.getSourceAsMap().get(MASTER_KEY); @@ -300,6 +301,7 @@ void initializeMasterKeyIfAbsent() { try (ThreadContext.StoredContext context = client.threadPool().getThreadContext().stashContext()) { GetRequest getRequest = new GetRequest(CONFIG_INDEX).id(MASTER_KEY); client.get(getRequest, ActionListener.wrap(response -> { + context.restore(); if (response.isExists()) { this.masterKey = (String) response.getSourceAsMap().get(MASTER_KEY); } else { diff --git a/src/main/java/org/opensearch/flowframework/workflow/RegisterModelGroupStep.java b/src/main/java/org/opensearch/flowframework/workflow/RegisterModelGroupStep.java index 1a2fcebe9..c0d6c701c 100644 --- a/src/main/java/org/opensearch/flowframework/workflow/RegisterModelGroupStep.java +++ b/src/main/java/org/opensearch/flowframework/workflow/RegisterModelGroupStep.java @@ -28,6 +28,7 @@ import java.util.Collections; import java.util.List; +import java.util.Locale; import java.util.Map; import java.util.Set; @@ -142,7 +143,7 @@ public void onFailure(Exception ex) { String modelGroupName = (String) inputs.get(NAME_FIELD); String description = (String) inputs.get(DESCRIPTION_FIELD); List backendRoles = getBackendRoles(inputs); - AccessMode modelAccessMode = (AccessMode) inputs.get(MODEL_ACCESS_MODE); + AccessMode modelAccessMode = AccessMode.from(((String) inputs.get(MODEL_ACCESS_MODE)).toLowerCase(Locale.ROOT)); Boolean isAddAllBackendRoles = inputs.containsKey(ADD_ALL_BACKEND_ROLES) ? Booleans.parseBoolean(inputs.get(ADD_ALL_BACKEND_ROLES).toString()) : null; From 6377acc5410d6f1ae7418f9d6186adb6f865ec96 Mon Sep 17 00:00:00 2001 From: Amit Galitzky Date: Wed, 1 May 2024 09:00:10 -0700 Subject: [PATCH 2/3] adding change log, fix casting Signed-off-by: Amit Galitzky --- CHANGELOG.md | 1 + .../opensearch-flow-framework.release-notes-2.14.0.0.md | 2 ++ .../flowframework/workflow/RegisterModelGroupStep.java | 5 ++++- 3 files changed, 7 insertions(+), 1 deletion(-) diff --git a/CHANGELOG.md b/CHANGELOG.md index b7c653e8c..987b3c945 100644 --- a/CHANGELOG.md +++ b/CHANGELOG.md @@ -25,6 +25,7 @@ Inspired from [Keep a Changelog](https://keepachangelog.com/en/1.1.0/) - Hide user and credential field from search response ([#680](https://github.com/opensearch-project/flow-framework/pull/680)) - Throw the correct error message in status API for WorkflowSteps ([#676](https://github.com/opensearch-project/flow-framework/pull/676)) - Delete workflow state when template is deleted and no resources exist ([#689](https://github.com/opensearch-project/flow-framework/pull/689)) +- Fixing model group parsing and restoring context ([#695] (https://github.com/opensearch-project/flow-framework/pull/695)) ### Infrastructure - Switch macos runner to macos-13 from macos-latest since macos-latest is now arm64 ([#686](https://github.com/opensearch-project/flow-framework/pull/686)) diff --git a/release-notes/opensearch-flow-framework.release-notes-2.14.0.0.md b/release-notes/opensearch-flow-framework.release-notes-2.14.0.0.md index ee08077c6..92a15d60f 100644 --- a/release-notes/opensearch-flow-framework.release-notes-2.14.0.0.md +++ b/release-notes/opensearch-flow-framework.release-notes-2.14.0.0.md @@ -13,6 +13,8 @@ Compatible with OpenSearch 2.14.0 - Hide user and credential field from search response ([#680](https://github.com/opensearch-project/flow-framework/pull/680)) - Throw the correct error message in status API for WorkflowSteps ([#676](https://github.com/opensearch-project/flow-framework/pull/676)) - Delete workflow state when template is deleted and no resources exist ([#689](https://github.com/opensearch-project/flow-framework/pull/689)) +- Fixing model group parsing and restoring context ([#695] (https://github.com/opensearch-project/flow-framework/pull/695)) + ### Infrastructure - Switch macOS runner to macos-13 from macos-latest since macos-latest is now arm64 ([#686](https://github.com/opensearch-project/flow-framework/pull/686)) diff --git a/src/main/java/org/opensearch/flowframework/workflow/RegisterModelGroupStep.java b/src/main/java/org/opensearch/flowframework/workflow/RegisterModelGroupStep.java index c0d6c701c..53a06d79f 100644 --- a/src/main/java/org/opensearch/flowframework/workflow/RegisterModelGroupStep.java +++ b/src/main/java/org/opensearch/flowframework/workflow/RegisterModelGroupStep.java @@ -143,7 +143,10 @@ public void onFailure(Exception ex) { String modelGroupName = (String) inputs.get(NAME_FIELD); String description = (String) inputs.get(DESCRIPTION_FIELD); List backendRoles = getBackendRoles(inputs); - AccessMode modelAccessMode = AccessMode.from(((String) inputs.get(MODEL_ACCESS_MODE)).toLowerCase(Locale.ROOT)); + AccessMode modelAccessMode = null; + if (inputs.get(MODEL_ACCESS_MODE) != null) { + modelAccessMode = AccessMode.from((inputs.get(MODEL_ACCESS_MODE)).toString().toLowerCase(Locale.ROOT)); + } Boolean isAddAllBackendRoles = inputs.containsKey(ADD_ALL_BACKEND_ROLES) ? Booleans.parseBoolean(inputs.get(ADD_ALL_BACKEND_ROLES).toString()) : null; From e34d8a420b510d02e8b2e79d7399eef76a41e890 Mon Sep 17 00:00:00 2001 From: Amit Galitzky Date: Wed, 1 May 2024 11:22:42 -0700 Subject: [PATCH 3/3] changed to contains key Signed-off-by: Amit Galitzky --- .../flowframework/workflow/RegisterModelGroupStep.java | 2 +- 1 file changed, 1 insertion(+), 1 deletion(-) diff --git a/src/main/java/org/opensearch/flowframework/workflow/RegisterModelGroupStep.java b/src/main/java/org/opensearch/flowframework/workflow/RegisterModelGroupStep.java index 53a06d79f..9771d7c77 100644 --- a/src/main/java/org/opensearch/flowframework/workflow/RegisterModelGroupStep.java +++ b/src/main/java/org/opensearch/flowframework/workflow/RegisterModelGroupStep.java @@ -144,7 +144,7 @@ public void onFailure(Exception ex) { String description = (String) inputs.get(DESCRIPTION_FIELD); List backendRoles = getBackendRoles(inputs); AccessMode modelAccessMode = null; - if (inputs.get(MODEL_ACCESS_MODE) != null) { + if (inputs.containsKey(MODEL_ACCESS_MODE)) { modelAccessMode = AccessMode.from((inputs.get(MODEL_ACCESS_MODE)).toString().toLowerCase(Locale.ROOT)); } Boolean isAddAllBackendRoles = inputs.containsKey(ADD_ALL_BACKEND_ROLES)