Skip to content

fix(test): Stabilize flaky test by awaiting model deployment #1617

New issue

Have a question about this project? Sign up for a free GitHub account to open an issue and contact its maintainers and the community.

Sign up for GitHub

By clicking “Sign up for GitHub”, you agree to our terms of service and privacy statement. We’ll occasionally send you account related emails.

Already on GitHub? Sign in to your account

Open
wants to merge 1 commit into
base: main
Choose a base branch
from
+36 −30
Open

fix(test): Stabilize flaky test by awaiting model deployment #1617

Show file tree
Hide file tree
Changes from all commits
Commits
Show all changes
1 commit
Select commit
File filter

Filter by extension

Filter by extension
Conversations
Failed to load comments.
Loading
Jump to
Jump to file
Failed to load files.
Loading
Diff view
Diff view
1 change: 1 addition & 0 deletions CHANGELOG.md
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -10,6 +10,7 @@ The format is based on [Keep a Changelog](https://keepachangelog.com/en/1.0.0/),
### Enhancements

### Bug Fixes
- Fix flaky test in `NeuralQueryEnricherProcessorIT` by waiting for model deployment ([#1617](https://github.com/opensearch-project/neural-search/pull/1617))

### Infrastructure

Expand Down
51 changes: 27 additions & 24 deletions src/test/java/org/opensearch/neuralsearch/processor/NeuralQueryEnricherProcessorIT.java
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -57,10 +57,10 @@ public void testNeuralQueryEnricherProcessor_whenNoModelIdPassed_thenSuccess() {
createPipelineProcessor(modelId, ingest_pipeline, ProcessorType.TEXT_EMBEDDING);
updateIndexSettings(index, Settings.builder().put("index.search.default_pipeline", search_pipeline));
NeuralQueryBuilder neuralQueryBuilder = NeuralQueryBuilder.builder()
.fieldName(TEST_KNN_VECTOR_FIELD_NAME_1)
.queryText("Hello World")
.k(1)
.build();
.fieldName(TEST_KNN_VECTOR_FIELD_NAME_1)
Copy link
Collaborator

@heemin32 heemin32 Oct 15, 2025

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

Could you remove all these unnecessary changes?

Copy link
Member

@owaiskazi19 owaiskazi19 Oct 15, 2025

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

Can run ./gradlew spotlessApply to apply formatting

.queryText("Hello World")
.k(1)
.build();
Map<String, Object> response = search(index, neuralQueryBuilder, 2);
assertFalse(response.isEmpty());
}
Expand Down Expand Up @@ -88,9 +88,11 @@ public void testNeuralQueryEnricherProcessor_whenGetEmptyQueryBody_thenSuccess()
updateIndexSettings(index, Settings.builder().put("index.search.default_pipeline", search_pipeline));
Request request = new Request("POST", "/" + index + "/_search");
Response response = client().performRequest(request);
assertEquals(request.getEndpoint() + ": failed", RestStatus.OK, RestStatus.fromCode(response.getStatusLine().getStatusCode()));
assertEquals(request.getEndpoint() + ": failed", RestStatus.OK,
RestStatus.fromCode(response.getStatusLine().getStatusCode()));
String responseBody = EntityUtils.toString(response.getEntity());
Map<String, Object> responseInMap = XContentHelper.convertToMap(XContentType.JSON.xContent(), responseBody, false);
Map<String, Object> responseInMap = XContentHelper.convertToMap(XContentType.JSON.xContent(), responseBody,
false);
assertFalse(responseInMap.isEmpty());
assertEquals(3, ((Map) responseInMap.get("hits")).size());
}
Expand All @@ -104,10 +106,10 @@ public void testNeuralQueryEnricherProcessor_whenHybridQueryBuilderAndNoModelIdP
createPipelineProcessor(modelId, ingest_pipeline, ProcessorType.TEXT_EMBEDDING);
updateIndexSettings(index, Settings.builder().put("index.search.default_pipeline", search_pipeline));
NeuralQueryBuilder neuralQueryBuilder = NeuralQueryBuilder.builder()
.fieldName(TEST_KNN_VECTOR_FIELD_NAME_1)
.queryText("Hello World")
.k(1)
.build();
.fieldName(TEST_KNN_VECTOR_FIELD_NAME_1)
.queryText("Hello World")
.k(1)
.build();
HybridQueryBuilder hybridQueryBuilder = new HybridQueryBuilder();
hybridQueryBuilder.add(neuralQueryBuilder);
Map<String, Object> response = search(index, hybridQueryBuilder, 2);
Expand All @@ -126,10 +128,10 @@ public void testNeuralQueryEnricherProcessor_whenNoModelIdPassed_statsEnabled_th
createPipelineProcessor(modelId, ingest_pipeline, ProcessorType.TEXT_EMBEDDING);
updateIndexSettings(index, Settings.builder().put("index.search.default_pipeline", search_pipeline));
NeuralQueryBuilder neuralQueryBuilder = NeuralQueryBuilder.builder()
.fieldName(TEST_KNN_VECTOR_FIELD_NAME_1)
.queryText("Hello World")
.k(1)
.build();
.fieldName(TEST_KNN_VECTOR_FIELD_NAME_1)
.queryText("Hello World")
.k(1)
.build();
Map<String, Object> response = search(index, neuralQueryBuilder, 2);
assertFalse(response.isEmpty());

Expand All @@ -142,7 +144,8 @@ public void testNeuralQueryEnricherProcessor_whenNoModelIdPassed_statsEnabled_th
Map<String, Object> allNodesStats = parseAggregatedNodeStatsResponse(responseBody);

// Parse json to get stats
assertEquals(2, getNestedValue(allNodesStats, EventStatName.NEURAL_QUERY_ENRICHER_PROCESSOR_EXECUTIONS.getFullPath()));
assertEquals(2,
getNestedValue(allNodesStats, EventStatName.NEURAL_QUERY_ENRICHER_PROCESSOR_EXECUTIONS.getFullPath()));
assertEquals(1, getNestedValue(stats, InfoStatName.NEURAL_QUERY_ENRICHER_PROCESSORS.getFullPath()));

// Reset stats
Expand All @@ -153,21 +156,21 @@ public void testNeuralQueryEnricherProcessor_whenNoModelIdPassed_statsEnabled_th
private void initializeIndexIfNotExist(String indexName) {
if (indexName.equals(NeuralQueryEnricherProcessorIT.index) && !indexExists(indexName)) {
prepareKnnIndex(
indexName,
Collections.singletonList(new KNNFieldConfig(TEST_KNN_VECTOR_FIELD_NAME_1, TEST_DIMENSION, TEST_SPACE_TYPE))
);
indexName,
Collections.singletonList(
new KNNFieldConfig(TEST_KNN_VECTOR_FIELD_NAME_1, TEST_DIMENSION, TEST_SPACE_TYPE)));
addKnnDoc(
indexName,
"1",
Collections.singletonList(TEST_KNN_VECTOR_FIELD_NAME_1),
Collections.singletonList(Floats.asList(testVector).toArray())
);
indexName,
"1",
Collections.singletonList(TEST_KNN_VECTOR_FIELD_NAME_1),
Collections.singletonList(Floats.asList(testVector).toArray()));
assertEquals(1, getDocCount(indexName));
}

if (sparseIndex.equals(indexName) && !indexExists(indexName)) {
prepareSparseEncodingIndex(indexName, List.of(TEST_RANK_FEATURES_FIELD_NAME_1));
addSparseEncodingDoc(indexName, "1", List.of(TEST_RANK_FEATURES_FIELD_NAME_1), List.of(Map.of("hi", 1.0f, "hello", 1.1f)));
addSparseEncodingDoc(indexName, "1", List.of(TEST_RANK_FEATURES_FIELD_NAME_1),
List.of(Map.of("hi", 1.0f, "hello", 1.1f)));
assertEquals(1, getDocCount(indexName));
}
}
Expand Down
14 changes: 8 additions & 6 deletions src/testFixtures/java/org/opensearch/neuralsearch/BaseNeuralSearchIT.java
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -320,8 +320,7 @@ private void doLoadAndWaitForModelToBeReady(String modelId) throws Exception {
logger.info("Model is already deployed. Skip loading.");
return;
}
loadModel(modelId);
waitForModelToBeReady(modelId);
loadAndWaitForModelToBeReady(modelId);
}

