diff --git a/test/e2e/framework/gpu_util.go b/test/e2e/framework/gpu_util.go index d0ff9798f19..2cd9f33ee77 100644 --- a/test/e2e/framework/gpu_util.go +++ b/test/e2e/framework/gpu_util.go @@ -32,7 +32,7 @@ const ( // TODO: Parametrize it by making it a feature in TestFramework. // so we can override the daemonset in other setups (non COS). // GPUDevicePluginDSYAML is the official Google Device Plugin Daemonset NVIDIA GPU manifest for GKE - GPUDevicePluginDSYAML = "https://raw.githubusercontent.com/GoogleCloudPlatform/container-engine-accelerators/master/device-plugin-daemonset.yaml" + GPUDevicePluginDSYAML = "https://raw.githubusercontent.com/kubernetes/kubernetes/master/cluster/addons/device-plugins/nvidia-gpu/daemonset.yaml" ) // TODO make this generic and not linked to COS only @@ -61,8 +61,8 @@ func NVIDIADevicePlugin(ns string) *v1.Pod { Spec: ds.Spec.Template.Spec, } - // Remove NVIDIA drivers installation - p.Spec.InitContainers = []v1.Container{} + // Remove node affinity + p.Spec.Affinity = nil return p } diff --git a/test/e2e/scheduling/nvidia-gpus.go b/test/e2e/scheduling/nvidia-gpus.go index 5deeea25f7e..44055575f96 100644 --- a/test/e2e/scheduling/nvidia-gpus.go +++ b/test/e2e/scheduling/nvidia-gpus.go @@ -163,7 +163,7 @@ func testNvidiaGPUsOnCOS(f *framework.Framework) { framework.Logf("Cluster is running on COS. Proceeding with test") if f.BaseName == "device-plugin-gpus" { - dsYamlUrl = framework.GPUDevicePluginDSYAML + dsYamlUrl = "https://raw.githubusercontent.com/GoogleCloudPlatform/container-engine-accelerators/master/daemonset.yaml" gpuResourceName = framework.NVIDIAGPUResourceName podCreationFunc = makeCudaAdditionDevicePluginTestPod } else {