add configurable pod runtime

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

internal/common/config.default.yaml

@@ -64,6 +64,7 @@ worker:
         minFreeMemory: 1Gi
     nvidia:
       gpuType: A40
+      runtime: nvidia
       jobSpec:
         nodeSelector: {}
       poolSizing:

internal/scheduler/pool_k8s.go

@@ -188,7 +188,8 @@ func (wpc *KubernetesWorkerPoolController) createWorkerJob(workerId string, cpu
 		workerMemory = wpc.config.Worker.DefaultWorkerMemoryRequest
 	}
 
-	if gpuType != "" {
+	workerPool, _ := wpc.GetWorkerPoolConfig()
+	if gpuType != "" && workerPool.Runtime == "nvidia" {
 		resourceRequests[corev1.ResourceName("nvidia.com/gpu")] = *resource.NewQuantity(1, resource.DecimalSI)
 	}
 
@@ -243,6 +244,10 @@ func (wpc *KubernetesWorkerPoolController) createWorkerJob(workerId string, cpu
 		},
 	}
 
+	if workerPool.Runtime != "" {
+		podTemplate.Spec.RuntimeClassName = ptr.To(workerPool.Runtime)
+	}
+
 	ttl := int32(30)
 	job := &batchv1.Job{
 		ObjectMeta: metav1.ObjectMeta{

internal/types/config.go

@@ -125,6 +125,7 @@ type WorkerConfig struct {
 
 type WorkerPoolConfig struct {
 	GPUType    string                            `key:"gpuType"`
+	Runtime    string                            `key:"runtime"`
 	JobSpec    WorkerPoolJobSpecConfig           `key:"jobSpec"`
 	PoolSizing WorkerPoolJobSpecPoolSizingConfig `key:"poolSizing"`
 }