-
Notifications
You must be signed in to change notification settings - Fork 24.9k
Commit
This commit does not belong to any branch on this repository, and may belong to a fork outside of the repository.
Create an annotation when a model snapshot is stored (#53783)
- Loading branch information
1 parent
c20dc02
commit 7c5c912
Showing
19 changed files
with
422 additions
and
80 deletions.
There are no files selected for viewing
This file contains bidirectional Unicode text that may be interpreted or compiled differently than what appears below. To review, open the file in an editor that reveals hidden Unicode characters.
Learn more about bidirectional Unicode characters
This file contains bidirectional Unicode text that may be interpreted or compiled differently than what appears below. To review, open the file in an editor that reveals hidden Unicode characters.
Learn more about bidirectional Unicode characters
67 changes: 67 additions & 0 deletions
67
...n/core/src/main/java/org/elasticsearch/xpack/core/ml/annotations/AnnotationPersister.java
This file contains bidirectional Unicode text that may be interpreted or compiled differently than what appears below. To review, open the file in an editor that reveals hidden Unicode characters.
Learn more about bidirectional Unicode characters
| Original file line number | Diff line number | Diff line change |
|---|---|---|
| @@ -0,0 +1,67 @@ | ||
| /* | ||
| * Copyright Elasticsearch B.V. and/or licensed to Elasticsearch B.V. under one | ||
| * or more contributor license agreements. Licensed under the Elastic License; | ||
| * you may not use this file except in compliance with the Elastic License. | ||
| */ | ||
| package org.elasticsearch.xpack.core.ml.annotations; | ||
|
|
||
| import org.apache.logging.log4j.LogManager; | ||
| import org.apache.logging.log4j.Logger; | ||
| import org.elasticsearch.action.index.IndexRequest; | ||
| import org.elasticsearch.action.index.IndexResponse; | ||
| import org.elasticsearch.client.Client; | ||
| import org.elasticsearch.common.Nullable; | ||
| import org.elasticsearch.common.collect.Tuple; | ||
| import org.elasticsearch.common.util.concurrent.ThreadContext; | ||
| import org.elasticsearch.common.xcontent.ToXContent; | ||
| import org.elasticsearch.common.xcontent.XContentBuilder; | ||
| import org.elasticsearch.common.xcontent.XContentFactory; | ||
| import org.elasticsearch.xpack.core.common.notifications.AbstractAuditor; | ||
|
|
||
| import java.io.IOException; | ||
| import java.util.Objects; | ||
|
|
||
| import static org.elasticsearch.xpack.core.ClientHelper.ML_ORIGIN; | ||
|
|
||
| /** | ||
| * Persists annotations to Elasticsearch index. | ||
| */ | ||
| public class AnnotationPersister { | ||
|
|
||
| private static final Logger logger = LogManager.getLogger(AnnotationPersister.class); | ||
|
|
||
| private final Client client; | ||
| private final AbstractAuditor auditor; | ||
|
|
||
| public AnnotationPersister(Client client, AbstractAuditor auditor) { | ||
| this.client = Objects.requireNonNull(client); | ||
| this.auditor = Objects.requireNonNull(auditor); | ||
| } | ||
|
|
||
| /** | ||
| * Persists the given annotation to annotations index. | ||
| * | ||
| * @param annotationId existing annotation id. If {@code null}, a new annotation will be created and id will be assigned automatically | ||
| * @param annotation annotation to be persisted | ||
| * @param errorMessage error message to report when annotation fails to be persisted | ||
| * @return tuple of the form (annotation id, annotation object) | ||
| */ | ||
| public Tuple<String, Annotation> persistAnnotation(@Nullable String annotationId, Annotation annotation, String errorMessage) { | ||
| Objects.requireNonNull(annotation); | ||
| try (XContentBuilder xContentBuilder = annotation.toXContent(XContentFactory.jsonBuilder(), ToXContent.EMPTY_PARAMS)) { | ||
| IndexRequest indexRequest = | ||
| new IndexRequest(AnnotationIndex.WRITE_ALIAS_NAME) | ||
| .id(annotationId) | ||
| .source(xContentBuilder); | ||
| try (ThreadContext.StoredContext ignore = client.threadPool().getThreadContext().stashWithOrigin(ML_ORIGIN)) { | ||
| IndexResponse response = client.index(indexRequest).actionGet(); | ||
| return Tuple.tuple(response.getId(), annotation); | ||
| } | ||
| } catch (IOException ex) { | ||
| String jobId = annotation.getJobId(); | ||
| logger.error(errorMessage, ex); | ||
| auditor.error(jobId, errorMessage); | ||
| return null; | ||
| } | ||
| } | ||
| } |
This file contains bidirectional Unicode text that may be interpreted or compiled differently than what appears below. To review, open the file in an editor that reveals hidden Unicode characters.
Learn more about bidirectional Unicode characters
This file contains bidirectional Unicode text that may be interpreted or compiled differently than what appears below. To review, open the file in an editor that reveals hidden Unicode characters.
Learn more about bidirectional Unicode characters
119 changes: 119 additions & 0 deletions
119
...e/src/test/java/org/elasticsearch/xpack/core/ml/annotations/AnnotationPersisterTests.java
This file contains bidirectional Unicode text that may be interpreted or compiled differently than what appears below. To review, open the file in an editor that reveals hidden Unicode characters.
Learn more about bidirectional Unicode characters
| Original file line number | Diff line number | Diff line change |
|---|---|---|
| @@ -0,0 +1,119 @@ | ||
| /* | ||
| * Copyright Elasticsearch B.V. and/or licensed to Elasticsearch B.V. under one | ||
| * or more contributor license agreements. Licensed under the Elastic License; | ||
| * you may not use this file except in compliance with the Elastic License. | ||
| */ | ||
| package org.elasticsearch.xpack.core.ml.annotations; | ||
|
|
||
| import org.elasticsearch.action.ActionFuture; | ||
| import org.elasticsearch.action.DocWriteRequest; | ||
| import org.elasticsearch.action.index.IndexRequest; | ||
| import org.elasticsearch.action.index.IndexResponse; | ||
| import org.elasticsearch.client.Client; | ||
| import org.elasticsearch.common.bytes.BytesReference; | ||
| import org.elasticsearch.common.collect.Tuple; | ||
| import org.elasticsearch.common.settings.Settings; | ||
| import org.elasticsearch.common.util.concurrent.ThreadContext; | ||
| import org.elasticsearch.common.xcontent.XContentParser; | ||
| import org.elasticsearch.test.ESTestCase; | ||
| import org.elasticsearch.threadpool.ThreadPool; | ||
| import org.elasticsearch.xpack.core.common.notifications.AbstractAuditor; | ||
| import org.junit.Before; | ||
| import org.mockito.ArgumentCaptor; | ||
| import org.mockito.InOrder; | ||
|
|
||
| import java.io.IOException; | ||
|
|
||
| import static org.elasticsearch.common.collect.Tuple.tuple; | ||
| import static org.elasticsearch.common.xcontent.json.JsonXContent.jsonXContent; | ||
| import static org.hamcrest.Matchers.equalTo; | ||
| import static org.hamcrest.Matchers.is; | ||
| import static org.hamcrest.Matchers.nullValue; | ||
| import static org.mockito.Matchers.any; | ||
| import static org.mockito.Mockito.doReturn; | ||
| import static org.mockito.Mockito.inOrder; | ||
| import static org.mockito.Mockito.mock; | ||
| import static org.mockito.Mockito.verifyNoMoreInteractions; | ||
| import static org.mockito.Mockito.when; | ||
|
|
||
| public class AnnotationPersisterTests extends ESTestCase { | ||
|
|
||
| private static final String ANNOTATION_ID = "existing_annotation_id"; | ||
| private static final String ERROR_MESSAGE = "an error occurred while persisting annotation"; | ||
|
|
||
| private Client client; | ||
| private AbstractAuditor auditor; | ||
| private IndexResponse indexResponse; | ||
|
|
||
| private ArgumentCaptor<IndexRequest> indexRequestCaptor; | ||
|
|
||
| @Before | ||
| public void setUpMocks() { | ||
| ThreadContext threadContext = new ThreadContext(Settings.EMPTY); | ||
| ThreadPool threadPool = mock(ThreadPool.class); | ||
| when(threadPool.getThreadContext()).thenReturn(threadContext); | ||
| client = mock(Client.class); | ||
| when(client.threadPool()).thenReturn(threadPool); | ||
|
|
||
| auditor = mock(AbstractAuditor.class); | ||
|
|
||
| indexResponse = mock(IndexResponse.class); | ||
| when(indexResponse.getId()).thenReturn(ANNOTATION_ID); | ||
|
|
||
| indexRequestCaptor = ArgumentCaptor.forClass(IndexRequest.class); | ||
| } | ||
|
|
||
| public void testPersistAnnotation_Create() throws IOException { | ||
| doReturn(instantFuture(indexResponse)).when(client).index(any()); | ||
|
|
||
| AnnotationPersister persister = new AnnotationPersister(client, auditor); | ||
| Annotation annotation = AnnotationTests.randomAnnotation(); | ||
| Tuple<String, Annotation> result = persister.persistAnnotation(null, annotation, ERROR_MESSAGE); | ||
| assertThat(result, is(equalTo(tuple(ANNOTATION_ID, annotation)))); | ||
|
|
||
| InOrder inOrder = inOrder(client); | ||
| inOrder.verify(client).threadPool(); | ||
| inOrder.verify(client).index(indexRequestCaptor.capture()); | ||
| verifyNoMoreInteractions(client, auditor); | ||
|
|
||
| IndexRequest indexRequest = indexRequestCaptor.getValue(); | ||
|
|
||
| assertThat(indexRequest.index(), is(equalTo(AnnotationIndex.WRITE_ALIAS_NAME))); | ||
| assertThat(indexRequest.id(), is(nullValue())); | ||
| assertThat(parseAnnotation(indexRequest.source()), is(equalTo(annotation))); | ||
| assertThat(indexRequest.opType(), equalTo(DocWriteRequest.OpType.INDEX)); | ||
| } | ||
|
|
||
| public void testPersistAnnotation_Update() throws IOException { | ||
| doReturn(instantFuture(indexResponse)).when(client).index(any()); | ||
|
|
||
| AnnotationPersister persister = new AnnotationPersister(client, auditor); | ||
| Annotation annotation = AnnotationTests.randomAnnotation(); | ||
| Tuple<String, Annotation> result = persister.persistAnnotation(ANNOTATION_ID, annotation, ERROR_MESSAGE); | ||
| assertThat(result, is(equalTo(tuple(ANNOTATION_ID, annotation)))); | ||
|
|
||
| InOrder inOrder = inOrder(client); | ||
| inOrder.verify(client).threadPool(); | ||
| inOrder.verify(client).index(indexRequestCaptor.capture()); | ||
| verifyNoMoreInteractions(client, auditor); | ||
|
|
||
| IndexRequest indexRequest = indexRequestCaptor.getValue(); | ||
| assertThat(indexRequest.index(), is(equalTo(AnnotationIndex.WRITE_ALIAS_NAME))); | ||
| assertThat(indexRequest.id(), is(equalTo(ANNOTATION_ID))); | ||
| assertThat(parseAnnotation(indexRequest.source()), is(equalTo(annotation))); | ||
| assertThat(indexRequest.opType(), equalTo(DocWriteRequest.OpType.INDEX)); | ||
| } | ||
|
|
||
| @SuppressWarnings("unchecked") | ||
| private static <T> ActionFuture<T> instantFuture(T response) { | ||
| ActionFuture future = mock(ActionFuture.class); | ||
| when(future.actionGet()).thenReturn(response); | ||
| return future; | ||
| } | ||
|
|
||
| private Annotation parseAnnotation(BytesReference source) throws IOException { | ||
| try (XContentParser parser = createParser(jsonXContent, source)) { | ||
| return Annotation.PARSER.parse(parser, null); | ||
| } | ||
| } | ||
| } |
This file contains bidirectional Unicode text that may be interpreted or compiled differently than what appears below. To review, open the file in an editor that reveals hidden Unicode characters.
Learn more about bidirectional Unicode characters
This file contains bidirectional Unicode text that may be interpreted or compiled differently than what appears below. To review, open the file in an editor that reveals hidden Unicode characters.
Learn more about bidirectional Unicode characters
This file contains bidirectional Unicode text that may be interpreted or compiled differently than what appears below. To review, open the file in an editor that reveals hidden Unicode characters.
Learn more about bidirectional Unicode characters
Oops, something went wrong.