Add fake compute domain controller

Author: enoodle

Dockerfile

@@ -42,6 +42,11 @@ COPY ./cmd/mig-faker/ ./cmd/mig-faker/
 COPY ./internal/ ./internal/
 RUN --mount=type=cache,target=/root/.cache/go-build make build OS=$TARGETOS ARCH=$TARGETARCH COMPONENTS=mig-faker
 
+FROM common-builder AS compute-domain-controller-builder
+COPY ./cmd/compute-domain-controller/ ./cmd/compute-domain-controller/
+COPY ./pkg/compute-domain/ ./pkg/compute-domain/
+RUN --mount=type=cache,target=/root/.cache/go-build make build OS=$TARGETOS ARCH=$TARGETARCH COMPONENTS=compute-domain-controller
+
 FROM common-builder AS dra-plugin-gpu-builder
 COPY ./cmd/dra-plugin-gpu/ ./cmd/dra-plugin-gpu/
 COPY ./internal/dra-plugin-gpu/ ./internal/dra-plugin-gpu/
@@ -101,3 +106,7 @@ ENTRYPOINT ["/bin/dra-plugin-gpu"]
 FROM ubuntu AS kwok-dra-plugin
 COPY --from=kwok-dra-plugin-builder /go/src/github.com/run-ai/fake-gpu-operator/bin/kwok-dra-plugin /bin/
 ENTRYPOINT ["/bin/kwok-dra-plugin"]
+
+FROM ubuntu AS compute-domain-controller
+COPY --from=compute-domain-controller-builder /go/src/github.com/run-ai/fake-gpu-operator/bin/compute-domain-controller /bin/
+ENTRYPOINT ["/bin/compute-domain-controller"]

Makefile

@@ -1,5 +1,5 @@
 BUILD_DIR=$(shell pwd)/bin
-COMPONENTS?=device-plugin dra-plugin-gpu status-updater kwok-gpu-device-plugin kwok-dra-plugin status-exporter status-exporter-kwok topology-server mig-faker
+COMPONENTS?=device-plugin dra-plugin-gpu status-updater kwok-gpu-device-plugin kwok-dra-plugin status-exporter status-exporter-kwok topology-server mig-faker compute-domain-controller
 
 DOCKER_REPO_BASE=ghcr.io/run-ai/fake-gpu-operator
 DOCKER_TAG?=0.0.0-dev

cmd/compute-domain-controller/computedomain_controller.go

