[Feature][APIServer] Support decimal memory values in KubeRay APIServer (#3956)

* [Add] Newly added support for configuring calculation templates to handle decimal numbers

Signed-off-by: daiping8 <[email protected]>

* [Update] Modified memory field type to integer and added optional memory unit in ComputeTemplate

Signed-off-by: daiping8 <[email protected]>

* [Update] Added `MemoryUnit` field to `ComputeTemplate` in E2E tests to support memory unit specification

Signed-off-by: daiping8 <[email protected]>

* [Update] Added memory unit field `CompTemplateMemUnitForE2E` to enhance E2E test configurations

Signed-off-by: daiping8 <[email protected]>

* [Update] Added `memory_unit` field in `converter_test.go` for extended E2E test coverage

Signed-off-by: daiping8 <[email protected]>

* [Update] Applied default `memory_unit` handling in `converter.go` and updated E2E tests for consistency

Signed-off-by: daiping8 <[email protected]>

* [Update] Adjusted E2E test to handle default `MemoryUnit` in `ComputeTemplate` and ensure compatibility

Signed-off-by: daiping8 <[email protected]>

* [Update] Replaced protobuf cloning with explicit struct initialization in E2E test to enhance clarity and removed unused import

Signed-off-by: daiping8 <[email protected]>

* [Update] Added logging for expected and actual templates in E2E test to improve debugging clarity

Signed-off-by: daiping8 <[email protected]>

* [Update] Removed default `MemoryUnit` assignment logic from `converter.go` and cleaned up related tests and configurations

Signed-off-by: daiping8 <[email protected]>

---------

Signed-off-by: daiping8 <[email protected]>

Author: daiping8

apiserver/pkg/model/converter.go

@@ -412,6 +412,7 @@ func FromKubeToAPIComputeTemplate(configMap *corev1.ConfigMap) *api.ComputeTempl
 	runtime.Namespace = configMap.Namespace
 	runtime.Cpu = uint32(cpu)
 	runtime.Memory = uint32(memory)
+	runtime.MemoryUnit = configMap.Data["memory_unit"]
 	runtime.Gpu = uint32(gpu)
 	runtime.GpuAccelerator = configMap.Data["gpu_accelerator"]
 

apiserver/pkg/model/converter_test.go

@@ -141,6 +141,7 @@ var configMapWithoutTolerations = corev1.ConfigMap{
 		"extended_resources": "{\"vpc.amazonaws.com/efa\": 32}",
 		"name":               "head-node-template",
 		"namespace":          "max",
+		"memory_unit":        "Gi",
 	},
 }
 
@@ -154,6 +155,20 @@ var configMapWithTolerations = corev1.ConfigMap{
 		"name":               "head-node-template",
 		"namespace":          "max",
 		"tolerations":        "[{\"key\":\"blah1\",\"operator\":\"Exists\",\"effect\":\"NoExecute\"}]",
+		"memory_unit":        "Gi",
+	},
+}
+
+var configMapWithMemoryUnit = corev1.ConfigMap{
+	Data: map[string]string{
+		"cpu":                "4",
+		"gpu":                "0",
+		"gpu_accelerator":    "",
+		"memory":             "8192",
+		"extended_resources": "{\"vpc.amazonaws.com/efa\": 32}",
+		"name":               "head-node-template",
+		"namespace":          "max",
+		"memory_unit":        "Mi",
 	},
 }
 
@@ -674,6 +689,11 @@ func TestPopulateTemplate(t *testing.T) {
 		template.ExtendedResources,
 		"Extended resources mismatch",
 	)
