Draft: GPU Mutating Webhook

cmd/gpu-mutating-webhook/admission_controller.go

@@ -0,0 +1,145 @@
+/**
+# Copyright 2024 NVIDIA CORPORATION
+#
+# Licensed under the Apache License, Version 2.0 (the "License");
+# you may not use this file except in compliance with the License.
+# You may obtain a copy of the License at
+#
+#     http://www.apache.org/licenses/LICENSE-2.0
+#
+# Unless required by applicable law or agreed to in writing, software
+# distributed under the License is distributed on an "AS IS" BASIS,
+# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+# See the License for the specific language governing permissions and
+# limitations under the License.
+**/
+
+package main
+
+import (
+	"encoding/json"
+	"errors"
+	"fmt"
+	"io/ioutil"
+	"log"
+	"net/http"
+
+	admissionv1 "k8s.io/api/admission/v1"
+	metav1 "k8s.io/apimachinery/pkg/apis/meta/v1"
+	"k8s.io/apimachinery/pkg/runtime"
+	"k8s.io/apimachinery/pkg/runtime/serializer"
+)
+
+const (
+	jsonContentType = `application/json`
+)
+
+var (
+	universalDeserializer = serializer.NewCodecFactory(runtime.NewScheme()).UniversalDeserializer()
+)
+
+type patchOperation struct {
+	Op    string      `json:"op"`
+	Path  string      `json:"path"`
+	Value interface{} `json:"value,omitempty"`
+}
+
+type admitFunc func(*admissionv1.AdmissionRequest) ([]patchOperation, error)
+
+// Swati: skip nvidia-dra-driver-gpu ns as well
+func isKubeNamespace(ns string) bool {
+	return (ns == metav1.NamespacePublic || ns == metav1.NamespaceSystem)
+}
+
+func doServeAdmitFunc(w http.ResponseWriter, r *http.Request, admit admitFunc) ([]byte, error) {
+	// Request validation. Only handle POST requests with a body and json content type.
+	if r.Method != http.MethodPost {
+		w.WriteHeader(http.StatusMethodNotAllowed)
+		return nil, fmt.Errorf("invalid method %s, only POST is allowed", r.Method)
+	}
+
+	body, err := ioutil.ReadAll(r.Body)
+	if err != nil {
+		w.WriteHeader(http.StatusBadRequest)
+		return nil, fmt.Errorf("could not read request body: %v", err)
+	}
+
+	if ct := r.Header.Get("Content-Type"); ct != jsonContentType {
+		w.WriteHeader(http.StatusBadRequest)
+		return nil, fmt.Errorf("unsupported content type %s, only %s is supported", ct, jsonContentType)
+	}
+
+	// Parse the AdmissionReview request.
+	var admissionReviewReq admissionv1.AdmissionReview
+	if _, _, err := universalDeserializer.Decode(body, nil, &admissionReviewReq); err != nil {
+		w.WriteHeader(http.StatusBadRequest)
+		return nil, fmt.Errorf("could not deserialize AdmissionReview: %v", err)
+	} else if admissionReviewReq.Request == nil {
+		w.WriteHeader(http.StatusBadRequest)
+		return nil, errors.New("malformed admission review: Request is nil")
+	}
+
+	// Build the response
+	admissionReviewResp := admissionv1.AdmissionReview{
+		TypeMeta: admissionReviewReq.TypeMeta,
+		Response: &admissionv1.AdmissionResponse{
+			UID: admissionReviewReq.Request.UID,
+		},
+	}
+
+	// Skip k8s namespaces
+	var patchOps []patchOperation
+	if !isKubeNamespace(admissionReviewReq.Request.Namespace) {
+		patchOps, err = admit(admissionReviewReq.Request)
+	}
+
+	if err != nil {
+		admissionReviewResp.Response.Allowed = false
+		admissionReviewResp.Response.Result = &metav1.Status{
+			Message: err.Error(),
+		}
+	} else {
+		patchBytes, err := json.Marshal(patchOps)
+		if err != nil {
+			w.WriteHeader(http.StatusInternalServerError)
+			return nil, fmt.Errorf("could not marshal JSON patch: %v", err)
+		}
+		admissionReviewResp.Response.Allowed = true
+		admissionReviewResp.Response.Patch = patchBytes
+
+		pt := admissionv1.PatchTypeJSONPatch
+		admissionReviewResp.Response.PatchType = &pt
+	}
+
+	respBytes, err := json.Marshal(admissionReviewResp)
+	if err != nil {
+		return nil, fmt.Errorf("could not marshal AdmissionReview response: %v", err)
+	}
+	return respBytes, nil
+}
+
+// serveAdmitFunc is a wrapper that handles HTTP, calls doServeAdmitFunc, and writes the result.
+func serveAdmitFunc(w http.ResponseWriter, r *http.Request, admit admitFunc) {
+	log.Print("Handling webhook request ...")
+
+	respBytes, err := doServeAdmitFunc(w, r, admit)
+	if err != nil {
+		log.Printf("Error handling webhook request: %v", err)
+		w.WriteHeader(http.StatusInternalServerError)
+		_, _ = w.Write([]byte(err.Error()))
+		return
+	}
+
+	log.Print("Webhook request handled successfully")
+	_, writeErr := w.Write(respBytes)
+	if writeErr != nil {
+		log.Printf("Could not write response: %v", writeErr)
+	}
+}
+
+// admitFuncHandler converts an admitFunc into an http.Handler
+func admitFuncHandler(admit admitFunc) http.Handler {
+	return http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) {
+		serveAdmitFunc(w, r, admit)
+	})
+}

