Merge pull request #24600 from vespa-engine/freva/aws-host-flavor-flag

Allow overriding host flavor

3 files changed, 62 insertions, 1 deletions

diff --git a/flags/src/main/java/com/yahoo/vespa/flags/PermanentFlags.java b/flags/src/main/java/com/yahoo/vespa/flags/PermanentFlags.java
index 7766faf28c7..f68f6816ec8 100644
--- a/flags/src/main/java/com/yahoo/vespa/flags/PermanentFlags.java
+++ b/flags/src/main/java/com/yahoo/vespa/flags/PermanentFlags.java
@@ -69,6 +69,12 @@ public class PermanentFlags {
                     "node resources of the host, the maximum number of containers, etc.",
             "Takes effect on next iteration of DynamicProvisioningMaintainer.");
 
+    public static final UnboundStringFlag HOST_FLAVOR = defineStringFlag(
+            "host-flavor", "",
+            "Specifies the Vespa flavor name that the hosts of the matching nodes should have.",
+            "Takes effect on next deployment (including internal redeployment).",
+            APPLICATION_ID, CLUSTER_TYPE);
+
     public static final UnboundBooleanFlag SKIP_MAINTENANCE_DEPLOYMENT = defineFeatureFlag(
             "node-repository-skip-maintenance-deployment", false,
             "Whether PeriodicApplicationMaintainer should skip deployment for an application",
diff --git a/node-repository/src/main/java/com/yahoo/vespa/hosted/provision/provisioning/NodeAllocation.java b/node-repository/src/main/java/com/yahoo/vespa/hosted/provision/provisioning/NodeAllocation.java
index 9557354725b..ff24f627aa4 100644
--- a/node-repository/src/main/java/com/yahoo/vespa/hosted/provision/provisioning/NodeAllocation.java
+++ b/node-repository/src/main/java/com/yahoo/vespa/hosted/provision/provisioning/NodeAllocation.java
@@ -9,6 +9,8 @@ import com.yahoo.config.provision.NodeResources;
 import com.yahoo.config.provision.NodeType;
 import com.yahoo.config.provision.SystemName;
 import com.yahoo.net.HostName;
+import com.yahoo.vespa.flags.FetchVector;
+import com.yahoo.vespa.flags.PermanentFlags;
 import com.yahoo.vespa.hosted.provision.Node;
 import com.yahoo.vespa.hosted.provision.NodeList;
 import com.yahoo.vespa.hosted.provision.NodeRepository;
@@ -81,6 +83,7 @@ class NodeAllocation {
 
     private final NodeRepository nodeRepository;
     private final NodeResourceLimits nodeResourceLimits;
+    private final Optional<String> requiredHostFlavor;
 
     NodeAllocation(NodesAndHosts<? extends NodeList> allNodesAndHosts, ApplicationId application, ClusterSpec cluster, NodeSpec requestedNodes,
                    Supplier<Integer> nextIndex, NodeRepository nodeRepository) {
@@ -90,7 +93,12 @@ class NodeAllocation {
         this.requestedNodes = requestedNodes;
         this.nextIndex = nextIndex;
         this.nodeRepository = nodeRepository;
-        nodeResourceLimits = new NodeResourceLimits(nodeRepository);
+        this.nodeResourceLimits = new NodeResourceLimits(nodeRepository);
+        this.requiredHostFlavor = Optional.of(PermanentFlags.HOST_FLAVOR.bindTo(nodeRepository.flagSource())
+                        .with(FetchVector.Dimension.APPLICATION_ID, application.serializedForm())
+                        .with(FetchVector.Dimension.CLUSTER_TYPE, cluster.type().name())
+                        .value())
+                .filter(s -> !s.isBlank());
     }
 
     /**
@@ -115,6 +123,7 @@ class NodeAllocation {
                 if ( candidate.state() == Node.State.active && allocation.removable()) continue; // don't accept; causes removal
                 if ( candidate.state() == Node.State.active && candidate.wantToFail()) continue; // don't accept; causes failing
                 if ( indexes.contains(membership.index())) continue; // duplicate index (just to be sure)
+                if ( requiredHostFlavor.isPresent() && ! candidate.parent.map(node -> node.flavor().name()).equals(requiredHostFlavor)) continue;
 
                 boolean resizeable = requestedNodes.considerRetiring() && candidate.isResizable;
                 boolean acceptToRetire = acceptToRetire(candidate);
diff --git a/node-repository/src/test/java/com/yahoo/vespa/hosted/provision/provisioning/DynamicProvisioningTest.java b/node-repository/src/test/java/com/yahoo/vespa/hosted/provision/provisioning/DynamicProvisioningTest.java
index 3570e56e196..94b525cfe3c 100644
--- a/node-repository/src/test/java/com/yahoo/vespa/hosted/provision/provisioning/DynamicProvisioningTest.java
+++ b/node-repository/src/test/java/com/yahoo/vespa/hosted/provision/provisioning/DynamicProvisioningTest.java
@@ -11,12 +11,15 @@ import com.yahoo.config.provision.Environment;
 import com.yahoo.config.provision.Flavor;
 import com.yahoo.config.provision.HostSpec;
 import com.yahoo.config.provision.NodeResources;
+import com.yahoo.config.provision.NodeResources.Architecture;
 import com.yahoo.config.provision.NodeResources.DiskSpeed;
 import com.yahoo.config.provision.NodeResources.StorageType;
 import com.yahoo.config.provision.NodeType;
 import com.yahoo.config.provision.RegionName;
 import com.yahoo.config.provision.SystemName;
 import com.yahoo.config.provision.Zone;
+import com.yahoo.vespa.flags.InMemoryFlagSource;
+import com.yahoo.vespa.flags.PermanentFlags;
 import com.yahoo.vespa.hosted.provision.Node;
 import com.yahoo.vespa.hosted.provision.NodeList;
 import com.yahoo.vespa.hosted.provision.node.Agent;
@@ -260,6 +263,49 @@ public class DynamicProvisioningTest {
     }
 
     @Test
+    public void migrates_nodes_on_host_flavor_flag_change() {
+        InMemoryFlagSource flagSource = new InMemoryFlagSource();
+        List<Flavor> flavors = List.of(new Flavor("x86", new NodeResources(1, 4, 50, 0.1, fast, local, Architecture.x86_64)),
+                                       new Flavor("arm", new NodeResources(1, 4, 50, 0.1, fast, local, Architecture.arm64)));
+        MockHostProvisioner hostProvisioner = new MockHostProvisioner(flavors);
+        ProvisioningTester tester = new ProvisioningTester.Builder().zone(zone)
+                .flavors(flavors)
+                .hostProvisioner(hostProvisioner)
+                .resourcesCalculator(0, 0)
+                .nameResolver(nameResolver)
+                .flagSource(flagSource)
+                .build();
+
+        ApplicationId app = ProvisioningTester.applicationId();
+        ClusterSpec cluster = ClusterSpec.request(ClusterSpec.Type.content, new ClusterSpec.Id("cluster1")).vespaVersion("8").build();
+        Capacity capacity = Capacity.from(new ClusterResources(4, 2, new NodeResources(1, 4, 50, 0.1, DiskSpeed.any, StorageType.any, Architecture.any)));
+
+        hostProvisioner.overrideHostFlavor("x86");
+        tester.activate(app, cluster, capacity);
+        NodeList nodes = tester.nodeRepository().nodes().list();
+        nodes.forEach(n -> System.out.println(n.hostname() + " " + n.flavor().name()));
+        assertEquals(4, nodes.owner(app).state(Node.State.active).size());
+        assertEquals(Set.of("x86"), nodes.parentsOf(nodes.owner(app).state(Node.State.active)).stream().map(n -> n.flavor().name()).collect(Collectors.toSet()));
+
+        hostProvisioner.overrideHostFlavor("arm");
+        flagSource.withStringFlag(PermanentFlags.HOST_FLAVOR.id(), "arm");
+        tester.activate(app, cluster, capacity);
+        nodes = tester.nodeRepository().nodes().list();
+        assertEquals(4, nodes.owner(app).state(Node.State.inactive).size());
+        assertEquals(4, nodes.owner(app).state(Node.State.active).size());
+        assertEquals(Set.of("x86"), nodes.parentsOf(tester.getNodes(app, Node.State.inactive)).stream().map(n -> n.flavor().name()).collect(Collectors.toSet()));
+        assertEquals(Set.of("arm"), nodes.parentsOf(tester.getNodes(app, Node.State.active)).stream().map(n -> n.flavor().name()).collect(Collectors.toSet()));
+
+        flagSource.removeFlag(PermanentFlags.HOST_FLAVOR.id()); // Resetting flag does not moves the nodes back
+        tester.activate(app, cluster, capacity);
+        nodes = tester.nodeRepository().nodes().list();
+        assertEquals(4, nodes.owner(app).state(Node.State.inactive).size());
+        assertEquals(4, nodes.owner(app).state(Node.State.active).size());
+        assertEquals(Set.of("x86"), nodes.parentsOf(tester.getNodes(app, Node.State.inactive)).stream().map(n -> n.flavor().name()).collect(Collectors.toSet()));
+        assertEquals(Set.of("arm"), nodes.parentsOf(tester.getNodes(app, Node.State.active)).stream().map(n -> n.flavor().name()).collect(Collectors.toSet()));
+    }
+
+    @Test
     public void test_changing_limits() {
         int memoryTax = 3;
         List<Flavor> flavors = List.of(new Flavor("1x", new NodeResources(1, 10 - memoryTax, 100, 0.1, fast, remote)),