diff options
Diffstat (limited to 'node-repository/src/test/java/com/yahoo/vespa/hosted/provision/provisioning/ContainerImagesTest.java')
-rw-r--r-- | node-repository/src/test/java/com/yahoo/vespa/hosted/provision/provisioning/ContainerImagesTest.java | 38 |
1 files changed, 22 insertions, 16 deletions
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 217ead40b81..bb7ea52ca0e 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 @@ -5,14 +5,11 @@ 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.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; @@ -29,7 +26,8 @@ public class ContainerImagesTest { public void image_selection() { 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)); + DockerImage gpuImage = DockerImage.fromString("registry.example.com/vespa/tenant-gpu"); + ContainerImages images = new ContainerImages(defaultImage, Optional.of(tenantImage), Optional.of(gpuImage)); assertEquals(defaultImage, images.get(node(NodeType.confighost))); // For preload purposes assertEquals(defaultImage, images.get(node(NodeType.config))); @@ -40,32 +38,40 @@ public class ContainerImagesTest { assertEquals(defaultImage, images.get(node(NodeType.proxyhost))); // For preload purposes assertEquals(defaultImage, images.get(node(NodeType.proxy))); + // Choose GPU when node has GPU resources + assertEquals(gpuImage, images.get(node(NodeType.tenant, null, true))); + // 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()); + images = new ContainerImages(defaultImage, Optional.empty(), Optional.of(gpuImage)); assertEquals(defaultImage, images.get(node(NodeType.host))); assertEquals(defaultImage, images.get(node(NodeType.tenant))); } private static Node node(NodeType type) { - return node(type, null); + return node(type, null, false); } 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 node(type, requested, false); + } + + private static Node node(NodeType type, DockerImage requested, boolean gpu) { + NodeResources resources = new NodeResources(4, 8, 100, 0.3); + if (gpu) { + resources = resources.with(new NodeResources.GpuResources(1, 16)); } + Node.Builder b = Node.reserve(Set.of("::1"), type + "1", "parent1", resources, type); + b.allocation(new Allocation(ApplicationId.defaultId(), + ClusterMembership.from("container/id1/4/37", + Version.fromString("1.2.3"), + Optional.ofNullable(requested)), + resources, + Generation.initial(), + false)); return b.build(); } |