diff --git a/plugin/src/main/java/org/opensearch/ml/cluster/MLCommonsClusterEventListener.java b/plugin/src/main/java/org/opensearch/ml/cluster/MLCommonsClusterEventListener.java index a01a657f2b..7540bf742c 100644 --- a/plugin/src/main/java/org/opensearch/ml/cluster/MLCommonsClusterEventListener.java +++ b/plugin/src/main/java/org/opensearch/ml/cluster/MLCommonsClusterEventListener.java @@ -83,34 +83,38 @@ public void clusterChanged(ClusterChangedEvent event) { } /* - * 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. + * The stats collector and memory retention jobs live in the `.plugins-ml-jobs` index (introduced in 3.1, + * replacing `.ml_commons_task_polling_job`). Indexing a job document auto-creates that index, and creating + * an index while a rolling upgrade is in flight strands its replicas: the primary is allocated to the + * new-version node that issued the write, and replicas can never be assigned to nodes older than the + * primary's node, so the cluster stays yellow until enough nodes are upgraded. + * + * To avoid that, only touch the jobs index when it is safe: + * - the index already exists (writing a document to an existing index cannot strand a replica), or + * - every node in the cluster is on at least this node's version, so a newly created index's replicas + * are allocatable anywhere. + * While a rolling upgrade is in flight neither holds; the jobs are deferred, not skipped — when the last + * old node leaves, the resulting cluster state change re-runs this listener and the jobs are created then. */ - for (DiscoveryNode node : state.nodes()) { - if (node.isDataNode() && node.getVersion().onOrAfter(Version.V_3_1_0)) { - if (mlFeatureEnabledSetting.isMetricCollectionEnabled() - && mlFeatureEnabledSetting.isStaticMetricCollectionEnabled() - && !clusterService.state().getMetadata().hasIndex(ML_JOBS_INDEX) - && !this.startedStatsJob) { - mlTaskManager.indexStatsCollectorJob(true); - // using this variable in case if same node has a cluster state change event and the state is not updated yet - this.startedStatsJob = true; - } - - if (mlFeatureEnabledSetting.isAgenticMemoryEnabled() - && !MLCommonsSettings.ML_COMMONS_MULTI_TENANCY_ENABLED.get(clusterService.getSettings()) - && !this.startedMemoryRetentionJob) { - int intervalHours = MLCommonsSettings.ML_COMMONS_MEMORY_RETENTION_JOB_INTERVAL_HOURS.get(clusterService.getSettings()); - mlTaskManager.indexMemoryRetentionJob(intervalHours); - this.startedMemoryRetentionJob = true; - } + boolean jobsIndexExists = state.getMetadata().hasIndex(ML_JOBS_INDEX); + boolean noOlderNodes = state.nodes().getMinNodeVersion().onOrAfter(Version.CURRENT); + if (jobsIndexExists || noOlderNodes) { + if (mlFeatureEnabledSetting.isMetricCollectionEnabled() + && mlFeatureEnabledSetting.isStaticMetricCollectionEnabled() + && !jobsIndexExists + && !this.startedStatsJob) { + mlTaskManager.indexStatsCollectorJob(true); + // using this variable in case if same node has a cluster state change event and the state is not updated yet + this.startedStatsJob = true; + } - break; + if (mlFeatureEnabledSetting.isAgenticMemoryEnabled() + && mlFeatureEnabledSetting.isMemoryRetentionEnabled() + && !MLCommonsSettings.ML_COMMONS_MULTI_TENANCY_ENABLED.get(clusterService.getSettings()) + && !this.startedMemoryRetentionJob) { + int intervalHours = MLCommonsSettings.ML_COMMONS_MEMORY_RETENTION_JOB_INTERVAL_HOURS.get(clusterService.getSettings()); + mlTaskManager.indexMemoryRetentionJob(intervalHours); + this.startedMemoryRetentionJob = true; } } } diff --git a/plugin/src/test/java/org/opensearch/ml/cluster/MLCommonsClusterEventListenerTests.java b/plugin/src/test/java/org/opensearch/ml/cluster/MLCommonsClusterEventListenerTests.java index d2623dd13e..8998b5a9f1 100644 --- a/plugin/src/test/java/org/opensearch/ml/cluster/MLCommonsClusterEventListenerTests.java +++ b/plugin/src/test/java/org/opensearch/ml/cluster/MLCommonsClusterEventListenerTests.java @@ -73,9 +73,8 @@ public void setup() { ); } - public void testClusterChanged_WithV31DataNode_MetricCollectionEnabled() { - DiscoveryNode dataNode = createDataNode(Version.V_3_1_0); - setupClusterState(dataNode, false); + public void testClusterChanged_AllNodesCurrent_MetricCollectionEnabled() { + setupClusterState(false, createDataNode("n1", Version.CURRENT)); when(mlFeatureEnabledSetting.isMetricCollectionEnabled()).thenReturn(true); when(mlFeatureEnabledSetting.isStaticMetricCollectionEnabled()).thenReturn(true); @@ -85,33 +84,55 @@ public void testClusterChanged_WithV31DataNode_MetricCollectionEnabled() { verify(mlTaskManager).indexStatsCollectorJob(true); } - public void testClusterChanged_WithPreV31DataNode_NoJobsStarted() { - DiscoveryNode dataNode = createDataNode(Version.V_3_0_0); - setupClusterState(dataNode, false); + public void testClusterChanged_MixedVersionCluster_NoIndex_NoJobsStarted() { + // rolling upgrade in flight: one upgraded node, one old node, jobs index absent. + // Creating the index now would strand its replica on version-allocation rules, + // so neither job may be created. + setupClusterState(false, createDataNode("new", Version.CURRENT), createDataNode("old", Version.V_3_1_0)); when(mlFeatureEnabledSetting.isMetricCollectionEnabled()).thenReturn(true); when(mlFeatureEnabledSetting.isStaticMetricCollectionEnabled()).thenReturn(true); + when(mlFeatureEnabledSetting.isAgenticMemoryEnabled()).thenReturn(true); + when(mlFeatureEnabledSetting.isMemoryRetentionEnabled()).thenReturn(true); listener.clusterChanged(event); verify(mlTaskManager, never()).indexStatsCollectorJob(anyBoolean()); + verify(mlTaskManager, never()).indexMemoryRetentionJob(anyInt()); } - public void testClusterChanged_WithPostV31DataNode_JobsStarted() { - DiscoveryNode dataNode = createDataNode(Version.V_3_2_0); - setupClusterState(dataNode, false); + public void testClusterChanged_MixedVersionCluster_IndexExists_RetentionJobStarted() { + // writing a job document into an existing index cannot strand a replica, + // so a mixed-version cluster must not block it + setupClusterState(true, createDataNode("new", Version.CURRENT), createDataNode("old", Version.V_3_1_0)); + when(clusterService.getSettings()).thenReturn(org.opensearch.common.settings.Settings.EMPTY); - when(mlFeatureEnabledSetting.isMetricCollectionEnabled()).thenReturn(true); - when(mlFeatureEnabledSetting.isStaticMetricCollectionEnabled()).thenReturn(true); + when(mlFeatureEnabledSetting.isAgenticMemoryEnabled()).thenReturn(true); + when(mlFeatureEnabledSetting.isMemoryRetentionEnabled()).thenReturn(true); listener.clusterChanged(event); - verify(mlTaskManager).indexStatsCollectorJob(true); + verify(mlTaskManager).indexMemoryRetentionJob(24); } - public void testClusterChanged_IndexAlreadyPresent_JobNotStarted() { - DiscoveryNode dataNode = createDataNode(Version.V_3_1_0); - setupClusterState(dataNode, true); + public void testClusterChanged_JobsDeferredUntilUpgradeCompletes() { + when(clusterService.getSettings()).thenReturn(org.opensearch.common.settings.Settings.EMPTY); + when(mlFeatureEnabledSetting.isAgenticMemoryEnabled()).thenReturn(true); + when(mlFeatureEnabledSetting.isMemoryRetentionEnabled()).thenReturn(true); + + // first event: upgrade in flight — deferred + setupClusterState(false, createDataNode("new", Version.CURRENT), createDataNode("old", Version.V_3_1_0)); + listener.clusterChanged(event); + verify(mlTaskManager, never()).indexMemoryRetentionJob(anyInt()); + + // second event: last old node upgraded — job created + setupClusterState(false, createDataNode("new", Version.CURRENT), createDataNode("old", Version.CURRENT)); + listener.clusterChanged(event); + verify(mlTaskManager).indexMemoryRetentionJob(24); + } + + public void testClusterChanged_IndexAlreadyPresent_StatsJobNotStarted() { + setupClusterState(true, createDataNode("n1", Version.CURRENT)); when(mlFeatureEnabledSetting.isMetricCollectionEnabled()).thenReturn(true); when(mlFeatureEnabledSetting.isStaticMetricCollectionEnabled()).thenReturn(true); @@ -122,24 +143,36 @@ public void testClusterChanged_IndexAlreadyPresent_JobNotStarted() { } public void testClusterChanged_MemoryRetentionJobStarted() { - DiscoveryNode dataNode = createDataNode(Version.V_3_1_0); - setupClusterState(dataNode, false); + setupClusterState(false, createDataNode("n1", Version.CURRENT)); when(clusterService.getSettings()).thenReturn(org.opensearch.common.settings.Settings.EMPTY); when(mlFeatureEnabledSetting.isAgenticMemoryEnabled()).thenReturn(true); + when(mlFeatureEnabledSetting.isMemoryRetentionEnabled()).thenReturn(true); listener.clusterChanged(event); verify(mlTaskManager).indexMemoryRetentionJob(24); } + public void testClusterChanged_MemoryRetentionJobNotStarted_WhenRetentionDisabled() { + setupClusterState(false, createDataNode("n1", Version.CURRENT)); + when(clusterService.getSettings()).thenReturn(org.opensearch.common.settings.Settings.EMPTY); + + when(mlFeatureEnabledSetting.isAgenticMemoryEnabled()).thenReturn(true); + when(mlFeatureEnabledSetting.isMemoryRetentionEnabled()).thenReturn(false); + + listener.clusterChanged(event); + + verify(mlTaskManager, never()).indexMemoryRetentionJob(anyInt()); + } + public void testClusterChanged_MemoryRetentionJobNotStarted_WhenMultiTenancyEnabled() { - DiscoveryNode dataNode = createDataNode(Version.V_3_1_0); - setupClusterState(dataNode, false); + setupClusterState(false, createDataNode("n1", Version.CURRENT)); when(clusterService.getSettings()) .thenReturn(org.opensearch.common.settings.Settings.builder().put("plugins.ml_commons.multi_tenancy_enabled", true).build()); when(mlFeatureEnabledSetting.isAgenticMemoryEnabled()).thenReturn(true); + when(mlFeatureEnabledSetting.isMemoryRetentionEnabled()).thenReturn(true); listener.clusterChanged(event); @@ -147,8 +180,7 @@ public void testClusterChanged_MemoryRetentionJobNotStarted_WhenMultiTenancyEnab } public void testClusterChanged_MemoryRetentionJobNotStarted_WhenAgenticMemoryDisabled() { - DiscoveryNode dataNode = createDataNode(Version.V_3_1_0); - setupClusterState(dataNode, false); + setupClusterState(false, createDataNode("n1", Version.CURRENT)); when(clusterService.getSettings()).thenReturn(org.opensearch.common.settings.Settings.EMPTY); when(mlFeatureEnabledSetting.isAgenticMemoryEnabled()).thenReturn(false); @@ -158,10 +190,10 @@ public void testClusterChanged_MemoryRetentionJobNotStarted_WhenAgenticMemoryDis verify(mlTaskManager, never()).indexMemoryRetentionJob(anyInt()); } - private DiscoveryNode createDataNode(Version version) { + private DiscoveryNode createDataNode(String id, Version version) { return new DiscoveryNode( - "dataNode", - "dataNodeId", + id, + id + "Id", buildNewFakeTransportAddress(), Collections.emptyMap(), Collections.singleton(DiscoveryNodeRole.DATA_ROLE), @@ -169,8 +201,12 @@ private DiscoveryNode createDataNode(Version version) { ); } - private void setupClusterState(DiscoveryNode node, boolean hasMLJobsIndex) { - DiscoveryNodes nodes = DiscoveryNodes.builder().add(node).build(); + private void setupClusterState(boolean hasMLJobsIndex, DiscoveryNode... discoveryNodes) { + DiscoveryNodes.Builder builder = DiscoveryNodes.builder(); + for (DiscoveryNode node : discoveryNodes) { + builder.add(node); + } + DiscoveryNodes nodes = builder.build(); when(event.state()).thenReturn(clusterState); when(event.previousState()).thenReturn(clusterState);