Round-trip multi-role JobSets + k8s helper consolidation

.github/workflows/dryrun.yml

@@ -8,7 +8,9 @@ on:
   pull_request:
     paths:
       - 'internal/emitter/**'
+      - 'internal/parser/**'
       - 'internal/splat/**'
+      - 'internal/k8senc/**'
       - 'internal/*/types.go'
       - 'scripts/dryrun/**'
       - '.github/workflows/dryrun.yml'

internal/jobset/tosplat.go

@@ -0,0 +1,92 @@
+package jobset
+
+import (
+	corev1 "k8s.io/api/core/v1"
+
+	"github.com/InsightSoftmax/BAMMM/internal/k8senc"
+	"github.com/InsightSoftmax/BAMMM/internal/splat"
+)
+
+// DefaultImage is the placeholder image the Kueue/YuniKorn JobSet emitters use
+// for script and executable tasks that carry no source container. The parser
+// uses it to recognize that such a container is really an inlined script.
+const DefaultImage = "ubuntu:22.04"
+
+// ToTasks converts a JobSet's replicatedJobs into SPLAT tasks, plus the shared
+// job-level volumes and the maximum in-pod retry (BackoffLimit) observed. It is
+// the inverse of the Kueue and YuniKorn JobSet emitters; scheduler-specific
+// metadata (queue labels, gang annotations) stays with the caller.
+func ToTasks(js *JobSet) (tasks []splat.Task, volumes []splat.Volume, maxRetries int) {
+	seen := map[string]splat.Volume{}
+	for i := range js.Spec.ReplicatedJobs {
+		rj := &js.Spec.ReplicatedJobs[i]
+		spec := &rj.Template.Spec
+		if bl := spec.BackoffLimit; bl != nil && int(*bl) > maxRetries {
+			maxRetries = int(*bl)
+		}
+
+		task := splat.Task{Name: rj.Name, Replicas: replicasOf(spec.Parallelism, spec.Completions)}
+		pod := &spec.Template.Spec
+		if len(pod.Containers) > 0 {
+			c := pod.Containers[0]
+			task.Resources = k8senc.ResourcesFromContainer(&c)
+			task.Execution = executionOf(&c)
+			k8senc.VolumesFromPod(&c, pod, seen)
+		}
+		task.Placement = placementOf(pod)
+		tasks = append(tasks, task)
+	}
+	return tasks, k8senc.SortVolumes(seen), maxRetries
+}
+
+// replicasOf recovers a task's replica count from the wrapped Job's parallelism
+// (the emitters set parallelism/completions only when replicas exceed 1).
+func replicasOf(parallelism, completions *int32) int {
+	switch {
+	case parallelism != nil && *parallelism > 0:
+		return int(*parallelism)
+	case completions != nil && *completions > 0:
+		return int(*completions)
+	default:
+		return 1
+	}
+}
+
+// executionOf inverts the emitters' container: an inlined "/bin/bash -c <script>"
+// on the placeholder image round-trips back to a script; anything else stays a
+// container execution.
+func executionOf(c *corev1.Container) *splat.Execution {
+	e := &splat.Execution{}
+	if c.Image == DefaultImage && len(c.Command) == 2 &&
+		c.Command[0] == "/bin/bash" && c.Command[1] == "-c" && len(c.Args) == 1 {
+		e.Script = c.Args[0]
+	} else {
+		e.Container = &splat.ContainerExecution{Image: c.Image, Command: c.Command, Args: c.Args}
+	}
+	if c.WorkingDir != "" {
+		e.WorkingDir = c.WorkingDir
+	}
+	if env := k8senc.EnvMap(c.Env); len(env) > 0 {
+		e.Environment.Vars = env
+	}
+	return e
+}
+
+// placementOf recovers node selectors and tolerations, returning nil when the
+// pod pins neither.
+func placementOf(pod *corev1.PodSpec) *splat.Placement {
+	pl := &splat.Placement{}
+	set := false
+	if len(pod.NodeSelector) > 0 {
+		pl.NodeSelector = pod.NodeSelector
+		set = true
+	}
+	if len(pod.Tolerations) > 0 {
+		pl.Tolerations = k8senc.Tolerations(pod.Tolerations)
+		set = true
+	}
+	if !set {
+		return nil
+	}
+	return pl
+}

internal/k8senc/k8senc.go

@@ -184,6 +184,66 @@ func AttachVolumes(pod *corev1.PodSpec, c *corev1.Container, volumes []splat.Vol
 	}
 }
 
