feat(wanda): Add params field for env var validation and discovery

Add `params` field to wanda spec for declaring allowed values of
environment variables used in templated fields (name, froms). This
enables:
- Strict validation: reject builds where env var values don't match
  declared params
- Dependency discovery: specs with params are indexed under all their
  expanded names (e.g., base$PY with params [3.10, 3.11] is indexed as
  both base3.10 and base3.11), enabling dependency resolution without
  requiring all env vars to be set

Validation runs before variable expansion in buildDepGraph, so errors
reference the original $VARNAME.

Topic: wanda-params
Relative: wanda-deps

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

Author: andrew-anyscale

wanda/deps.go

@@ -69,6 +69,11 @@ func (g *depGraph) loadSpec(specPath string, isRoot bool) error {
 	if err != nil {
 		return fmt.Errorf("parse %s: %w", specPath, err)
 	}
+
+	if err := spec.ValidateParams(g.lookup); err != nil {
+		return fmt.Errorf("%s: %w", specPath, err)
+	}
+
 	spec = spec.expandVar(g.lookup)
 
 	if err := checkUnexpandedVars(spec, specPath); err != nil {
@@ -331,7 +336,8 @@ type discovered struct {
 }
 
 // discoverSpecs scans searchRoot for *.wanda.yaml files and builds a name index.
-// Names are expanded using the provided lookup function.
+// Names are expanded using declared params first, then the lookup function.
+// Specs with params will have all param combinations indexed.
 // Returns an error if two specs expand to the same name.
 func discoverSpecs(searchRoot string, lookup lookupFunc) (specIndex, error) {
 	index := make(specIndex)
@@ -356,13 +362,25 @@ func discoverSpecs(searchRoot string, lookup lookupFunc) (specIndex, error) {
 					outCh <- discovered{skipped: true}
 					continue
 				}
-				spec = spec.expandVar(lookup)
 
-				if vars := findUnexpandedVars(spec.Name); len(vars) > 0 {
+				// Use params to enumerate all possible names.
+				// For vars without params, $VARNAME is preserved.
+				names := spec.ExpandedNames()
+				sentAny := false
+				for _, name := range names {
+					// If still has unexpanded vars, try env lookup
+					if vars := findUnexpandedVars(name); len(vars) > 0 {
+						name = expandVar(name, lookup)
+						if vars := findUnexpandedVars(name); len(vars) > 0 {
+							continue // Can't fully expand, skip this name
+						}
+					}
+					outCh <- discovered{name: name, path: path}
+					sentAny = true
+				}
+				if !sentAny {
 					outCh <- discovered{skipped: true}
-					continue
 				}
-				outCh <- discovered{name: spec.Name, path: path}
 			}
 		}()
 	}

wanda/deps_test.go

@@ -515,6 +515,70 @@ func TestDiscoverSpecs_WithVariables(t *testing.T) {
 	}
 }
 
+func TestDiscoverSpecs_WithParams(t *testing.T) {
+	tmpDir := t.TempDir()
+
+	// Spec with params - should be indexed under all expanded names
+	writeSpec(t, tmpDir, "base.wanda.yaml", strings.Join([]string{
+		"name: base$PY",
+		"params:",
+		"  PY:",
+		"    - '3.10'",
+		"    - '3.11'",
+		"    - '3.12'",
+		"dockerfile: Dockerfile",
+	}, "\n"))
+
+	// No env vars needed - params provide the values
+	index, err := discoverSpecs(tmpDir, noopLookup)
+	if err != nil {
+		t.Fatalf("discoverSpecs: %v", err)
+	}
+
+	// All three expanded names should be indexed
+	for _, name := range []string{"base3.10", "base3.11", "base3.12"} {
+		if _, ok := index[name]; !ok {
+			t.Errorf("index missing %q, got: %v", name, index)
+		}
+	}
+
+	// All should point to the same spec file
+	path := index["base3.10"]
+	if index["base3.11"] != path || index["base3.12"] != path {
+		t.Errorf("all names should map to same path, got: %v", index)
+	}
+}
+
+func TestDiscoverSpecs_ParamsAndEnvFallback(t *testing.T) {
+	tmpDir := t.TempDir()
+
+	// Spec with partial params - one var has params, one needs env
+	writeSpec(t, tmpDir, "base.wanda.yaml", strings.Join([]string{
+		"name: base$PY-$ARCH",
+		"params:",
+		"  PY:",
+		"    - '3.10'",
+		"dockerfile: Dockerfile",
+	}, "\n"))
+
+	lookup := func(key string) (string, bool) {
+		if key == "ARCH" {
+			return "amd64", true
+		}
+		return "", false
+	}
+
+	index, err := discoverSpecs(tmpDir, lookup)
+	if err != nil {
+		t.Fatalf("discoverSpecs: %v", err)
+	}
+
+	// Should be indexed as base3.10-amd64
+	if _, ok := index["base3.10-amd64"]; !ok {
+		t.Errorf("index missing base3.10-amd64, got: %v", index)
+	}
+}
+
 func TestBuildDepGraph_Discovery(t *testing.T) {
 	tmpDir := t.TempDir()
 
@@ -643,3 +707,101 @@ func TestBuildDepGraph_TransitiveDeps(t *testing.T) {
 		t.Error("expected c in graph (transitive dep)")
 	}
 }
