Merge branch 'main' into CVE-48734

brianf-aws · web-flow · commit 23afd3b8a9eb · 2025-06-18T21:42:51.000-07:00
Signed-off-by: Brian Flores &lt;iflorbri@amazon.com&gt;
diff --git a/common/src/main/java/org/opensearch/ml/common/CommonValue.java b/common/src/main/java/org/opensearch/ml/common/CommonValue.java
@@ -65,6 +65,7 @@ public class CommonValue {
     public static final String ML_MEMORY_MESSAGE_INDEX_MAPPING_PATH = "index-mappings/ml_memory_message.json";
     public static final String ML_MCP_SESSION_MANAGEMENT_INDEX_MAPPING_PATH = "index-mappings/ml_mcp_session_management.json";
     public static final String ML_MCP_TOOLS_INDEX_MAPPING_PATH = "index-mappings/ml_mcp_tools.json";
+    public static final String ML_JOBS_INDEX_MAPPING_PATH = "index-mappings/ml_jobs.json";
 
     // Calculate Versions independently of OpenSearch core version
     public static final Version VERSION_2_11_0 = Version.fromString("2.11.0");
diff --git a/common/src/main/java/org/opensearch/ml/common/MLIndex.java b/common/src/main/java/org/opensearch/ml/common/MLIndex.java
@@ -15,6 +15,8 @@
 import static org.opensearch.ml.common.CommonValue.ML_CONNECTOR_INDEX_MAPPING_PATH;
 import static org.opensearch.ml.common.CommonValue.ML_CONTROLLER_INDEX;
 import static org.opensearch.ml.common.CommonValue.ML_CONTROLLER_INDEX_MAPPING_PATH;
+import static org.opensearch.ml.common.CommonValue.ML_JOBS_INDEX;
+import static org.opensearch.ml.common.CommonValue.ML_JOBS_INDEX_MAPPING_PATH;
 import static org.opensearch.ml.common.CommonValue.ML_MCP_SESSION_MANAGEMENT_INDEX_MAPPING_PATH;
 import static org.opensearch.ml.common.CommonValue.ML_MCP_TOOLS_INDEX_MAPPING_PATH;
 import static org.opensearch.ml.common.CommonValue.ML_MEMORY_MESSAGE_INDEX;
@@ -44,7 +46,8 @@ public enum MLIndex {
     MEMORY_META(ML_MEMORY_META_INDEX, false, ML_MEMORY_META_INDEX_MAPPING_PATH),
     MEMORY_MESSAGE(ML_MEMORY_MESSAGE_INDEX, false, ML_MEMORY_MESSAGE_INDEX_MAPPING_PATH),
     MCP_SESSION_MANAGEMENT(MCP_SESSION_MANAGEMENT_INDEX, false, ML_MCP_SESSION_MANAGEMENT_INDEX_MAPPING_PATH),
-    MCP_TOOLS(MCP_TOOLS_INDEX, false, ML_MCP_TOOLS_INDEX_MAPPING_PATH);
+    MCP_TOOLS(MCP_TOOLS_INDEX, false, ML_MCP_TOOLS_INDEX_MAPPING_PATH),
+    JOBS(ML_JOBS_INDEX, false, ML_JOBS_INDEX_MAPPING_PATH);
 
     private final String indexName;
     // whether we use an alias for the index
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
@@ -345,9 +345,9 @@ private MLCommonsSettings() {}
 
     // 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);
+        .boolSetting("plugins.ml_commons.metrics_collection_enabled", false, Setting.Property.NodeScope, Setting.Property.Final);
 
     // 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);
+        .boolSetting("plugins.ml_commons.metrics_static_collection_enabled", false, Setting.Property.NodeScope, Setting.Property.Final);
 }
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
@@ -88,12 +88,6 @@ 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);
     }
 
     /**
diff --git a/common/src/main/resources/index-mappings/ml_jobs.json b/common/src/main/resources/index-mappings/ml_jobs.json
@@ -0,0 +1,48 @@
+{
+  "_meta": {
+    "schema_version": 1
+  },
+  "properties": {
+    "name": {
+      "type": "keyword"
+    },
+    "enabled": {
+      "type": "boolean"
+    },
+    "enabled_time": {
+      "type": "date",
+      "format": "strict_date_time||epoch_millis"
+    },
+    "last_update_time": {
+      "type": "date",
+      "format": "strict_date_time||epoch_millis"
+    },
+    "schedule": {
+      "properties": {
+        "interval": {
+          "properties": {
+            "start_time": {
+              "type": "date",
+              "format": "strict_date_time||epoch_millis"
+            },
+            "period": {
+              "type": "integer"
+            },
+            "unit": {
+              "type": "keyword"
+            }
+          }
+        }
+      }
+    },
+    "lock_duration_seconds": {
+      "type": "long"
+    },
+    "jitter": {
+      "type": "double"
+    },
+    "type": {
+      "type": "keyword"
+    }
+  }
+}
diff --git a/ml-algorithms/src/main/java/org/opensearch/ml/engine/indices/MLIndicesHandler.java b/ml-algorithms/src/main/java/org/opensearch/ml/engine/indices/MLIndicesHandler.java
@@ -92,6 +92,10 @@ public void initMLMcpToolsIndex(ActionListener<Boolean> listener) {
         initMLIndexIfAbsent(MLIndex.MCP_TOOLS, listener);
     }
 
+    public void initMLJobsIndex(ActionListener<Boolean> listener) {
+        initMLIndexIfAbsent(MLIndex.JOBS, listener);
+    }
+
     public void initMLAgentIndex(ActionListener<Boolean> listener) {
         initMLIndexIfAbsent(MLIndex.AGENT, listener);
     }
diff --git a/ml-algorithms/src/test/java/org/opensearch/ml/engine/indices/MLIndicesHandlerTest.java b/ml-algorithms/src/test/java/org/opensearch/ml/engine/indices/MLIndicesHandlerTest.java
@@ -11,6 +11,7 @@
 import static org.mockito.Mockito.when;
 import static org.opensearch.ml.common.CommonValue.META;
 import static org.opensearch.ml.common.CommonValue.ML_AGENT_INDEX;
+import static org.opensearch.ml.common.CommonValue.ML_JOBS_INDEX;
 import static org.opensearch.ml.common.CommonValue.ML_MEMORY_MESSAGE_INDEX;
 import static org.opensearch.ml.common.CommonValue.ML_MEMORY_META_INDEX;
 import static org.opensearch.ml.common.CommonValue.SCHEMA_VERSION_FIELD;
@@ -216,4 +217,36 @@ public void initMLConnectorIndex_ResourceAlreadyExistsException_RaceCondition()
         verify(listener).onResponse(argumentCaptor.capture());
         assertEquals(true, argumentCaptor.getValue());
     }
+
+    @Test
+    public void initMLJobsIndex() {
+        ActionListener<Boolean> listener = mock(ActionListener.class);
+        doAnswer(invocation -> {
+            ActionListener<AcknowledgedResponse> actionListener = invocation.getArgument(1);
+            actionListener.onResponse(new AcknowledgedResponse(true));
+            return null;
+        }).when(indicesAdminClient).putMapping(any(), any());
+        ArgumentCaptor<Boolean> argumentCaptor = ArgumentCaptor.forClass(Boolean.class);
+        indicesHandler.initMLJobsIndex(listener);
+
+        verify(listener).onResponse(argumentCaptor.capture());
+        assertEquals(true, argumentCaptor.getValue());
+    }
+
+    @Test
+    public void initMLJobsIndexNoIndex() {
+        ActionListener<Boolean> listener = mock(ActionListener.class);
+        when(metadata.hasIndex(anyString())).thenReturn(false);
+        doAnswer(invocation -> {
+            ActionListener<CreateIndexResponse> actionListener = invocation.getArgument(1);
+            actionListener.onResponse(new CreateIndexResponse(true, true, ML_JOBS_INDEX));
+            return null;
+        }).when(indicesAdminClient).create(any(), any());
+        ArgumentCaptor<Boolean> argumentCaptor = ArgumentCaptor.forClass(Boolean.class);
+        indicesHandler.initMLJobsIndex(listener);
+
+        verify(indicesAdminClient).create(isA(CreateIndexRequest.class), any());
+        verify(listener).onResponse(argumentCaptor.capture());
+        assertEquals(true, argumentCaptor.getValue());
+    }
 }
diff --git a/plugin/src/main/java/org/opensearch/ml/cluster/MLCommonsClusterEventListener.java b/plugin/src/main/java/org/opensearch/ml/cluster/MLCommonsClusterEventListener.java
@@ -5,12 +5,14 @@
 
 package org.opensearch.ml.cluster;
 
+import static org.opensearch.ml.common.CommonValue.TASK_POLLING_JOB_INDEX;
 import static org.opensearch.ml.common.settings.MLCommonsSettings.ML_COMMONS_MONITORING_REQUEST_COUNT;
 
 import java.util.List;
 import java.util.Set;
 import java.util.stream.Collectors;
 
+import org.opensearch.Version;
 import org.opensearch.cluster.ClusterChangedEvent;
 import org.opensearch.cluster.ClusterState;
 import org.opensearch.cluster.ClusterStateListener;
@@ -19,9 +21,11 @@
 import org.opensearch.cluster.service.ClusterService;
 import org.opensearch.common.settings.Settings;
 import org.opensearch.ml.autoredeploy.MLModelAutoReDeployer;
+import org.opensearch.ml.common.settings.MLFeatureEnabledSetting;
 import org.opensearch.ml.model.MLModelCacheHelper;
 import org.opensearch.ml.model.MLModelManager;
 import org.opensearch.ml.task.MLTaskManager;
+import org.opensearch.transport.client.Client;
 
 import lombok.extern.log4j.Log4j2;
 
@@ -32,22 +36,27 @@ public class MLCommonsClusterEventListener implements ClusterStateListener {
     private final MLModelManager mlModelManager;
     private final MLTaskManager mlTaskManager;
     private final MLModelCacheHelper modelCacheHelper;
-
     private final MLModelAutoReDeployer mlModelAutoReDeployer;
+    private final Client client;
+    private final MLFeatureEnabledSetting mlFeatureEnabledSetting;
 
     public MLCommonsClusterEventListener(
         ClusterService clusterService,
         MLModelManager mlModelManager,
         MLTaskManager mlTaskManager,
         MLModelCacheHelper modelCacheHelper,
-        MLModelAutoReDeployer mlModelAutoReDeployer
+        MLModelAutoReDeployer mlModelAutoReDeployer,
+        Client client,
+        MLFeatureEnabledSetting mlFeatureEnabledSetting
     ) {
         this.clusterService = clusterService;
         this.clusterService.addListener(this);
         this.mlModelManager = mlModelManager;
         this.mlTaskManager = mlTaskManager;
         this.modelCacheHelper = modelCacheHelper;
         this.mlModelAutoReDeployer = mlModelAutoReDeployer;
+        this.client = client;
+        this.mlFeatureEnabledSetting = mlFeatureEnabledSetting;
     }
 
     @Override
@@ -69,5 +78,28 @@ public void clusterChanged(ClusterChangedEvent event) {
             List<String> addedNodesIds = delta.addedNodes().stream().map(DiscoveryNode::getId).collect(Collectors.toList());
             mlModelAutoReDeployer.buildAutoReloadArrangement(addedNodesIds, state.getNodes().getClusterManagerNodeId());
         }
+
+        /*
+         * In version 3.1, a new index `.plugins-ml-jobs` replaces the old `.ml_commons_task_polling_job` index for the job scheduler.
+         * Version 3.1 also introduces a stats collector job that should run at startup if the relevant settings are enabled.
+         * When upgrading from 3.0 to 3.1, we need to ensure the new `.plugins-ml-jobs` index is created if either:
+         *   - The stats collector job is enabled, or
+         *   - The batch polling task job was already running.
+         * To avoid issues during blue/green or rolling upgrades, we wait for a data node running 3.1 or later before creating the new jobs index and starting the jobs.
+         * The following logic implements this behavior.
+         */
+        for (DiscoveryNode node : state.nodes()) {
+            if (node.isDataNode() && Version.V_3_1_0.onOrAfter(node.getVersion())) {
+                if (mlFeatureEnabledSetting.isMetricCollectionEnabled() && mlFeatureEnabledSetting.isStaticMetricCollectionEnabled()) {
+                    mlTaskManager.startStatsCollectorJob();
+                }
+
+                if (clusterService.state().getMetadata().hasIndex(TASK_POLLING_JOB_INDEX)) {
+                    mlTaskManager.startTaskPollingJob();
+                }
+
+                break;
+            }
+        }
     }
 }
