Skip to content

Commit 491c488

Browse files
daiping8daiping8
authored
[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]>
1 parent b5c6b0a commit 491c488

File tree

10 files changed

+238
-162
lines changed

10 files changed

+238
-162
lines changed

‎apiserver/pkg/model/converter.go‎

Lines changed: 1 addition & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -412,6 +412,7 @@ func FromKubeToAPIComputeTemplate(configMap *corev1.ConfigMap) *api.ComputeTempl
412412
runtime.Namespace = configMap.Namespace
413413
runtime.Cpu = uint32(cpu)
414414
runtime.Memory = uint32(memory)
415+
runtime.MemoryUnit = configMap.Data["memory_unit"]
415416
runtime.Gpu = uint32(gpu)
416417
runtime.GpuAccelerator = configMap.Data["gpu_accelerator"]
417418

‎apiserver/pkg/model/converter_test.go‎

Lines changed: 20 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -141,6 +141,7 @@ var configMapWithoutTolerations = corev1.ConfigMap{
141141
"extended_resources": "{\"vpc.amazonaws.com/efa\": 32}",
142142
"name": "head-node-template",
143143
"namespace": "max",
144+
"memory_unit": "Gi",
144145
},
145146
}
146147

@@ -154,6 +155,20 @@ var configMapWithTolerations = corev1.ConfigMap{
154155
"name": "head-node-template",
155156
"namespace": "max",
156157
"tolerations": "[{\"key\":\"blah1\",\"operator\":\"Exists\",\"effect\":\"NoExecute\"}]",
158+
"memory_unit": "Gi",
159+
},
160+
}
161+
162+
var configMapWithMemoryUnit = corev1.ConfigMap{
163+
Data: map[string]string{
164+
"cpu": "4",
165+
"gpu": "0",
166+
"gpu_accelerator": "",
167+
"memory": "8192",
168+
"extended_resources": "{\"vpc.amazonaws.com/efa\": 32}",
169+
"name": "head-node-template",
170+
"namespace": "max",
171+
"memory_unit": "Mi",
157172
},
158173
}
159174

@@ -674,6 +689,11 @@ func TestPopulateTemplate(t *testing.T) {
674689
template.ExtendedResources,
675690
"Extended resources mismatch",
676691
)
692+
693+
template = FromKubeToAPIComputeTemplate(&configMapWithMemoryUnit)
694+
assert.Equal(t, uint32(8192), template.Memory, "Memory mismatch")
695+
// test memory unit MiB
696+
assert.Equal(t, "Mi", template.MemoryUnit, "Memory Unit mismatch")
677697
}
678698

