Add transport action for model inference

src/main/java/com/amazon/opendistroforelasticsearch/ad/transport/EntityResultAction.java

@@ -0,0 +1,29 @@
+/*
+ * Copyright 2020 Amazon.com, Inc. or its affiliates. All Rights Reserved.
+ *
+ * Licensed under the Apache License, Version 2.0 (the "License").
+ * You may not use this file except in compliance with the License.
+ * A copy of the License is located at
+ *
+ *     http://www.apache.org/licenses/LICENSE-2.0
+ *
+ * or in the "license" file accompanying this file. This file is distributed
+ * on an "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either
+ * express or implied. See the License for the specific language governing
+ * permissions and limitations under the License.
+ */
+
+package com.amazon.opendistroforelasticsearch.ad.transport;
+
+import org.elasticsearch.action.ActionType;
+import org.elasticsearch.action.support.master.AcknowledgedResponse;
+
+public class EntityResultAction extends ActionType<AcknowledgedResponse> {
+    public static final EntityResultAction INSTANCE = new EntityResultAction();
+    public static final String NAME = "cluster:admin/ad/entity/result";
+
+    private EntityResultAction() {
+        super(NAME, AcknowledgedResponse::new);
+    }
+
+}

src/main/java/com/amazon/opendistroforelasticsearch/ad/transport/EntityResultRequest.java

@@ -0,0 +1,110 @@
+/*
+ * Copyright 2020 Amazon.com, Inc. or its affiliates. All Rights Reserved.
+ *
+ * Licensed under the Apache License, Version 2.0 (the "License").
+ * You may not use this file except in compliance with the License.
+ * A copy of the License is located at
+ *
+ *     http://www.apache.org/licenses/LICENSE-2.0
+ *
+ * or in the "license" file accompanying this file. This file is distributed
+ * on an "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either
+ * express or implied. See the License for the specific language governing
+ * permissions and limitations under the License.
+ */
+
+package com.amazon.opendistroforelasticsearch.ad.transport;
+
+import static org.elasticsearch.action.ValidateActions.addValidationError;
+
+import java.io.IOException;
+import java.util.Locale;
+import java.util.Map;
+
+import org.elasticsearch.action.ActionRequest;
+import org.elasticsearch.action.ActionRequestValidationException;
+import org.elasticsearch.common.Strings;
+import org.elasticsearch.common.io.stream.StreamInput;
+import org.elasticsearch.common.io.stream.StreamOutput;
+import org.elasticsearch.common.xcontent.ToXContentObject;
+import org.elasticsearch.common.xcontent.XContentBuilder;
+
+import com.amazon.opendistroforelasticsearch.ad.constant.CommonErrorMessages;
+import com.amazon.opendistroforelasticsearch.ad.constant.CommonMessageAttributes;
+
+public class EntityResultRequest extends ActionRequest implements ToXContentObject {
+
+    private String detectorId;
+    private Map<String, double[]> entities;
+    private long start;
+    private long end;
+
+    public EntityResultRequest(StreamInput in) throws IOException {
+        super(in);
+        this.detectorId = in.readString();
+        this.entities = in.readMap(StreamInput::readString, StreamInput::readDoubleArray);
+        this.start = in.readLong();
+        this.end = in.readLong();
+    }
+
+    public EntityResultRequest(String detectorId, Map<String, double[]> entities, long start, long end) {
+        super();
+        this.detectorId = detectorId;
+        this.entities = entities;
+        this.start = start;
+        this.end = end;
+    }
+
+    public String getDetectorId() {
+        return this.detectorId;
+    }
+
+    public Map<String, double[]> getEntities() {
+        return this.entities;
+    }
+
+    public long getStart() {
+        return this.start;
+    }
+
+    public long getEnd() {
+        return this.end;
+    }
+
+    @Override
+    public void writeTo(StreamOutput out) throws IOException {
+        super.writeTo(out);
+        out.writeString(this.detectorId);
+        out.writeMap(this.entities, StreamOutput::writeString, StreamOutput::writeDoubleArray);
+        out.writeLong(this.start);
+        out.writeLong(this.end);
+    }
+
+    @Override
+    public ActionRequestValidationException validate() {
+        ActionRequestValidationException validationException = null;
+        if (Strings.isEmpty(detectorId)) {
+            validationException = addValidationError(CommonErrorMessages.AD_ID_MISSING_MSG, validationException);
+        }
+        if (start <= 0 || end <= 0 || start > end) {
+            validationException = addValidationError(
+                String.format(Locale.ROOT, "%s: start %d, end %d", CommonErrorMessages.INVALID_TIMESTAMP_ERR_MSG, start, end),
+                validationException
+            );
+        }
+        return validationException;
+    }
+
+    @Override
+    public XContentBuilder toXContent(XContentBuilder builder, Params params) throws IOException {
+        builder.startObject();
+        builder.field(CommonMessageAttributes.ID_JSON_KEY, detectorId);
+        builder.field(CommonMessageAttributes.START_JSON_KEY, start);
+        builder.field(CommonMessageAttributes.END_JSON_KEY, end);
+        for (String entity : entities.keySet()) {
+            builder.field(entity, entities.get(entity));
+        }
+        builder.endObject();
+        return builder;
+    }
+}

