X-Git-Url: https://git.arvados.org/arvados.git/blobdiff_plain/b05ec24843655e162c8c3207e1695debdca9725e..565612fd40474044e2afaa4fcb993c8c0197ca8e:/lib/dispatchcloud/node_size_test.go

diff --git a/lib/dispatchcloud/node_size_test.go b/lib/dispatchcloud/node_size_test.go
index cdcf4033fc..eb3648e8ac 100644
--- a/lib/dispatchcloud/node_size_test.go
+++ b/lib/dispatchcloud/node_size_test.go
@@ -154,69 +154,83 @@ func (*NodeSizeSuite) TestChooseGPU(c *check.C) {
 		"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"},
+		"cheap_gpu":      {Price: 2.0, RAM: 2000000000, VCPUs: 4, Scratch: 2 * GiB, Name: "cheap_gpu", CUDA: arvados.CUDAFeatures{DeviceCount: 1, HardwareCapability: "8.0", DriverVersion: "10.0"}},
+		"invalid_gpu":    {Price: 1.9, RAM: 2000000000, VCPUs: 4, Scratch: 2 * GiB, Name: "invalid_gpu", CUDA: arvados.CUDAFeatures{DeviceCount: 1, HardwareCapability: "12.0.12", DriverVersion: "12.0.12"}},
+		"non_gpu":        {Price: 1.1, RAM: 2000000000, VCPUs: 4, Scratch: 2 * GiB, Name: "non_gpu"},
 	}
-	best, err := ChooseInstanceType(&arvados.Cluster{InstanceTypes: menu}, &arvados.Container{
-		Mounts: map[string]arvados.Mount{
-			"/tmp": {Kind: "tmp", Capacity: 2 * int64(GiB)},
+
+	type GPUTestCase struct {
+		CUDA             arvados.CUDARuntimeConstraints
+		SelectedInstance string
+	}
+	cases := []GPUTestCase{
+		GPUTestCase{
+			CUDA: arvados.CUDARuntimeConstraints{
+				DeviceCount:        1,
+				HardwareCapability: "9.0",
+				DriverVersion:      "11.0",
+			},
+			SelectedInstance: "best",
 		},
-		RuntimeConstraints: arvados.RuntimeConstraints{
-			VCPUs:                  2,
-			RAM:                    987654321,
-			KeepCacheRAM:           123456789,
-			CUDADeviceCount:        1,
-			CUDAHardwareCapability: "9.0",
-			CUDADriverVersion:      "11.0",
+		GPUTestCase{
+			CUDA: arvados.CUDARuntimeConstraints{
+				DeviceCount:        2,
+				HardwareCapability: "9.0",
+				DriverVersion:      "11.0",
+			},
+			SelectedInstance: "costly",
 		},
-	})
-	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)},
+		GPUTestCase{
+			CUDA: arvados.CUDARuntimeConstraints{
+				DeviceCount:        1,
+				HardwareCapability: "8.0",
+				DriverVersion:      "11.0",
+			},
+			SelectedInstance: "low_capability",
 		},
-		RuntimeConstraints: arvados.RuntimeConstraints{
-			VCPUs:                  2,
-			RAM:                    987654321,
-			KeepCacheRAM:           123456789,
-			CUDADeviceCount:        2,
-			CUDAHardwareCapability: "9.0",
-			CUDADriverVersion:      "11.0",
+		GPUTestCase{
+			CUDA: arvados.CUDARuntimeConstraints{
+				DeviceCount:        1,
+				HardwareCapability: "9.0",
+				DriverVersion:      "10.0",
+			},
+			SelectedInstance: "low_driver",
 		},
-	})
-	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)},
+		GPUTestCase{
+			CUDA: arvados.CUDARuntimeConstraints{
+				DeviceCount:        1,
+				HardwareCapability: "",
+				DriverVersion:      "10.0",
+			},
+			SelectedInstance: "",
 		},
-		RuntimeConstraints: arvados.RuntimeConstraints{
-			VCPUs:                  2,
-			RAM:                    987654321,
-			KeepCacheRAM:           123456789,
-			CUDADeviceCount:        1,
-			CUDAHardwareCapability: "8.0",
-			CUDADriverVersion:      "11.0",
+		GPUTestCase{
+			CUDA: arvados.CUDARuntimeConstraints{
+				DeviceCount:        0,
+				HardwareCapability: "9.0",
+				DriverVersion:      "11.0",
+			},
+			SelectedInstance: "non_gpu",
 		},
-	})
-	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")
+	}
 
+	for _, tc := range cases {
+		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,
+				CUDA:         tc.CUDA,
+			},
+		})
+		if best.Name != "" {
+			c.Check(err, check.IsNil)
+			c.Check(best.Name, check.Equals, tc.SelectedInstance)
+		} else {
+			c.Check(err, check.Not(check.IsNil))
+		}
+	}
 }