feat: add settings to control metric collection

Signed-off-by: Pavan Yekbote <[email protected]>

Author: pyek-bot

common/src/main/java/org/opensearch/ml/common/settings/MLCommonsSettings.java

@@ -342,4 +342,12 @@ private MLCommonsSettings() {}
     /** This setting sets the remote metadata service name */
     public static final Setting<String> REMOTE_METADATA_SERVICE_NAME = Setting
         .simpleString("plugins.ml_commons." + REMOTE_METADATA_SERVICE_NAME_KEY, Setting.Property.NodeScope, Setting.Property.Final);
+
+    // Feature flag for enabling telemetry metric collection via metrics framework
+    public static final Setting<Boolean> ML_COMMONS_METRIC_COLLECTION_ENABLED = Setting
+        .boolSetting("plugins.ml_commons.metrics_collection_enabled", false, Setting.Property.NodeScope, Setting.Property.Dynamic);
+
+    // Feature flag for enabling telemetry static metric collection job -- MLStatsJobProcessor
+    public static final Setting<Boolean> ML_COMMONS_STATIC_METRIC_COLLECTION_ENABLED = Setting
+        .boolSetting("plugins.ml_commons.metrics_static_collection_enabled", false, Setting.Property.NodeScope, Setting.Property.Dynamic);
 }

common/src/main/java/org/opensearch/ml/common/settings/MLFeatureEnabledSetting.java

@@ -11,10 +11,13 @@
 import static org.opensearch.ml.common.settings.MLCommonsSettings.ML_COMMONS_CONNECTOR_PRIVATE_IP_ENABLED;
 import static org.opensearch.ml.common.settings.MLCommonsSettings.ML_COMMONS_CONTROLLER_ENABLED;
 import static org.opensearch.ml.common.settings.MLCommonsSettings.ML_COMMONS_LOCAL_MODEL_ENABLED;
+import static org.opensearch.ml.common.settings.MLCommonsSettings.ML_COMMONS_METRIC_COLLECTION_ENABLED;
 import static org.opensearch.ml.common.settings.MLCommonsSettings.ML_COMMONS_MULTI_TENANCY_ENABLED;
 import static org.opensearch.ml.common.settings.MLCommonsSettings.ML_COMMONS_OFFLINE_BATCH_INFERENCE_ENABLED;
 import static org.opensearch.ml.common.settings.MLCommonsSettings.ML_COMMONS_OFFLINE_BATCH_INGESTION_ENABLED;
+import static org.opensearch.ml.common.settings.MLCommonsSettings.ML_COMMONS_RAG_PIPELINE_FEATURE_ENABLED;
 import static org.opensearch.ml.common.settings.MLCommonsSettings.ML_COMMONS_REMOTE_INFERENCE_ENABLED;
+import static org.opensearch.ml.common.settings.MLCommonsSettings.ML_COMMONS_STATIC_METRIC_COLLECTION_ENABLED;
 
 import java.util.ArrayList;
 import java.util.List;
