18321: Incorporate CUDA request into picking a node type

[arvados.git] / lib / dispatchcloud / node_size_test.go

diff --git a/lib/dispatchcloud/node_size_test.go b/lib/dispatchcloud/node_size_test.go
index abd292cbaf1278c7ca147e7da7352d74d0ad39d0..cdcf4033fcc629317ef0f6edddf799778bf50e5b 100644 (file)
--- a/lib/dispatchcloud/node_size_test.go
+++ b/lib/dispatchcloud/node_size_test.go
@@ -147,3 +147,76 @@ func (*NodeSizeSuite) TestScratchForDockerImage(c *check.C) {
        // Short manifest will return 0
        c.Check(n, check.Equals, int64(0))
 }
+
+func (*NodeSizeSuite) TestChooseGPU(c *check.C) {
+       menu := map[string]arvados.InstanceType{
+               "costly":         {Price: 4.4, RAM: 4000000000, VCPUs: 8, Scratch: 2 * GiB, Name: "costly", CUDA: arvados.CUDAFeatures{DeviceCount: 2, HardwareCapability: "9.0", DriverVersion: "11.0"}},
+               "low_capability": {Price: 2.1, RAM: 2000000000, VCPUs: 4, Scratch: 2 * GiB, Name: "low_capability", CUDA: arvados.CUDAFeatures{DeviceCount: 1, HardwareCapability: "8.0", DriverVersion: "11.0"}},
+               "best":           {Price: 2.2, RAM: 2000000000, VCPUs: 4, Scratch: 2 * GiB, Name: "best", CUDA: arvados.CUDAFeatures{DeviceCount: 1, HardwareCapability: "9.0", DriverVersion: "11.0"}},
+               "low_driver":     {Price: 2.1, RAM: 2000000000, VCPUs: 4, Scratch: 2 * GiB, Name: "low_driver", CUDA: arvados.CUDAFeatures{DeviceCount: 1, HardwareCapability: "9.0", DriverVersion: "10.0"}},
+               "small":          {Price: 1.1, RAM: 1000000000, VCPUs: 2, Scratch: 2 * GiB, Name: "small"},
+       }
+       best, err := ChooseInstanceType(&arvados.Cluster{InstanceTypes: menu}, &arvados.Container{
+               Mounts: map[string]arvados.Mount{
+                       "/tmp": {Kind: "tmp", Capacity: 2 * int64(GiB)},
+               },
+               RuntimeConstraints: arvados.RuntimeConstraints{
+                       VCPUs:                  2,
+                       RAM:                    987654321,
+                       KeepCacheRAM:           123456789,
+                       CUDADeviceCount:        1,
+                       CUDAHardwareCapability: "9.0",
+                       CUDADriverVersion:      "11.0",
+               },
+       })
+       c.Check(err, check.IsNil)
+       c.Check(best.Name, check.Equals, "best")
+       c.Check(best.RAM >= 1234567890, check.Equals, true)
+       c.Check(best.VCPUs >= 2, check.Equals, true)
+       c.Check(best.CUDA.DeviceCount >= 1, check.Equals, true)
+       c.Check(best.CUDA.DriverVersion, check.Equals, "11.0")
+       c.Check(best.CUDA.HardwareCapability, check.Equals, "9.0")
+
+       best, err = ChooseInstanceType(&arvados.Cluster{InstanceTypes: menu}, &arvados.Container{
+               Mounts: map[string]arvados.Mount{
+                       "/tmp": {Kind: "tmp", Capacity: 2 * int64(GiB)},
+               },
+               RuntimeConstraints: arvados.RuntimeConstraints{
+                       VCPUs:                  2,
+                       RAM:                    987654321,
+                       KeepCacheRAM:           123456789,
+                       CUDADeviceCount:        2,
+                       CUDAHardwareCapability: "9.0",
+                       CUDADriverVersion:      "11.0",
+               },
+       })
+       c.Check(err, check.IsNil)
+       c.Check(best.Name, check.Equals, "costly")
+       c.Check(best.RAM >= 1234567890, check.Equals, true)
+       c.Check(best.VCPUs >= 2, check.Equals, true)
+       c.Check(best.CUDA.DeviceCount >= 2, check.Equals, true)
+       c.Check(best.CUDA.DriverVersion, check.Equals, "11.0")
+       c.Check(best.CUDA.HardwareCapability, check.Equals, "9.0")
+
+       best, err = ChooseInstanceType(&arvados.Cluster{InstanceTypes: menu}, &arvados.Container{
+               Mounts: map[string]arvados.Mount{
+                       "/tmp": {Kind: "tmp", Capacity: 2 * int64(GiB)},
+               },
+               RuntimeConstraints: arvados.RuntimeConstraints{
+                       VCPUs:                  2,
+                       RAM:                    987654321,
+                       KeepCacheRAM:           123456789,
+                       CUDADeviceCount:        1,
+                       CUDAHardwareCapability: "8.0",
+                       CUDADriverVersion:      "11.0",
+               },
+       })
+       c.Check(err, check.IsNil)
+       c.Check(best.Name, check.Equals, "low_capability")
+       c.Check(best.RAM >= 1234567890, check.Equals, true)
+       c.Check(best.VCPUs >= 2, check.Equals, true)
+       c.Check(best.CUDA.DeviceCount >= 1, check.Equals, true)
+       c.Check(best.CUDA.DriverVersion, check.Equals, "11.0")
+       c.Check(best.CUDA.HardwareCapability, check.Equals, "8.0")
+
+}