+
+	template = FromKubeToAPIComputeTemplate(&configMapWithMemoryUnit)
+	assert.Equal(t, uint32(8192), template.Memory, "Memory mismatch")
+	// test memory unit MiB
+	assert.Equal(t, "Mi", template.MemoryUnit, "Memory Unit mismatch")
 }
 
 func tolerationToString(toleration *api.PodToleration) string {

apiserver/pkg/util/cluster.go

@@ -179,7 +179,8 @@ func buildHeadPodTemplate(imageVersion string, envs *api.EnvironmentVariables, s
 
 	// calculate resources
 	cpu := fmt.Sprint(computeRuntime.GetCpu())
-	memory := fmt.Sprintf("%d%s", computeRuntime.GetMemory(), "Gi")
+	memoryUnit := computeRuntime.GetMemoryUnit()
+	memory := fmt.Sprintf("%d%s", computeRuntime.GetMemory(), memoryUnit)
 
 	// build volume and volumeMounts
 	volMounts := buildVolumeMounts(spec.Volumes)
@@ -433,7 +434,8 @@ func buildWorkerPodTemplate(imageVersion string, envs *api.EnvironmentVariables,
 
 	// calculate resources
 	cpu := fmt.Sprint(computeRuntime.GetCpu())
-	memory := fmt.Sprintf("%d%s", computeRuntime.GetMemory(), "Gi")
+	memoryUnit := computeRuntime.GetMemoryUnit()
+	memory := fmt.Sprintf("%d%s", computeRuntime.GetMemory(), memoryUnit)
 
 	// build volume and volumeMounts
 	volMounts := buildVolumeMounts(spec.Volumes)
@@ -847,13 +849,25 @@ func NewComputeTemplate(runtime *api.ComputeTemplate) (*corev1.ConfigMap, error)
 	if err != nil {
 		return nil, fmt.Errorf("failed to marshal extended resources: %w", err)
 	}
+	memoryUnit := "Gi"
+	if runtime.MemoryUnit != "" {
+		memoryUnit = runtime.MemoryUnit
+	}
+
+	memory := strconv.FormatUint(uint64(runtime.Memory), 10)
+	quantity := memory + memoryUnit
+
+	if _, err := resource.ParseQuantity(quantity); err != nil {
+		return nil, fmt.Errorf("invalid memory quantity %q: %w", quantity, err)
+	}
 
 	// Create data map
 	dmap := map[string]string{
 		"name":               runtime.Name,
 		"namespace":          runtime.Namespace,
 		"cpu":                strconv.FormatUint(uint64(runtime.Cpu), 10),
-		"memory":             strconv.FormatUint(uint64(runtime.Memory), 10),
+		"memory":             memory,
+		"memory_unit":        memoryUnit,
 		"gpu":                strconv.FormatUint(uint64(runtime.Gpu), 10),
 		"gpu_accelerator":    runtime.GpuAccelerator,
 		"extended_resources": string(extendedResourcesJSON),

apiserver/pkg/util/cluster_test.go

@@ -270,6 +270,7 @@ var templateWorker = api.ComputeTemplate{
 			Effect:   "NoExecute",
 		},
 	},
+	MemoryUnit: "Gi",
 }
 
 var expectedToleration = corev1.Toleration{

apiserver/test/e2e/config_server_e2e_test.go

@@ -123,7 +123,26 @@ func TestCreateTemplate(t *testing.T) {
 			} else {
 				require.NoError(t, err, "No error expected")
 				require.Nil(t, actualRPCStatus, "No RPC status expected")
-				require.Truef(t, reflect.DeepEqual(tc.Input.ComputeTemplate, actualTemplate), "Equal templates expected")
+				if tc.Input.ComputeTemplate.MemoryUnit == "" {
+					require.Equal(t, "Gi", actualTemplate.MemoryUnit, "Default MemoryUnit should be Gi")
+					// Copy tc.Input.ComputeTemplate to the expected template with the default MemoryUnit
+					expected := &api.ComputeTemplate{
+						Name:              tc.Input.ComputeTemplate.Name,
+						Cpu:               tc.Input.ComputeTemplate.Cpu,
+						Memory:            tc.Input.ComputeTemplate.Memory,
+						Namespace:         tc.Input.ComputeTemplate.Namespace,
+						Gpu:               tc.Input.ComputeTemplate.Gpu,
+						GpuAccelerator:    tc.Input.ComputeTemplate.GpuAccelerator,
+						Tolerations:       tc.Input.ComputeTemplate.Tolerations,
+						ExtendedResources: tc.Input.ComputeTemplate.ExtendedResources,
+						MemoryUnit:        "Gi",
+					}
+					t.Logf("Expected template: %+v", expected)
+					t.Logf("Actual template: %+v", actualTemplate)
+					require.Truef(t, reflect.DeepEqual(expected, actualTemplate), "Equal templates expected (with default MemoryUnit)")
+				} else {
+					require.Truef(t, reflect.DeepEqual(tc.Input.ComputeTemplate, actualTemplate), "Equal templates expected")
+				}
 			}
 		})
 	}

proto/config.proto

@@ -130,6 +130,8 @@ message ComputeTemplate {
   repeated PodToleration tolerations = 7;
   // Optional. Name and number of the extended resources
   map<string, uint32> extended_resources = 8;
+  // Optional. The unit of memory, default is "Gi" for memory
+  string memory_unit = 9;
 }
 
 // This service is not implemented.