src/main/java/com/amazon/opendistroforelasticsearch/ad/transport/EntityResultTransportAction.java

@@ -0,0 +1,181 @@
+/*
+ * Copyright 2020 Amazon.com, Inc. or its affiliates. All Rights Reserved.
+ *
+ * Licensed under the Apache License, Version 2.0 (the "License").
+ * You may not use this file except in compliance with the License.
+ * A copy of the License is located at
+ *
+ *     http://www.apache.org/licenses/LICENSE-2.0
+ *
+ * or in the "license" file accompanying this file. This file is distributed
+ * on an "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either
+ * express or implied. See the License for the specific language governing
+ * permissions and limitations under the License.
+ */
+
+package com.amazon.opendistroforelasticsearch.ad.transport;
+
+import static com.amazon.opendistroforelasticsearch.ad.settings.AnomalyDetectorSettings.COOLDOWN_MINUTES;
+
+import java.time.Clock;
+import java.time.Duration;
+import java.time.Instant;
+import java.util.Arrays;
+import java.util.Map.Entry;
+import java.util.Optional;
+
+import org.apache.logging.log4j.LogManager;
+import org.apache.logging.log4j.Logger;
+import org.elasticsearch.action.ActionListener;
+import org.elasticsearch.action.support.ActionFilters;
+import org.elasticsearch.action.support.HandledTransportAction;
+import org.elasticsearch.action.support.master.AcknowledgedResponse;
+import org.elasticsearch.common.inject.Inject;
+import org.elasticsearch.common.settings.Settings;
+import org.elasticsearch.tasks.Task;
+import org.elasticsearch.transport.TransportService;
+
+import com.amazon.opendistroforelasticsearch.ad.NodeStateManager;
+import com.amazon.opendistroforelasticsearch.ad.breaker.ADCircuitBreakerService;
+import com.amazon.opendistroforelasticsearch.ad.caching.CacheProvider;
+import com.amazon.opendistroforelasticsearch.ad.caching.EntityCache;
+import com.amazon.opendistroforelasticsearch.ad.common.exception.EndRunException;
+import com.amazon.opendistroforelasticsearch.ad.common.exception.LimitExceededException;
+import com.amazon.opendistroforelasticsearch.ad.constant.CommonErrorMessages;
+import com.amazon.opendistroforelasticsearch.ad.ml.CheckpointDao;
+import com.amazon.opendistroforelasticsearch.ad.ml.EntityModel;
+import com.amazon.opendistroforelasticsearch.ad.ml.ModelManager;
+import com.amazon.opendistroforelasticsearch.ad.ml.ModelState;
+import com.amazon.opendistroforelasticsearch.ad.ml.ThresholdingResult;
+import com.amazon.opendistroforelasticsearch.ad.model.AnomalyDetector;
+import com.amazon.opendistroforelasticsearch.ad.model.AnomalyResult;
+import com.amazon.opendistroforelasticsearch.ad.model.Entity;
+import com.amazon.opendistroforelasticsearch.ad.settings.AnomalyDetectorSettings;
+import com.amazon.opendistroforelasticsearch.ad.transport.handler.MultitiEntityResultHandler;
+import com.amazon.opendistroforelasticsearch.ad.util.ParseUtils;
+
+public class EntityResultTransportAction extends HandledTransportAction<EntityResultRequest, AcknowledgedResponse> {
+
+    private static final Logger LOG = LogManager.getLogger(EntityResultTransportAction.class);
+    private ModelManager manager;
+    private ADCircuitBreakerService adCircuitBreakerService;
+    private MultitiEntityResultHandler anomalyResultHandler;
+    private CheckpointDao checkpointDao;
+    private EntityCache cache;
+    private final NodeStateManager stateManager;
+    private final int coolDownMinutes;
+    private final Clock clock;
+
+    @Inject
+    public EntityResultTransportAction(
+        ActionFilters actionFilters,
+        TransportService transportService,
+        ModelManager manager,
+        ADCircuitBreakerService adCircuitBreakerService,
+        MultitiEntityResultHandler anomalyResultHandler,
+        CheckpointDao checkpointDao,
+        CacheProvider entityCache,
+        NodeStateManager stateManager,
+        Settings settings,
+        Clock clock
+    ) {
+        super(EntityResultAction.NAME, transportService, actionFilters, EntityResultRequest::new);
+        this.manager = manager;
+        this.adCircuitBreakerService = adCircuitBreakerService;
+        this.anomalyResultHandler = anomalyResultHandler;
+        this.checkpointDao = checkpointDao;
+        this.cache = entityCache;
+        this.stateManager = stateManager;
+        this.coolDownMinutes = (int) (COOLDOWN_MINUTES.get(settings).getMinutes());
+        this.clock = clock;
+    }
+
+    @Override
+    protected void doExecute(Task task, EntityResultRequest request, ActionListener<AcknowledgedResponse> listener) {
+        if (adCircuitBreakerService.isOpen()) {
+            listener.onFailure(new LimitExceededException(request.getDetectorId(), CommonErrorMessages.MEMORY_CIRCUIT_BROKEN_ERR_MSG));
+            return;
+        }
+
+        try {
+            String detectorId = request.getDetectorId();
+            stateManager.getAnomalyDetector(detectorId, onGetDetector(listener, detectorId, request));
+        } catch (Exception exception) {
+            LOG.error("fail to get entity's anomaly grade", exception);
+            listener.onFailure(exception);
+        }
+
+    }
+
+    private ActionListener<Optional<AnomalyDetector>> onGetDetector(
+        ActionListener<AcknowledgedResponse> listener,
+        String detectorId,
+        EntityResultRequest request
+    ) {
+        return ActionListener.wrap(detectorOptional -> {
+            if (!detectorOptional.isPresent()) {
+                listener.onFailure(new EndRunException(detectorId, "AnomalyDetector is not available.", true));
+                return;
+            }
+
+            AnomalyDetector detector = detectorOptional.get();
+            // we only support 1 categorical field now
+            String categoricalField = detector.getCategoryField().get(0);
+
+            ADResultBulkRequest currentBulkRequest = new ADResultBulkRequest();
+            // index pressure is high. Only save anomalies
+            boolean onlySaveAnomalies = stateManager
+                .getLastIndexThrottledTime()
+                .plus(Duration.ofMinutes(coolDownMinutes))
+                .isAfter(clock.instant());
+
+            for (Entry<String, double[]> entity : request.getEntities().entrySet()) {
+                String entityName = entity.getKey();
+                // For ES, the limit of the document ID is 512 bytes.
+                // truncate to 256 characters if too long since we are using it as part of document id.
+                if (entityName.length() > AnomalyDetectorSettings.MAX_ENTITY_LENGTH) {
+                    continue;
+                }
+
+                double[] datapoint = entity.getValue();
+                String modelId = manager.getEntityModelId(detectorId, entityName);
+                ModelState<EntityModel> entityModel = cache.get(modelId, detector, datapoint, entityName);
+                if (entityModel == null) {
+                    // cache miss
+                    continue;
+                }
+                ThresholdingResult result = manager.getAnomalyResultForEntity(detectorId, datapoint, entityName, entityModel, modelId);
+                // result.getRcfScore() = 0 means the model is not initialized
+                // result.getGrade() = 0 means it is not an anomaly
+                // So many EsRejectedExecutionException if we write no matter what
+                if (result.getRcfScore() > 0 && (!onlySaveAnomalies || result.getGrade() > 0)) {
+                    this.anomalyResultHandler
+                        .write(
+                            new AnomalyResult(
+                                detectorId,
+                                result.getRcfScore(),
+                                result.getGrade(),
+                                result.getConfidence(),
+                                ParseUtils.getFeatureData(datapoint, detector),
+                                Instant.ofEpochMilli(request.getStart()),
+                                Instant.ofEpochMilli(request.getEnd()),
+                                Instant.now(),
+                                Instant.now(),
+                                null,
+                                Arrays.asList(new Entity(categoricalField, entityName))
+                            ),
+                            currentBulkRequest
+                        );
+                }
+            }
+            this.anomalyResultHandler.flush(currentBulkRequest, detectorId);
+            // bulk all accumulated checkpoint requests
+            this.checkpointDao.flush();
+
+            listener.onResponse(new AcknowledgedResponse(true));
+        }, exception -> {
+            LOG.error("fail to get entity's anomaly grade", exception);
+            listener.onFailure(exception);
+        });
+    }
+}