diff --git a/plugin/src/main/java/org/opensearch/ml/cluster/MLCommonsClusterManagerEventListener.java b/plugin/src/main/java/org/opensearch/ml/cluster/MLCommonsClusterManagerEventListener.java
@@ -8,28 +8,18 @@
 import static org.opensearch.ml.common.settings.MLCommonsSettings.ML_COMMONS_SYNC_UP_JOB_INTERVAL_IN_SECONDS;
 import static org.opensearch.ml.plugin.MachineLearningPlugin.GENERAL_THREAD_POOL;
 
-import java.io.IOException;
-import java.time.Instant;
-import java.time.temporal.ChronoUnit;
 import java.util.List;
 
-import org.opensearch.action.index.IndexRequest;
-import org.opensearch.action.support.WriteRequest;
 import org.opensearch.cluster.LocalNodeClusterManagerListener;
 import org.opensearch.cluster.service.ClusterService;
 import org.opensearch.common.lifecycle.LifecycleListener;
 import org.opensearch.common.settings.Settings;
 import org.opensearch.common.unit.TimeValue;
-import org.opensearch.common.xcontent.json.JsonXContent;
 import org.opensearch.core.action.ActionListener;
-import org.opensearch.jobscheduler.spi.schedule.IntervalSchedule;
 import org.opensearch.ml.autoredeploy.MLModelAutoReDeployer;
