Stop reading container images from ZK

14 files changed, 89 insertions, 140 deletions

diff --git a/node-repository/src/main/java/com/yahoo/vespa/hosted/provision/NodeRepository.java b/node-repository/src/main/java/com/yahoo/vespa/hosted/provision/NodeRepository.java
index 67de210127d..7f0c52a5aca 100644
--- a/node-repository/src/main/java/com/yahoo/vespa/hosted/provision/NodeRepository.java
+++ b/node-repository/src/main/java/com/yahoo/vespa/hosted/provision/NodeRepository.java
@@ -80,6 +80,7 @@ public class NodeRepository extends AbstractComponent {
              zone,
              new DnsNameResolver(),
              DockerImage.fromString(config.containerImage()),
+             Optional.of(config.tenantContainerImage()).filter(s -> !s.isEmpty()).map(DockerImage::fromString),
              flagSource,
              metricsDb,
              config.useCuratorClientCache(),
@@ -98,6 +99,7 @@ public class NodeRepository extends AbstractComponent {
                           Zone zone,
                           NameResolver nameResolver,
                           DockerImage containerImage,
+                          Optional<DockerImage> tenantContainerImage,
                           FlagSource flagSource,
                           MetricsDb metricsDb,
                           boolean useCuratorClientCache,
@@ -118,7 +120,7 @@ public class NodeRepository extends AbstractComponent {
         this.osVersions = new OsVersions(this);
         this.infrastructureVersions = new InfrastructureVersions(db);
         this.firmwareChecks = new FirmwareChecks(db, clock);
-        this.containerImages = new ContainerImages(db, containerImage);
+        this.containerImages = new ContainerImages(containerImage, tenantContainerImage);
         this.archiveUris = new ArchiveUris(db);
         this.jobControl = new JobControl(new JobControlFlags(db, flagSource));
         this.applications = new Applications(db);
diff --git a/node-repository/src/main/java/com/yahoo/vespa/hosted/provision/persistence/CuratorDatabaseClient.java b/node-repository/src/main/java/com/yahoo/vespa/hosted/provision/persistence/CuratorDatabaseClient.java
index 45aaedd9550..0a05d9d5bf4 100644
--- a/node-repository/src/main/java/com/yahoo/vespa/hosted/provision/persistence/CuratorDatabaseClient.java
+++ b/node-repository/src/main/java/com/yahoo/vespa/hosted/provision/persistence/CuratorDatabaseClient.java
@@ -6,7 +6,6 @@ import com.yahoo.component.Version;
 import com.yahoo.config.provision.ApplicationId;
 import com.yahoo.config.provision.ApplicationLockException;
 import com.yahoo.config.provision.ApplicationTransaction;
-import com.yahoo.config.provision.DockerImage;
 import com.yahoo.config.provision.HostName;
 import com.yahoo.config.provision.NodeFlavors;
 import com.yahoo.config.provision.NodeType;
@@ -67,7 +66,7 @@ public class CuratorDatabaseClient {
     private static final Path inactiveJobsPath = root.append("inactiveJobs");
     private static final Path infrastructureVersionsPath = root.append("infrastructureVersions");
     private static final Path osVersionsPath = root.append("osVersions");
-    private static final Path containerImagesPath = root.append("dockerImages");
+    private static final Path containerImagesPath = root.append("dockerImages"); // TODO(mpolden): Delete this path from ZK after 2021-11-01
     private static final Path firmwareCheckPath = root.append("firmwareCheck");
     private static final Path archiveUrisPath = root.append("archiveUris");
 
@@ -98,7 +97,6 @@ public class CuratorDatabaseClient {
         db.create(inactiveJobsPath);
         db.create(infrastructureVersionsPath);
         db.create(osVersionsPath);
-        db.create(containerImagesPath);
         db.create(firmwareCheckPath);
         db.create(archiveUrisPath);
         db.create(loadBalancersPath);
@@ -410,24 +408,6 @@ public class CuratorDatabaseClient {
         return db.lock(lockPath.append("osVersionsLock"), defaultLockTimeout);
     }
 
-    // Container images -----------------------------------------------------------
-
-    public Map<NodeType, DockerImage> readContainerImages() {
-        return read(containerImagesPath, NodeTypeContainerImagesSerializer::fromJson).orElseGet(TreeMap::new);
-    }
-
-    public void writeContainerImages(Map<NodeType, DockerImage> images) {
-        NestedTransaction transaction = new NestedTransaction();
-        CuratorTransaction curatorTransaction = db.newCuratorTransactionIn(transaction);
-        curatorTransaction.add(CuratorOperations.setData(containerImagesPath.getAbsolute(),
-                                                         NodeTypeContainerImagesSerializer.toJson(images)));
-        transaction.commit();
-    }
-
-    public Lock lockContainerImages() {
-        return db.lock(lockPath.append("dockerImagesLock"), defaultLockTimeout);
-    }
-
     // Firmware checks -----------------------------------------------------------
 
     /** Stores the instant after which a firmware check is required, or clears any outstanding ones if empty is given. */
diff --git a/node-repository/src/main/java/com/yahoo/vespa/hosted/provision/provisioning/ContainerImages.java b/node-repository/src/main/java/com/yahoo/vespa/hosted/provision/provisioning/ContainerImages.java
index cd9303d7f6e..a3227061e2c 100644
--- a/node-repository/src/main/java/com/yahoo/vespa/hosted/provision/provisioning/ContainerImages.java
+++ b/node-repository/src/main/java/com/yahoo/vespa/hosted/provision/provisioning/ContainerImages.java
@@ -3,73 +3,49 @@ package com.yahoo.vespa.hosted.provision.provisioning;
 
 import com.yahoo.config.provision.DockerImage;
 import com.yahoo.config.provision.NodeType;
-import com.yahoo.lang.CachedSupplier;
-import com.yahoo.vespa.curator.Lock;
-import com.yahoo.vespa.hosted.provision.persistence.CuratorDatabaseClient;
+import com.yahoo.vespa.hosted.provision.Node;
 
-import java.time.Duration;
-import java.util.Collections;
-import java.util.Map;
+import java.util.Objects;
 import java.util.Optional;
-import java.util.logging.Logger;
 
 /**
- * Multi-thread safe class to get and set container images for given node types. Images are stored in ZooKeeper so that
- * nodes receive the same image from all config servers.
+ * This class decides the container image to use for a given node. Two sources are considered, in the following order:
+ *
+ * 1. Requested image (from node allocation, this is set by either a feature flag or through services.xml)
+ * 2. Default image, specified in the node repository config file
+ *
+ * Independent of source, the registry part of the image is rewritten to match the one set in the node repository config
+ * file.
  *
  * @author freva
+ * @author mpolden
  */
 public class ContainerImages {
 
-    private static final Duration cacheTtl = Duration.ofMinutes(1);
-    private static final Logger log = Logger.getLogger(ContainerImages.class.getName());
-
-    private final CuratorDatabaseClient db;
     private final DockerImage defaultImage;
+    private final Optional<DockerImage> tenantImage;
 
-    /**
-     * The container image is read on every request to /nodes/v2/node/[fqdn]. Cache current images to avoid
-     * unnecessary ZK reads. When images change, some nodes may need to wait for TTL until they see the new image,
-     * this is fine.
-     */
-    private final CachedSupplier<Map<NodeType, DockerImage>> images;
-
-    public ContainerImages(CuratorDatabaseClient db, DockerImage defaultImage) {
-        this.db = db;
-        this.defaultImage = defaultImage;
-        this.images = new CachedSupplier<>(() -> Collections.unmodifiableMap(db.readContainerImages()), cacheTtl);
-    }
-
-    /** Returns the current images for each node type */
-    public Map<NodeType, DockerImage> getImages() {
-        return images.get();
+    public ContainerImages(DockerImage defaultImage, Optional<DockerImage> tenantContainerImage) {
+        this.defaultImage = Objects.requireNonNull(defaultImage);
+        this.tenantImage = Objects.requireNonNull(tenantContainerImage);
     }
 
-    /** Returns the container image to use for given node type */
-    public DockerImage imageFor(NodeType type) {
-        NodeType typeToUseForLookup = type.isHost() ? type.childNodeType() : type;
-        DockerImage image = getImages().get(typeToUseForLookup);
-        if (image == null) {
+    /** Returns the container image to use for given node */
+    public DockerImage get(Node node) {
+        Optional<DockerImage> requestedImage = node.allocation()
+                                                   .flatMap(allocation -> allocation.membership().cluster().dockerImageRepo());
+        NodeType nodeType = node.type().isHost() ? node.type().childNodeType() : node.type();
+        final DockerImage image;
+        if (requestedImage.isPresent()) {
+            image = requestedImage.get();
+        } else if (nodeType == NodeType.tenant) {
+            image = tenantImage.orElse(defaultImage);
+        } else {
             image = defaultImage;
         }
         return rewriteRegistry(image);
     }
 
-    /** Set the docker image for nodes of given type */
-    public void setImage(NodeType nodeType, Optional<DockerImage> image) {
-        if (nodeType.isHost()) {
-            throw new IllegalArgumentException("Setting container image for " + nodeType + " nodes is unsupported");
-        }
-        try (Lock lock = db.lockContainerImages()) {
-            Map<NodeType, DockerImage> images = db.readContainerImages();
-            image.ifPresentOrElse(img -> images.put(nodeType, img),
-                                  () -> images.remove(nodeType));
-            db.writeContainerImages(images);
-            this.images.invalidate(); // Throw away current cache
-            log.info("Set container image for " + nodeType + " nodes to " + image.map(DockerImage::asString).orElse(null));
-        }
-    }
-
     /** Rewrite the registry part of given image, using this zone's default image */
     private DockerImage rewriteRegistry(DockerImage image) {
         return image.withRegistry(defaultImage.registry());
diff --git a/node-repository/src/main/java/com/yahoo/vespa/hosted/provision/restapi/NodesResponse.java b/node-repository/src/main/java/com/yahoo/vespa/hosted/provision/restapi/NodesResponse.java
index df6e24aa8a3..02b426ed6fc 100644
--- a/node-repository/src/main/java/com/yahoo/vespa/hosted/provision/restapi/NodesResponse.java
+++ b/node-repository/src/main/java/com/yahoo/vespa/hosted/provision/restapi/NodesResponse.java
@@ -146,8 +146,7 @@ class NodesResponse extends SlimeJsonResponse {
             toSlime(allocation.membership(), object.setObject("membership"));
             object.setLong("restartGeneration", allocation.restartGeneration().wanted());
             object.setLong("currentRestartGeneration", allocation.restartGeneration().current());
-            object.setString("wantedDockerImage", allocation.membership().cluster().dockerImage()
-                    .orElseGet(() -> nodeRepository.containerImages().imageFor(node.type()).withTag(allocation.membership().cluster().vespaVersion()).asString()));
+            object.setString("wantedDockerImage", nodeRepository.containerImages().get(node).withTag(allocation.membership().cluster().vespaVersion()).asString());
             object.setString("wantedVespaVersion", allocation.membership().cluster().vespaVersion().toFullString());
             NodeResourcesSerializer.toSlime(allocation.requestedResources(), object.setObject("requestedResources"));
             allocation.networkPorts().ifPresent(ports -> NetworkPortsSerializer.toSlime(ports, object.setArray("networkPorts")));
@@ -214,7 +213,7 @@ class NodesResponse extends SlimeJsonResponse {
                    .or(() -> Optional.of(node)
                                      .filter(n -> n.flavor().getType() != Flavor.Type.DOCKER_CONTAINER)
                                      .flatMap(n -> n.status().vespaVersion()
-                                                    .map(version -> nodeRepository.containerImages().imageFor(n.type()).withTag(version))));
+                                                    .map(version -> nodeRepository.containerImages().get(n).withTag(version))));
     }
 
     private void ipAddressesToSlime(Set<String> ipAddresses, Cursor array) {
diff --git a/node-repository/src/main/java/com/yahoo/vespa/hosted/provision/restapi/NodesV2ApiHandler.java b/node-repository/src/main/java/com/yahoo/vespa/hosted/provision/restapi/NodesV2ApiHandler.java
index fa89e2e58f0..1163890addf 100644
--- a/node-repository/src/main/java/com/yahoo/vespa/hosted/provision/restapi/NodesV2ApiHandler.java
+++ b/node-repository/src/main/java/com/yahoo/vespa/hosted/provision/restapi/NodesV2ApiHandler.java
@@ -3,7 +3,6 @@ package com.yahoo.vespa.hosted.provision.restapi;
 
 import com.yahoo.component.Version;
 import com.yahoo.config.provision.ApplicationId;
-import com.yahoo.config.provision.DockerImage;
 import com.yahoo.config.provision.Flavor;
 import com.yahoo.config.provision.HostFilter;
 import com.yahoo.config.provision.NodeFlavors;
@@ -358,12 +357,11 @@ public class NodesV2ApiHandler extends LoggingRequestHandler {
 
     private MessageResponse setTargetVersions(String nodeTypeS, Inspector inspector) {
         NodeType nodeType = NodeType.valueOf(nodeTypeS.toLowerCase());
-        List<String> messageParts = new ArrayList<>(4);
+        List<String> messageParts = new ArrayList<>();
 
         boolean force = inspector.field("force").asBool();
         Inspector versionField = inspector.field("version");
         Inspector osVersionField = inspector.field("osVersion");
-        Inspector containerImageField = inspector.field("dockerImage");
         Inspector upgradeBudgetField = inspector.field("upgradeBudget");
 
         if (versionField.valid()) {
@@ -396,16 +394,8 @@ public class NodesV2ApiHandler extends LoggingRequestHandler {
             }
         }
 
-        if (containerImageField.valid()) {
-            Optional<DockerImage> dockerImage = Optional.of(containerImageField.asString())
-                    .filter(s -> !s.isEmpty())
-                    .map(DockerImage::fromString);
-            nodeRepository.containerImages().setImage(nodeType, dockerImage);
-            messageParts.add("container image to " + dockerImage.map(DockerImage::asString).orElse(null));
-        }
-
         if (messageParts.isEmpty()) {
-            throw new IllegalArgumentException("At least one of 'version', 'osVersion' or 'dockerImage' must be set");
+            throw new IllegalArgumentException("At least one of 'version' or 'osVersion' must be set");
         }
 
         return new MessageResponse("Set " + String.join(", ", messageParts) +
diff --git a/node-repository/src/main/java/com/yahoo/vespa/hosted/provision/restapi/UpgradeResponse.java b/node-repository/src/main/java/com/yahoo/vespa/hosted/provision/restapi/UpgradeResponse.java
index cc5d7745d2f..e03edecff5e 100644
--- a/node-repository/src/main/java/com/yahoo/vespa/hosted/provision/restapi/UpgradeResponse.java
+++ b/node-repository/src/main/java/com/yahoo/vespa/hosted/provision/restapi/UpgradeResponse.java
@@ -21,13 +21,11 @@ public class UpgradeResponse extends HttpResponse {
 
     private final InfrastructureVersions infrastructureVersions;
     private final OsVersions osVersions;
-    private final ContainerImages containerImages;
 
     public UpgradeResponse(InfrastructureVersions infrastructureVersions, OsVersions osVersions, ContainerImages containerImages) {
         super(200);
         this.infrastructureVersions = infrastructureVersions;
         this.osVersions = osVersions;
-        this.containerImages = containerImages;
     }
 
     @Override
@@ -41,9 +39,7 @@ public class UpgradeResponse extends HttpResponse {
         Cursor osVersionsObject = root.setObject("osVersions");
         osVersions.readChange().targets().forEach((nodeType, target) -> osVersionsObject.setString(nodeType.name(), target.version().toFullString()));
 
-
-        Cursor dockerImagesObject = root.setObject("dockerImages");
-        containerImages.getImages().forEach((nodeType, image) -> dockerImagesObject.setString(nodeType.name(), image.asString()));
+        root.setObject("dockerImages"); // Unused, but present to avoid breaking API
 
         new JsonFormat(true).encode(stream, slime);
     }
diff --git a/node-repository/src/main/java/com/yahoo/vespa/hosted/provision/testutils/MockNodeRepository.java b/node-repository/src/main/java/com/yahoo/vespa/hosted/provision/testutils/MockNodeRepository.java
index 00937fe22ce..1a2d5294aa5 100644
--- a/node-repository/src/main/java/com/yahoo/vespa/hosted/provision/testutils/MockNodeRepository.java
+++ b/node-repository/src/main/java/com/yahoo/vespa/hosted/provision/testutils/MockNodeRepository.java
@@ -67,6 +67,7 @@ public class MockNodeRepository extends NodeRepository {
               Zone.defaultZone(),
               new MockNameResolver().mockAnyLookup(),
               DockerImage.fromString("docker-registry.domain.tld:8080/dist/vespa"),
+              Optional.empty(),
               new InMemoryFlagSource(),
               new MemoryMetricsDb(Clock.fixed(Instant.ofEpochMilli(123), ZoneId.of("Z"))),
               true,
diff --git a/node-repository/src/test/java/com/yahoo/vespa/hosted/provision/NodeRepositoryTester.java b/node-repository/src/test/java/com/yahoo/vespa/hosted/provision/NodeRepositoryTester.java
index aaccd7b31cc..b391292884f 100644
--- a/node-repository/src/test/java/com/yahoo/vespa/hosted/provision/NodeRepositoryTester.java
+++ b/node-repository/src/test/java/com/yahoo/vespa/hosted/provision/NodeRepositoryTester.java
@@ -44,6 +44,7 @@ public class NodeRepositoryTester {
                                             Zone.defaultZone(),
                                             new MockNameResolver().mockAnyLookup(),
                                             DockerImage.fromString("docker-registry.domain.tld:8080/dist/vespa"),
+                                            Optional.empty(),
                                             new InMemoryFlagSource(),
                                             new MemoryMetricsDb(clock),
                                             true,
diff --git a/node-repository/src/test/java/com/yahoo/vespa/hosted/provision/maintenance/CapacityCheckerTester.java b/node-repository/src/test/java/com/yahoo/vespa/hosted/provision/maintenance/CapacityCheckerTester.java
index 7e1def4b754..b10b4c2958a 100644
--- a/node-repository/src/test/java/com/yahoo/vespa/hosted/provision/maintenance/CapacityCheckerTester.java
+++ b/node-repository/src/test/java/com/yahoo/vespa/hosted/provision/maintenance/CapacityCheckerTester.java
@@ -70,6 +70,7 @@ public class CapacityCheckerTester {
                                             zone,
                                             new MockNameResolver().mockAnyLookup(),
                                             DockerImage.fromString("docker-registry.domain.tld:8080/dist/vespa"),
+                                            Optional.empty(),
                                             new InMemoryFlagSource(),
                                             new MemoryMetricsDb(clock),
                                             true,
diff --git a/node-repository/src/test/java/com/yahoo/vespa/hosted/provision/maintenance/SpareCapacityMaintainerTest.java b/node-repository/src/test/java/com/yahoo/vespa/hosted/provision/maintenance/SpareCapacityMaintainerTest.java
index 7a42c3a11f0..373bfe20162 100644
--- a/node-repository/src/test/java/com/yahoo/vespa/hosted/provision/maintenance/SpareCapacityMaintainerTest.java
+++ b/node-repository/src/test/java/com/yahoo/vespa/hosted/provision/maintenance/SpareCapacityMaintainerTest.java
@@ -32,6 +32,7 @@ import java.time.Duration;
 import java.time.Instant;
 import java.util.ArrayList;
 import java.util.List;
+import java.util.Optional;
 import java.util.Set;
 import java.util.stream.Collectors;
 import java.util.stream.IntStream;
@@ -263,6 +264,7 @@ public class SpareCapacityMaintainerTest {
                                                 new Zone(Environment.prod, RegionName.from("us-east-3")),
                                                 new MockNameResolver().mockAnyLookup(),
                                                 DockerImage.fromString("docker-registry.domain.tld:8080/dist/vespa"),
+                                                Optional.empty(),
                                                 new InMemoryFlagSource(),
                                                 new MemoryMetricsDb(clock),
                                                 true,
diff --git a/node-repository/src/test/java/com/yahoo/vespa/hosted/provision/provisioning/ContainerImagesTest.java b/node-repository/src/test/java/com/yahoo/vespa/hosted/provision/provisioning/ContainerImagesTest.java
index 6f624ade97d..924787a1ca6 100644
--- a/node-repository/src/test/java/com/yahoo/vespa/hosted/provision/provisioning/ContainerImagesTest.java
+++ b/node-repository/src/test/java/com/yahoo/vespa/hosted/provision/provisioning/ContainerImagesTest.java
@@ -1,13 +1,22 @@
 // Copyright 2020 Oath Inc. Licensed under the terms of the Apache 2.0 license. See LICENSE in the project root.
 package com.yahoo.vespa.hosted.provision.provisioning;
 
+import com.yahoo.component.Version;
+import com.yahoo.config.provision.ApplicationId;
+import com.yahoo.config.provision.ClusterMembership;
 import com.yahoo.config.provision.DockerImage;
+import com.yahoo.config.provision.Flavor;
 import com.yahoo.config.provision.NodeResources;
 import com.yahoo.config.provision.NodeType;
-import com.yahoo.vespa.flags.InMemoryFlagSource;
+import com.yahoo.vespa.hosted.provision.Node;
+import com.yahoo.vespa.hosted.provision.node.Allocation;
+import com.yahoo.vespa.hosted.provision.node.Generation;
+import com.yahoo.vespa.hosted.provision.node.IP;
+import com.yahoo.vespa.hosted.provision.testutils.MockNodeFlavors;
 import org.junit.Test;
 
 import java.util.Optional;
+import java.util.Set;
 
 import static org.junit.Assert.assertEquals;
 
@@ -18,34 +27,46 @@ public class ContainerImagesTest {
 
     @Test
     public void image_selection() {
-        var flagSource = new InMemoryFlagSource();
-        var tester = new ProvisioningTester.Builder().flagSource(flagSource).build();
-
-        var proxyImage = DockerImage.fromString("docker-registry.domain.tld:8080/dist/proxy");
-        tester.nodeRepository().containerImages().setImage(NodeType.proxy, Optional.of(proxyImage));
-
-        // Host uses tenant default image (for preload purposes)
-        var defaultImage = DockerImage.fromString("docker-registry.domain.tld:8080/dist/vespa");
-        var hosts = tester.makeReadyNodes(2, "default", NodeType.host);
-        tester.activateTenantHosts();
-        for (var host : hosts) {
-            assertEquals(defaultImage, tester.nodeRepository().containerImages().imageFor(host.type()));
-        }
+        DockerImage defaultImage = DockerImage.fromString("registry.example.com/vespa/default");
+        DockerImage tenantImage = DockerImage.fromString("registry.example.com/vespa/tenant");
+        ContainerImages images = new ContainerImages(defaultImage, Optional.of(tenantImage));
 
-        // Tenant node uses tenant default image
-        var resources = new NodeResources(2, 8, 50, 1);
-        for (var host : hosts) {
-            var nodes = tester.makeReadyChildren(2, resources, host.hostname());
-            for (var node : nodes) {
-                assertEquals(defaultImage, tester.nodeRepository().containerImages().imageFor(node.type()));
-            }
-        }
+        assertEquals(defaultImage, images.get(node(NodeType.confighost)));  // For preload purposes
+        assertEquals(defaultImage, images.get(node(NodeType.config)));
+
+        assertEquals(tenantImage, images.get(node(NodeType.host))); // For preload purposes
+        assertEquals(tenantImage, images.get(node(NodeType.tenant)));
+
+        assertEquals(defaultImage, images.get(node(NodeType.proxyhost))); // For preload purposes
+        assertEquals(defaultImage, images.get(node(NodeType.proxy)));
+
+        // Tenant node requesting a special image
+        DockerImage requested = DockerImage.fromString("registry.example.com/vespa/special");
+        assertEquals(requested, images.get(node(NodeType.tenant, requested)));
+
+        // When there is no custom tenant image, the default one is used
+        images = new ContainerImages(defaultImage, Optional.empty());
+        assertEquals(defaultImage, images.get(node(NodeType.host)));
+        assertEquals(defaultImage, images.get(node(NodeType.tenant)));
+    }
+
+    private static Node node(NodeType type) {
+        return node(type, null);
+    }
 
-        // Proxy host uses image used by child nodes (proxy nodes), which is overridden in this case (for preload purposes)
-        var proxyHosts = tester.makeReadyNodes(2, "default", NodeType.proxyhost, 1);
-        for (var host : proxyHosts) {
-            assertEquals(proxyImage, tester.nodeRepository().containerImages().imageFor(host.type()));
+    private static Node node(NodeType type, DockerImage requested) {
+        Flavor flavor = new MockNodeFlavors().getFlavorOrThrow("default");
+        Node.Builder b = Node.create(type + "1", new IP.Config(Set.of(), Set.of()), type + "1.example.com", flavor, type);
+        if (requested != null) {
+            b.allocation(new Allocation(ApplicationId.defaultId(),
+                                        ClusterMembership.from("container/id1/4/37",
+                                                               Version.fromString("1.2.3"),
+                                                               Optional.of(requested)),
+                                        NodeResources.unspecified(),
+                                        Generation.initial(),
+                                        false));
         }
+        return b.build();
     }
 
 }
diff --git a/node-repository/src/test/java/com/yahoo/vespa/hosted/provision/provisioning/ProvisioningTester.java b/node-repository/src/test/java/com/yahoo/vespa/hosted/provision/provisioning/ProvisioningTester.java
index 4c9597d00bd..bf5f6f1738f 100644
--- a/node-repository/src/test/java/com/yahoo/vespa/hosted/provision/provisioning/ProvisioningTester.java
+++ b/node-repository/src/test/java/com/yahoo/vespa/hosted/provision/provisioning/ProvisioningTester.java
@@ -111,6 +111,7 @@ public class ProvisioningTester {
                                                  zone,
                                                  nameResolver,
                                                  containerImage,
+                                                 Optional.empty(),
                                                  flagSource,
                                                  new MemoryMetricsDb(clock),
                                                  true,
diff --git a/node-repository/src/test/java/com/yahoo/vespa/hosted/provision/restapi/NodesV2ApiTest.java b/node-repository/src/test/java/com/yahoo/vespa/hosted/provision/restapi/NodesV2ApiTest.java
index 1ab0efb42bd..30f4705812d 100644
--- a/node-repository/src/test/java/com/yahoo/vespa/hosted/provision/restapi/NodesV2ApiTest.java
+++ b/node-repository/src/test/java/com/yahoo/vespa/hosted/provision/restapi/NodesV2ApiTest.java
@@ -228,7 +228,6 @@ public class NodesV2ApiTest {
                                    Utf8.toBytes("{\"modelName\": null}"), Request.Method.PATCH),
                        "{\"message\":\"Updated dockerhost1.yahoo.com\"}");
         tester.assertPartialResponse(new Request("http://localhost:8080/nodes/v2/node/dockerhost1.yahoo.com"), "modelName", false);
-        tester.container().handleRequest((new Request("http://localhost:8080/nodes/v2/upgrade/tenant", Utf8.toBytes("{\"dockerImage\": \"ignored-registry.example.com/my/image\"}"), Request.Method.PATCH)));
 
         ((OrchestratorMock) tester.container().components().getComponent(OrchestratorMock.class.getName()))
                 .suspend(new HostName("host4.yahoo.com"));
@@ -679,7 +678,7 @@ public class NodesV2ApiTest {
                                           Utf8.toBytes("{}"),
                                           Request.Method.PATCH),
                               400,
-                              "{\"error-code\":\"BAD_REQUEST\",\"message\":\"At least one of 'version', 'osVersion' or 'dockerImage' must be set\"}");
+                              "{\"error-code\":\"BAD_REQUEST\",\"message\":\"At least one of 'version' or 'osVersion' must be set\"}");
 
         // Downgrade without force fails
         tester.assertResponse(new Request("http://localhost:8080/nodes/v2/upgrade/confighost",
@@ -761,26 +760,6 @@ public class NodesV2ApiTest {
                                           Request.Method.PATCH),
                               200,
                               "{\"message\":\"Set osVersion to null for nodes of type confighost\"}");
-
-        // Set container image for config and tenant
-        assertResponse(new Request("http://localhost:8080/nodes/v2/upgrade/tenant",
-                        Utf8.toBytes("{\"dockerImage\": \"my-repo.my-domain.example:1234/repo/tenant\"}"),
-                        Request.Method.PATCH),
-                "{\"message\":\"Set container image to my-repo.my-domain.example:1234/repo/tenant for nodes of type tenant\"}");
-        assertResponse(new Request("http://localhost:8080/nodes/v2/upgrade/config",
-                        Utf8.toBytes("{\"dockerImage\": \"my-repo.my-domain.example:1234/repo/image\"}"),
-                        Request.Method.PATCH),
-                "{\"message\":\"Set container image to my-repo.my-domain.example:1234/repo/image for nodes of type config\"}");
-
-        assertResponse(new Request("http://localhost:8080/nodes/v2/upgrade/"),
-                "{\"versions\":{\"config\":\"6.123.456\",\"confighost\":\"6.124.42\",\"controller\":\"6.123.456\"},\"osVersions\":{\"host\":\"7.5.2\"},\"dockerImages\":{\"tenant\":\"my-repo.my-domain.example:1234/repo/tenant\",\"config\":\"my-repo.my-domain.example:1234/repo/image\"}}");
-
-        // Cannot set container image for non docker node type
-        tester.assertResponse(new Request("http://localhost:8080/nodes/v2/upgrade/confighost",
-                                          Utf8.toBytes("{\"dockerImage\": \"my-repo.my-domain.example:1234/repo/image\"}"),
-                                          Request.Method.PATCH),
-                              400,
-                              "{\"error-code\":\"BAD_REQUEST\",\"message\":\"Setting container image for confighost nodes is unsupported\"}");
     }
 
     @Test
diff --git a/node-repository/src/test/java/com/yahoo/vespa/hosted/provision/restapi/responses/node4-after-changes.json b/node-repository/src/test/java/com/yahoo/vespa/hosted/provision/restapi/responses/node4-after-changes.json
index 6f25254b420..7770227f5f8 100644
--- a/node-repository/src/test/java/com/yahoo/vespa/hosted/provision/restapi/responses/node4-after-changes.json
+++ b/node-repository/src/test/java/com/yahoo/vespa/hosted/provision/restapi/responses/node4-after-changes.json
@@ -25,7 +25,7 @@
   },
   "restartGeneration": 1,
   "currentRestartGeneration": 1,
-  "wantedDockerImage": "docker-registry.domain.tld:8080/my/image:6.42.0",
+  "wantedDockerImage": "docker-registry.domain.tld:8080/dist/vespa:6.42.0",
   "wantedVespaVersion": "6.42.0",
   "requestedResources": { "vcpu":1.0, "memoryGb":4.0, "diskGb":100.0, "bandwidthGbps":1.0, "diskSpeed":"fast", "storageType":"any" },
   "orchestratorStatus": "ALLOWED_TO_BE_DOWN",