Skip to content

Commit 4ce2802

Browse files
feat(raycicmd): add matrix expansion types and selector resolution
Rayci will expand matrix into individual Buildkite steps, allowing for fine-grained depends_on targeting of specific matrix combinations. ## Matrix Formats Supported Simple array (anonymous dimension): matrix: ["3.10", "3.11", "3.12"] Named dimensions: matrix: setup: python: ["3.10", "3.11"] cuda: ["12.1.1", "12.8.1"] ## Key Generation Expanded keys follow deterministic format: {base-key}-{dim1}{val1}-{dim2}{val2} Examples: ray-build + python=3.11, cuda=12.1.1 → ray-build-cuda1211-python311 ray-build + ["linux", "darwin"] → ray-build-linux, ray-build-darwin ## Selector Syntax for depends_on Simple reference (waits for ALL instances via meta-step): depends_on: ray-build Selector syntax (waits for matching instances only): depends_on: - key: ray-build matrix: cuda: "12.1.1" # matches all python versions with cuda 12.1.1 Multi-value selector: depends_on: - key: ray-build matrix: python: ["3.10", "3.11"] # matches specific python versions Topic: matrix-expand Signed-off-by: andrew <andrew@anyscale.com>
1 parent 6165f02 commit 4ce2802

File tree

2 files changed

+835
-0
lines changed

2 files changed

+835
-0
lines changed

raycicmd/matrix_expand.go

