diff --git a/test/e2e/scheduling/nvidia-gpus.go b/test/e2e/scheduling/nvidia-gpus.go index b0b4c62e7c9..fc337a67eff 100644 --- a/test/e2e/scheduling/nvidia-gpus.go +++ b/test/e2e/scheduling/nvidia-gpus.go @@ -18,7 +18,6 @@ package scheduling import ( "os" - "strings" "time" "k8s.io/api/core/v1" @@ -35,7 +34,6 @@ import ( const ( testPodNamePrefix = "nvidia-gpu-" - cosOSImage = "Container-Optimized OS from Google" // Nvidia driver installation can take upwards of 5 minutes. driverInstallTimeout = 10 * time.Minute ) @@ -69,15 +67,12 @@ func makeCudaAdditionDevicePluginTestPod() *v1.Pod { return testPod } -func isClusterRunningCOS(f *framework.Framework) bool { +func logOSImages(f *framework.Framework) { nodeList, err := f.ClientSet.CoreV1().Nodes().List(metav1.ListOptions{}) framework.ExpectNoError(err, "getting node list") for _, node := range nodeList.Items { - if !strings.Contains(node.Status.NodeInfo.OSImage, cosOSImage) { - return false - } + framework.Logf("OS Image: %v", node.Status.NodeInfo.OSImage) } - return true } func areGPUsAvailableOnAllSchedulableNodes(f *framework.Framework) bool { @@ -111,14 +106,7 @@ func getGPUsAvailable(f *framework.Framework) int64 { } func SetupNVIDIAGPUNode(f *framework.Framework, setupResourceGatherer bool) *framework.ContainerResourceGatherer { - // Skip the test if the base image is not COS. - // TODO: Add support for other base images. - // CUDA apps require host mounts which is not portable across base images (yet). - framework.Logf("Checking base image") - if !isClusterRunningCOS(f) { - Skip("Nvidia GPU tests are supproted only on Container Optimized OS image currently") - } - framework.Logf("Cluster is running on COS. Proceeding with test") + logOSImages(f) dsYamlUrlFromEnv := os.Getenv("NVIDIA_DRIVER_INSTALLER_DAEMONSET") if dsYamlUrlFromEnv != "" { @@ -163,7 +151,7 @@ func SetupNVIDIAGPUNode(f *framework.Framework, setupResourceGatherer bool) *fra return rsgather } -func testNvidiaGPUsOnCOS(f *framework.Framework) { +func testNvidiaGPUs(f *framework.Framework) { rsgather := SetupNVIDIAGPUNode(f, true) framework.Logf("Creating as many pods as there are Nvidia GPUs and have the pods run a CUDA app") podList := []*v1.Pod{} @@ -186,7 +174,7 @@ func testNvidiaGPUsOnCOS(f *framework.Framework) { var _ = SIGDescribe("[Feature:GPUDevicePlugin]", func() { f := framework.NewDefaultFramework("device-plugin-gpus") - It("run Nvidia GPU Device Plugin tests on Container Optimized OS only", func() { - testNvidiaGPUsOnCOS(f) + It("run Nvidia GPU Device Plugin tests", func() { + testNvidiaGPUs(f) }) })