+// VolumesFromPod recovers job-level SPLAT volumes from a container's mounts
+// paired with the pod's volume sources, writing into seen (keyed by name so
+// identical per-role volumes dedupe across a multi-role job). Shared by the
+// Volcano parser and the JobSet (Kueue/YuniKorn) parsers.
+func VolumesFromPod(c *corev1.Container, pod *corev1.PodSpec, seen map[string]splat.Volume) {
+	sources := map[string]corev1.Volume{}
+	for _, v := range pod.Volumes {
+		sources[v.Name] = v
+	}
+	for _, m := range c.VolumeMounts {
+		if _, ok := seen[m.Name]; ok {
+			continue
+		}
+		v := splat.Volume{Name: m.Name, MountPath: m.MountPath, ReadOnly: m.ReadOnly}
+		if src, ok := sources[m.Name]; ok {
+			switch {
+			case src.PersistentVolumeClaim != nil:
+				v.PVC = src.PersistentVolumeClaim.ClaimName
+			case src.ConfigMap != nil:
+				v.ConfigMap = src.ConfigMap.Name
+			case src.Secret != nil:
+				v.Secret = src.Secret.SecretName
+			case src.HostPath != nil:
+				v.HostPath = src.HostPath.Path
+			case src.EmptyDir != nil:
+				v.EmptyDir = true
+			}
+		}
+		seen[m.Name] = v
+	}
+}
+
+// SortVolumes returns the volumes collected in seen, ordered by name.
+func SortVolumes(seen map[string]splat.Volume) []splat.Volume {
+	if len(seen) == 0 {
+		return nil
+	}
+	names := make([]string, 0, len(seen))
+	for n := range seen {
+		names = append(names, n)
+	}
+	sort.Strings(names)
+	out := make([]splat.Volume, 0, len(names))
+	for _, n := range names {
+		out = append(out, seen[n])
+	}
+	return out
+}
+
+// Tolerations converts core/v1 pod tolerations into the loosely-typed slice
+// SPLAT stores in Placement.Tolerations (which round-trips back via ConvertVia).
+// Shared by the Volcano and Armada parsers.
+func Tolerations(ts []corev1.Toleration) []interface{} {
+	out := make([]interface{}, 0, len(ts))
+	for i := range ts {
+		out = append(out, ts[i])
+	}
+	return out
+}
+
 // ConvertVia re-materializes a loosely-typed value (e.g. map[string]interface{}
 // from a YAML round-trip, or a concrete struct) into dst by marshaling through
 // YAML.

internal/parser/armada/parser.go

@@ -7,7 +7,6 @@ import (
 	"math"
 	"strconv"
 
-	corev1 "k8s.io/api/core/v1"
 	"sigs.k8s.io/yaml"
 
 	armadatypes "github.com/InsightSoftmax/BAMMM/internal/armada"
@@ -96,7 +95,7 @@ func taskFromJob(aj *armadatypes.Job, index int) splat.Task {
 	task.Execution = exec
 
 	if len(aj.PodSpec.Tolerations) > 0 {
-		task.Placement = &splat.Placement{Tolerations: tolerations(aj.PodSpec.Tolerations)}
+		task.Placement = &splat.Placement{Tolerations: k8senc.Tolerations(aj.PodSpec.Tolerations)}
 	}
 	return task
 }
@@ -183,11 +182,3 @@ func applyExtensions(job *splat.Job, req *armadatypes.Request, rawPriority float
 		job.Spec.Extensions.Armada = ext
 	}
 }
-
-func tolerations(ts []corev1.Toleration) []interface{} {
-	out := make([]interface{}, 0, len(ts))
-	for i := range ts {
-		out = append(out, ts[i])
-	}
-	return out
-}

internal/parser/kueue/parser.go

@@ -10,6 +10,7 @@ import (
 	corev1 "k8s.io/api/core/v1"
 	"sigs.k8s.io/yaml"
 
+	"github.com/InsightSoftmax/BAMMM/internal/jobset"
 	"github.com/InsightSoftmax/BAMMM/internal/k8senc"
 	"github.com/InsightSoftmax/BAMMM/internal/parser"
 	"github.com/InsightSoftmax/BAMMM/internal/splat"
@@ -31,6 +32,18 @@ const queueNameLabel = "kueue.x-k8s.io/queue-name"
 func Parse(data []byte) (*splat.Job, error) {
 	docs := k8senc.SplitYAMLDocs(data)
 
+	// A multi-role job is emitted as a JobSet; if one is present it is the whole
+	// job, so parse it and return before looking for a single batch/v1 Job.
+	for _, doc := range docs {
+		if k8senc.DocumentKind(doc) == jobset.Kind {
+			var js jobset.JobSet
+			if err := yaml.Unmarshal(doc, &js); err != nil {
+				return nil, fmt.Errorf("kueue: unmarshal JobSet: %w", err)
+			}
+			return jobFromJobSet(&js), nil
+		}
+	}
+
 	var k8sJob *batchv1.Job
 	configMaps := map[string]map[string]string{} // name -> data
 
@@ -73,6 +86,25 @@ func Parse(data []byte) (*splat.Job, error) {
 	return job, nil
 }
 
+// jobFromJobSet builds a multi-role SPLAT job from a Kueue-admitted JobSet.
+func jobFromJobSet(js *jobset.JobSet) *splat.Job {
+	job := &splat.Job{APIVersion: splat.APIVersion, Kind: splat.Kind}
+	job.Metadata.Name = js.Name
+	job.Metadata.Annotations = map[string]string{"bammm.io/source-format": "kueue"}
+	if js.Namespace != "" && js.Namespace != "default" {
+		job.Metadata.Annotations["bammm.io/namespace"] = js.Namespace
+	}
+	applyLabels(job, js.Labels)
+
+	tasks, volumes, maxRetries := jobset.ToTasks(js)
+	job.Spec.Tasks = tasks
+	job.Spec.Volumes = volumes
+	if maxRetries > 0 {
+		job.Spec.Lifecycle.MaxRetries = maxRetries
+	}
+	return job
+}
+
 func applyLabels(job *splat.Job, labels map[string]string) {
 	for k, v := range labels {
 		switch k {