src/test/java/com/amazon/opendistroforelasticsearch/ad/AbstractADTest.java

@@ -16,13 +16,18 @@
 package com.amazon.opendistroforelasticsearch.ad;
 
 import static org.hamcrest.Matchers.containsString;
+import static org.mockito.ArgumentMatchers.any;
+import static org.mockito.Mockito.doAnswer;
+import static org.mockito.Mockito.mock;
+import static org.mockito.Mockito.when;
 
 import java.util.ArrayList;
 import java.util.Arrays;
 import java.util.Collections;
 import java.util.HashMap;
 import java.util.List;
 import java.util.Map;
+import java.util.concurrent.ExecutorService;
 import java.util.concurrent.TimeUnit;
 import java.util.regex.Matcher;
 import java.util.regex.Pattern;
@@ -34,8 +39,11 @@
 import org.apache.logging.log4j.core.appender.AbstractAppender;
 import org.apache.logging.log4j.core.layout.PatternLayout;
 import org.apache.logging.log4j.util.StackLocatorUtil;
+import org.elasticsearch.Version;
 import org.elasticsearch.action.ActionResponse;
 import org.elasticsearch.action.support.PlainActionFuture;
+import org.elasticsearch.cluster.metadata.AliasMetadata;
+import org.elasticsearch.cluster.metadata.IndexMetadata;
 import org.elasticsearch.common.bytes.BytesReference;
 import org.elasticsearch.common.settings.Settings;
 import org.elasticsearch.common.xcontent.NamedXContentRegistry;
