Skip to content

[Backport 3.8] [Memory Retention] Fix yellow cluster during rolling upgrade #4940

New issue

Have a question about this project? Sign up for a free GitHub account to open an issue and contact its maintainers and the community.

Sign up for GitHub

By clicking “Sign up for GitHub”, you agree to our terms of service and privacy statement. We’ll occasionally send you account related emails.

Already on GitHub? Sign in to your account

Merged
rithinpullela merged 1 commit into from
Jul 28, 2026
+92 −52
Merged

[Backport 3.8] [Memory Retention] Fix yellow cluster during rolling upgrade #4940

Show file tree
Hide file tree
Changes from all commits
Commits
Show all changes
1 commit
Select commit
File filter

Filter by extension

Filter by extension
Conversations
Failed to load comments.
Loading
Jump to
Jump to file
Failed to load files.
Loading
Diff view
Diff view
56 changes: 30 additions & 26 deletions plugin/src/main/java/org/opensearch/ml/cluster/MLCommonsClusterEventListener.java
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -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;
}
}
}
Expand Down
88 changes: 62 additions & 26 deletions plugin/src/test/java/org/opensearch/ml/cluster/MLCommonsClusterEventListenerTests.java
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -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);
Expand All @@ -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);
Expand All @@ -122,33 +143,44 @@ 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);

verify(mlTaskManager, never()).indexMemoryRetentionJob(anyInt());
}

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);
Expand All @@ -158,19 +190,23 @@ 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),
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);
Expand Down
Loading