Merge branch 'main' into multi-tenancy-index

brianf-aws · web-flow · commit a78bdefd5d46 · 2025-08-11T21:27:03.000-07:00
Signed-off-by: Brian Flores &lt;iflorbri@amazon.com&gt;
diff --git a/common/src/main/java/org/opensearch/ml/common/settings/MLCommonsSettings.java b/common/src/main/java/org/opensearch/ml/common/settings/MLCommonsSettings.java
@@ -361,7 +361,7 @@ private MLCommonsSettings() {}
 
     // 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.Final);
+        .boolSetting("plugins.ml_commons.metrics_static_collection_enabled", false, Setting.Property.NodeScope, Setting.Property.Dynamic);
 
     // Feature flag for Agentic memory APIs
     public static final Setting<Boolean> ML_COMMONS_AGENTIC_MEMORY_ENABLED = Setting
diff --git a/common/src/main/java/org/opensearch/ml/common/settings/MLFeatureEnabledSetting.java b/common/src/main/java/org/opensearch/ml/common/settings/MLFeatureEnabledSetting.java
@@ -106,6 +106,12 @@ public MLFeatureEnabledSetting(ClusterService clusterService, Settings settings)
         clusterService.getClusterSettings().addSettingsUpdateConsumer(ML_COMMONS_AGENTIC_SEARCH_ENABLED, it -> isAgenticSearchEnabled = it);
         clusterService.getClusterSettings().addSettingsUpdateConsumer(ML_COMMONS_MCP_CONNECTOR_ENABLED, it -> isMcpConnectorEnabled = it);
         clusterService.getClusterSettings().addSettingsUpdateConsumer(ML_COMMONS_AGENTIC_MEMORY_ENABLED, it -> isAgenticMemoryEnabled = it);
+        clusterService.getClusterSettings().addSettingsUpdateConsumer(ML_COMMONS_STATIC_METRIC_COLLECTION_ENABLED, it -> {
+            isStaticMetricCollectionEnabled = it;
+            for (SettingsChangeListener listener : listeners) {
+                listener.onStaticMetricCollectionEnabledChanged(it);
+            }
+        });
     }
 
     /**
diff --git a/common/src/main/java/org/opensearch/ml/common/settings/SettingsChangeListener.java b/common/src/main/java/org/opensearch/ml/common/settings/SettingsChangeListener.java
@@ -18,5 +18,20 @@ public interface SettingsChangeListener {
      *                    <li><code>false</code> if multi-tenancy is disabled</li>
      *                  </ul>
      */
-    void onMultiTenancyEnabledChanged(boolean isEnabled);
+    default void onMultiTenancyEnabledChanged(boolean isEnabled) {
+        // do nothing
+    }
+
+    /**
+     * Callback method that gets triggered when the static metric collection setting changes.
+     *
+     * @param isEnabled A boolean value indicating the new state of the static metric collection setting:
+     *                  <ul>
+     *                    <li><code>true</code> if static metric collection is enabled</li>
+     *                    <li><code>false</code> if static metric collection is disabled</li>
+     *                  </ul>
+     */
+    default void onStaticMetricCollectionEnabledChanged(boolean isEnabled) {
+        // do nothing
+    }
 }
diff --git a/common/src/test/java/org/opensearch/ml/common/settings/MLFeatureEnabledSettingTests.java b/common/src/test/java/org/opensearch/ml/common/settings/MLFeatureEnabledSettingTests.java
@@ -170,4 +170,19 @@ public void testAgenticMemoryCanBeDisabled() {
         MLFeatureEnabledSetting setting = new MLFeatureEnabledSetting(mockClusterService, settings);
         assertFalse(setting.isAgenticMemoryEnabled());
     }
+
+    @Test
+    public void testStaticMetricCollectionSettingChangeNotifiesListeners() {
+        Settings settings = Settings.builder().put("plugins.ml_commons.metrics_static_collection_enabled", false).build();
+
+        MLFeatureEnabledSetting setting = new MLFeatureEnabledSetting(mockClusterService, settings);
+
+        SettingsChangeListener mockListener = mock(SettingsChangeListener.class);
+        setting.addListener(mockListener);
+
+        mockClusterSettings.applySettings(Settings.builder().put("plugins.ml_commons.metrics_static_collection_enabled", true).build());
+
+        verify(mockListener).onStaticMetricCollectionEnabledChanged(true);
+        assertTrue(setting.isStaticMetricCollectionEnabled());
+    }
 }
