feat(telemetry aware plugin): adding counters and implementing telemetry aware plugin

Author: pyek-bot

plugin/src/main/java/org/opensearch/ml/plugin/MachineLearningPlugin.java

@@ -55,6 +55,9 @@
 import org.opensearch.env.Environment;
 import org.opensearch.env.NodeEnvironment;
 import org.opensearch.indices.SystemIndexDescriptor;
+import org.opensearch.jobscheduler.spi.JobSchedulerExtension;
+import org.opensearch.jobscheduler.spi.ScheduledJobParser;
+import org.opensearch.jobscheduler.spi.ScheduledJobRunner;
 import org.opensearch.ml.action.agents.DeleteAgentTransportAction;
 import org.opensearch.ml.action.agents.GetAgentTransportAction;
 import org.opensearch.ml.action.agents.TransportRegisterAgentAction;
@@ -111,6 +114,7 @@
 import org.opensearch.ml.cluster.DiscoveryNodeHelper;
 import org.opensearch.ml.cluster.MLCommonsClusterEventListener;
 import org.opensearch.ml.cluster.MLCommonsClusterManagerEventListener;
+import org.opensearch.ml.common.CommonValue;
 import org.opensearch.ml.common.FunctionName;
 import org.opensearch.ml.common.input.execute.anomalylocalization.AnomalyLocalizationInput;
 import org.opensearch.ml.common.input.execute.metricscorrelation.MetricsCorrelationInput;
@@ -199,7 +203,8 @@
 import org.opensearch.ml.engine.utils.AgentModelsSearcher;
 import org.opensearch.ml.helper.ConnectorAccessControlHelper;
 import org.opensearch.ml.helper.ModelAccessControlHelper;
-import org.opensearch.ml.jobs.MLBatchTaskUpdateJobRunner;
+import org.opensearch.ml.jobs.MLJobParameter;
+import org.opensearch.ml.jobs.MLJobRunner;
 import org.opensearch.ml.memory.ConversationalMemoryHandler;
 import org.opensearch.ml.memory.action.conversation.CreateConversationAction;
 import org.opensearch.ml.memory.action.conversation.CreateConversationTransportAction;
@@ -293,6 +298,7 @@
 import org.opensearch.ml.stats.MLNodeLevelStat;
 import org.opensearch.ml.stats.MLStat;
 import org.opensearch.ml.stats.MLStats;
+import org.opensearch.ml.stats.otel.counters.MLOperationalMetricsCounter;
 import org.opensearch.ml.stats.suppliers.CounterSupplier;
 import org.opensearch.ml.stats.suppliers.IndexStatusSupplier;
 import org.opensearch.ml.task.MLExecuteTaskRunner;
@@ -313,6 +319,7 @@
 import org.opensearch.plugins.SearchPipelinePlugin;
 import org.opensearch.plugins.SearchPlugin;
 import org.opensearch.plugins.SystemIndexPlugin;
+import org.opensearch.plugins.TelemetryAwarePlugin;
 import org.opensearch.remote.metadata.client.SdkClient;
 import org.opensearch.remote.metadata.client.impl.SdkClientFactory;
 import org.opensearch.repositories.RepositoriesService;
@@ -326,6 +333,8 @@
 import org.opensearch.searchpipelines.questionanswering.generative.GenerativeQARequestProcessor;
 import org.opensearch.searchpipelines.questionanswering.generative.GenerativeQAResponseProcessor;
 import org.opensearch.searchpipelines.questionanswering.generative.ext.GenerativeQAParamExtBuilder;
+import org.opensearch.telemetry.metrics.MetricsRegistry;
+import org.opensearch.telemetry.tracing.Tracer;
 import org.opensearch.threadpool.ExecutorBuilder;
 import org.opensearch.threadpool.FixedExecutorBuilder;
 import org.opensearch.threadpool.ThreadPool;
@@ -343,7 +352,9 @@ public class MachineLearningPlugin extends Plugin
         SearchPipelinePlugin,
         ExtensiblePlugin,
         IngestPlugin,