cmd/gpu-mutating-webhook/main.go

@@ -0,0 +1,195 @@
+/**
+# Copyright 2025 NVIDIA CORPORATION
+#
+# Licensed under the Apache License, Version 2.0 (the "License");
+# you may not use this file except in compliance with the License.
+# You may obtain a copy of the License at
+#
+#     http://www.apache.org/licenses/LICENSE-2.0
+#
+# Unless required by applicable law or agreed to in writing, software
+# distributed under the License is distributed on an "AS IS" BASIS,
+# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+# See the License for the specific language governing permissions and
+# limitations under the License.
+**/
+
+package main
+
+import (
+	"fmt"
+	"log"
+	"net/http"
+	"path/filepath"
+	"strings"
+
+	admissionv1 "k8s.io/api/admission/v1"
+	corev1 "k8s.io/api/core/v1"
+	metav1 "k8s.io/apimachinery/pkg/apis/meta/v1"
+	"k8s.io/klog/v2"
+)
+
+const (
+	tlsDir          = `/etc/webhook/tls`
+	tlsCertFile     = `tls.crt`
+	tlsKeyFile      = `tls.key`
+	gpuResourceName = "nvidia.com/gpu"
+	gpuClaimName    = "nvidia-gpu-resourceclaim"
+	gpuTemplateName = "nvidia-gpu-resourceclaim-template"
+)
+
+var (
+	podResource = metav1.GroupVersionResource{Version: "v1", Resource: "pods"}
+)
+
+func applyGPUMutation(req *admissionv1.AdmissionRequest) ([]patchOperation, error) {
+	// Only mutate Pod CREATE
+	// Swati: may be add UPDATE
+	if req.Resource != podResource || req.Operation != admissionv1.Create {
+		klog.Infof("skip mutation for %v/%v", req.Resource, req.Operation)
+		return nil, nil
+	}
+
+	var pod corev1.Pod
+	if _, _, err := universalDeserializer.Decode(req.Object.Raw, nil, &pod); err != nil {
+		klog.Errorf("failed to decode Pod: %v", err)
+		return nil, fmt.Errorf("could not deserialize pod: %w", err)
+	}
+
+	key := escapeJSONPointer(gpuResourceName)
+	var patches []patchOperation
+	var ctrGPUResourceClaims []string
+
+	// Iterate on all containers and check for "nvidia.com/gpu" limits
+	// using the logic described here for prefering limits over requests
+	// GPUs are only supposed to be specified in the limits section, meaning
+	// - can specify GPU limits without specifying requests. limit will be used as request value by default
+	// - can specify GPU in both limits and requests but they must be equal
+	// - cannot specify GPU requests without specifying limits
+	// refer: https://kubernetes.io/docs/tasks/manage-gpus/scheduling-gpus/#using-device-plugins
+	for ci, ctr := range pod.Spec.Containers {
+		ctrName := ctr.Name
+		limitCount, limitOk := ctr.Resources.Limits[gpuResourceName]
+
+		// skip if no GPUs in limits
+		if !limitOk || limitCount.Value() < 1 {
+			continue
+		}
+		gpuCount := limitCount.Value()
+
+		// check any GPUs in requests
+		// it must be equal to limits
+		if reqCount, reqOK := ctr.Resources.Requests[gpuResourceName]; reqOK {
+			if reqCount.Value() != gpuCount {
+				klog.Warningf("container[%q]: gpu request (%d) != limit (%d), skipping mutation", ctrName, reqCount.Value(), gpuCount)
+				continue
+			}
+			reqPatch := removeResourceRequest(ci, "requests", key)
+			patches = append(patches, reqPatch)
+			klog.Infof("removed container[%q].Resources.Requests: %v", ctrName, reqPatch)
+		}
+		limitPatch := removeResourceRequest(ci, "limits", key)
+		patches = append(patches, limitPatch)
+		klog.Infof("removed container[%q].Resources.Limits: %v", ctrName, limitPatch)
+
+		// ensure container-claims slice exists
+		// this is JSON way to first creating the field if it does not exist and append later with "-"
+		if len(ctr.Resources.Claims) == 0 {
+			createPatch := createClaimPatch(fmt.Sprintf("/spec/containers/%d/resources/claims", ci))
+			patches = append(patches, createPatch)
+			klog.Infof("created container[%q] empty claims array: %v", ctrName, createPatch)
+		}
+
+		// append one claim per GPU
+		for i := int64(0); i < gpuCount; i++ {
+			claimName := fmt.Sprintf("%s-%d", gpuClaimName, i)
+			ctrGPUResourceClaims = append(ctrGPUResourceClaims, claimName)
+			appendPatch := appendClaimPatch(
+				fmt.Sprintf("/spec/containers/%d/resources/claims", ci),
+				map[string]string{"name": claimName},
+			)
+			patches = append(patches, appendPatch)
+			klog.Infof("added to container[%q].Resources.Claims: %v", ctrName, appendPatch)
+		}
+	}
+
+	// Add claims pod-level
+	podName := pod.Name
+	if len(ctrGPUResourceClaims) > 0 {
+		// ensure pod-claims slice exists
+		if len(pod.Spec.ResourceClaims) == 0 {
+			createPatch := createClaimPatch("/spec/resourceClaims")
+			patches = append(patches, createPatch)
+			klog.Infof("created pod[%q] empty claims array: %v", podName, createPatch)
+		}
+
+		// append each container GPU claim at pod-level
+		for _, name := range ctrGPUResourceClaims {
+			appendPatch := appendClaimPatch(
+				"/spec/resourceClaims",
+				map[string]string{
+					"name":                      name,
+					"resourceClaimTemplateName": gpuTemplateName,
+				},
+			)
+			patches = append(patches, appendPatch)
+			klog.Infof("added ResourceClaim %q (template=%q) to %q: %v", name, gpuTemplateName, podName, appendPatch)
+		}
+	}
+
+	return patches, nil
+}
+
+// escapeJSONPointer replace "/" with "~1"
+// refer: https://github.com/json-patch/json-patch-tests/issues/42
+// needed for "nvidia.com/gpu". otherwise JSON will treat "/" as a path delimiter and treat "gpu" as new field
+func escapeJSONPointer(s string) string {
+	return strings.ReplaceAll(s, "/", "~1")
+}
+
+// removeResourceRequest removes either .resources.requests or .resources.limits
+func removeResourceRequest(ci int, field, key string) patchOperation {
+	return patchOperation{
+		Op:   "remove",
+		Path: fmt.Sprintf("/spec/containers/%d/resources/%s/%s", ci, field, key),
+	}
+}
+
+// createClaimPatch creates an empty slice at the given path
+func createClaimPatch(path string) patchOperation {
+	return patchOperation{
+		Op:    "add",
+		Path:  path,
+		Value: []map[string]string{},
+	}
+}
+
+// appendClaimPatch appends to the slice at path
+// "-" is JSON way to inserting at the end of the array when no index is specified.
+// refer: https://datatracker.ietf.org/doc/html/rfc6902
+func appendClaimPatch(path string, entry map[string]string) patchOperation {
+	return patchOperation{
+		Op:    "add",
+		Path:  path + "/-",
+		Value: entry,
+	}
+}
+
+func main() {
+	certPath := filepath.Join(tlsDir, tlsCertFile)
+	keyPath := filepath.Join(tlsDir, tlsKeyFile)
+
+	mux := http.NewServeMux()
+	mux.Handle("/mutate", admitFuncHandler(applyGPUMutation))
+
+	server := &http.Server{
+		Addr:    ":8443",
+		Handler: mux,
+	}
+
+	if err := server.ListenAndServeTLS(certPath, keyPath); err != nil {
+		// Swati: need better error handling here
+		log.Fatalf("Failed to start server: %v", err)
+	}
+	klog.Infof("Started gpu-mutating-webhook server at %s", server.Addr)
+}

deployments/container/Dockerfile

@@ -67,4 +67,5 @@ RUN mkdir /licenses && mv /NGC-DL-CONTAINER-LICENSE /licenses/NGC-DL-CONTAINER-L
 COPY --from=build /artifacts/compute-domain-controller     /usr/bin/compute-domain-controller
 COPY --from=build /artifacts/compute-domain-kubelet-plugin /usr/bin/compute-domain-kubelet-plugin
 COPY --from=build /artifacts/gpu-kubelet-plugin            /usr/bin/gpu-kubelet-plugin
+COPY --from=build /artifacts/gpu-mutating-webhook          /usr/bin/gpu-mutating-webhook
 COPY --from=build /build/templates                         /templates