Skip to content

Commit 4eb5dff

Browse files
benwtrentbenwtrent
authored
[ML] mute inference_rescorer yaml tests and add single node rest tests (#97179)
There is a common issue with the ML test cleanup code where we grab model stats before taking up a cleaning action. However, in grabbing those stats the search fails because the index was just recently created. This moves the yaml tests as they existed (pretty much line for line) into a single node rest test. I also mute the yaml tests instead of simply deleting them (as I would prefer having those). related to: #80703
1 parent d8e22b7 commit 4eb5dff

File tree

3 files changed

+207
-0
lines changed

3 files changed

+207
-0
lines changed

‎x-pack/plugin/ml/qa/ml-with-security/build.gradle

+1
Original file line numberDiff line numberDiff line change
@@ -181,6 +181,7 @@ tasks.named("yamlRestTest").configure {
181181
'ml/inference_crud/Test put nlp model config with vocabulary set',
182182
'ml/inference_crud/Test put model model aliases with nlp model',
183183
'ml/inference_processor/Test create processor with missing mandatory fields',
184+
'ml/inference_rescore/Test rescore with missing model',
184185
'ml/inference_stats_crud/Test get stats given missing trained model',
185186
'ml/inference_stats_crud/Test get stats given expression without matches and allow_no_match is false',
186187
'ml/jobs_crud/Test cannot create job with model snapshot id set',

‎x-pack/plugin/ml/qa/single-node-tests/src/javaRestTest/java/org/elasticsearch/xpack/ml/integration/InferenceRescorerIT.java

+194
Original file line numberDiff line numberDiff line change
@@ -0,0 +1,194 @@
1+
/*
2+
* Copyright Elasticsearch B.V. and/or licensed to Elasticsearch B.V. under one
3+
* or more contributor license agreements. Licensed under the Elastic License
4+
* 2.0; you may not use this file except in compliance with the Elastic License
5+
* 2.0.
6+
*/
7+
8+
package org.elasticsearch.xpack.ml.integration;
9+
10+
import org.elasticsearch.client.Request;
11+
import org.elasticsearch.client.Response;
12+
import org.elasticsearch.common.settings.Settings;
13+
import org.elasticsearch.common.xcontent.support.XContentMapValues;
14+
import org.junit.Before;
15+
16+
import java.io.IOException;
17+
import java.util.List;
18+
19+
import static org.hamcrest.Matchers.equalTo;
20+
21+
public class InferenceRescorerIT extends InferenceTestCase {
22+
23+
private static final String MODEL_ID = "ltr-model";
24+
private static final String INDEX_NAME = "store";
25+
26+
@Before
27+
public void setupModelAndData() throws IOException {
28+
putRegressionModel(MODEL_ID, """
29+
{
30+
"description": "super complex model for tests",
31+
"input": {"field_names": ["cost", "product"]},
32+
"inference_config": {
33+
"regression": {
34+
}
35+
},
36+
"definition": {
37+
"preprocessors" : [{
38+
"one_hot_encoding": {
39+
"field": "product",
40+
"hot_map": {
41+
"TV": "type_tv",
42+
"VCR": "type_vcr",
43+
"Laptop": "type_laptop"
44+
}
45+
}
46+
}],
47+
"trained_model": {
48+
"ensemble": {
49+
"feature_names": ["cost", "type_tv", "type_vcr", "type_laptop"],
50+
"target_type": "regression",
51+
"trained_models": [
52+
{
53+
"tree": {
54+
"feature_names": [
55+
"cost"
56+
],
57+
"tree_structure": [
58+
{
59+
"node_index": 0,
60+
"split_feature": 0,
61+
"split_gain": 12,
62+
"threshold": 400,
63+
"decision_type": "lte",
64+
"default_left": true,
65+
"left_child": 1,
66+
"right_child": 2
67+
},
68+
{
69+
"node_index": 1,
70+
"leaf_value": 5.0
71+
},
72+
{
73+
"node_index": 2,
74+
"leaf_value": 2.0
75+
}
76+
],
77+
"target_type": "regression"
78+
}
79+
},
80+
{
81+
"tree": {
82+
"feature_names": [
83+
"type_tv"
84+
],
85+
"tree_structure": [
86+
{
87+
"node_index": 0,
88+
"split_feature": 0,
89+
"split_gain": 12,
90+
"threshold": 1,
91+
"decision_type": "lt",
92+
"default_left": true,
93+
"left_child": 1,
94+
"right_child": 2
95+
},
96+
{
97+
"node_index": 1,
98+
"leaf_value": 1.0
99+
},
100+
{
101+
"node_index": 2,
102+
"leaf_value": 12.0
103+
}
104+
],
105+
"target_type": "regression"
106+
}
107+
}
108+
]
109+
}
110+
}
111+
}
112+
}""");
113+
createIndex(INDEX_NAME, Settings.EMPTY, """
114+
"properties":{
115+
"product":{"type": "keyword"},
116+
"cost":{"type": "integer"}}""");
117+
indexData("{ \"product\": \"TV\", \"cost\": 300}");
118+
indexData("{ \"product\": \"TV\", \"cost\": 400}");
119+
indexData("{ \"product\": \"TV\", \"cost\": 600}");
120+
indexData("{ \"product\": \"VCR\", \"cost\": 15}");
121+
indexData("{ \"product\": \"VCR\", \"cost\": 350}");
122+
indexData("{ \"product\": \"VCR\", \"cost\": 580}");
123+
indexData("{ \"product\": \"Laptop\", \"cost\": 100}");
124+
indexData("{ \"product\": \"Laptop\", \"cost\": 300}");
125+
indexData("{ \"product\": \"Laptop\", \"cost\": 500}");
126+
adminClient().performRequest(new Request("POST", INDEX_NAME + "/_refresh"));
127+
}
128+
129+
public void testInferenceRescore() throws Exception {
130+
Request request = new Request("GET", "store/_search?size=3");
131+
request.setJsonEntity("""
132+
{
133+
"rescore": {
134+
"window_size": 10,
135+
"inference": { "model_id": "ltr-model" }
136+
}
137+
}""");
138+
assertHitScores(client().performRequest(request), List.of(17.0, 17.0, 14.0));
139+
request.setJsonEntity("""
140+
{
141+
"query": {"term": {"product": "Laptop"}},
142+
"rescore": {
143+
"window_size": 10,
144+
"inference": { "model_id": "ltr-model" }
145+
}
146+
}""");
147+
assertHitScores(client().performRequest(request), List.of(6.0, 6.0, 3.0));
148+
}
149+
150+
public void testInferenceRescoreSmallWindow() throws Exception {
151+
Request request = new Request("GET", "store/_search?size=5");
152+
request.setJsonEntity("""
153+
{
154+
"rescore": {
155+
"window_size": 2,
156+
"inference": { "model_id": "ltr-model" }
157+
}
158+
}""");
159+
assertHitScores(client().performRequest(request), List.of(17.0, 17.0, 1.0, 1.0, 1.0));
160+
}
161+
162+
public void testInferenceRescorerWithChainedRescorers() throws IOException {
163+
Request request = new Request("GET", "store/_search?size=5");
164+
request.setJsonEntity("""
165+
{
166+
"rescore": [
167+
{
168+
"window_size": 4,
169+
"query": { "rescore_query":{ "script_score": {"query": {"match_all": {}}, "script": {"source": "return 4"}}}}
170+
},
171+
{
172+
"window_size": 3,
173+
"inference": { "model_id": "ltr-model" }
174+
},
175+
{
176+
"window_size": 2,
177+
"query": { "rescore_query": { "script_score": {"query": {"match_all": {}}, "script": {"source": "return 20"}}}}
178+
}
179+
]
180+
}""");
181+
assertHitScores(client().performRequest(request), List.of(37.0, 37.0, 14.0, 5.0, 1.0));
182+
}
183+
184+
private void indexData(String data) throws IOException {
185+
Request request = new Request("POST", INDEX_NAME + "/_doc");
186+
request.setJsonEntity(data);
187+
client().performRequest(request);
188+
}
189+
190+
@SuppressWarnings("unchecked")
191+
private static void assertHitScores(Response response, List<Double> expectedScores) throws IOException {
192+
assertThat((List<Double>) XContentMapValues.extractValue("hits.hits._score", responseAsMap(response)), equalTo(expectedScores));
193+
}
194+
}

‎x-pack/plugin/src/yamlRestTest/resources/rest-api-spec/test/ml/inference_rescore.yml

+12
Original file line numberDiff line numberDiff line change
@@ -134,6 +134,10 @@ setup:
134134
135135
---
136136
"Test rescore with stored model":
137+
- skip:
138+
version: all
139+
reason: "@AwaitsFix https://github.com/elastic/elasticsearch/issues/80703"
140+
137141
- do:
138142
search:
139143
index: store
@@ -166,6 +170,10 @@ setup:
166170
- match: { hits.hits.2._score: 3.0 }
167171
---
168172
"Test rescore with stored model and smaller window_size":
173+
- skip:
174+
version: all
175+
reason: "@AwaitsFix https://github.com/elastic/elasticsearch/issues/80703"
176+
169177
- do:
170178
search:
171179
index: store
@@ -184,6 +192,10 @@ setup:
184192
- match: { hits.hits.4._score: 1.0 }
185193
---
186194
"Test rescore with stored model and chained rescorers":
195+
- skip:
196+
version: all
197+
reason: "@AwaitsFix https://github.com/elastic/elasticsearch/issues/80703"
198+
187199
- do:
188200
search:
189201
index: store

0 commit comments

Comments
 (0)