use nvidia gpu only for gpu worker pods

beam-cloud · Jan 22, 2024 · 014dddd · 014dddd
1 parent d3440f2
commit 014dddd
Show file tree

Hide file tree

Showing 3 changed files with 6 additions and 119 deletions.
diff --git a/docker/Dockerfile.k3d b/docker/Dockerfile.k3d
@@ -26,7 +26,6 @@ rm -rf /var/lib/apt/lists/* /var/log/*
 EOT
 
 COPY --from=k3s /bin /bin
-COPY docker/config.toml.tmpl /var/lib/rancher/k3s/agent/etc/containerd/config.toml.tmpl
 
 VOLUME /var/lib/kubelet
 VOLUME /var/lib/rancher/k3s

diff --git a/docker/config.toml.tmpl b/docker/config.toml.tmpl
diff --git a/internal/scheduler/pool_k8s.go b/internal/scheduler/pool_k8s.go
@@ -188,7 +188,9 @@ func (wpc *KubernetesWorkerPoolController) createWorkerJob(workerId string, cpu
 		workerMemory = wpc.config.Worker.DefaultWorkerMemoryRequest
 	}
 
+	runtimeClassName := ""
 	if gpuType != "" {
+		runtimeClassName = "nvidia"
 		resourceRequests[corev1.ResourceName("nvidia.com/gpu")] = *resource.NewQuantity(1, resource.DecimalSI)
 	}
 
@@ -243,6 +245,10 @@ func (wpc *KubernetesWorkerPoolController) createWorkerJob(workerId string, cpu
 		},
 	}
 
+	if runtimeClassName != "" {
+		podTemplate.Spec.RuntimeClassName = ptr.To(runtimeClassName)
+	}
+
 	ttl := int32(30)
 	job := &batchv1.Job{
 		ObjectMeta: metav1.ObjectMeta{