679699
func tolerationToString(toleration *api.PodToleration) string {

‎apiserver/pkg/util/cluster.go‎

Lines changed: 17 additions & 3 deletions
Original file line numberDiff line numberDiff line change
@@ -179,7 +179,8 @@ func buildHeadPodTemplate(imageVersion string, envs *api.EnvironmentVariables, s
179179

180180
// calculate resources
181181
cpu := fmt.Sprint(computeRuntime.GetCpu())
182-
memory := fmt.Sprintf("%d%s", computeRuntime.GetMemory(), "Gi")
182+
memoryUnit := computeRuntime.GetMemoryUnit()
183+
memory := fmt.Sprintf("%d%s", computeRuntime.GetMemory(), memoryUnit)
183184

184185
// build volume and volumeMounts
185186
volMounts := buildVolumeMounts(spec.Volumes)
@@ -433,7 +434,8 @@ func buildWorkerPodTemplate(imageVersion string, envs *api.EnvironmentVariables,
433434

434435
// calculate resources
435436
cpu := fmt.Sprint(computeRuntime.GetCpu())
436-
memory := fmt.Sprintf("%d%s", computeRuntime.GetMemory(), "Gi")
437+
memoryUnit := computeRuntime.GetMemoryUnit()
438+
memory := fmt.Sprintf("%d%s", computeRuntime.GetMemory(), memoryUnit)
437439

438440
// build volume and volumeMounts
439441
volMounts := buildVolumeMounts(spec.Volumes)
@@ -847,13 +849,25 @@ func NewComputeTemplate(runtime *api.ComputeTemplate) (*corev1.ConfigMap, error)
847849
if err != nil {
848850
return nil, fmt.Errorf("failed to marshal extended resources: %w", err)
849851
}
852+
memoryUnit := "Gi"
853+
if runtime.MemoryUnit != "" {
854+
memoryUnit = runtime.MemoryUnit
855+
}
856+
857+
memory := strconv.FormatUint(uint64(runtime.Memory), 10)
858+
quantity := memory + memoryUnit
859+
860+
if _, err := resource.ParseQuantity(quantity); err != nil {
861+
return nil, fmt.Errorf("invalid memory quantity %q: %w", quantity, err)
862+
}
850863

851864
// Create data map
852865
dmap := map[string]string{
853866
"name": runtime.Name,
854867
"namespace": runtime.Namespace,
855868
"cpu": strconv.FormatUint(uint64(runtime.Cpu), 10),
856-
"memory": strconv.FormatUint(uint64(runtime.Memory), 10),
869+
"memory": memory,
870+
"memory_unit": memoryUnit,
857871
"gpu": strconv.FormatUint(uint64(runtime.Gpu), 10),
858872
"gpu_accelerator": runtime.GpuAccelerator,
859873
"extended_resources": string(extendedResourcesJSON),

‎apiserver/pkg/util/cluster_test.go‎

Lines changed: 1 addition & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -270,6 +270,7 @@ var templateWorker = api.ComputeTemplate{
270270
Effect: "NoExecute",
271271
},
272272
},
273+
MemoryUnit: "Gi",
273274
}
274275

275276
var expectedToleration = corev1.Toleration{

‎apiserver/test/e2e/config_server_e2e_test.go‎

Lines changed: 20 additions & 1 deletion
Original file line numberDiff line numberDiff line change
@@ -123,7 +123,26 @@ func TestCreateTemplate(t *testing.T) {
123123
} else {
124124
require.NoError(t, err, "No error expected")
125125
require.Nil(t, actualRPCStatus, "No RPC status expected")
126-
require.Truef(t, reflect.DeepEqual(tc.Input.ComputeTemplate, actualTemplate), "Equal templates expected")
126+
if tc.Input.ComputeTemplate.MemoryUnit == "" {
127+
require.Equal(t, "Gi", actualTemplate.MemoryUnit, "Default MemoryUnit should be Gi")
128+
// Copy tc.Input.ComputeTemplate to the expected template with the default MemoryUnit
129+
expected := &api.ComputeTemplate{
130+
Name: tc.Input.ComputeTemplate.Name,
131+
Cpu: tc.Input.ComputeTemplate.Cpu,
132+
Memory: tc.Input.ComputeTemplate.Memory,
133+
Namespace: tc.Input.ComputeTemplate.Namespace,
134+
Gpu: tc.Input.ComputeTemplate.Gpu,
135+
GpuAccelerator: tc.Input.ComputeTemplate.GpuAccelerator,
136+
Tolerations: tc.Input.ComputeTemplate.Tolerations,
137+
ExtendedResources: tc.Input.ComputeTemplate.ExtendedResources,
138+
MemoryUnit: "Gi",
139+
}
140+
t.Logf("Expected template: %+v", expected)
141+
t.Logf("Actual template: %+v", actualTemplate)
142+
require.Truef(t, reflect.DeepEqual(expected, actualTemplate), "Equal templates expected (with default MemoryUnit)")
143+
} else {
144+
require.Truef(t, reflect.DeepEqual(tc.Input.ComputeTemplate, actualTemplate), "Equal templates expected")
145+
}
127146
}
128147
})
129148
}

‎proto/config.proto‎

Lines changed: 2 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -130,6 +130,8 @@ message ComputeTemplate {
130130
repeated PodToleration tolerations = 7;
131131
// Optional. Name and number of the extended resources
132132
map<string, uint32> extended_resources = 8;
133+
// Optional. The unit of memory, default is "Gi" for memory
134+
string memory_unit = 9;
133135
}
134136

135137
// This service is not implemented.

‎proto/go_client/cluster.pb.go‎

Lines changed: 2 additions & 2 deletions
Some generated files are not rendered by default. Learn more about customizing how changed files appear on GitHub.

0 commit comments

Comments
 (0)