-import org.opensearch.ml.common.CommonValue;
 import org.opensearch.ml.common.settings.MLFeatureEnabledSetting;
 import org.opensearch.ml.engine.encryptor.Encryptor;
 import org.opensearch.ml.engine.indices.MLIndicesHandler;
-import org.opensearch.ml.jobs.MLJobParameter;
-import org.opensearch.ml.jobs.MLJobType;
 import org.opensearch.remote.metadata.client.SdkClient;
 import org.opensearch.threadpool.Scheduler;
 import org.opensearch.threadpool.ThreadPool;
@@ -105,40 +95,6 @@ public void onClusterManager() {
                 TimeValue.timeValueSeconds(jobInterval),
                 GENERAL_THREAD_POOL
             );
-        startStatsCollectorJob();
-    }
-
-    private void startStatsCollectorJob() {
-        try {
-            int intervalInMinutes = 5;
-            Long lockDurationSeconds = 20L;
-
-            MLJobParameter jobParameter = new MLJobParameter(
-                MLJobType.STATS_COLLECTOR.name(),
-                new IntervalSchedule(Instant.now(), intervalInMinutes, ChronoUnit.MINUTES),
-                lockDurationSeconds,
-                null,
-                MLJobType.STATS_COLLECTOR
-            );
-
-            IndexRequest indexRequest = new IndexRequest()
-                .index(CommonValue.ML_JOBS_INDEX)
-                .id(MLJobType.STATS_COLLECTOR.name())
-                .source(jobParameter.toXContent(JsonXContent.contentBuilder(), null))
-                .setRefreshPolicy(WriteRequest.RefreshPolicy.IMMEDIATE);
-
-            client
-                .index(
-                    indexRequest,
-                    ActionListener
-                        .wrap(
-                            r -> log.info("Indexed ml stats collection job successfully"),
-                            e -> log.error("Failed to index stats collection job", e)
-                        )
-                );
-        } catch (IOException e) {
-            log.error("Failed to index stats collection job", e);
-        }
     }
 
     private void startSyncModelRoutingCron() {
diff --git a/plugin/src/main/java/org/opensearch/ml/jobs/processors/MLJobProcessor.java b/plugin/src/main/java/org/opensearch/ml/jobs/processors/MLJobProcessor.java
@@ -14,7 +14,6 @@
 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;
 
@@ -36,7 +35,8 @@ public MLJobProcessor(ClusterService clusterService, Client client, ThreadPool t
 
     public void process(ScheduledJobParameter scheduledJobParameter, JobExecutionContext jobExecutionContext, boolean isProcessorEnabled) {
         if (!isProcessorEnabled) {
-            throw new MLException(scheduledJobParameter.getName() + " not enabled.");
+            log.warn("{} not enabled.", scheduledJobParameter.getName());
+            return;
         }
 
         process(scheduledJobParameter, jobExecutionContext);
diff --git a/plugin/src/main/java/org/opensearch/ml/plugin/MachineLearningPlugin.java b/plugin/src/main/java/org/opensearch/ml/plugin/MachineLearningPlugin.java
@@ -5,10 +5,13 @@
 
 package org.opensearch.ml.plugin;
 
+import static org.opensearch.ml.common.CommonValue.MCP_SESSION_MANAGEMENT_INDEX;
+import static org.opensearch.ml.common.CommonValue.MCP_TOOLS_INDEX;
 import static org.opensearch.ml.common.CommonValue.ML_AGENT_INDEX;
 import static org.opensearch.ml.common.CommonValue.ML_CONFIG_INDEX;
 import static org.opensearch.ml.common.CommonValue.ML_CONNECTOR_INDEX;
 import static org.opensearch.ml.common.CommonValue.ML_CONTROLLER_INDEX;
+import static org.opensearch.ml.common.CommonValue.ML_JOBS_INDEX;
 import static org.opensearch.ml.common.CommonValue.ML_MEMORY_MESSAGE_INDEX;
 import static org.opensearch.ml.common.CommonValue.ML_MEMORY_META_INDEX;
 import static org.opensearch.ml.common.CommonValue.ML_MODEL_GROUP_INDEX;
@@ -773,7 +776,9 @@ public Collection<Object> createComponents(
             mlModelManager,
             mlTaskManager,
             modelCacheHelper,
-            mlModelAutoRedeployer
+            mlModelAutoRedeployer,
+            client,
+            mlFeatureEnabledSetting
         );
         MLCommonsClusterManagerEventListener clusterManagerEventListener = new MLCommonsClusterManagerEventListener(
             clusterService,
@@ -791,8 +796,11 @@ public Collection<Object> createComponents(
         MLJobRunner
             .getInstance()
             .initialize(clusterService, threadPool, client, sdkClient, connectorAccessControlHelper, mlFeatureEnabledSetting);
-        MLOperationalMetricsCounter.initialize(clusterService.getClusterName().toString(), metricsRegistry, mlFeatureEnabledSetting);
-        MLAdoptionMetricsCounter.initialize(clusterService.getClusterName().toString(), metricsRegistry, mlFeatureEnabledSetting);
+
+        if (mlFeatureEnabledSetting.isMetricCollectionEnabled()) {
+            MLOperationalMetricsCounter.initialize(clusterService.getClusterName().toString(), metricsRegistry, mlFeatureEnabledSetting);
+            MLAdoptionMetricsCounter.initialize(clusterService.getClusterName().toString(), metricsRegistry, mlFeatureEnabledSetting);
+        }
 
         mcpToolsHelper = new McpToolsHelper(client, threadPool, toolFactoryWrapper);
         McpAsyncServerHolder.init(mlIndicesHandler, mcpToolsHelper);
@@ -1267,6 +1275,9 @@ public Collection<SystemIndexDescriptor> getSystemIndexDescriptors(Settings sett
         systemIndexDescriptors.add(new SystemIndexDescriptor(ML_MEMORY_META_INDEX, "ML Commons Memory Meta Index"));
         systemIndexDescriptors.add(new SystemIndexDescriptor(ML_MEMORY_MESSAGE_INDEX, "ML Commons Memory Message Index"));
         systemIndexDescriptors.add(new SystemIndexDescriptor(ML_STOP_WORDS_INDEX, "ML Commons Stop Words Index"));
+        systemIndexDescriptors.add(new SystemIndexDescriptor(MCP_SESSION_MANAGEMENT_INDEX, "ML Commons MCP session management Index"));
+        systemIndexDescriptors.add(new SystemIndexDescriptor(MCP_TOOLS_INDEX, "ML Commons MCP tools Index"));
+        systemIndexDescriptors.add(new SystemIndexDescriptor(ML_JOBS_INDEX, "ML Commons Jobs Index"));
         return systemIndexDescriptors;
     }
 
diff --git a/plugin/src/main/java/org/opensearch/ml/task/MLTaskManager.java b/plugin/src/main/java/org/opensearch/ml/task/MLTaskManager.java
diff --git a/plugin/src/test/java/org/opensearch/ml/task/MLTaskManagerTests.java b/plugin/src/test/java/org/opensearch/ml/task/MLTaskManagerTests.java

Original file line number	Diff line number	Diff line change
`@@ -92,6 +92,10 @@ public void initMLMcpToolsIndex(ActionListener<Boolean> listener) {`
`92`	`92`	`initMLIndexIfAbsent(MLIndex.MCP_TOOLS, listener);`
`93`	`93`	`}`
`94`	`94`
	`95`	`+ public void initMLJobsIndex(ActionListener<Boolean> listener) {`
	`96`	`+ initMLIndexIfAbsent(MLIndex.JOBS, listener);`
	`97`	`+ }`
	`98`	`+`
`95`	`99`	`public void initMLAgentIndex(ActionListener<Boolean> listener) {`
`96`	`100`	`initMLIndexIfAbsent(MLIndex.AGENT, listener);`
`97`	`101`	`}`