diff --git a/internal/compose/compose.go b/internal/compose/compose.go index d5fa75e..4936e50 100644 --- a/internal/compose/compose.go +++ b/internal/compose/compose.go @@ -9,6 +9,8 @@ package compose import ( "fmt" + "maps" + "slices" "strings" "github.com/Masterminds/semver/v3" @@ -377,7 +379,12 @@ func applyDefaults(site, derived, base values.Tree, in Input) error { } if len(in.GPUNodeSelectorTerms) > 0 { if _, set := values.Get(base, requiredNodeTerms); !set { - if err := values.Set(derived, requiredNodeTerms, in.GPUNodeSelectorTerms); err != nil { + // The form's own terms replace a derived list, so they carry the pin instead. + layer, terms := derived, in.GPUNodeSelectorTerms + if have, _ := values.Get(in.Overrides, requiredNodeTerms); len(asList(have)) > 0 { + layer, terms = in.Overrides, AndTerms(asList(have), terms) + } + if err := values.Set(layer, requiredNodeTerms, terms); err != nil { return err } } @@ -385,6 +392,66 @@ func applyDefaults(site, derived, base values.Tree, in Input) error { return nil } +// AndTerms is the nodeSelectorTerms that require both lists: terms are ORed +// and a term's matchExpressions ANDed, so each pair of terms becomes one. An +// In expression of pin replaces the same key's In in a term. A term that is +// not a map is kept as is. +func AndTerms(terms, pin []any) []any { + if len(terms) == 0 { + return pin + } + var out []any + for _, t := range terms { + tm, ok := asTree(t) + if !ok { + out = append(out, t) + continue + } + for _, p := range pin { + pm, ok := asTree(p) + if !ok { + continue + } + exprs := slices.Clone(asList(tm["matchExpressions"])) + for _, e := range asList(pm["matchExpressions"]) { + exprs = putExpr(exprs, e) + } + term := maps.Clone(tm) + term["matchExpressions"] = exprs + out = append(out, term) + } + } + return out +} + +func putExpr(exprs []any, e any) []any { + em, ok := asTree(e) + if ok && em["operator"] == "In" { + for i, x := range exprs { + if xm, ok := asTree(x); ok && xm["key"] == em["key"] && xm["operator"] == "In" { + exprs[i] = e + return exprs + } + } + } + return append(exprs, e) +} + +func asList(v any) []any { + l, _ := v.([]any) + return l +} + +func asTree(v any) (values.Tree, bool) { + switch m := v.(type) { + case values.Tree: + return m, true + case map[string]any: + return m, true + } + return nil, false +} + // applyInputs puts each input's answer in the derived layer, over the variant's // values and under the form's overrides: the site's, else the input's default. func applyInputs(derived, site values.Tree, in Input) error { diff --git a/internal/compose/compose_test.go b/internal/compose/compose_test.go index 2a5dbf4..4c15bc8 100644 --- a/internal/compose/compose_test.go +++ b/internal/compose/compose_test.go @@ -1,6 +1,7 @@ package compose import ( + "reflect" "strings" "testing" @@ -636,11 +637,12 @@ func TestComposeGPUResource(t *testing.T) { } // Product affinity is derived, so the catalog's or the site's own required -// node affinity stands, and the form still overrides it. +// node affinity stands; the form's own terms carry it instead. func TestComposeGPUNodeSelectorTerms(t *testing.T) { const path = "affinity.nodeAffinity.requiredDuringSchedulingIgnoredDuringExecution.nodeSelectorTerms" - terms := []any{values.Tree{"matchExpressions": []any{values.Tree{"key": "k", "operator": "In", "values": []any{"p"}}}}} - theirs := []any{"theirs"} + pin := values.Tree{"key": "k", "operator": "In", "values": []any{"p"}} + terms := []any{values.Tree{"matchExpressions": []any{pin}}} + theirs := values.Tree{"key": "storage", "operator": "In", "values": []any{"true"}} in := testInput() in.GPUNodeSelectorTerms = terms @@ -656,7 +658,7 @@ func TestComposeGPUNodeSelectorTerms(t *testing.T) { in := testInput() in.GPUNodeSelectorTerms = terms tree := values.Tree{} - if err := values.Set(tree, path, theirs); err != nil { + if err := values.Set(tree, path, []any{values.Tree{"matchExpressions": []any{theirs}}}); err != nil { t.Fatal(err) } if layer == "catalog" { @@ -671,10 +673,42 @@ func TestComposeGPUNodeSelectorTerms(t *testing.T) { if got := p.LayerOf(path); got != layer { t.Errorf("%s set terms, but they are %q's", layer, got) } - if _, ok := p.Layers[values.LayerDerived]; ok && layer == "catalog" { - if _, set := values.Get(p.Layers[values.LayerDerived], path); set { - t.Error("derived must yield to the catalog's own terms") - } + if _, set := values.Get(p.Layers[values.LayerDerived], path); set { + t.Errorf("%s set terms, but derived set them too", layer) + } + got, _ := values.Get(p.Values(), path) + want := []any{values.Tree{"matchExpressions": []any{theirs}}} + if layer == "form" { + want = []any{values.Tree{"matchExpressions": []any{theirs, pin}}} + } + if !reflect.DeepEqual(got, want) { + t.Errorf("%s set terms: got %v, want %v", layer, got, want) } } } + +func TestAndTerms(t *testing.T) { + expr := func(key, op string, v ...any) values.Tree { + return values.Tree{"key": key, "operator": op, "values": v} + } + term := func(e ...any) values.Tree { return values.Tree{"matchExpressions": e} } + pin := []any{term(expr("gpu", "In", "H800"))} + + if got := AndTerms(nil, pin); !reflect.DeepEqual(got, pin) { + t.Errorf("no terms: got %v, want the pin", got) + } + terms := []any{ + term(expr("storage", "In", "true")), + term(expr("gpu", "In", "H100"), expr("gpu", "NotIn", "H800-bad")), + } + want := []any{ + term(expr("storage", "In", "true"), expr("gpu", "In", "H800")), + term(expr("gpu", "In", "H800"), expr("gpu", "NotIn", "H800-bad")), + } + if got := AndTerms(terms, pin); !reflect.DeepEqual(got, want) { + t.Errorf("got %v, want %v", got, want) + } + if n := len(terms[0].(values.Tree)["matchExpressions"].([]any)); n != 1 { + t.Errorf("the terms passed in were changed: %v", terms) + } +} diff --git a/internal/server/deploy.go b/internal/server/deploy.go index 6e94fc8..e04b600 100644 --- a/internal/server/deploy.go +++ b/internal/server/deploy.go @@ -274,8 +274,11 @@ func (s *Server) compose(ctx context.Context, req planRequest) (*plan.Plan, erro case len(terms) == 0: warnings = append(warnings, fmt.Sprintf("no node label to match GPU products %s on for %s, so the deploy is not pinned to them -- list it with a labelKey in the site profile's nodes.gpus", strings.Join(products, ", "), gpuRes)) case asked: + // ANDed into the form's own terms rather than over them, so a term it set still holds. const path = "affinity.nodeAffinity.requiredDuringSchedulingIgnoredDuringExecution.nodeSelectorTerms" - if err := values.Set(overrides, path, terms); err != nil { + have, _ := values.Get(overrides, path) + list, _ := have.([]any) + if err := values.Set(overrides, path, compose.AndTerms(list, terms)); err != nil { return nil, err } default: diff --git a/internal/server/gpuselect_test.go b/internal/server/gpuselect_test.go index 793876f..82269e9 100644 --- a/internal/server/gpuselect_test.go +++ b/internal/server/gpuselect_test.go @@ -62,6 +62,34 @@ func TestGPUProductsKeepAPreferredAffinityRule(t *testing.T) { } } +// The form's own required term, a preset's say, holds beside the products: +// terms are ORed, so the product expression goes into it rather than beside it. +func TestGPUProductsAreANDedIntoTheFormsTerms(t *testing.T) { + srv, _ := deployServer(t, true) + storage := map[string]any{"key": "storage", "operator": "In", "values": []any{"true"}} + code, p := post(t, srv, "/api/plans", map[string]any{ + "model": "qwen3.6-35b-a3b", "serviceId": "r", + "gpuProducts": []string{"NVIDIA-H800"}, + "overrides": map[string]any{ + "affinity": map[string]any{"nodeAffinity": map[string]any{"requiredDuringSchedulingIgnoredDuringExecution": map[string]any{ + "nodeSelectorTerms": []any{map[string]any{"matchExpressions": []any{storage}}}, + }}}, + }, + }) + if code != 200 { + t.Fatalf("status %d: %v", code, p) + } + na := planValues(p)["affinity"].(map[string]any)["nodeAffinity"].(map[string]any) + terms := na["requiredDuringSchedulingIgnoredDuringExecution"].(map[string]any)["nodeSelectorTerms"].([]any) + want := []any{map[string]any{"matchExpressions": []any{ + storage, + map[string]any{"key": "nvidia.com/gpu.product", "operator": "In", "values": []any{"NVIDIA-H800"}}, + }}} + if !reflect.DeepEqual(terms, want) { + t.Fatalf("got %v, want %v", terms, want) + } +} + func matchExpr(t *testing.T, p map[string]any) map[string]any { t.Helper() na := planValues(p)["affinity"].(map[string]any)["nodeAffinity"].(map[string]any)