From 810ebfa59458d2458fd2a41f5ddf3f00d3fe316a Mon Sep 17 00:00:00 2001 From: "github-actions[bot]" Date: Wed, 1 May 2024 18:29:48 +0000 Subject: [PATCH] Fixing model group parsing and restoring context (#695) * adding context.restore when missed and fixing model group parsing Signed-off-by: Amit Galitzky * adding change log, fix casting Signed-off-by: Amit Galitzky * changed to contains key Signed-off-by: Amit Galitzky --------- Signed-off-by: Amit Galitzky (cherry picked from commit 28ea3ddeab3ce20ddf58412d859b1e99ae716cfa) Signed-off-by: github-actions[bot] --- CHANGELOG.md | 1 + .../opensearch-flow-framework.release-notes-2.14.0.0.md | 2 ++ .../transport/CreateWorkflowTransportAction.java | 1 + .../org/opensearch/flowframework/util/EncryptorUtils.java | 4 +++- .../flowframework/workflow/RegisterModelGroupStep.java | 6 +++++- 5 files changed, 12 insertions(+), 2 deletions(-) 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/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..9771d7c77 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,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) inputs.get(MODEL_ACCESS_MODE); + AccessMode modelAccessMode = 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) ? Booleans.parseBoolean(inputs.get(ADD_ALL_BACKEND_ROLES).toString()) : null;