Skip to content

Commit

Permalink
Fixing model group parsing and restoring context (#695)
Browse files Browse the repository at this point in the history
* adding context.restore when missed and fixing model group parsing

Signed-off-by: Amit Galitzky <[email protected]>

* adding change log, fix casting

Signed-off-by: Amit Galitzky <[email protected]>

* changed to contains key

Signed-off-by: Amit Galitzky <[email protected]>

---------

Signed-off-by: Amit Galitzky <[email protected]>
(cherry picked from commit 28ea3dd)
Signed-off-by: github-actions[bot] <github-actions[bot]@users.noreply.github.com>
  • Loading branch information
github-actions[bot] committed May 1, 2024
1 parent 99c5ca1 commit 810ebfa
Show file tree
Hide file tree
Showing 5 changed files with 12 additions and 2 deletions.
1 change: 1 addition & 0 deletions CHANGELOG.md
Original file line number Diff line number Diff line change
Expand Up @@ -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))
Expand Down
Original file line number Diff line number Diff line change
Expand Up @@ -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))
Expand Down
Original file line number Diff line number Diff line change
Expand Up @@ -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";
Expand Down
Original file line number Diff line number Diff line change
Expand Up @@ -250,7 +250,6 @@ public void initializeMasterKey(ActionListener<Boolean> listener) {

GetRequest getRequest = new GetRequest(CONFIG_INDEX).id(MASTER_KEY);
client.get(getRequest, ActionListener.wrap(getResponse -> {

if (!getResponse.isExists()) {

// Generate new key and index
Expand All @@ -260,6 +259,7 @@ public void initializeMasterKey(ActionListener<Boolean> listener) {
.setRefreshPolicy(WriteRequest.RefreshPolicy.IMMEDIATE);

client.index(masterKeyIndexRequest, ActionListener.wrap(indexResponse -> {
context.restore();

Check warning on line 262 in src/main/java/org/opensearch/flowframework/util/EncryptorUtils.java

View check run for this annotation

Codecov / codecov/patch

src/main/java/org/opensearch/flowframework/util/EncryptorUtils.java#L262

Added line #L262 was not covered by tests
// Set generated key to master
logger.info("Config has been initialized successfully");
this.masterKey = generatedKey;
Expand All @@ -270,6 +270,7 @@ public void initializeMasterKey(ActionListener<Boolean> listener) {
}));

} else {
context.restore();

Check warning on line 273 in src/main/java/org/opensearch/flowframework/util/EncryptorUtils.java

View check run for this annotation

Codecov / codecov/patch

src/main/java/org/opensearch/flowframework/util/EncryptorUtils.java#L273

Added line #L273 was not covered by tests
// Set existing key to master
logger.info("Config has already been initialized");
this.masterKey = (String) getResponse.getSourceAsMap().get(MASTER_KEY);
Expand Down Expand Up @@ -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 {
Expand Down
Original file line number Diff line number Diff line change
Expand Up @@ -28,6 +28,7 @@

import java.util.Collections;
import java.util.List;
import java.util.Locale;
import java.util.Map;
import java.util.Set;

Expand Down Expand Up @@ -142,7 +143,10 @@ public void onFailure(Exception ex) {
String modelGroupName = (String) inputs.get(NAME_FIELD);
String description = (String) inputs.get(DESCRIPTION_FIELD);
List<String> 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;
Expand Down

0 comments on commit 810ebfa

Please sign in to comment.