[raycicmd] Apply array adjustments during expansion (#464)

Modify expandSingleArrayStep to parse adjustments from inside the array
definition and apply them in two passes: additions are appended first,
then skipped elements are filtered out. Store the final element list on
arrayConfig so depends_on resolution uses post-adjustment elements,
ensuring skipped elements are excluded and added elements are included.
Update array-test.rayci.yaml with adjustment examples.

Topic: apply-adjustments
Relative: parse-adjustments
Signed-off-by: andrew <andrew@anyscale.com>

Signed-off-by: andrew <andrew@anyscale.com>

Author: andrew-anyscale

.buildkite/array-test.Dockerfile

@@ -0,0 +1,26 @@
+# syntax=docker/dockerfile:1.3-labs
+FROM ubuntu:22.04
+
+ARG PYTHON_VERSION
+ARG CUDA_VERSION
+
+# Validate that only expected combinations are built.
+RUN <<'EOF'
+#!/bin/bash
+set -euo pipefail
+
+valid=false
+if [[ "$PYTHON_VERSION" = "3.12" && "$CUDA_VERSION" = "13.0.0-cudnn" ]]; then
+  valid=true
+fi
+if [[ "$PYTHON_VERSION" = "3.11" && "$CUDA_VERSION" = "12.8.1-cudnn" ]]; then
+  valid=true
+fi
+
+if [[ "$valid" != "true" ]]; then
+  echo "ERROR: unexpected combination py${PYTHON_VERSION}+cu${CUDA_VERSION}"
+  exit 1
+fi
+
+echo "py${PYTHON_VERSION}+cu${CUDA_VERSION}" > /image-info.txt
+EOF

.buildkite/array-test.rayci.yaml

@@ -8,11 +8,11 @@ steps:
     key: array-step
     command: |
       echo "Running os={{array.os}} variant={{array.variant}}"
-      if [ "{{array.os}}" != "linux" ] && [ "{{array.os}}" != "macos" ]; then
+      if [[ "{{array.os}}" != "linux" && "{{array.os}}" != "macos" ]]; then
         echo "ERROR: unexpected os={{array.os}}"
         exit 1
       fi
-      if [ "{{array.variant}}" != "1" ]; then
+      if [[ "{{array.variant}}" != "1" ]]; then
         echo "ERROR: unexpected variant={{array.variant}}"
         exit 1
       fi
@@ -21,6 +21,35 @@ steps:
       os: ["linux", "windows", "macos"]
       variant: ["1", "2"]
 
+  - name: array-cuda-build
+    label: "wanda: cuda py{{array.python}} cu{{array.cuda}}"
+    wanda: .buildkite/array-test.wanda.yaml
+    env:
+      PYTHON_VERSION: "{{array.python}}"
+      CUDA_VERSION: "{{array.cuda}}"
+    depends_on: forge
+    array:
+      python: ["3.12"]
+      cuda: ["13.0.0-cudnn"]
+      adjustments:
+        - with:
+            python: "3.11"
+            cuda: "12.8.1-cudnn"
+
+  # Verify that the expected images were built
+  - label: "Verify cuda images"
+    key: verify-cuda-images
+    command: |
+      tag="${RAYCI_WORK_REPO}:${RAYCI_BUILD_ID}"
+      expected="array-test-py3.12-cu13.0.0-cudnn array-test-py3.11-cu12.8.1-cudnn"
+      for name in $expected; do
+        echo "Checking ${tag}-${name} ..."
+        docker manifest inspect "${tag}-${name}"
+      done
+    depends_on:
+      - array-cuda-build
+      - forge
+
   # Depends on only the linux-1 and macos-1 variants
   - label: "Depends check only"
     key: depends-on-variants
@@ -35,4 +64,6 @@ steps:
     tags: always
     key: final-validation
     command: echo "Trigger check"
-    depends_on: depends-on-variants
+    depends_on:
+      - depends-on-variants
+      - verify-cuda-images

.buildkite/array-test.wanda.yaml

@@ -0,0 +1,7 @@
+name: "array-test-py$PYTHON_VERSION-cu$CUDA_VERSION"
+froms:
+  - ubuntu:22.04
+dockerfile: .buildkite/array-test.Dockerfile
+build_args:
+  - PYTHON_VERSION
+  - CUDA_VERSION

raycicmd/array_config.go

@@ -7,9 +7,11 @@ import (
 	"strings"
 )
 
-// arrayConfig represents a parsed array definition.
+// arrayConfig represents a parsed array definition. After expansion,
+// elements holds the final list (post-adjustments).
 type arrayConfig struct {
-	dims map[string][]string // dimension name -> values
+	dims     map[string][]string // dimension name -> values
+	elements []*arrayElement     // populated during expansion
 }
 
 // arrayElement is one combination from expand().

raycicmd/array_expand.go

@@ -2,6 +2,7 @@ package raycicmd
 
 import (
 	"fmt"
+	"maps"
 	"slices"
 )
 
@@ -113,6 +114,27 @@ func expandSingleArrayStep(
 		)
 	}
 
