diff --git a/x-pack/plugin/ml/src/main/java/org/elasticsearch/xpack/ml/MachineLearningExtension.java b/x-pack/plugin/ml/src/main/java/org/elasticsearch/xpack/ml/MachineLearningExtension.java
index ad1d406d0cbbe..b0e65d524036c 100644
--- a/x-pack/plugin/ml/src/main/java/org/elasticsearch/xpack/ml/MachineLearningExtension.java
+++ b/x-pack/plugin/ml/src/main/java/org/elasticsearch/xpack/ml/MachineLearningExtension.java
@@ -10,6 +10,7 @@
 import org.elasticsearch.common.settings.ClusterSettings;
 import org.elasticsearch.common.settings.Settings;
 import org.elasticsearch.xpack.ml.autoscaling.AbstractNodeAvailabilityZoneMapper;
+import org.elasticsearch.xpack.ml.autoscaling.NodeRealAvailabilityZoneMapper;
 
 public interface MachineLearningExtension {
 
@@ -29,5 +30,7 @@ default String[] getAnalyticsDestIndexAllowedSettings() {
         return DefaultMachineLearningExtension.ANALYTICS_DEST_INDEX_ALLOWED_SETTINGS;
     }
 
-    AbstractNodeAvailabilityZoneMapper getNodeAvailabilityZoneMapper(Settings settings, ClusterSettings clusterSettings);
+    default AbstractNodeAvailabilityZoneMapper getNodeAvailabilityZoneMapper(Settings settings, ClusterSettings clusterSettings) {
+        return new NodeRealAvailabilityZoneMapper(settings, clusterSettings);
+    }
 }
diff --git a/x-pack/plugin/ml/src/main/java/org/elasticsearch/xpack/ml/autoscaling/NodeFakeAvailabilityZoneMapper.java b/x-pack/plugin/ml/src/main/java/org/elasticsearch/xpack/ml/autoscaling/NodeFakeAvailabilityZoneMapper.java
index 163e9cae2b251..618db1da71ab8 100644
--- a/x-pack/plugin/ml/src/main/java/org/elasticsearch/xpack/ml/autoscaling/NodeFakeAvailabilityZoneMapper.java
+++ b/x-pack/plugin/ml/src/main/java/org/elasticsearch/xpack/ml/autoscaling/NodeFakeAvailabilityZoneMapper.java
@@ -16,7 +16,6 @@
 import org.elasticsearch.common.settings.ClusterSettings;
 import org.elasticsearch.common.settings.Settings;
 
-import java.util.ArrayList;
 import java.util.Collection;
 import java.util.HashMap;
 import java.util.List;
@@ -46,19 +45,16 @@ public NodeFakeAvailabilityZoneMapper(Settings settings, ClusterSettings cluster
      *         cluster.
      */
     public NodesByAvailabilityZone buildNodesByAvailabilityZone(DiscoveryNodes discoveryNodes) {
-        return doBuildNodesByAvailabilityZone(discoveryNodes);
-    }
-
-    private static NodesByAvailabilityZone doBuildNodesByAvailabilityZone(DiscoveryNodes discoveryNodes) {
         Collection<DiscoveryNode> nodes = discoveryNodes.getNodes().values();
 
         Map<List<String>, Collection<DiscoveryNode>> allNodesByAvailabilityZone = new HashMap<>();
         Map<List<String>, Collection<DiscoveryNode>> mlNodesByAvailabilityZone = new HashMap<>();
         for (DiscoveryNode node : nodes) {
             List<String> nodeIdValues = List.of(node.getId());
-            allNodesByAvailabilityZone.computeIfAbsent(nodeIdValues, k -> new ArrayList<>()).add(node);
+            List<DiscoveryNode> nodeList = List.of(node);
+            allNodesByAvailabilityZone.put(nodeIdValues, nodeList);
             if (node.getRoles().contains(DiscoveryNodeRole.ML_ROLE)) {
-                mlNodesByAvailabilityZone.computeIfAbsent(nodeIdValues, k -> new ArrayList<>()).add(node);
+                mlNodesByAvailabilityZone.put(nodeIdValues, nodeList);
             }
         }
         return new NodesByAvailabilityZone(Map.copyOf(allNodesByAvailabilityZone), Map.copyOf(mlNodesByAvailabilityZone));