X-Git-Url: https://git.arvados.org/arvados.git/blobdiff_plain/cbf93e8d897448dbd52369afe89fef2392140ff1..09cbdc3074b3f1e69c9c537875146f6da0a6ed8f:/lib/dispatchcloud/node_size_test.go

diff --git a/lib/dispatchcloud/node_size_test.go b/lib/dispatchcloud/node_size_test.go
index eef86f7477..86bfbec7b6 100644
--- a/lib/dispatchcloud/node_size_test.go
+++ b/lib/dispatchcloud/node_size_test.go
@@ -5,7 +5,7 @@
 package dispatchcloud
 
 import (
-	"git.curoverse.com/arvados.git/sdk/go/arvados"
+	"git.arvados.org/arvados.git/sdk/go/arvados"
 	check "gopkg.in/check.v1"
 )
 
@@ -73,8 +73,17 @@ func (*NodeSizeSuite) TestChoose(c *check.C) {
 			"best":   {Price: 3.3, RAM: 4000000000, VCPUs: 4, Scratch: 2 * GiB, Name: "best"},
 			"costly": {Price: 4.4, RAM: 4000000000, VCPUs: 8, Scratch: 2 * GiB, Name: "costly"},
 		},
+		{
+			"small":  {Price: 1.1, RAM: 1000000000, VCPUs: 2, Scratch: GiB, Name: "small"},
+			"nearly": {Price: 2.2, RAM: 1200000000, VCPUs: 4, Scratch: 2 * GiB, Name: "nearly"},
+			"best":   {Price: 3.3, RAM: 4000000000, VCPUs: 4, Scratch: 2 * GiB, Name: "best"},
+			"costly": {Price: 4.4, RAM: 4000000000, VCPUs: 8, Scratch: 2 * GiB, Name: "costly"},
+		},
 	} {
-		best, err := ChooseInstanceType(&arvados.Cluster{InstanceTypes: menu}, &arvados.Container{
+		best, err := ChooseInstanceType(&arvados.Cluster{InstanceTypes: menu, Containers: arvados.ContainersConfig{
+			LocalKeepBlobBuffersPerVCPU: 1,
+			ReserveExtraRAM:             268435456,
+		}}, &arvados.Container{
 			Mounts: map[string]arvados.Mount{
 				"/tmp": {Kind: "tmp", Capacity: 2 * int64(GiB)},
 			},
@@ -92,7 +101,30 @@ func (*NodeSizeSuite) TestChoose(c *check.C) {
 	}
 }
 
-func (*NodeSizeSuite) TestChoosePreemptable(c *check.C) {
+func (*NodeSizeSuite) TestChooseWithBlobBuffersOverhead(c *check.C) {
+	menu := map[string]arvados.InstanceType{
+		"nearly": {Price: 2.2, RAM: 4000000000, VCPUs: 4, Scratch: 2 * GiB, Name: "small"},
+		"best":   {Price: 3.3, RAM: 8000000000, VCPUs: 4, Scratch: 2 * GiB, Name: "best"},
+		"costly": {Price: 4.4, RAM: 12000000000, VCPUs: 8, Scratch: 2 * GiB, Name: "costly"},
+	}
+	best, err := ChooseInstanceType(&arvados.Cluster{InstanceTypes: menu, Containers: arvados.ContainersConfig{
+		LocalKeepBlobBuffersPerVCPU: 16, // 1 GiB per vcpu => 2 GiB
+		ReserveExtraRAM:             268435456,
+	}}, &arvados.Container{
+		Mounts: map[string]arvados.Mount{
+			"/tmp": {Kind: "tmp", Capacity: 2 * int64(GiB)},
+		},
+		RuntimeConstraints: arvados.RuntimeConstraints{
+			VCPUs:        2,
+			RAM:          987654321,
+			KeepCacheRAM: 123456789,
+		},
+	})
+	c.Check(err, check.IsNil)
+	c.Check(best.Name, check.Equals, "best")
+}
+
+func (*NodeSizeSuite) TestChoosePreemptible(c *check.C) {
 	menu := map[string]arvados.InstanceType{
 		"costly":      {Price: 4.4, RAM: 4000000000, VCPUs: 8, Scratch: 2 * GiB, Preemptible: true, Name: "costly"},
 		"almost best": {Price: 2.2, RAM: 2000000000, VCPUs: 4, Scratch: 2 * GiB, Name: "almost best"},
@@ -141,3 +173,90 @@ 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"}},
+		"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"},
+	}
+
+	type GPUTestCase struct {
+		CUDA             arvados.CUDARuntimeConstraints
+		SelectedInstance string
+	}
+	cases := []GPUTestCase{
+		GPUTestCase{
+			CUDA: arvados.CUDARuntimeConstraints{
+				DeviceCount:        1,
+				HardwareCapability: "9.0",
+				DriverVersion:      "11.0",
+			},
+			SelectedInstance: "best",
+		},
+		GPUTestCase{
+			CUDA: arvados.CUDARuntimeConstraints{
+				DeviceCount:        2,
+				HardwareCapability: "9.0",
+				DriverVersion:      "11.0",
+			},
+			SelectedInstance: "costly",
+		},
+		GPUTestCase{
+			CUDA: arvados.CUDARuntimeConstraints{
+				DeviceCount:        1,
+				HardwareCapability: "8.0",
+				DriverVersion:      "11.0",
+			},
+			SelectedInstance: "low_capability",
+		},
+		GPUTestCase{
+			CUDA: arvados.CUDARuntimeConstraints{
+				DeviceCount:        1,
+				HardwareCapability: "9.0",
+				DriverVersion:      "10.0",
+			},
+			SelectedInstance: "low_driver",
+		},
+		GPUTestCase{
+			CUDA: arvados.CUDARuntimeConstraints{
+				DeviceCount:        1,
+				HardwareCapability: "",
+				DriverVersion:      "10.0",
+			},
+			SelectedInstance: "",
+		},
+		GPUTestCase{
+			CUDA: arvados.CUDARuntimeConstraints{
+				DeviceCount:        0,
+				HardwareCapability: "9.0",
+				DriverVersion:      "11.0",
+			},
+			SelectedInstance: "non_gpu",
+		},
+	}
+
+	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))
+		}
+	}
+}