[raycicmd] Implicit dimension matching for array-to-array depends_on

When array step A depends on array step B via a plain string reference and both share array dimension names (e.g. "python", "cuda"), each expanded variant of A now depends only on the B variant with matching values on those shared dimensions, instead of all variants of B.

This unblocks parallelism in the pipeline DAG: py3.11 build no longer waits on py3.10 or py3.12 builds of its dependency.

Behavior:
- Overlapping dims: constrain to matching values (new)
- No overlapping dims: fan out to all variants (unchanged)
- Non-array deps: pass through as-is (unchanged)
- Explicit selector filter: takes precedence over implicit matching (unchanged)
- Opt-out: use map selector form `- step-name:` to force all-variants fanout

Topic: implicit-dim-matching
Relative: array-validation-guards
Labels: draft
Signed-off-by: andrew <andrew@anyscale.com>

Author: andrew-anyscale

raycicmd/array_depends.go

@@ -7,8 +7,9 @@ import (
 
 // arraySelector represents a dependency selector with optional array filter.
 type arraySelector struct {
-	key    string              // step identifier: "key" for command steps, "name" for wanda steps
-	filter map[string][]string // dimension -> allowed values (nil = all instances)
+	key      string              // step identifier
+	filter   map[string][]string // dimension constraints; nil triggers implicit matching or fan-out
+	matchAll bool                // true = depend on all variants (e.g. `- step: all`)
 }
 
 // parseArrayDependsOn parses a depends_on field into a list of selectors.
@@ -20,6 +21,8 @@ type arraySelector struct {
 //	depends_on:                       # selector with array filter
 //	  - ray-build:
 //	      python: "3.11"
+//	depends_on:                       # explicit all-variants
+//	  - ray-build: all
 func parseArrayDependsOn(v any) ([]*arraySelector, error) {
 	if v == nil {
 		return nil, nil
@@ -73,16 +76,28 @@ func parseArraySelectorMap(m map[string]any) (*arraySelector, error) {
 
 	sel := &arraySelector{key: stepName}
 
-	if filterVal != nil {
-		filterMap, ok := filterVal.(map[string]any)
-		if !ok {
-			return nil, fmt.Errorf("selector %q filter must be a map, got %T", stepName, filterVal)
+	switch fv := filterVal.(type) {
+	case nil:
+		// `- step-name:` with no value; treat same as plain string.
+	case string:
+		if fv != "all" {
+			return nil, fmt.Errorf(
+				"selector %q: string value must be \"all\", got %q",
+				stepName, fv,
+			)
 		}
-		parsed, err := parseSelectorFilter(filterMap)
+		sel.matchAll = true
+	case map[string]any:
+		parsed, err := parseSelectorFilter(fv)
 		if err != nil {
 			return nil, fmt.Errorf("in selector for %q: %w", stepName, err)
 		}
 		sel.filter = parsed
+	default:
+		return nil, fmt.Errorf(
+			"selector %q filter must be a map or \"all\", got %T",
+			stepName, filterVal,
+		)
 	}
 
 	return sel, nil
@@ -108,9 +123,15 @@ func parseSelectorFilter(m map[string]any) (map[string][]string, error) {
 }
 
 // resolveDependsOn resolves a depends_on field to concrete step keys.
-// Plain references to array steps fan out to all instances.
-// Selector references resolve to only matching instances.
-func resolveDependsOn(dependsOn any, configs map[string]*arrayConfig) ([]string, error) {
+// When currentElem is non-nil and a plain reference targets an array
+// step with overlapping dimensions, resolution is constrained to the
+// variant matching the current element's values on those dimensions.
+// Explicit selector filters take precedence over implicit matching.
+func resolveDependsOn(
+	dependsOn any,
+	configs map[string]*arrayConfig,
+	currentElem *arrayElement,
+) ([]string, error) {
 	selectors, err := parseArrayDependsOn(dependsOn)
 	if err != nil {
 		return nil, err
@@ -120,7 +141,7 @@ func resolveDependsOn(dependsOn any, configs map[string]*arrayConfig) ([]string,
 	for _, sel := range selectors {
 		cfg, isArray := configs[sel.key]
 		if !isArray {
-			if sel.filter != nil {
+			if sel.filter != nil || sel.matchAll {
 				return nil, fmt.Errorf(
 					"cannot use array selector on non-array step %q",
 					sel.key,
@@ -129,8 +150,31 @@ func resolveDependsOn(dependsOn any, configs map[string]*arrayConfig) ([]string,
 			result = append(result, sel.key)
 			continue
 		}
+
+		// For plain references (no explicit filter/matchAll) from an
+		// array step, build an implicit filter on overlapping dims.
+		// When there is no overlap, sel.filter stays nil and
+		// resolveArraySelector matches all elements (fan-out).
+		implicit := false
+		if sel.filter == nil && !sel.matchAll && currentElem != nil {
+			if f := implicitDimFilter(currentElem, cfg); f != nil {
+				sel = &arraySelector{
+					key:    sel.key,
+					filter: f,
+				}
+				implicit = true
+			}
+		}
+
 		matches, err := resolveArraySelector(sel, cfg)
 		if err != nil {
+			if implicit {
+				return nil, fmt.Errorf(
+					"implicit dimension matching on %q: %w; "+
+						"use `- %s: all` to depend on all variants",
+					sel.key, err, sel.key,
+				)
+			}
 			return nil, err
 		}
 		result = append(result, matches...)
@@ -139,6 +183,27 @@ func resolveDependsOn(dependsOn any, configs map[string]*arrayConfig) ([]string,
 	return result, nil
 }
 
+// implicitDimFilter builds a selector filter from the overlapping
+// dimensions between the current element and the target config.
+// Returns nil if there is no overlap, causing the caller to fall
+// through to fan-out-to-all behavior.
+func implicitDimFilter(
+	current *arrayElement, target *arrayConfig,
+) map[string][]string {
+	var filter map[string][]string
+	for dim := range target.dims {
+		val, ok := current.values[dim]
+		if !ok {
+			continue
+		}
+		if filter == nil {
+			filter = make(map[string][]string)
+		}
+		filter[dim] = []string{val}
+	}
+	return filter
+}
+
 // resolveArraySelector resolves a selector against an array config
 // to concrete step keys, using the final element list (post-adjustments).
 func resolveArraySelector(sel *arraySelector, cfg *arrayConfig) ([]string, error) {

raycicmd/array_expand.go

@@ -10,10 +10,13 @@ import (
 // depends_on references to point to the expanded keys.
 func expandArraySteps(gs []*pipelineGroup) error {
 	configs := make(map[string]*arrayConfig)
+	// elems maps each expanded step key to its arrayElement,
+	// used for implicit dimension matching in Pass 2.
+	elems := make(map[string]*arrayElement)
 
 	// Pass 1: expand array steps into resolvedSteps.
 	for _, g := range gs {
-		if err := g.buildResolvedSteps(configs); err != nil {
+		if err := g.buildResolvedSteps(configs, elems); err != nil {
 			return fmt.Errorf("expand arrays: %w", err)
 		}
 	}
@@ -25,8 +28,9 @@ func expandArraySteps(gs []*pipelineGroup) error {
 			if !ok {
 				continue
 			}
+			currentElem := elems[stepKey(rs.src)]
 			resolved, err := resolveDependsOn(
-				dependsOn, configs,
+				dependsOn, configs, currentElem,
 			)
 			if err != nil {
 				return fmt.Errorf(
@@ -65,6 +69,7 @@ func expandArraySteps(gs []*pipelineGroup) error {
 // array steps into multiple entries.
 func (g *pipelineGroup) buildResolvedSteps(
 	configs map[string]*arrayConfig,
+	elems map[string]*arrayElement,
 ) error {
 	var result []*resolvedStep
 
@@ -95,8 +100,10 @@ func (g *pipelineGroup) buildResolvedSteps(
 		if err != nil {
 			return err
 		}
-		for _, es := range expanded {
+		cfg := configs[baseKey]
+		for j, es := range expanded {
 			result = append(result, &resolvedStep{src: es})
+			elems[stepKey(es)] = cfg.elements[j]
 		}
 	}