-        SystemIndexPlugin {
+        SystemIndexPlugin,
+        TelemetryAwarePlugin,
+        JobSchedulerExtension {
     public static final String ML_THREAD_POOL_PREFIX = "thread_pool.ml_commons.";
     public static final String GENERAL_THREAD_POOL = "opensearch_ml_general";
     public static final String SDK_CLIENT_THREAD_POOL = "opensearch_ml_sdkclient";
@@ -356,6 +367,8 @@ public class MachineLearningPlugin extends Plugin
     public static final String DEPLOY_THREAD_POOL = "opensearch_ml_deploy";
     public static final String ML_BASE_URI = "/_plugins/_ml";
 
+    public static final String ML_COMMONS_JOBS_TYPE = "opensearch_ml_commons_jobs";
+
     private MLStats mlStats;
     private MLModelCacheHelper modelCacheHelper;
     private MLTaskManager mlTaskManager;
@@ -398,11 +411,13 @@ public class MachineLearningPlugin extends Plugin
     private ScriptService scriptService;
     private Encryptor encryptor;
 
-    public MachineLearningPlugin(Settings settings) {
-        // Handle this here as this feature is tied to Search/Query API, not to a ml-common API
-        // and as such, it can't be lazy-loaded when a ml-commons API is invoked.
-        this.ragSearchPipelineEnabled = MLCommonsSettings.ML_COMMONS_RAG_PIPELINE_FEATURE_ENABLED.get(settings);
-    }
+//    public MachineLearningPlugin(Settings settings) {
+//        // Handle this here as this feature is tied to Search/Query API, not to a ml-common API
+//        // and as such, it can't be lazy-loaded when a ml-commons API is invoked.
+//        this.ragSearchPipelineEnabled = MLCommonsSettings.ML_COMMONS_RAG_PIPELINE_FEATURE_ENABLED.get(settings);
+//    }
+
+    public MachineLearningPlugin() {}
 
     @Override
     public List<ActionHandler<? extends ActionRequest, ? extends ActionResponse>> getActions() {
@@ -485,7 +500,9 @@ public Collection<Object> createComponents(
         NodeEnvironment nodeEnvironment,
         NamedWriteableRegistry namedWriteableRegistry,
         IndexNameExpressionResolver indexNameExpressionResolver,
-        Supplier<RepositoriesService> repositoriesServiceSupplier
+        Supplier<RepositoriesService> repositoriesServiceSupplier,
+        Tracer tracer,
+        MetricsRegistry metricsRegistry
     ) {
         this.indexUtils = new IndexUtils(client, clusterService);
         this.client = client;
@@ -723,7 +740,11 @@ public Collection<Object> createComponents(
             .getClusterSettings()
             .addSettingsUpdateConsumer(MLCommonsSettings.ML_COMMONS_RAG_PIPELINE_FEATURE_ENABLED, it -> ragSearchPipelineEnabled = it);
 
-        MLBatchTaskUpdateJobRunner.getJobRunnerInstance().initialize(clusterService, threadPool, client);
+        MLJobRunner.getInstance().initialize(clusterService, threadPool, client);
+
+        // todo: add setting
+         MLOperationalMetricsCounter.initialize(clusterService.getClusterName().toString(), metricsRegistry);
+        // MLAdoptionMetricsCounter.initialize(clusterService.getClusterName().toString(), metricsRegistry);
 
         return ImmutableList
             .of(
@@ -1174,4 +1195,24 @@ public Collection<SystemIndexDescriptor> getSystemIndexDescriptors(Settings sett
         systemIndexDescriptors.add(new SystemIndexDescriptor(ML_STOP_WORDS_INDEX, "ML Commons Stop Words Index"));
         return systemIndexDescriptors;
     }
+
+    @Override
+    public String getJobType() {
+        return ML_COMMONS_JOBS_TYPE;
+    }
+
+    @Override
+    public String getJobIndex() {
+        return CommonValue.ML_JOBS_INDEX;
+    }
+
+    @Override
+    public ScheduledJobRunner getJobRunner() {
+        return MLJobRunner.getInstance();
+    }
+
+    @Override
+    public ScheduledJobParser getJobParser() {
+        return (parser, id, jobDocVersion) -> MLJobParameter.parse(parser);
+    }
 }

plugin/src/main/java/org/opensearch/ml/stats/otel/counters/AbstractMLMetricsCounter.java

@@ -0,0 +1,38 @@
+package org.opensearch.ml.stats.otel.counters;
+
+import java.util.Map;
+import java.util.concurrent.ConcurrentHashMap;
+import java.util.stream.Stream;
+
+import org.opensearch.telemetry.metrics.Counter;
+import org.opensearch.telemetry.metrics.MetricsRegistry;
+import org.opensearch.telemetry.metrics.tags.Tags;
+
+public abstract class AbstractMLMetricsCounter<T extends Enum<T>> {
+    private static final String PREFIX = "ml.commons.";
+    private static final String UNIT = "1";
+    private static final String CLUSTER_NAME_TAG = "cluster_name";
+
+    protected final String clusterName;
+    protected final MetricsRegistry metricsRegistry;
+    protected final Map<T, Counter> metricCounterMap;
+
+    protected AbstractMLMetricsCounter(String clusterName, MetricsRegistry metricsRegistry, Class<T> metricClass) {
+        this.clusterName = clusterName;
+        this.metricsRegistry = metricsRegistry;
+        this.metricCounterMap = new ConcurrentHashMap<>();
+        Stream.of(metricClass.getEnumConstants()).forEach(metric -> metricCounterMap.computeIfAbsent(metric, this::createMetricCounter));
+    }
+
+    public void incrementCounter(T metric, Tags customTags) {
+        Counter counter = metricCounterMap.computeIfAbsent(metric, this::createMetricCounter);
+        Tags metricsTags = (customTags == null ? Tags.create() : customTags).addTag(CLUSTER_NAME_TAG, clusterName);
+        counter.add(1, metricsTags);
+    }
+
+    private Counter createMetricCounter(T metric) {
+        return metricsRegistry.createCounter(PREFIX + metric.name(), getMetricDescription(metric), UNIT);
+    }
+
+    protected abstract String getMetricDescription(T metric);
+}

plugin/src/main/java/org/opensearch/ml/stats/otel/counters/MLAdoptionMetricsCounter.java

@@ -0,0 +1,29 @@
+package org.opensearch.ml.stats.otel.counters;
+
+import org.opensearch.ml.stats.otel.metrics.AdoptionMetric;
+import org.opensearch.telemetry.metrics.MetricsRegistry;
+
+public class MLAdoptionMetricsCounter extends AbstractMLMetricsCounter<AdoptionMetric> {
+
+    private static MLAdoptionMetricsCounter instance;
+
+    private MLAdoptionMetricsCounter(String clusterName, MetricsRegistry metricsRegistry) {
+        super(clusterName, metricsRegistry, AdoptionMetric.class);
+    }
+
+    public static synchronized void initialize(String clusterName, MetricsRegistry metricsRegistry) {
+        instance = new MLAdoptionMetricsCounter(clusterName, metricsRegistry);
+    }
+
+    public static synchronized MLAdoptionMetricsCounter getInstance() {
+        if (instance == null) {
+            throw new IllegalStateException("MLAdoptionMetricsCounter is not initialized. Call initialize() first.");
+        }
+        return instance;
+    }
+
+    @Override
+    protected String getMetricDescription(AdoptionMetric metric) {
+        return metric.getDescription();
+    }
+}

plugin/src/main/java/org/opensearch/ml/stats/otel/counters/MLOperationalMetricsCounter.java

@@ -0,0 +1,30 @@
+package org.opensearch.ml.stats.otel.counters;
+
+import org.opensearch.ml.stats.otel.metrics.OperationalMetric;
+import org.opensearch.telemetry.metrics.MetricsRegistry;
+
+public class MLOperationalMetricsCounter extends AbstractMLMetricsCounter<OperationalMetric> {
+
+    private static MLOperationalMetricsCounter instance;
+
+    private MLOperationalMetricsCounter(String clusterName, MetricsRegistry metricsRegistry) {
+        super(clusterName, metricsRegistry, OperationalMetric.class);
+    }
+
+    public static synchronized void initialize(String clusterName, MetricsRegistry metricsRegistry) {
+        instance = new MLOperationalMetricsCounter(clusterName, metricsRegistry);
+    }
+
+    public static synchronized MLOperationalMetricsCounter getInstance() {
+        if (instance == null) {
+            throw new IllegalStateException("MLOperationalMetricsCounter is not initialized. Call initialize() first.");
+        }
+
+        return instance;
+    }
+
+    @Override
+    protected String getMetricDescription(OperationalMetric metric) {
+        return metric.getDescription();
+    }
+}

plugin/src/main/java/org/opensearch/ml/stats/otel/metrics/AdoptionMetric.java

@@ -0,0 +1,15 @@
+package org.opensearch.ml.stats.otel.metrics;
+
+import lombok.Getter;
+
+@Getter
+public enum AdoptionMetric {
+    MODEL_COUNT("Number of models created"),
+    CONNECTOR_COUNT("Number of connectors created");
+
+    private final String description;
+
+    AdoptionMetric(String description) {
+        this.description = description;
+    }
+}

plugin/src/main/java/org/opensearch/ml/stats/otel/metrics/OperationalMetric.java

@@ -0,0 +1,15 @@
+package org.opensearch.ml.stats.otel.metrics;
+
+import lombok.Getter;
+
+@Getter
+public enum OperationalMetric {
+    MODEL_PREDICT_COUNT("Total number of predict calls made"),
+    MODEL_PREDICT_LATENCY("Latency for model predict");
+
+    private final String description;
+
+    OperationalMetric(String description) {
+        this.description = description;
+    }
+}

plugin/src/main/java/org/opensearch/ml/task/MLPredictTaskRunner.java

@@ -6,8 +6,8 @@
 package org.opensearch.ml.task;
 
 import static org.opensearch.core.xcontent.XContentParserUtils.ensureExpectedToken;
+import static org.opensearch.ml.common.CommonValue.ML_JOBS_INDEX;
 import static org.opensearch.ml.common.CommonValue.ML_MODEL_INDEX;
-import static org.opensearch.ml.common.CommonValue.TASK_POLLING_JOB_INDEX;
 import static org.opensearch.ml.common.MLModel.ALGORITHM_FIELD;
 import static org.opensearch.ml.common.utils.StringUtils.getErrorMessage;
 import static org.opensearch.ml.permission.AccessController.checkUserPermissions;
@@ -73,7 +73,10 @@
 import org.opensearch.ml.stats.MLActionLevelStat;
 import org.opensearch.ml.stats.MLNodeLevelStat;
 import org.opensearch.ml.stats.MLStats;
+import org.opensearch.ml.stats.otel.counters.MLOperationalMetricsCounter;
+import org.opensearch.ml.stats.otel.metrics.OperationalMetric;
 import org.opensearch.ml.utils.MLNodeUtils;
+import org.opensearch.telemetry.metrics.tags.Tags;
 import org.opensearch.threadpool.ThreadPool;
 import org.opensearch.transport.TransportResponseHandler;
 import org.opensearch.transport.TransportService;
@@ -433,7 +436,8 @@ private void runPredict(
                                                     remoteJob
                                                 );
 
-                                                if (!clusterService.state().metadata().indices().containsKey(TASK_POLLING_JOB_INDEX)) {
+                                                // todo: logic for starting the job
+                                                if (!clusterService.state().metadata().indices().containsKey(ML_JOBS_INDEX)) {
                                                     mlTaskManager.startTaskPollingJob();
                                                 }
 
@@ -459,6 +463,7 @@ private void runPredict(
                             } else {
                                 handleAsyncMLTaskComplete(mlTask);
                                 mlModelManager.trackPredictDuration(modelId, startTime);
+                                MLOperationalMetricsCounter.getInstance().incrementCounter(OperationalMetric.MODEL_PREDICT_COUNT, Tags.create().addTag("MODEL_ID", modelId));
                                 internalListener.onResponse(output);
                             }
                         }, e -> handlePredictFailure(mlTask, internalListener, e, false, modelId, actionName));