@@ -318,4 +326,40 @@ public HttpRequest releaseAndCopy() {
 
         }, null);
     }
+
+    protected boolean areEqualWithArrayValue(Map<String, double[]> first, Map<String, double[]> second) {
+        if (first.size() != second.size()) {
+            return false;
+        }
+
+        return first.entrySet().stream().allMatch(e -> Arrays.equals(e.getValue(), second.get(e.getKey())));
+    }
+
+    protected IndexMetadata indexMeta(String name, long creationDate, String... aliases) {
+        IndexMetadata.Builder builder = IndexMetadata
+            .builder(name)
+            .settings(
+                Settings
+                    .builder()
+                    .put("index.number_of_shards", 1)
+                    .put("index.number_of_replicas", 1)
+                    .put("index.version.created", Version.CURRENT.id)
+            );
+        builder.creationDate(creationDate);
+        for (String alias : aliases) {
+            builder.putAlias(AliasMetadata.builder(alias).build());
+        }
+        return builder.build();
+    }
+
+    protected void setUpADThreadPool(ThreadPool mockThreadPool) {
+        ExecutorService executorService = mock(ExecutorService.class);
+
+        when(mockThreadPool.executor(AnomalyDetectorPlugin.AD_THREAD_POOL_NAME)).thenReturn(executorService);
+        doAnswer(invocation -> {
+            Runnable runnable = invocation.getArgument(0);
+            runnable.run();
+            return null;
+        }).when(executorService).execute(any(Runnable.class));
+    }
 }