From df6f0f8322d220c83ffc8f60882352c573128f35 Mon Sep 17 00:00:00 2001 From: AirFunnyB Date: Tue, 28 Jul 2026 01:57:15 -0600 Subject: [PATCH] Defer .plugins-ml-jobs creation until rolling upgrade completes; gate retention job on retention_enabled (#4937) MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit Creating the jobs index from the first upgraded node strands its replica (replicas cannot be allocated to nodes older than the primary's node), leaving the cluster yellow for the whole rolling-upgrade window and failing downstream BWC tests. Only touch the jobs index when it already exists or when every node in the cluster is on at least this node's version; otherwise defer — the cluster state change from the last old node leaving re-runs the listener. Also require plugins.ml_commons.memory.retention_enabled (opt-in, default false) before creating the memory retention job. Signed-off-by: Erfan Ballew Co-authored-by: Erfan Ballew (cherry picked from commit bc8e96f3082da8928c83cdb807c9463edbec57b0) Signed-off-by: opensearch-ci-bot --- .../MLCommonsClusterEventListener.java | 56 ++++++------ .../MLCommonsClusterEventListenerTests.java | 88 +++++++++++++------ 2 files changed, 92 insertions(+), 52 deletions(-) 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);