+	// Parse and apply adjustments if present.
+	var adjustments []*arrayAdjustment
+	if arrayMap, ok := arrayDef.(map[string]any); ok {
+		if adjDef, ok := arrayMap["adjustments"]; ok {
+			adjustments, err = parseArrayAdjustments(adjDef)
+			if err != nil {
+				return nil, fmt.Errorf(
+					"step %q: %w", baseKey, err,
+				)
+			}
+		}
+	}
+
+	elements, err = applyAdjustments(elements, adjustments, cfg)
+	if err != nil {
+		return nil, fmt.Errorf(
+			"step %q: %w", baseKey, err,
+		)
+	}
+
+	cfg.elements = elements
 	configs[baseKey] = cfg
 
 	var result []map[string]any
@@ -137,6 +159,78 @@ func expandSingleArrayStep(
 	return result, nil
 }
 
+// applyAdjustments processes adjustments in two passes: first appends
+// additions, then filters out skipped elements.
+func applyAdjustments(
+	elements []*arrayElement,
+	adjustments []*arrayAdjustment,
+	cfg *arrayConfig,
+) ([]*arrayElement, error) {
+	// Pass 1: append additions.
+	for _, adj := range adjustments {
+		if adj.skip {
+			continue
+		}
+		for dim := range cfg.dims {
+			if _, ok := adj.with[dim]; !ok {
+				return nil, fmt.Errorf(
+					"addition adjustment must specify all "+
+						"dimensions; missing %q", dim,
+				)
+			}
+		}
+		elements = append(elements, &arrayElement{
+			values: maps.Clone(adj.with),
+		})
+	}
+
+	// Pass 2: collect skip adjustments, validate, then remove.
+	var skipAdjs []*arrayAdjustment
+	for _, adj := range adjustments {
+		if adj.skip {
+			skipAdjs = append(skipAdjs, adj)
+		}
+	}
+	for _, adj := range skipAdjs {
+		if !hasElement(elements, adj.with) {
+			return nil, fmt.Errorf(
+				"skip adjustment with=%v matches no element",
+				adj.with,
+			)
+		}
+	}
+	elements = slices.DeleteFunc(elements, func(e *arrayElement) bool {
+		for _, adj := range skipAdjs {
+			if elemMatchesWith(e, adj.with) {
+				return true
+			}
+		}
+		return false
+	})
+
+	return elements, nil
+}
+
+func hasElement(
+	elements []*arrayElement, with map[string]string,
+) bool {
+	for _, elem := range elements {
+		if elemMatchesWith(elem, with) {
+			return true
+		}
+	}
+	return false
+}
+
+func elemMatchesWith(elem *arrayElement, with map[string]string) bool {
+	for dim, val := range with {
+		if elem.values[dim] != val {
+			return false
+		}
+	}
+	return true
+}
+
 // arraySelector represents a dependency selector with optional array filter.
 type arraySelector struct {
 	key    string              // step identifier: "key" for command steps, "name" for wanda steps
@@ -271,7 +365,8 @@ func resolveDependsOn(dependsOn any, configs map[string]*arrayConfig) ([]string,
 	return result, nil
 }
 
-// resolveArraySelector resolves a selector against an array config to concrete step keys.
+// 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) {
 	for dim := range sel.filter {
 		if _, ok := cfg.dims[dim]; !ok {
@@ -282,9 +377,8 @@ func resolveArraySelector(sel *arraySelector, cfg *arrayConfig) ([]string, error
 		}
 	}
 
-	elements := cfg.expand()
 	var matches []string
-	for _, elem := range elements {
+	for _, elem := range cfg.elements {
 		if elem.matchesSelector(sel) {
 			matches = append(matches, elem.generateKey(sel.key))
 		}