+
+func TestBuildDepGraph_ParamsValidation(t *testing.T) {
+	tmpDir := t.TempDir()
+
+	writeSpec(t, tmpDir, "spec.wanda.yaml", strings.Join([]string{
+		"name: myimage$PY_VERSION",
+		"params:",
+		"  PY_VERSION:",
+		"    - '3.10'",
+		"    - '3.11'",
+		"dockerfile: Dockerfile",
+	}, "\n"))
+
+	t.Run("valid param value", func(t *testing.T) {
+		lookup := func(k string) (string, bool) {
+			if k == "PY_VERSION" {
+				return "3.10", true
+			}
+			return "", false
+		}
+		_, err := buildDepGraph(filepath.Join(tmpDir, "spec.wanda.yaml"), lookup, "")
+		if err != nil {
+			t.Errorf("unexpected error with valid param: %v", err)
+		}
+	})
+
+	t.Run("invalid param value", func(t *testing.T) {
+		lookup := func(k string) (string, bool) {
+			if k == "PY_VERSION" {
+				return "3.9", true
+			}
+			return "", false
+		}
+		_, err := buildDepGraph(filepath.Join(tmpDir, "spec.wanda.yaml"), lookup, "")
+		if err == nil {
+			t.Error("expected error for invalid param value")
+		}
+		if !strings.Contains(err.Error(), "3.9") {
+			t.Errorf("error should mention invalid value '3.9': %v", err)
+		}
+	})
+}
+
+func TestBuildDepGraph_DiscoveryWithParams(t *testing.T) {
+	tmpDir := t.TempDir()
+
+	if err := os.Mkdir(filepath.Join(tmpDir, ".git"), 0755); err != nil {
+		t.Fatal(err)
+	}
+
+	// Base spec with params - discoverable via params, loadable with env var
+	baseDir := filepath.Join(tmpDir, "base")
+	if err := os.MkdirAll(baseDir, 0755); err != nil {
+		t.Fatal(err)
+	}
+	writeSpec(t, baseDir, "base.wanda.yaml", strings.Join([]string{
+		"name: base$PY",
+		"params:",
+		"  PY:",
+		"    - '3.10'",
+		"    - '3.11'",
+		"dockerfile: Dockerfile",
+	}, "\n"))
+
+	// App spec depends on base3.10
+	appDir := filepath.Join(tmpDir, "app")
+	if err := os.MkdirAll(appDir, 0755); err != nil {
+		t.Fatal(err)
+	}
+	writeSpec(t, appDir, "app.wanda.yaml", strings.Join([]string{
+		"name: app",
+		`froms: ["cr.ray.io/rayproject/base3.10"]`,
+		"dockerfile: Dockerfile",
+	}, "\n"))
+
+	// Discovery finds base3.10 via params (no env var needed for discovery).
+	// Loading the spec requires env var to be set for expansion.
+	lookup := func(key string) (string, bool) {
+		if key == "PY" {
+			return "3.10", true
+		}
+		return "", false
+	}
+
+	graph, err := buildDepGraph(filepath.Join(appDir, "app.wanda.yaml"), lookup, testPrefix)
+	if err != nil {
+		t.Fatalf("buildDepGraph: %v", err)
+	}
+
+	// base3.10 was discovered via params and loaded with PY=3.10
+	if graph.Specs["base3.10"] == nil {
+		t.Error("expected base3.10 in graph")
+	}
+
+	if len(graph.Order) != 2 {
+		t.Errorf("Order has %d items, want 2", len(graph.Order))
+	}
+}

wanda/spec.go