Lines changed: 370 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -0,0 +1,370 @@
1+
package raycicmd
2+
3+
import (
4+
"fmt"
5+
"maps"
6+
"slices"
7+
"sort"
8+
"strings"
9+
)
10+
11+
// dimension is a named dimension in a matrix (e.g., "python", "cuda").
12+
type dimension string
13+
14+
// anonymousDim is used for simple array matrices: matrix: ["a", "b", "c"]
15+
const anonymousDim dimension = ""
16+
17+
// variant is a value within a dimension (e.g., "3.10", "12.1.1").
18+
type variant string
19+
20+
// matrixConfig represents an expanded step's matrix definition.
21+
type matrixConfig struct {
22+
Setup map[dimension][]variant
23+
}
24+
25+
// matrixInstance represents a single expanded matrix combination.
26+
type matrixInstance struct {
27+
Values map[dimension]variant
28+
}
29+
30+
// matrixSelector represents a dependency selector with optional matrix filter.
31+
type matrixSelector struct {
32+
Key string // base step key
33+
Matrix map[dimension][]variant // partial dimension constraints (nil = all), each dimension can match multiple values
34+
}
35+
36+
// parseMatrixConfig parses the matrix field from a step.
37+
// Named dimensions matrix
38+
//
39+
// matrix:
40+
// setup:
41+
// python: ["3.10", "3.11"]
42+
// cuda: ["12.1.1", "12.8.1"]
43+
//
44+
// Simple array matrix (no dimension name)
45+
//
46+
// matrix:
47+
// - "darwin"
48+
// - "Linux"
49+
// - "Windows"
50+
func parseMatrixConfig(v any) (*matrixConfig, error) {
51+
cfg := &matrixConfig{
52+
Setup: make(map[dimension][]variant),
53+
}
54+
55+
switch val := v.(type) {
56+
case []any:
57+
// Simple array: matrix: ["3.10", "3.11"]
58+
variants, err := anySliceToVariantSlice(val)
59+
if err != nil {
60+
return nil, fmt.Errorf("invalid matrix array: %w", err)
61+
}
62+
cfg.Setup[anonymousDim] = variants
63+
64+
case map[string]any:
65+
// Map with setup and optional adjustments
66+
if setup, ok := val["setup"]; ok {
67+
setupMap, ok := setup.(map[string]any)
68+
if !ok {
69+
return nil, fmt.Errorf("matrix.setup must be a map")
70+
}
71+
for dim, vals := range setupMap {
72+
valsSlice, ok := vals.([]any)
73+
if !ok {
74+
return nil, fmt.Errorf("matrix.setup.%s must be an array", dim)
75+
}
76+
variants, err := anySliceToVariantSlice(valsSlice)
77+
if err != nil {
78+
return nil, fmt.Errorf("invalid values for dimension %s: %w", dim, err)
79+
}
80+
cfg.Setup[dimension(dim)] = variants
81+
}
82+
} else {
83+
return nil, fmt.Errorf("matrix map must have 'setup' key")
84+
}
85+
86+
if _, ok := val["adjustments"]; ok {
87+
return nil, fmt.Errorf("matrix.adjustments is not supported")
88+
}
89+
90+
default:
91+
return nil, fmt.Errorf("matrix must be an array or map, got %T", v)
92+
}
93+
94+
return cfg, nil
95+
}
96+
97+
func (cfg *matrixConfig) sortedDimensions() []dimension {
98+
dims := make([]dimension, 0, len(cfg.Setup))
99+
for dim := range cfg.Setup {
100+
dims = append(dims, dim)
101+
}
102+
sort.Slice(dims, func(i, j int) bool { return dims[i] < dims[j] })
103+
return dims
104+
}
105+
106+
// expand generates all combinations from a matrix config.
107+
func (cfg *matrixConfig) expand() []*matrixInstance {
108+
if len(cfg.Setup) == 0 {
109+
return nil
110+
}
111+
112+
result := []*matrixInstance{{Values: make(map[dimension]variant)}}
113+
for _, dim := range cfg.sortedDimensions() {
114+
var expanded []*matrixInstance
115+
for _, inst := range result {
116+
for _, val := range cfg.Setup[dim] {
117+
newInst := &matrixInstance{Values: maps.Clone(inst.Values)}
118+
newInst.Values[dim] = val
119+
expanded = append(expanded, newInst)
120+
}
121+
}
122+
result = expanded
123+
}
124+
return result
125+
}
126+
127+
// generateKey creates the expanded key for an instance.
128+
// Format: {base-key}-{dim1}{val1}-{dim2}{val2} (dims sorted alphabetically)
129+
func (inst *matrixInstance) generateKey(baseKey string, cfg *matrixConfig) string {
130+
parts := []string{baseKey}
131+
for _, dim := range cfg.sortedDimensions() {
132+
parts = append(parts, sanitizeKeyPart(string(dim)+string(inst.Values[dim])))
133+
}
134+
return strings.Join(parts, "-")
135+
}
136+
137+
// sanitizeKeyPart removes invalid characters from a key part.
138+
// Buildkite keys allow: alphanumeric, underscores, dashes, colons.
139+
// We exclude dashes here since they're used as separators between parts.
140+
func sanitizeKeyPart(s string) string {
141+
return strings.Map(func(r rune) rune {
142+
if (r >= 'a' && r <= 'z') || (r >= 'A' && r <= 'Z') ||
143+
(r >= '0' && r <= '9') || r == '_' || r == ':' {
144+
return r
145+
}
146+
return -1 // drop character
147+
}, s)
148+
}
149+
150+
// generateTags creates auto-generated tags for matrix values.
151+
// These tags are used to tag the expanded steps and depends_on selectors.
152+
// Format: {dim1}-{val1}, {dim2}-{val2} (dims sorted alphabetically)
153+
func (inst *matrixInstance) generateTags() []string {
154+
var tags []string
155+
for dim, val := range inst.Values {
156+
if dim == anonymousDim {
157+
// Simple array: just use the value as tag (e.g., 3.10)
158+
tags = append(tags, string(val))
159+
} else {
160+
// Named dimension: use dimension-value format (e.g., python-3.11)
161+
tags = append(tags, string(dim)+"-"+string(val))
162+
}
163+
}
164+
sort.Strings(tags)
165+
return tags
166+
}
167+
168+
// substituteValues replaces {{matrix.X}} placeholders in any value.
169+
// It recursively traverses maps and slices, returning a new structure with
170+
// all string values substituted. Non-string leaf values are returned as-is.
171+
func (inst *matrixInstance) substituteValues(v any) any {
172+
switch val := v.(type) {
173+
case string:
174+
var replacerArgs []string
175+
for dim, dimVal := range inst.Values {
176+
if dim == anonymousDim {
177+
// Simple array: use {{matrix}}
178+
replacerArgs = append(replacerArgs, "{{matrix}}", string(dimVal))
179+
} else {
180+
// Named dimension: use {{matrix.dim}} (e.g., {{matrix.python}})
181+
replacerArgs = append(replacerArgs, "{{matrix."+string(dim)+"}}", string(dimVal))
182+
}
183+
}
184+
return strings.NewReplacer(replacerArgs...).Replace(val)
185+
186+
case map[string]any:
187+
result := make(map[string]any)
188+
for k, v := range val {
189+
result[k] = inst.substituteValues(v)
190+
}
191+
return result
192+
193+
case []any:
194+
result := make([]any, len(val))
195+
for i, v := range val {
196+
result[i] = inst.substituteValues(v)
197+
}
198+
return result
199+
200+
default:
201+
return v
202+
}
203+
}
204+
205+
// parseMatrixDependsOn parses a depends_on field into a list of selectors.
206+
//
207+
// Supported YAML formats:
208+
//
209+
// depends_on: step-key # single string
210+
// depends_on: [step-a, step-b] # string array
211+
// depends_on: # selector with matrix filter
212+
// - key: ray-build
213+
// matrix:
214+
// python: "3.11"
215+
func parseMatrixDependsOn(v any) ([]*matrixSelector, error) {
216+
if v == nil {
217+
return nil, nil
218+
}
219+
220+
switch val := v.(type) {
221+
case string:
222+
return []*matrixSelector{{Key: val}}, nil
223+
224+
case []string:
225+
var selectors []*matrixSelector
226+
for _, key := range val {
227+
selectors = append(selectors, &matrixSelector{Key: key})
228+
}
229+
return selectors, nil
230+
231+
case []any:
232+
var selectors []*matrixSelector
233+
for i, item := range val {
234+
switch itemVal := item.(type) {
235+
case string:
236+
selectors = append(selectors, &matrixSelector{Key: itemVal})
237+
238+
case map[string]any:
239+
sel, err := parseMatrixSelectorMap(itemVal)
240+
if err != nil {
241+
return nil, fmt.Errorf("depends_on[%d]: %w", i, err)
242+
}
243+
selectors = append(selectors, sel)
244+
245+
default:
246+
return nil, fmt.Errorf("depends_on[%d]: unexpected type %T", i, item)
247+
}
248+
}
249+
return selectors, nil
250+
251+
default:
252+
return nil, fmt.Errorf("depends_on must be string or array, got %T", v)
253+
}
254+
}
255+
256+
// parseMatrixSelectorMap parses a selector map into a matrixSelector.
257+
//
258+
// Example YAML:
259+
//
260+
// key: ray-build
261+
// matrix:
262+
// python: ["3.10", "3.11"]
263+
// cuda: "12.1.1"
264+
func parseMatrixSelectorMap(m map[string]any) (*matrixSelector, error) {
265+
sel := &matrixSelector{}
266+
267+
key, ok := m["key"]
268+
if !ok {
269+
return nil, fmt.Errorf("selector missing 'key' field")
270+
}
271+
keyStr, ok := key.(string)
272+
if !ok {
273+
return nil, fmt.Errorf("selector 'key' must be a string")
274+
}
275+
sel.Key = keyStr
276+
277+
if matrix, ok := m["matrix"]; ok {
278+
matrixMap, ok := matrix.(map[string]any)
279+
if !ok {
280+
return nil, fmt.Errorf("selector 'matrix' must be a map")
281+
}
282+
sel.Matrix = make(map[dimension][]variant)
283+
for k, v := range matrixMap {
284+
dim := dimension(k)
285+
switch val := v.(type) {
286+
case string:
287+
sel.Matrix[dim] = []variant{variant(val)}
288+
case []any:
289+
variants, err := anySliceToVariantSlice(val)
290+
if err != nil {
291+
return nil, fmt.Errorf("selector 'matrix.%s': %w", k, err)
292+
}
293+
sel.Matrix[dim] = variants
294+
case []string:
295+
variants := make([]variant, len(val))
296+
for i, s := range val {
297+
variants[i] = variant(s)
298+
}
299+
sel.Matrix[dim] = variants
300+
default:
301+
return nil, fmt.Errorf("selector 'matrix.%s' must be a string or array", k)
302+
}
303+
}
304+
}
305+
306+
return sel, nil
307+
}
308+
309+
// expand returns the expanded step keys that match this selector.
310+
//
311+
// For example, if ray-build expanded to [ray-build-python310, ray-build-python311]
312+
// and the selector filters to python: "3.11", this returns [ray-build-python311].
313+
func (sel *matrixSelector) expand(stepKeyToConfig map[string]*matrixConfig) ([]string, error) {
314+
cfg, ok := stepKeyToConfig[sel.Key]
315+
if !ok {
316+
return []string{sel.Key}, nil
317+
}
318+
319+
if sel.Matrix == nil {
320+
return []string{sel.Key}, nil
321+
}
322+
323+
for dim := range sel.Matrix {
324+
if _, ok := cfg.Setup[dim]; !ok {
325+
return nil, fmt.Errorf("selector dimension %q not found in matrix for %q", dim, sel.Key)
326+
}
327+
}
328+
329+
instances := cfg.expand()
330+
331+
var matches []string
332+
for _, inst := range instances {
333+
if inst.matches(sel) {
334+
matches = append(matches, inst.generateKey(sel.Key, cfg))
335+
}
336+
}
337+
338+
if len(matches) == 0 {
339+
return nil, fmt.Errorf("no matches for selector {key: %q, matrix: %v}", sel.Key, sel.Matrix)
340+
}
341+
342+
return matches, nil
343+
}
344+
345+
func (inst *matrixInstance) matches(sel *matrixSelector) bool {
346+
for dim, allowedVals := range sel.Matrix {
347+
if !slices.Contains(allowedVals, inst.Values[dim]) {
348+
return false
349+
}
350+
}
351+
return true
352+
}
353+
354+
// hasMatrixPlaceholder checks if a string contains any {{matrix...}} placeholder.
355+
func hasMatrixPlaceholder(s string) bool {
356+
return strings.Contains(s, "{{matrix")
357+
}
358+
359+
// anySliceToVariantSlice converts []any to []variant.
360+
func anySliceToVariantSlice(arr []any) ([]variant, error) {
361+
result := make([]variant, len(arr))
362+
for i, v := range arr {
363+
s, ok := v.(string)
364+
if !ok {
365+
return nil, fmt.Errorf("element %d is not a string: %T", i, v)
366+
}
367+
result[i] = variant(s)
368+
}
369+
return result, nil
370+
}

0 commit comments

Comments
 (0)