protected void loadAndWaitForModelToBeReady(String modelId) throws Exception {
Expand Down Expand Up @@ -370,9 +369,12 @@ protected void waitForModelToBeReady(String modelId) throws Exception {
throw new RuntimeException("Model " + modelId + " failed to be ready for inference after " + MAX_ATTEMPTS + " attempts");
}

/**
* LOADED is same as DEPLOYED but it is deprecated at OS2.7. As Neural Search is introduced after OS2.7, we could exclude that state.
*/
protected boolean isModelReadyForInference(@NonNull final String modelId) throws IOException, ParseException {
MLModelState state = getModelState(modelId);
return MLModelState.LOADED.equals(state) || MLModelState.DEPLOYED.equals(state);
return MLModelState.DEPLOYED.equals(state);
Copy link
Collaborator

@heemin32 heemin32 Oct 15, 2025

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

Could you add a comment on why we don't check for LOADED state?

LOADED is same as DEPLOYED but it is deprecated at OS2.7. As Neural Search is introduced after OS2.7, we could exclude that state.

}

/**
Expand All @@ -384,7 +386,7 @@ protected boolean isModelReadyForInference(@NonNull final String modelId) throws
protected String prepareModel() {
String requestBody = Files.readString(Path.of(classLoader.getResource("processor/UploadModelRequestBody.json").toURI()));
String modelId = registerModelGroupAndUploadModel(requestBody);
loadModel(modelId);
loadAndWaitForModelToBeReady(modelId);
return modelId;
}

Expand All @@ -399,14 +401,14 @@ protected String prepareSparseEncodingModel() {
Path.of(classLoader.getResource("processor/UploadSparseEncodingModelRequestBody.json").toURI())
);
String modelId = registerModelGroupAndUploadModel(requestBody);
loadModel(modelId);
loadAndWaitForModelToBeReady(modelId);
return modelId;
}

protected String prepareSemanticHighlightingLocalModel() throws Exception {
String requestBody = Files.readString(Path.of(classLoader.getResource("highlight/LocalQuestionAnsweringModel.json").toURI()));
String modelId = registerModelGroupAndUploadModel(requestBody);
loadModel(modelId);
loadAndWaitForModelToBeReady(modelId);
return modelId;
}

Expand Down