Skip to content

Commit f32c3be

Browse files
valeriy42valeriy42
committed
clean up
1 parent c40723a commit f32c3be

File tree

1 file changed

+74
-79
lines changed

1 file changed

+74
-79
lines changed

‎x-pack/plugin/ml/src/internalClusterTest/java/org/elasticsearch/xpack/ml/integration/UnusedStatsRemoverIT.java‎

Lines changed: 74 additions & 79 deletions
Original file line numberDiff line numberDiff line change
@@ -64,123 +64,118 @@ public void createComponents() {
6464
}
6565

6666
public void testRemoveUnusedStats() throws Exception {
67-
6867
String modelId = "model-with-stats";
6968
putDFA(modelId);
7069

71-
72-
70+
// Existing analytics and models
7371
indexStatDocument(new DataCounts("analytics-with-stats", 1, 1, 1), DataCounts.documentId("analytics-with-stats"));
74-
indexStatDocument(new DataCounts("missing-analytics-with-stats", 1, 1, 1), DataCounts.documentId("missing-analytics-with-stats"));
72+
indexStatDocument(new InferenceStats(1, 1, 1, 1, modelId, "test", Instant.now()), InferenceStats.docId(modelId, "test"));
7573
indexStatDocument(
7674
new InferenceStats(1, 1, 1, 1, TrainedModelProvider.MODELS_STORED_AS_RESOURCE.iterator().next(), "test", Instant.now()),
7775
InferenceStats.docId(TrainedModelProvider.MODELS_STORED_AS_RESOURCE.iterator().next(), "test")
7876
);
77+
78+
// Unused analytics/model stats
79+
indexStatDocument(new DataCounts("missing-analytics-with-stats", 1, 1, 1), DataCounts.documentId("missing-analytics-with-stats"));
7980
indexStatDocument(
8081
new InferenceStats(1, 1, 1, 1, "missing-model", "test", Instant.now()),
8182
InferenceStats.docId("missing-model", "test")
8283
);
83-
indexStatDocument(
84-
new InferenceStats(1, 1, 1, 1, modelId, "test", Instant.now()),
85-
InferenceStats.docId(modelId, "test")
86-
);
87-
client().admin().indices().prepareRefresh(MlStatsIndex.indexPattern()).get();
8884

89-
PlainActionFuture<Boolean> deletionListener = new PlainActionFuture<>();
90-
UnusedStatsRemover statsRemover = new UnusedStatsRemover(client, new TaskId("test", 0L));
91-
statsRemover.remove(10000.0f, deletionListener, () -> false);
92-
deletionListener.actionGet();
85+
refreshStatsIndex();
86+
runUnusedStatsRemover();
9387

94-
client().admin().indices().prepareRefresh(MlStatsIndex.indexPattern()).get();
88+
final String index = MlStatsIndex.TEMPLATE_NAME + "-000001";
9589

96-
final String initialStateIndex = MlStatsIndex.TEMPLATE_NAME + "-000001";
90+
// Validate expected docs
91+
assertDocExists(index, InferenceStats.docId(modelId, "test"));
92+
assertDocExists(index, DataCounts.documentId("analytics-with-stats"));
93+
assertDocExists(index, InferenceStats.docId(TrainedModelProvider.MODELS_STORED_AS_RESOURCE.iterator().next(), "test"));
9794

98-
// Make sure that stats that should exist still exist
99-
assertTrue(client().prepareGet(initialStateIndex, InferenceStats.docId(modelId, "test")).get().isExists());
100-
assertTrue(
101-
client().prepareGet(
102-
initialStateIndex,
103-
InferenceStats.docId(TrainedModelProvider.MODELS_STORED_AS_RESOURCE.iterator().next(), "test")
104-
).get().isExists()
95+
// Validate removed docs
96+
assertDocDoesNotExist(index, InferenceStats.docId("missing-model", "test"));
97+
assertDocDoesNotExist(index, DataCounts.documentId("missing-analytics-with-stats"));
98+
}
99+
100+
public void testRemovingUnusedStatsFromReadOnlyIndexShouldFailSilently() throws Exception {
101+
String modelId = "model-with-stats";
102+
putDFA(modelId);
103+
104+
indexStatDocument(
105+
new InferenceStats(1, 1, 1, 1, "missing-model", "test", Instant.now()),
106+
InferenceStats.docId("missing-model", "test")
105107
);
106-
assertTrue(client().prepareGet(initialStateIndex, DataCounts.documentId("analytics-with-stats")).get().isExists());
108+
makeIndexReadOnly();
109+
refreshStatsIndex();
110+
111+
runUnusedStatsRemover();
112+
refreshStatsIndex();
107113

108-
// make sure that unused stats were deleted
109-
assertFalse(client().prepareGet(initialStateIndex, DataCounts.documentId("missing-analytics-with-stats")).get().isExists());
110-
assertFalse(client().prepareGet(initialStateIndex, InferenceStats.docId("missing-model", "test")).get().isExists());
114+
final String index = MlStatsIndex.TEMPLATE_NAME + "-000001";
115+
assertDocExists(index, InferenceStats.docId("missing-model", "test")); // should still exist
111116
}
112117

113118
private void putDFA(String modelId) {
114119
prepareIndex("foo").setId("some-empty-doc").setSource("{}", XContentType.JSON).get();
115-
PutDataFrameAnalyticsAction.Request request = new PutDataFrameAnalyticsAction.Request(
120+
121+
PutDataFrameAnalyticsAction.Request analyticsRequest = new PutDataFrameAnalyticsAction.Request(
116122
new DataFrameAnalyticsConfig.Builder().setId("analytics-with-stats")
117123
.setModelMemoryLimit(ByteSizeValue.ofGb(1))
118124
.setSource(new DataFrameAnalyticsSource(new String[] { "foo" }, null, null, null))
119125
.setDest(new DataFrameAnalyticsDest("bar", null))
120126
.setAnalysis(new Regression("prediction"))
121127
.build()
122128
);
123-
client.execute(PutDataFrameAnalyticsAction.INSTANCE, request).actionGet();
124-
125-
client.execute(
126-
PutTrainedModelAction.INSTANCE,
127-
new PutTrainedModelAction.Request(
128-
TrainedModelConfig.builder()
129-
.setModelId(modelId)
130-
.setInferenceConfig(RegressionConfig.EMPTY_PARAMS)
131-
.setInput(new TrainedModelInput(Arrays.asList("foo", "bar")))
132-
.setParsedDefinition(
133-
new TrainedModelDefinition.Builder().setPreProcessors(Collections.emptyList())
134-
.setTrainedModel(
135-
Tree.builder()
136-
.setFeatureNames(Arrays.asList("foo", "bar"))
137-
.setRoot(TreeNode.builder(0).setLeafValue(42))
138-
.build()
139-
)
140-
)
141-
.validate(true)
142-
.build(),
143-
false
144-
)
145-
).actionGet();
129+
client.execute(PutDataFrameAnalyticsAction.INSTANCE, analyticsRequest).actionGet();
130+
131+
TrainedModelDefinition.Builder definition = new TrainedModelDefinition.Builder().setPreProcessors(Collections.emptyList())
132+
.setTrainedModel(
133+
Tree.builder().setFeatureNames(Arrays.asList("foo", "bar")).setRoot(TreeNode.builder(0).setLeafValue(42)).build()
134+
);
135+
136+
TrainedModelConfig modelConfig = TrainedModelConfig.builder()
137+
.setModelId(modelId)
138+
.setInferenceConfig(RegressionConfig.EMPTY_PARAMS)
139+
.setInput(new TrainedModelInput(Arrays.asList("foo", "bar")))
140+
.setParsedDefinition(definition)
141+
.validate(true)
142+
.build();
143+
144+
client.execute(PutTrainedModelAction.INSTANCE, new PutTrainedModelAction.Request(modelConfig, false)).actionGet();
146145
}
147146

148-
public void testRemovingUnusedStatsFromReadOnlyIndexShouldFailSilently() throws Exception {
149-
150-
String modelId = "model-with-stats";
151-
putDFA(modelId);
152-
153-
indexStatDocument(
154-
new InferenceStats(1, 1, 1, 1, "missing-model", "test", Instant.now()),
155-
InferenceStats.docId("missing-model", "test")
156-
);
157-
158-
// set index to read-only
159-
client().admin().indices().prepareUpdateSettings(MlStatsIndex.indexPattern()).setSettings(Settings.builder().put("index.blocks.write", true)).get();
147+
private void indexStatDocument(ToXContentObject object, String docId) throws Exception {
148+
IndexRequest doc = new IndexRequest(MlStatsIndex.writeAlias()).id(docId);
149+
try (XContentBuilder builder = XContentFactory.jsonBuilder()) {
150+
object.toXContent(builder, new ToXContent.MapParams(Collections.singletonMap(ToXContentParams.FOR_INTERNAL_STORAGE, "true")));
151+
doc.source(builder);
152+
client.index(doc).actionGet();
153+
}
154+
}
160155

156+
private void refreshStatsIndex() {
161157
client().admin().indices().prepareRefresh(MlStatsIndex.indexPattern()).get();
158+
}
162159

160+
private void runUnusedStatsRemover() {
163161
PlainActionFuture<Boolean> deletionListener = new PlainActionFuture<>();
164-
UnusedStatsRemover statsRemover = new UnusedStatsRemover(client, new TaskId("test", 0L));
165-
statsRemover.remove(10000.0f, deletionListener, () -> false);
162+
new UnusedStatsRemover(client, new TaskId("test", 0L)).remove(10000.0f, deletionListener, () -> false);
166163
deletionListener.actionGet();
164+
}
167165

168-
client().admin().indices().prepareRefresh(MlStatsIndex.indexPattern()).get();
169-
// make sure that unused stats are still there
170-
final String initialStateIndex = MlStatsIndex.TEMPLATE_NAME + "-000001";
171-
assertTrue(client().prepareGet(initialStateIndex, InferenceStats.docId("missing-model", "test")).get().isExists());
166+
private void makeIndexReadOnly() {
167+
client().admin()
168+
.indices()
169+
.prepareUpdateSettings(MlStatsIndex.indexPattern())
170+
.setSettings(Settings.builder().put("index.blocks.write", true))
171+
.get();
172172
}
173173

174-
private void indexStatDocument(ToXContentObject object, String docId) throws Exception {
175-
ToXContent.Params params = new ToXContent.MapParams(
176-
Collections.singletonMap(ToXContentParams.FOR_INTERNAL_STORAGE, Boolean.toString(true))
177-
);
178-
IndexRequest doc = new IndexRequest(MlStatsIndex.writeAlias());
179-
doc.id(docId);
180-
try (XContentBuilder builder = XContentFactory.jsonBuilder()) {
181-
object.toXContent(builder, params);
182-
doc.source(builder);
183-
client.index(doc).actionGet();
184-
}
174+
private void assertDocExists(String index, String docId) {
175+
assertTrue(client().prepareGet(index, docId).get().isExists());
176+
}
177+
178+
private void assertDocDoesNotExist(String index, String docId) {
179+
assertFalse(client().prepareGet(index, docId).get().isExists());
185180
}
186181
}

0 commit comments

Comments
 (0)