diff --git a/plugin/src/main/java/org/opensearch/ml/cluster/MLCommonsClusterEventListener.java b/plugin/src/main/java/org/opensearch/ml/cluster/MLCommonsClusterEventListener.java
@@ -91,7 +91,7 @@ public void clusterChanged(ClusterChangedEvent event) {
         for (DiscoveryNode node : state.nodes()) {
             if (node.isDataNode() && Version.V_3_1_0.onOrAfter(node.getVersion())) {
                 if (mlFeatureEnabledSetting.isMetricCollectionEnabled() && mlFeatureEnabledSetting.isStaticMetricCollectionEnabled()) {
-                    mlTaskManager.startStatsCollectorJob();
+                    mlTaskManager.indexStatsCollectorJob(true);
                 }
 
                 if (clusterService.state().getMetadata().hasIndex(TASK_POLLING_JOB_INDEX)) {
diff --git a/plugin/src/main/java/org/opensearch/ml/jobs/MLJobParameter.java b/plugin/src/main/java/org/opensearch/ml/jobs/MLJobParameter.java
@@ -49,14 +49,14 @@ public class MLJobParameter implements ScheduledJobParameter {
 
     public MLJobParameter() {}
 
-    public MLJobParameter(String name, Schedule schedule, Long lockDurationSeconds, Double jitter, MLJobType jobType) {
+    public MLJobParameter(String name, Schedule schedule, Long lockDurationSeconds, Double jitter, MLJobType jobType, boolean isEnabled) {
         this.jobName = name;
         this.schedule = schedule;
         this.lockDurationSeconds = lockDurationSeconds;
         this.jitter = jitter;
 
         Instant now = Instant.now();
-        this.isEnabled = true;
+        this.isEnabled = isEnabled;
         this.enabledTime = now;
         this.lastUpdateTime = now;
         this.jobType = jobType;
diff --git a/plugin/src/main/java/org/opensearch/ml/jobs/MLJobRunner.java b/plugin/src/main/java/org/opensearch/ml/jobs/MLJobRunner.java
@@ -93,6 +93,10 @@ public void runJob(ScheduledJobParameter scheduledJobParameter, JobExecutionCont
             throw new IllegalArgumentException("Job parameters is invalid.");
         }
 
+        if (!jobParameter.isEnabled()) {
+            throw new IllegalStateException(String.format("Attempted to run disabled job of type: %s", jobParameter.getJobType().name()));
+        }
+
         switch (jobParameter.getJobType()) {
             case STATS_COLLECTOR:
                 MLStatsJobProcessor
diff --git a/plugin/src/main/java/org/opensearch/ml/plugin/MachineLearningPlugin.java b/plugin/src/main/java/org/opensearch/ml/plugin/MachineLearningPlugin.java
@@ -600,6 +600,8 @@ public Collection<Object> createComponents(
         Path dataPath = environment.dataFiles()[0];
 
         mlFeatureEnabledSetting = new MLFeatureEnabledSetting(clusterService, settings);
+        mlFeatureEnabledSetting.addListener(mlTaskManager);
+
         mlIndicesHandler = new MLIndicesHandler(clusterService, client, mlFeatureEnabledSetting);
 
         SdkClient sdkClient = SdkClientFactory
@@ -666,6 +668,7 @@ public Collection<Object> createComponents(
         mlInputDatasetHandler = new MLInputDatasetHandler(client);
         modelAccessControlHelper = new ModelAccessControlHelper(clusterService, settings);
         connectorAccessControlHelper = new ConnectorAccessControlHelper(clusterService, settings);
+
         mlModelManager = new MLModelManager(
             clusterService,
             scriptService,
diff --git a/plugin/src/main/java/org/opensearch/ml/task/MLTaskManager.java b/plugin/src/main/java/org/opensearch/ml/task/MLTaskManager.java
@@ -51,6 +51,7 @@
 import org.opensearch.ml.common.exception.MLException;
 import org.opensearch.ml.common.exception.MLLimitExceededException;
 import org.opensearch.ml.common.exception.MLResourceNotFoundException;
+import org.opensearch.ml.common.settings.SettingsChangeListener;
 import org.opensearch.ml.engine.indices.MLIndicesHandler;
 import org.opensearch.ml.jobs.MLJobParameter;
 import org.opensearch.ml.jobs.MLJobType;
@@ -73,7 +74,7 @@
  * MLTaskManager is responsible for managing MLTask.
  */
 @Log4j2
-public class MLTaskManager {
+public class MLTaskManager implements SettingsChangeListener {
     public static int TASK_SEMAPHORE_TIMEOUT = 5000; // 5 seconds
     private final Map<String, MLTaskCache> taskCaches;
     private final Client client;
@@ -553,7 +554,8 @@ public void startTaskPollingJob() {
                 new IntervalSchedule(Instant.now(), 1, ChronoUnit.MINUTES),
                 20L,
                 null,
-                MLJobType.BATCH_TASK_UPDATE
+                MLJobType.BATCH_TASK_UPDATE,
+                true
             );
 
             IndexRequest indexRequest = new IndexRequest()
@@ -562,24 +564,27 @@ public void startTaskPollingJob() {
                 .source(jobParameter.toXContent(JsonXContent.contentBuilder(), null))
                 .setRefreshPolicy(WriteRequest.RefreshPolicy.IMMEDIATE);
 
-            startJob(indexRequest, MLJobType.BATCH_TASK_UPDATE, () -> this.taskPollingJobStarted = true);
+            indexJob(indexRequest, MLJobType.BATCH_TASK_UPDATE, () -> this.taskPollingJobStarted = true);
         } catch (IOException e) {
             log.error("Failed to index task polling job", e);
         }
     }
 
-    public void startStatsCollectorJob() {
-        if (statsCollectorJobStarted) {
-            return;
-        }
+    @Override
+    public void onStaticMetricCollectionEnabledChanged(boolean isEnabled) {
+        log.info("Static metric collection setting changed to: {}", isEnabled);
+        indexStatsCollectorJob(isEnabled);
+    }
 
+    public void indexStatsCollectorJob(boolean enabled) {
         try {
             MLJobParameter jobParameter = new MLJobParameter(
                 MLJobType.STATS_COLLECTOR.name(),
                 new IntervalSchedule(Instant.now(), 5, ChronoUnit.MINUTES),
                 60L,
                 null,
-                MLJobType.STATS_COLLECTOR
+                MLJobType.STATS_COLLECTOR,
+                enabled
             );
 
             IndexRequest indexRequest = new IndexRequest()
@@ -588,7 +593,7 @@ public void startStatsCollectorJob() {
                 .source(jobParameter.toXContent(JsonXContent.contentBuilder(), null))
                 .setRefreshPolicy(WriteRequest.RefreshPolicy.IMMEDIATE);
 
-            startJob(indexRequest, MLJobType.STATS_COLLECTOR, () -> this.statsCollectorJobStarted = true);
+            indexJob(indexRequest, MLJobType.STATS_COLLECTOR, () -> {});
         } catch (IOException e) {
             log.error("Failed to index stats collection job", e);
         }
@@ -601,7 +606,7 @@ public void startStatsCollectorJob() {
      * @param jobType the type of job being started
      * @param successCallback callback to execute on successful job indexing
      */
-    private void startJob(IndexRequest indexRequest, MLJobType jobType, Runnable successCallback) {
+    private void indexJob(IndexRequest indexRequest, MLJobType jobType, Runnable successCallback) {
         mlIndicesHandler.initMLJobsIndex(ActionListener.wrap(success -> {
             if (success) {
                 try (ThreadContext.StoredContext context = client.threadPool().getThreadContext().stashContext()) {
diff --git a/plugin/src/test/java/org/opensearch/ml/jobs/MLJobParameterTests.java b/plugin/src/test/java/org/opensearch/ml/jobs/MLJobParameterTests.java
@@ -35,7 +35,7 @@ public void setUp() {
         lockDurationSeconds = 20L;
         jitter = 0.5;
         jobType = null;
-        jobParameter = new MLJobParameter(jobName, schedule, lockDurationSeconds, jitter, jobType);
+        jobParameter = new MLJobParameter(jobName, schedule, lockDurationSeconds, jitter, jobType, true);
     }
 
     @Test
@@ -54,7 +54,7 @@ public void testToXContent() throws Exception {
     @Test
     public void testNullCase() throws IOException {
         String newJobName = "test-job";
-        MLJobParameter nullParameter = new MLJobParameter(newJobName, null, null, null, null);
+        MLJobParameter nullParameter = new MLJobParameter(newJobName, null, null, null, null, true);
         nullParameter.setLastUpdateTime(null);
         nullParameter.setEnabledTime(null);
 
@@ -64,6 +64,7 @@ public void testNullCase() throws IOException {
 
         assertTrue(jsonString.contains(newJobName));
         assertEquals(newJobName, nullParameter.getName());
+        assertTrue(nullParameter.isEnabled());
         assertNull(nullParameter.getSchedule());
         assertNull(nullParameter.getLockDurationSeconds());
         assertNull(nullParameter.getJitter());
diff --git a/plugin/src/test/java/org/opensearch/ml/jobs/MLJobRunnerTests.java b/plugin/src/test/java/org/opensearch/ml/jobs/MLJobRunnerTests.java
@@ -73,4 +73,11 @@ public void testRunJobWithNullJobType() {
         when(jobParameter.getJobType()).thenReturn(null);
         jobRunner.runJob(jobParameter, jobExecutionContext);
     }
+
+    @Test(expected = IllegalStateException.class)
+    public void testRunJobWithDisabledJob() {
+        when(jobParameter.isEnabled()).thenReturn(false);
+        when(jobParameter.getJobType()).thenReturn(MLJobType.STATS_COLLECTOR);
+        jobRunner.runJob(jobParameter, jobExecutionContext);
+    }
 }
diff --git a/plugin/src/test/java/org/opensearch/ml/task/MLTaskManagerTests.java b/plugin/src/test/java/org/opensearch/ml/task/MLTaskManagerTests.java
@@ -404,7 +404,7 @@ public void testStartStatsCollectorJob() throws IOException {
             return null;
         }).when(client).index(any(), any());
 
-        mlTaskManager.startStatsCollectorJob();
+        mlTaskManager.indexStatsCollectorJob(true);
 
         ArgumentCaptor<IndexRequest> indexRequestCaptor = ArgumentCaptor.forClass(IndexRequest.class);
         verify(client).index(indexRequestCaptor.capture(), any());
@@ -429,7 +429,7 @@ public void testStartStatsCollectorJob_IndexException() throws IOException {
             return null;
         }).when(client).index(any(), any());
 
-        mlTaskManager.startStatsCollectorJob();
+        mlTaskManager.indexStatsCollectorJob(true);
 
         verify(client).index(any(), any());
     }
@@ -469,4 +469,24 @@ public void testUpdateMLTaskDirectly_TaskDoneState() {
         mlTaskManager.updateMLTaskDirectly("task_id", updatedFields, listener);
         verify(listener).onResponse(any(UpdateResponse.class));
     }
+
+    public void testOnStaticMetricCollectionEnabledChanged() {
+        doAnswer(invocation -> {
+            ActionListener<Boolean> listener = invocation.getArgument(0);
+            listener.onResponse(true);
+            return null;
+        }).when(mlIndicesHandler).initMLJobsIndex(any());
+
+        doAnswer(invocation -> {
+            ActionListener<IndexResponse> listener = invocation.getArgument(1);
+            listener.onResponse(indexResponse);
+            return null;
+        }).when(client).index(any(), any());
+
+        mlTaskManager.onStaticMetricCollectionEnabledChanged(true);
+        verify(mlTaskManager).indexStatsCollectorJob(true);
+
+        mlTaskManager.onStaticMetricCollectionEnabledChanged(false);
+        verify(mlTaskManager).indexStatsCollectorJob(false);
+    }
 }

Original file line number	Diff line number	Diff line change
`@@ -91,7 +91,7 @@ public void clusterChanged(ClusterChangedEvent event) {`
`91`	`91`	`for (DiscoveryNode node : state.nodes()) {`
`92`	`92`	`if (node.isDataNode() && Version.V_3_1_0.onOrAfter(node.getVersion())) {`
`93`	`93`	`if (mlFeatureEnabledSetting.isMetricCollectionEnabled() && mlFeatureEnabledSetting.isStaticMetricCollectionEnabled()) {`
`94`		`- mlTaskManager.startStatsCollectorJob();`
	`94`	`+ mlTaskManager.indexStatsCollectorJob(true);`
`95`	`95`	`}`
`96`	`96`
`97`	`97`	`if (clusterService.state().getMetadata().hasIndex(TASK_POLLING_JOB_INDEX)) {`