fix: change logic for starting job to support single node cluster as well

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

Author: pyek-bot

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());
+    }
 }

plugin/src/main/java/org/opensearch/ml/cluster/MLCommonsClusterEventListener.java

@@ -75,23 +75,31 @@ public void clusterChanged(ClusterChangedEvent event) {
             Set<String> removedNodeIds = delta.removedNodes().stream().map(DiscoveryNode::getId).collect(Collectors.toSet());
             mlModelManager.removeWorkerNodes(removedNodeIds, false);
         } else if (delta.added()) {
-            for (DiscoveryNode node : delta.addedNodes()) {
-                // 3.1 introduces a new index for the job scheduler to track jobs
-                // the statsCollectorJob needs to be run when a cluster is started with the stats settings enabled
-                // As a result, we need to wait for a data node to come up before creating the new jobs index
-                if (node.isDataNode() && Version.V_3_1_0.onOrAfter(node.getVersion())) {
-                    if (mlFeatureEnabledSetting.isMetricCollectionEnabled() && mlFeatureEnabledSetting.isStaticMetricCollectionEnabled()) {
-                        mlTaskManager.startStatsCollectorJob();
-                    }
+            List<String> addedNodesIds = delta.addedNodes().stream().map(DiscoveryNode::getId).collect(Collectors.toList());
+            mlModelAutoReDeployer.buildAutoReloadArrangement(addedNodesIds, state.getNodes().getClusterManagerNodeId());
+        }
 
-                    if (clusterService.state().getMetadata().hasIndex(TASK_POLLING_JOB_INDEX)) {
-                        mlTaskManager.startTaskPollingJob();
-                    }
+        /*
+         * 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();
                 }
-            }
 
-            List<String> addedNodesIds = delta.addedNodes().stream().map(DiscoveryNode::getId).collect(Collectors.toList());
-            mlModelAutoReDeployer.buildAutoReloadArrangement(addedNodesIds, state.getNodes().getClusterManagerNodeId());
+                if (clusterService.state().getMetadata().hasIndex(TASK_POLLING_JOB_INDEX)) {
+                    mlTaskManager.startTaskPollingJob();
+                }
+
+                break;
+            }
         }
     }
 }

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

@@ -82,6 +82,7 @@ public class MLTaskManager {
     private final MLIndicesHandler mlIndicesHandler;
     private final Map<MLTaskType, AtomicInteger> runningTasksCount;
     private boolean taskPollingJobStarted;
+    private boolean statsCollectorJobStarted;
     public static final ImmutableSet<MLTaskState> TASK_DONE_STATES = ImmutableSet
         .of(MLTaskState.COMPLETED, MLTaskState.COMPLETED_WITH_ERROR, MLTaskState.FAILED, MLTaskState.CANCELLED);
 
@@ -546,74 +547,70 @@ public void startTaskPollingJob() {
             return;
         }
 
-        mlIndicesHandler.initMLJobsIndex(ActionListener.wrap(success -> {
-            if (success) {
-                String id = "ml_batch_task_polling_job";
-                String jobName = "poll_batch_jobs";
-                String interval = "1";
-                Long lockDurationSeconds = 20L;
-
-                MLJobParameter jobParameter = new MLJobParameter(
-                    jobName,
-                    new IntervalSchedule(Instant.now(), Integer.parseInt(interval), ChronoUnit.MINUTES),
-                    lockDurationSeconds,
-                    null,
-                    MLJobType.BATCH_TASK_UPDATE
-                );
-                IndexRequest indexRequest = new IndexRequest()
-                    .index(CommonValue.ML_JOBS_INDEX)
-                    .id(id)
-                    .source(jobParameter.toXContent(JsonXContent.contentBuilder(), null))
-                    .setRefreshPolicy(WriteRequest.RefreshPolicy.IMMEDIATE);
-
-                try (ThreadContext.StoredContext context = client.threadPool().getThreadContext().stashContext()) {
-                    client.index(indexRequest, ActionListener.runBefore(ActionListener.wrap(r -> {
-                        log.info("Indexed ml task polling job successfully");
-                        this.taskPollingJobStarted = true;
-                    }, e -> log.error("Failed to index task polling job", e)), context::restore));
-                }
-            }
-        }, e -> log.error("Failed to initialize ML jobs index", e)));
+        try {
+            MLJobParameter jobParameter = new MLJobParameter(
+                MLJobType.BATCH_TASK_UPDATE.name(),
+                new IntervalSchedule(Instant.now(), 1, ChronoUnit.MINUTES),
+                20L,
+                null,
+                MLJobType.BATCH_TASK_UPDATE
+            );
+
+            IndexRequest indexRequest = new IndexRequest()
+                .index(CommonValue.ML_JOBS_INDEX)
+                .id(MLJobType.BATCH_TASK_UPDATE.name())
+                .source(jobParameter.toXContent(JsonXContent.contentBuilder(), null))
+                .setRefreshPolicy(WriteRequest.RefreshPolicy.IMMEDIATE);
+
+            startJob(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;
+        }
+
+        try {
+            MLJobParameter jobParameter = new MLJobParameter(
+                MLJobType.STATS_COLLECTOR.name(),
+                new IntervalSchedule(Instant.now(), 5, ChronoUnit.MINUTES),
+                60L,
+                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);
+
+            startJob(indexRequest, MLJobType.STATS_COLLECTOR, () -> this.statsCollectorJobStarted = true);
+        } catch (IOException e) {
+            log.error("Failed to index stats collection job", e);
+        }
+    }
+
+    /**
+     * Start a job by indexing the job parameter to ML jobs index.
+     * 
+     * @param indexRequest the index request containing the job parameter
+     * @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) {
         mlIndicesHandler.initMLJobsIndex(ActionListener.wrap(success -> {
             if (success) {
-                try {
-                    int intervalInMinutes = 5;
-                    Long lockDurationSeconds = 60L;
-
-                    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);
-
-                    try (ThreadContext.StoredContext context = client.threadPool().getThreadContext().stashContext()) {
-                        client
-                            .index(
-                                indexRequest,
-                                ActionListener
-                                    .runBefore(
-                                        ActionListener
-                                            .wrap(
-                                                r -> log.info("Indexed ml stats collection job successfully"),
-                                                e -> log.error("Failed to index stats collection job", e)
-                                            ),
-                                        context::restore
-                                    )
-                            );
-                    }
-                } catch (IOException e) {
-                    log.error("Failed to index stats collection job", e);
+                try (ThreadContext.StoredContext context = client.threadPool().getThreadContext().stashContext()) {
+                    client.index(indexRequest, ActionListener.runBefore(ActionListener.wrap(r -> {
+                        log.info("Indexed {} successfully", jobType.name());
+                        if (successCallback != null) {
+                            successCallback.run();
+                        }
+                    }, e -> log.error("Failed to index {} job", jobType.name(), e)), context::restore));
                 }
             }
         }, e -> log.error("Failed to initialize ML jobs index", e)));