-
Notifications
You must be signed in to change notification settings - Fork 79
Define gpu memory utilization #502
New issue
Have a question about this project? Sign up for a free GitHub account to open an issue and contact its maintainers and the community.
By clicking “Sign up for GitHub”, you agree to our terms of service and privacy statement. We’ll occasionally send you account related emails.
Already on GitHub? Sign in to your account
Changes from all commits
File filter
Filter by extension
Conversations
Jump to
Diff view
Diff view
There are no files selected for viewing
| Original file line number | Diff line number | Diff line change |
|---|---|---|
|
|
@@ -128,6 +128,126 @@ func TestConfigureCmdThinkFlag(t *testing.T) { | |
| } | ||
| } | ||
|
|
||
| func TestConfigureCmdGPUMemoryUtilizationFlag(t *testing.T) { | ||
| // Create the configure command | ||
| cmd := newConfigureCmd() | ||
|
|
||
| // Verify the --gpu-memory-utilization flag exists | ||
| gpuMemFlag := cmd.Flags().Lookup("gpu-memory-utilization") | ||
| if gpuMemFlag == nil { | ||
| t.Fatal("--gpu-memory-utilization flag not found") | ||
| } | ||
|
|
||
| // Verify the default value is empty (nil pointer) | ||
| if gpuMemFlag.DefValue != "" { | ||
| t.Errorf("Expected default gpu-memory-utilization value to be '' (nil), got '%s'", gpuMemFlag.DefValue) | ||
| } | ||
|
|
||
| // Verify the flag type | ||
| if gpuMemFlag.Value.Type() != "float64" { | ||
| t.Errorf("Expected gpu-memory-utilization flag type to be 'float64', got '%s'", gpuMemFlag.Value.Type()) | ||
| } | ||
|
|
||
| // Test setting the flag value | ||
| err := cmd.Flags().Set("gpu-memory-utilization", "0.7") | ||
| if err != nil { | ||
| t.Errorf("Failed to set gpu-memory-utilization flag: %v", err) | ||
| } | ||
|
|
||
| // Verify the value was set | ||
| gpuMemValue := gpuMemFlag.Value.String() | ||
| if gpuMemValue != "0.7" { | ||
| t.Errorf("Expected gpu-memory-utilization flag value to be '0.7', got '%s'", gpuMemValue) | ||
| } | ||
| } | ||
|
|
||
| func TestGPUMemoryUtilizationBehavior(t *testing.T) { | ||
| // Helper to create float64 pointer | ||
| float64Ptr := func(f float64) *float64 { return &f } | ||
|
Contributor
There was a problem hiding this comment. Choose a reason for hiding this commentThe reason will be displayed to describe this comment to others. Learn more. |
||
|
|
||
| tests := []struct { | ||
| name string | ||
| gpuMemValue *float64 | ||
| expectError bool | ||
| expectGPUMemSet bool | ||
| expectedGPUMemUtil float64 | ||
| }{ | ||
| { | ||
| name: "default - not set (nil)", | ||
| gpuMemValue: nil, | ||
| expectError: false, | ||
| expectGPUMemSet: false, | ||
| }, | ||
| { | ||
| name: "valid value 0.5", | ||
| gpuMemValue: float64Ptr(0.5), | ||
| expectError: false, | ||
| expectGPUMemSet: true, | ||
| expectedGPUMemUtil: 0.5, | ||
| }, | ||
| { | ||
| name: "edge case 0.0", | ||
| gpuMemValue: float64Ptr(0.0), | ||
| expectError: false, | ||
| expectGPUMemSet: true, | ||
| expectedGPUMemUtil: 0.0, | ||
| }, | ||
| { | ||
| name: "edge case 1.0", | ||
| gpuMemValue: float64Ptr(1.0), | ||
| expectError: false, | ||
| expectGPUMemSet: true, | ||
| expectedGPUMemUtil: 1.0, | ||
| }, | ||
| { | ||
| name: "invalid - negative value", | ||
| gpuMemValue: float64Ptr(-0.1), | ||
| expectError: true, | ||
| }, | ||
| { | ||
| name: "invalid - value > 1.0", | ||
| gpuMemValue: float64Ptr(1.5), | ||
| expectError: true, | ||
| }, | ||
| } | ||
|
|
||
| for _, tt := range tests { | ||
| t.Run(tt.name, func(t *testing.T) { | ||
| flags := ConfigureFlags{ | ||
| GPUMemoryUtilization: tt.gpuMemValue, | ||
| } | ||
|
|
||
| req, err := flags.BuildConfigureRequest("test-model") | ||
|
|
||
| if tt.expectError { | ||
| if err == nil { | ||
| t.Fatal("Expected error but got none") | ||
| } | ||
| return | ||
| } | ||
|
|
||
| if err != nil { | ||
| t.Fatalf("Unexpected error: %v", err) | ||
| } | ||
|
|
||
| if tt.expectGPUMemSet { | ||
| // GPU memory utilization should be set | ||
| if req.VLLM == nil || req.VLLM.GPUMemoryUtilization == nil { | ||
| t.Fatal("Expected GPU memory utilization to be set") | ||
| } | ||
| if *req.VLLM.GPUMemoryUtilization != tt.expectedGPUMemUtil { | ||
| t.Errorf("Expected GPU memory utilization to be %f, got %f", tt.expectedGPUMemUtil, *req.VLLM.GPUMemoryUtilization) | ||
| } | ||
| } else { | ||
| // GPU memory utilization should NOT be set | ||
| if req.VLLM != nil && req.VLLM.GPUMemoryUtilization != nil { | ||
| t.Errorf("Expected GPU memory utilization to be nil when not set, got %f", *req.VLLM.GPUMemoryUtilization) | ||
| } | ||
| } | ||
| }) | ||
| } | ||
| } | ||
|
|
||
| func TestThinkFlagBehavior(t *testing.T) { | ||
| // Helper to create bool pointer | ||
| boolPtr := func(b bool) *bool { return &b } | ||
|
|
||
| Original file line number | Diff line number | Diff line change |
|---|---|---|
|
|
@@ -60,6 +60,15 @@ func (c *Config) GetArgs(bundle types.ModelBundle, socket string, mode inference | |
|
|
||
| // Add vLLM-specific arguments from backend config | ||
| if config != nil && config.VLLM != nil { | ||
| // Add GPU memory utilization if specified | ||
| if config.VLLM.GPUMemoryUtilization != nil { | ||
| utilization := *config.VLLM.GPUMemoryUtilization | ||
| if utilization < 0.0 || utilization > 1.0 { | ||
| return nil, fmt.Errorf("gpu-memory-utilization must be between 0.0 and 1.0, got %f", utilization) | ||
| } | ||
|
Comment on lines
+66
to
+68
Contributor
There was a problem hiding this comment. Choose a reason for hiding this commentThe reason will be displayed to describe this comment to others. Learn more. The validation logic for the |
||
| args = append(args, "--gpu-memory-utilization", strconv.FormatFloat(utilization, 'f', -1, 64)) | ||
| } | ||
|
|
||
| // Add HuggingFace overrides if specified | ||
| if len(config.VLLM.HFOverrides) > 0 { | ||
| hfOverridesJSON, err := json.Marshal(config.VLLM.HFOverrides) | ||
|
|
||
There was a problem hiding this comment.
Choose a reason for hiding this comment
The reason will be displayed to describe this comment to others. Learn more.
suggestion (testing): Add a test case for non-float input to the --gpu-memory-utilization flag to exercise the error path in Float64PtrValue.Set.
Since
Float64PtrValue.Setreturns an error whenstrconv.ParseFloatfails, please add a subtest that callscmd.Flags().Set("gpu-memory-utilization", "not-a-number")and asserts that an error is returned. This will exercise the error path and verify the custom flag’s parsing behavior end-to-end.