@@ -12,6 +12,11 @@ import (
 type Spec struct {
 	Name string `yaml:"name,omitempty"`
 
+	// Params declares allowed values for environment variables used in
+	// templated fields, currently only `name` and `froms` are supported.
+	// Keys are variable names (without $), values are lists of allowed values.
+	Params map[string][]string `yaml:"params,omitempty"`
+
 	Tags []string `yaml:"tags"`
 
 	// Inputs
@@ -120,6 +125,7 @@ func (s *Spec) expandVar(lookup lookupFunc) *Spec {
 	result := new(Spec)
 
 	result.Name = expandVar(s.Name, lookup)
+	result.Params = s.Params // Params are expanded by expandVarWithParams
 	result.Tags = stringsExpanVar(s.Tags, lookup)
 	result.Froms = stringsExpanVar(s.Froms, lookup)
 	result.Srcs = stringsExpanVar(s.Srcs, lookup)
@@ -129,3 +135,146 @@ func (s *Spec) expandVar(lookup lookupFunc) *Spec {
 
 	return result
 }
+
+// extractVarNames extracts all variable names from a string.
+// Uses the same parsing rules as expandVar.
+func extractVarNames(s string) []string {
+	var vars []string
+	inName := false
+	nameStart := 0
+
+	for i, r := range s {
+		if !inName {
+			if r == '$' {
+				inName = true
+				nameStart = i + 1
+			}
+		} else {
+			if r == '$' {
+				if nameStart == i {
+					// $$ escape sequence
+					inName = false
+					continue
+				}
+			}
+			if r >= 'a' && r <= 'z' || r >= 'A' && r <= 'Z' || r == '_' {
+				continue
+			}
+			if r >= '0' && r <= '9' && i > nameStart {
+				continue
+			}
+
+			if i > nameStart {
+				vars = append(vars, s[nameStart:i])
+			}
+			if r == '$' {
+				nameStart = i + 1
+			} else {
+				inName = false
+			}
+		}
+	}
+
+	if inName && len(s) > nameStart {
+		vars = append(vars, s[nameStart:])
+	}
+
+	return vars
+}
+
+// expandVarWithParams returns all possible expanded values for a string
+// using declared params. Generates the cartesian product when multiple
+// variables are present. Variables without params are preserved as $VARNAME.
+func expandVarWithParams(s string, params map[string][]string) []string {
+	vars := extractVarNames(s)
+	if len(vars) == 0 {
+		return []string{s}
+	}
+
+	// Find which vars have params
+	var paramVars []string
+	for _, v := range vars {
+		if _, ok := params[v]; ok {
+			paramVars = append(paramVars, v)
+		}
+	}
+
+	if len(paramVars) == 0 {
+		return []string{s}
+	}
+
+	// Generate all combinations
+	combinations := []map[string]string{{}}
+	for _, v := range paramVars {
+		var newCombinations []map[string]string
+		for _, combo := range combinations {
+			for _, val := range params[v] {
+				newCombo := make(map[string]string, len(combo)+1)
+				for k, v := range combo {
+					newCombo[k] = v
+				}
+				newCombo[v] = val
+				newCombinations = append(newCombinations, newCombo)
+			}
+		}
+		combinations = newCombinations
+	}
+
+	// Expand each combination
+	var results []string
+	for _, combo := range combinations {
+		expanded := expandVar(s, func(k string) (string, bool) {
+			v, ok := combo[k]
+			return v, ok
+		})
+		results = append(results, expanded)
+	}
+
+	return results
+}
+
+// ExpandedNames returns all possible names based on declared params.
+func (s *Spec) ExpandedNames() []string {
+	return expandVarWithParams(s.Name, s.Params)
+}
+
+// ExpandedFroms returns all possible froms based on declared params.
+// Results are deduplicated.
+func (s *Spec) ExpandedFroms() []string {
+	seen := make(map[string]bool)
+	var results []string
+	for _, from := range s.Froms {
+		for _, expanded := range expandVarWithParams(from, s.Params) {
+			if !seen[expanded] {
+				seen[expanded] = true
+				results = append(results, expanded)
+			}
+		}
+	}
+	return results
+}
+
+// ValidateParams validates that environment variable values match declared params.
+func (s *Spec) ValidateParams(lookup lookupFunc) error {
+	for varName, allowed := range s.Params {
+		val, ok := lookup(varName)
+		if !ok {
+			continue // Unset vars are handled by expandVar
+		}
+
+		valid := false
+		for _, a := range allowed {
+			if val == a {
+				valid = true
+				break
+			}
+		}
+		if !valid {
+			return fmt.Errorf(
+				"env var %s has value %q, but params only allow: %v",
+				varName, val, allowed,
+			)
+		}
+	}
+	return nil
+}