support resource claim cel builder

Author: wangqianqianjun

internal/constants/constants.go

@@ -74,8 +74,6 @@ const (
 	// In remote vGPU mode, selected workload is set by user with /workload annotation or generated by system
 	SelectedWorkloadAnnotation = Domain + "/selected-workload"
 
-	CELFilterExpressionAnnotation = Domain + "/cel-filter-expression"
-
 	WorkloadModeAnnotation = Domain + "/workload-mode"
 	WorkloadModeDynamic    = "dynamic"
 	WorkloadModeFixed      = "fixed"

internal/gpuallocator/gpuallocator.go

@@ -1355,7 +1355,7 @@ func (s *GpuAllocator) ComposeAllocationRequest(pod *v1.Pod) (*tfv1.AllocRequest
 		Limit:    gpuLimitResource,
 
 		DisableCELFilter:    disableCELFilter,
-		CELFilterExpression: pod.Annotations[constants.CELFilterExpressionAnnotation],
+		CELFilterExpression: pod.Annotations[constants.DRACelExpressionAnnotation],
 
 		Count:    uint(count),
 		GPUModel: pod.Annotations[constants.GPUModelAnnotation],

internal/webhook/v1/pod_dra.go

@@ -136,7 +136,7 @@ func (p *DRAProcessor) HandleDRAAdmission(ctx context.Context, pod *corev1.Pod,
 	return nil
 }
 
-// TODO: support more attributes for filtering
+// BuildCELSelector constructs a CEL expression for DRA device selection based on TensorFusion requirements
 func BuildCELSelector(pod *corev1.Pod, tfInfo *utils.TensorFusionInfo) (string, error) {
 	var conditions []string
 
@@ -154,6 +154,28 @@ func BuildCELSelector(pod *corev1.Pod, tfInfo *utils.TensorFusionInfo) (string,
 		conditions = append(conditions, fmt.Sprintf(`device.attributes["model"] == "%s"`, tfInfo.Profile.GPUModel))
 	}
 
+	// 3. GPU count requirement (important for multi-GPU workloads)
+	if tfInfo.Profile.GPUCount > 0 {
+		conditions = append(conditions, fmt.Sprintf(`int(device.attributes["gpu_count"]) >= %d`, tfInfo.Profile.GPUCount))
+	}
+
+	// 4. Pool name filter (for resource isolation and scheduling preferences)
+	if tfInfo.Profile.PoolName != "" {
+		conditions = append(conditions, fmt.Sprintf(`device.attributes["pool_name"] == "%s"`, tfInfo.Profile.PoolName))
+	}
+
+	// 5. Workload name filter (for workload-specific device assignment)
+	if tfInfo.WorkloadName != "" {
+		conditions = append(conditions, fmt.Sprintf(`device.attributes["workload_name"] == "%s"`, tfInfo.WorkloadName))
+		// Workload namespace is same as pod namespace in TensorFusion
+		conditions = append(conditions, fmt.Sprintf(`device.attributes["workload_namespace"] == "%s"`, pod.Namespace))
+	}
+
+	// 6. Pod namespace filter (for namespace-based device isolation)
+	if pod.Namespace != "" {
+		conditions = append(conditions, fmt.Sprintf(`device.attributes["pod_namespace"] == "%s"`, pod.Namespace))
+	}
+
 	// Return a basic condition if no specific requirements
 	if len(conditions) == 0 {
 		// Simple condition that should work with most DRA drivers

internal/webhook/v1/pod_webhook_dra_test.go

@@ -149,37 +149,137 @@ func TestDRAProcessor_HandleDRAAdmission(t *testing.T) {
 }
 
 func TestBuildCELSelector(t *testing.T) {
-	pod := &corev1.Pod{
-		ObjectMeta: metav1.ObjectMeta{
-			Name:      "test-pod",
-			Namespace: "test-namespace",
+	tests := []struct {
+		name                 string
+		pod                  *corev1.Pod
+		tfInfo               *utils.TensorFusionInfo
+		expectedConditions   []string
+		unexpectedConditions []string
+	}{
+		{
+			name: "Basic resource filters",
+			pod: &corev1.Pod{
+				ObjectMeta: metav1.ObjectMeta{
+					Name:      "test-pod",
+					Namespace: "test-namespace",
+				},
+			},
+			tfInfo: &utils.TensorFusionInfo{
+				Profile: &tfv1.WorkloadProfileSpec{
+					GPUCount: 2,
+					Resources: tfv1.Resources{
+						Requests: tfv1.Resource{
+							Tflops: resource.MustParse("20"),
+							Vram:   resource.MustParse("16Gi"),
+						},
+					},
+					GPUModel: "H100",
+				},
+			},
+			expectedConditions: []string{
+				`device.attributes["tflops"].quantity >= quantity("20")`,
+				`device.attributes["vram"].quantity >= quantity("16Gi")`,
+				`device.attributes["model"] == "H100"`,
+				`int(device.attributes["gpu_count"]) >= 2`,
+				`device.attributes["pod_namespace"] == "test-namespace"`,
+			},
 		},
-	}
-
-	tfInfo := &utils.TensorFusionInfo{
-		Profile: &tfv1.WorkloadProfileSpec{
-			GPUCount: 2,
-			Resources: tfv1.Resources{
-				Requests: tfv1.Resource{
-					Tflops: resource.MustParse("20"),
-					Vram:   resource.MustParse("16Gi"),
+		{
+			name: "All filters including pool and workload",
+			pod: &corev1.Pod{
+				ObjectMeta: metav1.ObjectMeta{
+					Name:      "test-pod",
+					Namespace: "production",
+				},
+			},
+			tfInfo: &utils.TensorFusionInfo{
+				Profile: &tfv1.WorkloadProfileSpec{
+					GPUCount: 1,
+					Resources: tfv1.Resources{
+						Requests: tfv1.Resource{
+							Tflops: resource.MustParse("10"),
+							Vram:   resource.MustParse("8Gi"),
+						},
+					},
+					GPUModel: "A100",
+					PoolName: "high-priority",
+				},
+				WorkloadName: "ml-training-job",
+			},
+			expectedConditions: []string{
+				`device.attributes["tflops"].quantity >= quantity("10")`,
+				`device.attributes["vram"].quantity >= quantity("8Gi")`,
+				`device.attributes["model"] == "A100"`,
+				`int(device.attributes["gpu_count"]) >= 1`,
+				`device.attributes["pool_name"] == "high-priority"`,
+				`device.attributes["workload_name"] == "ml-training-job"`,
+				`device.attributes["workload_namespace"] == "production"`,
+				`device.attributes["pod_namespace"] == "production"`,
+			},
+		},
+		{
+			name: "Zero resources fallback to default condition",
+			pod: &corev1.Pod{
+				ObjectMeta: metav1.ObjectMeta{
+					Name:      "test-pod",
+					Namespace: "default",
+				},
+			},
+			tfInfo: &utils.TensorFusionInfo{
+				Profile: &tfv1.WorkloadProfileSpec{
+					GPUCount: 0, // Zero count should not add condition
+					Resources: tfv1.Resources{
+						Requests: tfv1.Resource{
+							// Zero resources
+						},
+					},
+				},
+			},
+			expectedConditions: []string{
+				`device.attributes["pod_namespace"] == "default"`,
+			},
+		},
+		{
+			name: "Empty resources fallback to basic condition",
+			pod: &corev1.Pod{
+				ObjectMeta: metav1.ObjectMeta{
+					Name:      "test-pod",
+					Namespace: "",
+				},
+			},
+			tfInfo: &utils.TensorFusionInfo{
+				Profile: &tfv1.WorkloadProfileSpec{
+					// All empty/zero values
 				},
 			},
-			GPUModel: "H100",
+			expectedConditions: []string{
+				`device.attributes.exists("type")`,
+			},
 		},
 	}
 
-	celExpression, err := BuildCELSelector(pod, tfInfo)
-	require.NoError(t, err)
-	require.NotEmpty(t, celExpression)
+	for _, tt := range tests {
+		t.Run(tt.name, func(t *testing.T) {
+			celExpression, err := BuildCELSelector(tt.pod, tt.tfInfo)
+			require.NoError(t, err)
+			require.NotEmpty(t, celExpression)
+
+			// Verify expected conditions are present
+			for _, condition := range tt.expectedConditions {
+				assert.Contains(t, celExpression, condition, "Expected condition not found: %s", condition)
+			}
 
-	// Verify it contains the expected resource filters
-	assert.Contains(t, celExpression, `device.attributes["tflops"].quantity >= quantity("20")`)
-	assert.Contains(t, celExpression, `device.attributes["vram"].quantity >= quantity("16Gi")`)
-	assert.Contains(t, celExpression, `device.attributes["model"] == "H100"`)
+			// Verify unexpected conditions are not present
+			for _, condition := range tt.unexpectedConditions {
+				assert.NotContains(t, celExpression, condition, "Unexpected condition found: %s", condition)
+			}
 
-	// Verify conditions are combined with AND
-	assert.Contains(t, celExpression, " && ")
+			// Verify proper AND joining (unless it's the fallback condition)
+			if len(tt.expectedConditions) > 1 {
+				assert.Contains(t, celExpression, " && ", "Conditions should be joined with &&")
+			}
+		})
+	}
 }
 
 func TestHasDRAClaim(t *testing.T) {