Return the image for child node type for docker hosts

1 files changed, 8 insertions, 2 deletions

diff --git a/node-repository/src/main/java/com/yahoo/vespa/hosted/provision/restapi/v2/NodesResponse.java b/node-repository/src/main/java/com/yahoo/vespa/hosted/provision/restapi/v2/NodesResponse.java
index be99bb946d5..2ab916e6375 100644
--- a/node-repository/src/main/java/com/yahoo/vespa/hosted/provision/restapi/v2/NodesResponse.java
+++ b/node-repository/src/main/java/com/yahoo/vespa/hosted/provision/restapi/v2/NodesResponse.java
@@ -161,7 +161,7 @@ class NodesResponse extends HttpResponse {
             toSlime(allocation.membership(), object.setObject("membership"));
             object.setLong("restartGeneration", allocation.restartGeneration().wanted());
             object.setLong("currentRestartGeneration", allocation.restartGeneration().current());
-            object.setString("wantedDockerImage", nodeRepository.dockerImage(node.type()).withTag(allocation.membership().cluster().vespaVersion()).asString());
+            object.setString("wantedDockerImage", dockerImageFor(node.type()).withTag(allocation.membership().cluster().vespaVersion()).asString());
             object.setString("wantedVespaVersion", allocation.membership().cluster().vespaVersion().toFullString());
             allocation.networkPorts().ifPresent(ports -> NetworkPortsSerializer.toSlime(ports, object.setArray("networkPorts")));
             orchestrator.apply(new HostName(node.hostname()))
@@ -220,7 +220,13 @@ class NodesResponse extends HttpResponse {
                 .or(() -> Optional.of(node)
                         .filter(n -> n.flavor().getType() != Flavor.Type.DOCKER_CONTAINER)
                         .flatMap(n -> n.status().vespaVersion()
-                                .map(version -> nodeRepository.dockerImages().dockerImageFor(n.type()).withTag(version))));
+                                .map(version -> dockerImageFor(n.type()).withTag(version))));
+    }
+
+    // Docker hosts are not running in an image, but return the image of the node type running on it anyway,
+    // this allows the docker host to pre-download the (likely) image its node will run
+    private DockerImage dockerImageFor(NodeType nodeType) {
+        return nodeRepository.dockerImage(nodeType.isDockerHost() ? nodeType.childNodeType() : nodeType);
     }