diff --git a/pkg/devcontainer/config/config.go b/pkg/devcontainer/config/config.go index 2c85f6030..c799f0637 100644 --- a/pkg/devcontainer/config/config.go +++ b/pkg/devcontainer/config/config.go @@ -2,6 +2,7 @@ package config import ( "encoding/json" + "fmt" "path" "path/filepath" "strconv" @@ -311,27 +312,89 @@ type HostRequirements struct { // Amount of required disk space in bytes. Supports units tb, gb, mb and kb. Storage string `json:"storage,omitempty"` - // If GPU support should be enabled - GPU types.StrBool `json:"gpu,omitempty"` + // If GPU support should be enabled. Accepts bool, string ("optional"), or object with cores/memory. + GPU *GPURequirement `json:"gpu,omitempty"` +} + +// GPURequirement represents the gpu field in hostRequirements. +// It supports bool, string, and object formats per the devcontainer spec. +type GPURequirement struct { + Value string // "true", "false", or "optional" + Cores int // Object format, advisory + GPUMemory string // Object format, advisory +} + +// UnmarshalJSON handles bool, string, and object formats for the gpu field. +func (g *GPURequirement) UnmarshalJSON(data []byte) error { + var raw any + if err := json.Unmarshal(data, &raw); err != nil { + return err + } + switch v := raw.(type) { + case bool: + g.Value = strconv.FormatBool(v) + case string: + g.Value = v + case map[string]any: + g.parseObject(v) + default: + return fmt.Errorf("gpu: %w", types.ErrUnsupportedType) + } + return nil +} + +// MarshalJSON produces the canonical JSON format so that round-tripping +// through json.Marshal/Unmarshal (used by Convert for deep copies) is lossless. +func (g GPURequirement) MarshalJSON() ([]byte, error) { + if g.Cores > 0 || g.GPUMemory != "" { + obj := map[string]any{} + if g.Cores > 0 { + obj["cores"] = g.Cores + } + if g.GPUMemory != "" { + obj["memory"] = g.GPUMemory + } + return json.Marshal(obj) + } + b, err := strconv.ParseBool(g.Value) + if err == nil { + return json.Marshal(b) + } + return json.Marshal(g.Value) +} + +func (g *GPURequirement) parseObject(obj map[string]any) { + g.Value = "true" + if f, ok := obj["cores"].(float64); ok { + g.Cores = int(f) + } + if s, ok := obj["memory"].(string); ok { + g.GPUMemory = s + } } // ShouldEnableGPU determines if GPU should be enabled based on requirements and availability. func (h *HostRequirements) ShouldEnableGPU(gpuAvailable bool) (enable bool, warnIfMissing bool) { - if h == nil || h.GPU == "" { + if h == nil || h.GPU == nil { return false, false } - gpuValue := string(h.GPU) - if gpuValue == "optional" { + if h.GPU.Value == "optional" { return gpuAvailable, false } - required, err := h.GPU.Bool() + required, err := strconv.ParseBool(h.GPU.Value) if err != nil { return false, false } - return required && gpuAvailable, required && !gpuAvailable + if required && !gpuAvailable { + return false, true + } + if !required { + return false, false + } + return true, false } type PortAttribute struct { diff --git a/pkg/devcontainer/config/gpu_test.go b/pkg/devcontainer/config/gpu_test.go new file mode 100644 index 000000000..75db8df5f --- /dev/null +++ b/pkg/devcontainer/config/gpu_test.go @@ -0,0 +1,184 @@ +package config + +import ( + "encoding/json" + "testing" +) + +func TestGPURequirement_BoolAndString(t *testing.T) { + tests := []struct { + name string + input string + wantVal string + }{ + {"true", `{"gpu": true}`, "true"}, + {"false", `{"gpu": false}`, "false"}, + {"optional", `{"gpu": "optional"}`, "optional"}, + } + for _, tt := range tests { + t.Run(tt.name, func(t *testing.T) { + var hr HostRequirements + if err := json.Unmarshal([]byte(tt.input), &hr); err != nil { + t.Fatalf("unmarshal: %v", err) + } + if hr.GPU == nil { + t.Fatal("GPU is nil") + } + if hr.GPU.Value != tt.wantVal { + t.Errorf("Value = %q, want %q", hr.GPU.Value, tt.wantVal) + } + }) + } +} + +func TestGPURequirement_ObjectFormat(t *testing.T) { + tests := []struct { + name string + input string + wantVal string + wantCores int + wantMem string + }{ + { + "cores and memory", + `{"gpu": {"cores": 4, "memory": "8gb"}}`, + "true", 4, "8gb", + }, + { + "cores only", + `{"gpu": {"cores": 2}}`, + "true", 2, "", + }, + { + "memory only", + `{"gpu": {"memory": "16gb"}}`, + "true", 0, "16gb", + }, + { + "empty object", + `{"gpu": {}}`, + "true", 0, "", + }, + } + for _, tt := range tests { + t.Run(tt.name, func(t *testing.T) { + var hr HostRequirements + if err := json.Unmarshal([]byte(tt.input), &hr); err != nil { + t.Fatalf("unmarshal: %v", err) + } + if hr.GPU == nil { + t.Fatal("GPU is nil") + } + if hr.GPU.Value != tt.wantVal { + t.Errorf("Value = %q, want %q", hr.GPU.Value, tt.wantVal) + } + if hr.GPU.Cores != tt.wantCores { + t.Errorf("Cores = %d, want %d", hr.GPU.Cores, tt.wantCores) + } + if hr.GPU.GPUMemory != tt.wantMem { + t.Errorf("GPUMemory = %q, want %q", hr.GPU.GPUMemory, tt.wantMem) + } + }) + } +} + +func TestGPURequirement_InvalidType(t *testing.T) { + var hr HostRequirements + err := json.Unmarshal([]byte(`{"gpu": [1, 2, 3]}`), &hr) + if err == nil { + t.Fatal("expected error for array input, got nil") + } +} + +func TestGPURequirement_OmittedIsNil(t *testing.T) { + var hr HostRequirements + if err := json.Unmarshal([]byte(`{"cpus": 4}`), &hr); err != nil { + t.Fatalf("unmarshal: %v", err) + } + if hr.GPU != nil { + t.Errorf("GPU = %+v, want nil", hr.GPU) + } +} + +func TestShouldEnableGPU(t *testing.T) { + gpu := func(val string) *GPURequirement { return &GPURequirement{Value: val} } + + tests := []struct { + name string + hr *HostRequirements + available bool + wantEnable bool + wantWarning bool + }{ + {"nil GPU", &HostRequirements{}, false, false, false}, + {"true available", &HostRequirements{GPU: gpu("true")}, true, true, false}, + {"true unavailable", &HostRequirements{GPU: gpu("true")}, false, false, true}, + {"false", &HostRequirements{GPU: gpu("false")}, true, false, false}, + {"optional available", &HostRequirements{GPU: gpu("optional")}, true, true, false}, + {"optional unavailable", &HostRequirements{GPU: gpu("optional")}, false, false, false}, + { + "object available", + &HostRequirements{GPU: &GPURequirement{Value: "true", Cores: 4, GPUMemory: "8gb"}}, + true, true, false, + }, + { + "object unavailable", + &HostRequirements{GPU: &GPURequirement{Value: "true", Cores: 4, GPUMemory: "8gb"}}, + false, false, true, + }, + } + for _, tt := range tests { + t.Run(tt.name, func(t *testing.T) { + enable, warn := tt.hr.ShouldEnableGPU(tt.available) + if enable != tt.wantEnable { + t.Errorf("enable = %v, want %v", enable, tt.wantEnable) + } + if warn != tt.wantWarning { + t.Errorf("warnIfMissing = %v, want %v", warn, tt.wantWarning) + } + }) + } +} + +func TestShouldEnableGPU_NilReceiver(t *testing.T) { + var hr *HostRequirements + enable, warn := hr.ShouldEnableGPU(true) + if enable || warn { + t.Errorf("got (%v, %v), want (false, false)", enable, warn) + } +} + +func TestGPURequirement_RoundTrip(t *testing.T) { + tests := []struct { + name string + input string + }{ + {"bool true", `{"gpu":true}`}, + {"bool false", `{"gpu":false}`}, + {"string optional", `{"gpu":"optional"}`}, + {"object", `{"gpu":{"cores":4,"memory":"8gb"}}`}, + {"empty object", `{"gpu":{}}`}, + } + for _, tt := range tests { + t.Run(tt.name, func(t *testing.T) { + var hr1 HostRequirements + if err := json.Unmarshal([]byte(tt.input), &hr1); err != nil { + t.Fatalf("first unmarshal: %v", err) + } + data, err := json.Marshal(&hr1) + if err != nil { + t.Fatalf("marshal: %v", err) + } + var hr2 HostRequirements + if err := json.Unmarshal(data, &hr2); err != nil { + t.Fatalf("second unmarshal: %v", err) + } + if hr1.GPU == nil || hr2.GPU == nil { + t.Fatal("GPU is nil after round-trip") + } + if *hr1.GPU != *hr2.GPU { + t.Errorf("round-trip mismatch: %+v != %+v", *hr1.GPU, *hr2.GPU) + } + }) + } +}