@@ -0,0 +1,207 @@
+/*
+ * Copyright 2025 The Kubernetes Authors.
+ *
+ * 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 (
+	"context"
+	"fmt"
+
+	resourceapi "k8s.io/api/resource/v1"
+	apierrors "k8s.io/apimachinery/pkg/api/errors"
+	metav1 "k8s.io/apimachinery/pkg/apis/meta/v1"
+	"k8s.io/apimachinery/pkg/runtime"
+	ctrl "sigs.k8s.io/controller-runtime"
+	"sigs.k8s.io/controller-runtime/pkg/client"
+	"sigs.k8s.io/controller-runtime/pkg/controller/controllerutil"
+	"sigs.k8s.io/controller-runtime/pkg/log"
+
+	computedomainv1beta1 "github.com/NVIDIA/k8s-dra-driver-gpu/api/nvidia.com/resource/v1beta1"
+	"github.com/run-ai/fake-gpu-operator/pkg/compute-domain/consts"
+)
+
+const (
+	// DefaultComputeDomainAllocationMode is the default allocation mode when not specified
+	DefaultComputeDomainAllocationMode = "Single"
+)
+
+// ComputeDomainReconciler watches ComputeDomain resources and keeps the
+// associated ResourceClaimTemplates in sync.
+type ComputeDomainReconciler struct {
+	client.Client
+	Scheme *runtime.Scheme
+}
+
+//+kubebuilder:rbac:groups=resource.nvidia.com,resources=computedomains,verbs=get;list;watch;create;update;patch;delete
+//+kubebuilder:rbac:groups=resource.nvidia.com,resources=computedomains/status,verbs=get;update;patch
+//+kubebuilder:rbac:groups=resource.nvidia.com,resources=computedomains/finalizers,verbs=update
+//+kubebuilder:rbac:groups=resource.k8s.io,resources=resourceclaimtemplates,verbs=get;list;watch;create;update;patch;delete
+
+func (r *ComputeDomainReconciler) Reconcile(ctx context.Context, req ctrl.Request) (ctrl.Result, error) {
+	logger := log.FromContext(ctx)
+
+	domain := &computedomainv1beta1.ComputeDomain{}
+	if err := r.Get(ctx, req.NamespacedName, domain); err != nil {
+		return ctrl.Result{}, client.IgnoreNotFound(err)
+	}
+
+	if domain.DeletionTimestamp.IsZero() {
+		if err := r.ensureFinalizer(ctx, domain); err != nil {
+			return ctrl.Result{}, err
+		}
+		if err := r.ensureResourceClaimTemplates(ctx, domain); err != nil {
+			return ctrl.Result{}, err
+		}
+	} else {
+		if err := r.handleDeletion(ctx, domain); err != nil {
+			return ctrl.Result{}, err
+		}
+		return ctrl.Result{}, nil
+	}
+
+	logger.V(4).Info("reconciled ComputeDomain", "namespace", domain.Namespace, "name", domain.Name)
+	return ctrl.Result{}, nil
+}
+
+func (r *ComputeDomainReconciler) ensureFinalizer(ctx context.Context, domain *computedomainv1beta1.ComputeDomain) error {
+	if controllerutil.ContainsFinalizer(domain, consts.ComputeDomainFinalizer) {
+		return nil
+	}
+
+	controllerutil.AddFinalizer(domain, consts.ComputeDomainFinalizer)
+	return r.Update(ctx, domain)
+}
+
+func (r *ComputeDomainReconciler) handleDeletion(ctx context.Context, domain *computedomainv1beta1.ComputeDomain) error {
+	if !controllerutil.ContainsFinalizer(domain, consts.ComputeDomainFinalizer) {
+		return nil
+	}
+
+	if err := r.deleteResourceClaimTemplates(ctx, domain); err != nil {
+		return err
+	}
+
+	controllerutil.RemoveFinalizer(domain, consts.ComputeDomainFinalizer)
+	return r.Update(ctx, domain)
+}
+
+func (r *ComputeDomainReconciler) ensureResourceClaimTemplates(ctx context.Context, domain *computedomainv1beta1.ComputeDomain) error {
+	return r.ensureTemplate(ctx, domain, domain.Name, consts.ComputeDomainWorkloadDeviceClass, "workload")
+}
+
+func (r *ComputeDomainReconciler) getAllocationMode(domain *computedomainv1beta1.ComputeDomain) string {
+	if domain.Spec.Channel != nil && domain.Spec.Channel.AllocationMode != "" {
+		return domain.Spec.Channel.AllocationMode
+	}
+	return DefaultComputeDomainAllocationMode
+}
+
+func (r *ComputeDomainReconciler) ensureTemplate(
+	ctx context.Context,
+	domain *computedomainv1beta1.ComputeDomain,
+	name string,
+	deviceClass string,
+	templateType string,
+) error {
+	key := client.ObjectKey{Namespace: domain.Namespace, Name: name}
+	existing := &resourceapi.ResourceClaimTemplate{}
+	err := r.Get(ctx, key, existing)
+	if err == nil {
+		return nil
+	}
+	if !apierrors.IsNotFound(err) {
+		return err
+	}
+
+	template := &resourceapi.ResourceClaimTemplate{
+		ObjectMeta: metav1.ObjectMeta{
+			Name:      name,
+			Namespace: domain.Namespace,
+			Labels: map[string]string{
+				"resource.nvidia.com/computeDomain":       domain.Name,
+				"resource.nvidia.com/computeDomainTarget": templateType,
+			},
+			Finalizers: []string{
+				"resource.nvidia.com/computeDomain",
+			},
+		},
+		Spec: resourceapi.ResourceClaimTemplateSpec{
+			ObjectMeta: metav1.ObjectMeta{
+				Labels: map[string]string{
+					"nvidia.com/computeDomain": domain.Name,
+				},
+			},
+			Spec: resourceapi.ResourceClaimSpec{
+				Devices: resourceapi.DeviceClaim{
+					Config: []resourceapi.DeviceClaimConfiguration{
+						{
+							DeviceConfiguration: resourceapi.DeviceConfiguration{
+								Opaque: &resourceapi.OpaqueDeviceConfiguration{
+									Driver: consts.ComputeDomainDriverName,
+									Parameters: runtime.RawExtension{
+										Raw: []byte(fmt.Sprintf(`{
+											"allocationMode": "%s",
+											"apiVersion": "resource.nvidia.com/v1beta1",
+											"domainID": "%s",
+											"kind": "ComputeDomainChannelConfig"
+										}`, r.getAllocationMode(domain), domain.UID)),
+									},
+								},
+							},
+						},
+					},
+					Requests: []resourceapi.DeviceRequest{
+						{
+							Name: "channel",
+							Exactly: &resourceapi.ExactDeviceRequest{
+								AllocationMode:  resourceapi.DeviceAllocationModeExactCount,
+								Count:           1,
+								DeviceClassName: deviceClass,
+							},
+						},
+					},
+				},
+			},
+		},
+	}
+
+	if err := controllerutil.SetControllerReference(domain, template, r.Scheme); err != nil {
+		return err
+	}
+
+	return client.IgnoreAlreadyExists(r.Create(ctx, template))
+}
+
+func (r *ComputeDomainReconciler) deleteResourceClaimTemplates(ctx context.Context, domain *computedomainv1beta1.ComputeDomain) error {
+	template := &resourceapi.ResourceClaimTemplate{
+		ObjectMeta: metav1.ObjectMeta{
+			Name:      domain.Name,
+			Namespace: domain.Namespace,
+		},
+	}
+	if err := r.Delete(ctx, template); err != nil && !apierrors.IsNotFound(err) {
+		return err
+	}
+	return nil
+}
+
+// SetupWithManager wires the reconciler into the controller-runtime manager.
+func (r *ComputeDomainReconciler) SetupWithManager(mgr ctrl.Manager) error {
+	return ctrl.NewControllerManagedBy(mgr).
+		For(&computedomainv1beta1.ComputeDomain{}).
+		Owns(&resourceapi.ResourceClaimTemplate{}).
+		Complete(r)
+}

cmd/compute-domain-controller/computedomain_controller_test.go

@@ -0,0 +1,158 @@
+/*
+ * Copyright 2025 The Kubernetes Authors.
+ *
+ * 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 (
+	"context"
+	"testing"
+
+	"github.com/stretchr/testify/assert"
+	"github.com/stretchr/testify/require"
+
+	resourceapi "k8s.io/api/resource/v1"
+	metav1 "k8s.io/apimachinery/pkg/apis/meta/v1"
+	"k8s.io/apimachinery/pkg/runtime"
+	"k8s.io/apimachinery/pkg/types"
+	ctrl "sigs.k8s.io/controller-runtime"
+	"sigs.k8s.io/controller-runtime/pkg/client"
+	"sigs.k8s.io/controller-runtime/pkg/client/fake"
+
+	computedomainv1beta1 "github.com/NVIDIA/k8s-dra-driver-gpu/api/nvidia.com/resource/v1beta1"
+	"github.com/run-ai/fake-gpu-operator/pkg/compute-domain/consts"
+)
+
+func TestComputeDomainReconciler_Reconcile(t *testing.T) {
+	scheme := runtime.NewScheme()
+	_ = resourceapi.AddToScheme(scheme)
+	_ = computedomainv1beta1.AddToScheme(scheme)
+
+	tests := map[string]struct {
+		computeDomain            *computedomainv1beta1.ComputeDomain
+		expectedWorkloadTemplate bool
+		expectedFinalizer        bool
+	}{
+		"new ComputeDomain creates templates": {
+			computeDomain: &computedomainv1beta1.ComputeDomain{
+				ObjectMeta: metav1.ObjectMeta{
+					Name:      "test-domain",
+					Namespace: "default",
+					UID:       "test-uid",
+				},
+			},
+			expectedWorkloadTemplate: true,
+			expectedFinalizer:        true,
+		},
+	}
+
+	for name, test := range tests {
+		t.Run(name, func(t *testing.T) {
+			// Setup
+			objs := []client.Object{test.computeDomain}
+
+			fakeClient := fake.NewClientBuilder().
+				WithScheme(scheme).
+				WithObjects(objs...).
+				Build()
+
+			reconciler := &ComputeDomainReconciler{
+				Client: fakeClient,
+				Scheme: scheme,
+			}
+
+			// Execute
+			req := ctrl.Request{
+				NamespacedName: types.NamespacedName{
+					Name:      test.computeDomain.GetName(),
+					Namespace: test.computeDomain.GetNamespace(),
+				},
+			}
+
+			result, err := reconciler.Reconcile(context.Background(), req)
+
+			// Verify
+			require.NoError(t, err)
+			assert.Equal(t, ctrl.Result{}, result)
+
+			// Check ResourceClaimTemplates
+			if test.expectedWorkloadTemplate {
+				workloadTemplate := &resourceapi.ResourceClaimTemplate{}
+				err := fakeClient.Get(context.Background(), types.NamespacedName{
+					Name:      "test-domain",
+					Namespace: "default",
+				}, workloadTemplate)
+				assert.NoError(t, err)
+				// Check config section
+				assert.Len(t, workloadTemplate.Spec.Spec.Devices.Config, 1)
+				assert.Equal(t, consts.ComputeDomainDriverName, workloadTemplate.Spec.Spec.Devices.Config[0].Opaque.Driver)
+				// Check requests - only channel request expected
+				assert.Len(t, workloadTemplate.Spec.Spec.Devices.Requests, 1)
+				assert.Equal(t, "channel", workloadTemplate.Spec.Spec.Devices.Requests[0].Name)
+				assert.Equal(t, resourceapi.DeviceAllocationModeExactCount, workloadTemplate.Spec.Spec.Devices.Requests[0].Exactly.AllocationMode)
+				assert.Equal(t, int64(1), workloadTemplate.Spec.Spec.Devices.Requests[0].Exactly.Count)
+				assert.Equal(t, consts.ComputeDomainWorkloadDeviceClass, workloadTemplate.Spec.Spec.Devices.Requests[0].Exactly.DeviceClassName)
+				// Check labels
+				assert.Equal(t, test.computeDomain.GetName(), workloadTemplate.Labels["resource.nvidia.com/computeDomain"])
+				assert.Equal(t, "workload", workloadTemplate.Labels["resource.nvidia.com/computeDomainTarget"])
+				// Check labels copied into generated claims
+				assert.Equal(t, test.computeDomain.GetName(), workloadTemplate.Spec.Labels["nvidia.com/computeDomain"])
+				// Check finalizers
+				assert.Contains(t, workloadTemplate.Finalizers, "resource.nvidia.com/computeDomain")
+			}
+
+			// Check finalizer
+			updatedDomain := &computedomainv1beta1.ComputeDomain{}
+			err = fakeClient.Get(context.Background(), req.NamespacedName, updatedDomain)
+			require.NoError(t, err)
+
+			finalizers := updatedDomain.GetFinalizers()
+			hasFinalizer := false
+			for _, f := range finalizers {
+				if f == consts.ComputeDomainFinalizer {
+					hasFinalizer = true
+					break
+				}
+			}
+			assert.Equal(t, test.expectedFinalizer, hasFinalizer)
+		})
+	}
+}
+
+func TestComputeDomainReconciler_Reconcile_NotFound(t *testing.T) {
+	scheme := runtime.NewScheme()
+	_ = resourceapi.AddToScheme(scheme)
+	_ = computedomainv1beta1.AddToScheme(scheme)
+
+	fakeClient := fake.NewClientBuilder().
+		WithScheme(scheme).
+		Build()
+
+	reconciler := &ComputeDomainReconciler{
+		Client: fakeClient,
+		Scheme: scheme,
+	}
+
+	req := ctrl.Request{
+		NamespacedName: types.NamespacedName{
+			Name:      "non-existent",
+			Namespace: "default",
+		},
+	}
+
+	result, err := reconciler.Reconcile(context.Background(), req)
+	assert.NoError(t, err)
+	assert.Equal(t, ctrl.Result{}, result)
+}