@@ -42,6 +45,11 @@ public class MLFeatureEnabledSetting {
 
     private volatile Boolean isRagSearchPipelineEnabled;
 
+    // block any push
+    private volatile Boolean isMetricCollectionEnabled;
+    // block static push
+    private volatile Boolean isStaticMetricCollectionEnabled;
+
     private final List<SettingsChangeListener> listeners = new ArrayList<>();
 
     public MLFeatureEnabledSetting(ClusterService clusterService, Settings settings) {
@@ -53,6 +61,9 @@ public MLFeatureEnabledSetting(ClusterService clusterService, Settings settings)
         isBatchIngestionEnabled = ML_COMMONS_OFFLINE_BATCH_INGESTION_ENABLED.get(settings);
         isBatchInferenceEnabled = ML_COMMONS_OFFLINE_BATCH_INFERENCE_ENABLED.get(settings);
         isMultiTenancyEnabled = ML_COMMONS_MULTI_TENANCY_ENABLED.get(settings);
+        isRagSearchPipelineEnabled = ML_COMMONS_RAG_PIPELINE_FEATURE_ENABLED.get(settings);
+        isMetricCollectionEnabled = ML_COMMONS_METRIC_COLLECTION_ENABLED.get(settings);
+        isStaticMetricCollectionEnabled = ML_COMMONS_STATIC_METRIC_COLLECTION_ENABLED.get(settings);
 
         clusterService
             .getClusterSettings()
@@ -74,6 +85,12 @@ public MLFeatureEnabledSetting(ClusterService clusterService, Settings settings)
         clusterService
             .getClusterSettings()
             .addSettingsUpdateConsumer(MLCommonsSettings.ML_COMMONS_RAG_PIPELINE_FEATURE_ENABLED, it -> isRagSearchPipelineEnabled = it);
+        clusterService
+            .getClusterSettings()
+            .addSettingsUpdateConsumer(ML_COMMONS_METRIC_COLLECTION_ENABLED, it -> isMetricCollectionEnabled = it);
+        clusterService
+            .getClusterSettings()
+            .addSettingsUpdateConsumer(ML_COMMONS_STATIC_METRIC_COLLECTION_ENABLED, it -> isStaticMetricCollectionEnabled = it);
     }
 
     /**
@@ -148,6 +165,14 @@ public boolean isRagSearchPipelineEnabled() {
         return isRagSearchPipelineEnabled;
     }
 
+    public boolean isMetricCollectionEnabled() {
+        return isMetricCollectionEnabled;
+    }
+
+    public boolean isStaticMetricCollectionEnabled() {
+        return isStaticMetricCollectionEnabled;
+    }
+
     @VisibleForTesting
     public void notifyMultiTenancyListeners(boolean isEnabled) {
         for (SettingsChangeListener listener : listeners) {

plugin/src/main/java/org/opensearch/ml/jobs/MLJobRunner.java

@@ -9,6 +9,7 @@
 import org.opensearch.jobscheduler.spi.JobExecutionContext;
 import org.opensearch.jobscheduler.spi.ScheduledJobParameter;
 import org.opensearch.jobscheduler.spi.ScheduledJobRunner;
+import org.opensearch.ml.common.settings.MLFeatureEnabledSetting;
 import org.opensearch.ml.helper.ConnectorAccessControlHelper;
 import org.opensearch.ml.jobs.processors.MLBatchTaskUpdateProcessor;
 import org.opensearch.ml.jobs.processors.MLStatsJobProcessor;
@@ -52,6 +53,9 @@ public static MLJobRunner getInstance() {
     @Setter
     private ConnectorAccessControlHelper connectorAccessControlHelper;
 
+    @Setter
+    private MLFeatureEnabledSetting mlFeatureEnabledSetting;
+
     private boolean initialized;
 
     private MLJobRunner() {
@@ -63,14 +67,16 @@ public void initialize(
         final ThreadPool threadPool,
         final Client client,
         final SdkClient sdkClient,
-        final ConnectorAccessControlHelper connectorAccessControlHelper
+        final ConnectorAccessControlHelper connectorAccessControlHelper,
+        final MLFeatureEnabledSetting mlFeatureEnabledSetting
     ) {
         this.clusterService = clusterService;
         this.threadPool = threadPool;
         this.client = client;
         this.sdkClient = sdkClient;
         this.connectorAccessControlHelper = connectorAccessControlHelper;
         this.initialized = true;
+        this.mlFeatureEnabledSetting = mlFeatureEnabledSetting;
     }
 
     @Override
@@ -84,7 +90,7 @@ public void runJob(ScheduledJobParameter scheduledJobParameter, JobExecutionCont
             case STATS_COLLECTOR:
                 MLStatsJobProcessor
                     .getInstance(clusterService, client, threadPool, connectorAccessControlHelper, sdkClient)
-                    .process(jobParameter, jobExecutionContext);
+                    .process(jobParameter, jobExecutionContext, mlFeatureEnabledSetting.isStaticMetricCollectionEnabled());
                 break;
             case BATCH_TASK_UPDATE:
                 MLBatchTaskUpdateProcessor.getInstance(clusterService, client, threadPool).process(jobParameter, jobExecutionContext);

plugin/src/main/java/org/opensearch/ml/jobs/processors/MLJobProcessor.java

@@ -14,6 +14,7 @@
 import org.opensearch.jobscheduler.spi.JobExecutionContext;
 import org.opensearch.jobscheduler.spi.ScheduledJobParameter;
 import org.opensearch.jobscheduler.spi.utils.LockService;
+import org.opensearch.ml.common.exception.MLException;
 import org.opensearch.threadpool.ThreadPool;
 import org.opensearch.transport.client.Client;
 
@@ -33,6 +34,14 @@ public MLJobProcessor(ClusterService clusterService, Client client, ThreadPool t
 
     public abstract void run();
 
+    public void process(ScheduledJobParameter scheduledJobParameter, JobExecutionContext jobExecutionContext, boolean isProcessorEnabled) {
+        if (!isProcessorEnabled) {
+            throw new MLException(scheduledJobParameter.getName() + " not enabled.");
+        }
+
+        process(scheduledJobParameter, jobExecutionContext);
+    }
+
     public void process(ScheduledJobParameter scheduledJobParameter, JobExecutionContext jobExecutionContext) {
         final LockService lockService = jobExecutionContext.getLockService();
 

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

@@ -764,11 +764,11 @@ public Collection<Object> createComponents(
             mlFeatureEnabledSetting
         );
 
-        MLJobRunner.getInstance().initialize(clusterService, threadPool, client, sdkClient, connectorAccessControlHelper);
-
-        // todo: add setting
-        MLOperationalMetricsCounter.initialize(clusterService.getClusterName().toString(), metricsRegistry);
-        MLAdoptionMetricsCounter.initialize(clusterService.getClusterName().toString(), metricsRegistry);
+        MLJobRunner
+            .getInstance()
+            .initialize(clusterService, threadPool, client, sdkClient, connectorAccessControlHelper, mlFeatureEnabledSetting);
+        MLOperationalMetricsCounter.initialize(clusterService.getClusterName().toString(), metricsRegistry, mlFeatureEnabledSetting);
+        MLAdoptionMetricsCounter.initialize(clusterService.getClusterName().toString(), metricsRegistry, mlFeatureEnabledSetting);
 
         return ImmutableList
             .of(

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

@@ -9,27 +9,43 @@
 import java.util.concurrent.ConcurrentHashMap;
 import java.util.stream.Stream;
 
+import org.opensearch.ml.common.settings.MLFeatureEnabledSetting;
 import org.opensearch.telemetry.metrics.Counter;
 import org.opensearch.telemetry.metrics.MetricsRegistry;
 import org.opensearch.telemetry.metrics.tags.Tags;
 
+import lombok.extern.log4j.Log4j2;
+
+@Log4j2
 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";
 
+    private final MLFeatureEnabledSetting mlFeatureEnabledSetting;
+
     protected final String clusterName;
     protected final MetricsRegistry metricsRegistry;
     protected final Map<T, Counter> metricCounterMap;
 
-    protected AbstractMLMetricsCounter(String clusterName, MetricsRegistry metricsRegistry, Class<T> metricClass) {
+    protected AbstractMLMetricsCounter(
+        String clusterName,
+        MetricsRegistry metricsRegistry,
+        Class<T> metricClass,
+        MLFeatureEnabledSetting mlFeatureEnabledSetting
+    ) {
         this.clusterName = clusterName;
         this.metricsRegistry = metricsRegistry;
         this.metricCounterMap = new ConcurrentHashMap<>();
+        this.mlFeatureEnabledSetting = mlFeatureEnabledSetting;
         Stream.of(metricClass.getEnumConstants()).forEach(metric -> metricCounterMap.computeIfAbsent(metric, this::createMetricCounter));
     }
 
     public void incrementCounter(T metric, Tags customTags) {
+        if (!mlFeatureEnabledSetting.isMetricCollectionEnabled()) {
+            return;
+        }
+
         Counter counter = metricCounterMap.computeIfAbsent(metric, this::createMetricCounter);
         Tags metricsTags = (customTags == null ? Tags.create() : customTags).addTag(CLUSTER_NAME_TAG, clusterName);
         counter.add(1, metricsTags);

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

@@ -5,19 +5,24 @@
 
 package org.opensearch.ml.stats.otel.counters;
 
+import org.opensearch.ml.common.settings.MLFeatureEnabledSetting;
 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);
+    private MLAdoptionMetricsCounter(String clusterName, MetricsRegistry metricsRegistry, MLFeatureEnabledSetting mlFeatureEnabledSetting) {
+        super(clusterName, metricsRegistry, AdoptionMetric.class, mlFeatureEnabledSetting);
     }
 
-    public static synchronized void initialize(String clusterName, MetricsRegistry metricsRegistry) {
-        instance = new MLAdoptionMetricsCounter(clusterName, metricsRegistry);
+    public static synchronized void initialize(
+        String clusterName,
+        MetricsRegistry metricsRegistry,
+        MLFeatureEnabledSetting mlFeatureEnabledSetting
+    ) {
+        instance = new MLAdoptionMetricsCounter(clusterName, metricsRegistry, mlFeatureEnabledSetting);
     }
 
     public static synchronized MLAdoptionMetricsCounter getInstance() {

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

@@ -5,19 +5,28 @@
 
 package org.opensearch.ml.stats.otel.counters;
 
+import org.opensearch.ml.common.settings.MLFeatureEnabledSetting;
 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);
+    private MLOperationalMetricsCounter(
+        String clusterName,
+        MetricsRegistry metricsRegistry,
+        MLFeatureEnabledSetting mlFeatureEnabledSetting
+    ) {
+        super(clusterName, metricsRegistry, OperationalMetric.class, mlFeatureEnabledSetting);
     }
 
-    public static synchronized void initialize(String clusterName, MetricsRegistry metricsRegistry) {
-        instance = new MLOperationalMetricsCounter(clusterName, metricsRegistry);
+    public static synchronized void initialize(
+        String clusterName,
+        MetricsRegistry metricsRegistry,
+        MLFeatureEnabledSetting mlFeatureEnabledSetting
+    ) {
+        instance = new MLOperationalMetricsCounter(clusterName, metricsRegistry, mlFeatureEnabledSetting);
     }
 
     public static synchronized MLOperationalMetricsCounter getInstance() {