From 19c69850b83a7ec788e8de881d1067bab03421bf Mon Sep 17 00:00:00 2001
From: stackedsax <alex.scammon@gresearch.co.uk>
Date: Sat, 18 Jul 2026 09:35:18 -0700
Subject: [PATCH] Round-trip multi-role JobSets; consolidate k8s helpers

Kueue and YuniKorn emit multi-role jobs as a JobSet (one replicatedJob per
role), but the parsers only understood a single batch/v1 Job, so those jobs
could not round-trip back to SPLAT. Add JobSet parsing:

- internal/jobset.ToTasks inverts the emitters' replicatedJobs into SPLAT
  tasks (execution incl. inlined-script detection, resources, placement,
  replicas) plus shared job volumes and the max BackoffLimit.
- The Kueue and YuniKorn parsers detect kind: JobSet and delegate; YuniKorn
  additionally recovers gang scheduling from the task-groups annotation.

Tech debt folded in while here:
- Move the duplicated tolerations helper and the Volcano parser's volume
  recovery into k8senc (Tolerations, VolumesFromPod, SortVolumes); Volcano
  and the JobSet parser now share them.
- Trigger the dry-run workflow on parser/** and k8senc/** changes too, not
  just emitters.

Tests: JobSet parse for Kueue (container + inlined-script roles) and YuniKorn
(gang from task-groups), plus an emit->parse round-trip guarding against
emitter/parser drift.

Co-Authored-By: Claude Opus 4.8 (1M context) <noreply@anthropic.com>
Claude-Session: https://claude.ai/code/session_016PnTz6Zxqa4jHocK8kCbyx
---
 .github/workflows/dryrun.yml            |   2 +
 internal/jobset/tosplat.go              |  92 +++++++++++++
 internal/k8senc/k8senc.go               |  60 +++++++++
 internal/parser/armada/parser.go        |  11 +-
 internal/parser/kueue/parser.go         |  32 +++++
 internal/parser/kueue/parser_test.go    | 164 ++++++++++++++++++++++++
 internal/parser/volcano/parser.go       |  67 +---------
 internal/parser/yunikorn/parser.go      |  41 +++++-
 internal/parser/yunikorn/parser_test.go |  78 +++++++++++
 9 files changed, 471 insertions(+), 76 deletions(-)
 create mode 100644 internal/jobset/tosplat.go

diff --git a/.github/workflows/dryrun.yml b/.github/workflows/dryrun.yml
index 56803f1..ce31d57 100644
--- a/.github/workflows/dryrun.yml
+++ b/.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'
diff --git a/internal/jobset/tosplat.go b/internal/jobset/tosplat.go
new file mode 100644
index 0000000..d6ab688
--- /dev/null
+++ b/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
+}
diff --git a/internal/k8senc/k8senc.go b/internal/k8senc/k8senc.go
index 8b275cd..3c2ba68 100644
--- a/internal/k8senc/k8senc.go
+++ b/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.
diff --git a/internal/parser/armada/parser.go b/internal/parser/armada/parser.go
index b13593e..e835d74 100644
--- a/internal/parser/armada/parser.go
+++ b/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
-}
diff --git a/internal/parser/kueue/parser.go b/internal/parser/kueue/parser.go
index 2fcafd2..2a786db 100644
--- a/internal/parser/kueue/parser.go
+++ b/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 {
diff --git a/internal/parser/kueue/parser_test.go b/internal/parser/kueue/parser_test.go
index 608da90..ade9c82 100644
--- a/internal/parser/kueue/parser_test.go
+++ b/internal/parser/kueue/parser_test.go
@@ -5,9 +5,66 @@ import (
 	"testing"
 	"time"
 
+	kueueemit "github.com/InsightSoftmax/BAMMM/internal/emitter/kueue"
 	"github.com/InsightSoftmax/BAMMM/internal/parser/kueue"
+	"github.com/InsightSoftmax/BAMMM/internal/splat"
 )
 
+// TestRoundTrip_MultiRoleJobSet guards the emitter and parser against drift: a
+// multi-role job emitted as a JobSet must parse back to the same roles.
+func TestRoundTrip_MultiRoleJobSet(t *testing.T) {
+	orig := &splat.Job{
+		APIVersion: splat.APIVersion,
+		Kind:       splat.Kind,
+		Metadata:   splat.Metadata{Name: "sim"},
+		Spec: splat.Spec{
+			Schedule: splat.Schedule{Queue: "gpu"},
+			Tasks: []splat.Task{
+				{
+					Name:      "driver",
+					Replicas:  1,
+					Resources: &splat.Resources{CPUsPerTask: 4},
+					Execution: &splat.Execution{Container: &splat.ContainerExecution{Image: "sim:1", Command: []string{"drive"}}},
+				},
+				{
+					Name:      "compute",
+					Replicas:  4,
+					Resources: &splat.Resources{CPUsPerTask: 32, GPU: &splat.GPURequest{Count: 4}},
+					Execution: &splat.Execution{Container: &splat.ContainerExecution{Image: "sim:1", Command: []string{"work"}}},
+				},
+			},
+		},
+	}
+
+	out, err := kueueemit.Emit(orig)
+	if err != nil {
+		t.Fatalf("Emit: %v", err)
+	}
+	job, err := kueue.Parse(out)
+	if err != nil {
+		t.Fatalf("Parse: %v", err)
+	}
+
+	if job.Spec.Schedule.Queue != "gpu" {
+		t.Errorf("queue: got %q", job.Spec.Schedule.Queue)
+	}
+	if len(job.Spec.Tasks) != 2 {
+		t.Fatalf("tasks: got %d want 2", len(job.Spec.Tasks))
+	}
+	for i, want := range orig.Spec.Tasks {
+		got := job.Spec.Tasks[i]
+		if got.Name != want.Name || got.Replicas != want.Replicas {
+			t.Errorf("task %d: got %s/%d want %s/%d", i, got.Name, got.Replicas, want.Name, want.Replicas)
+		}
+		if got.Execution == nil || got.Execution.Container == nil || got.Execution.Container.Image != "sim:1" {
+			t.Errorf("task %d execution: %+v", i, got.Execution)
+		}
+	}
+	if job.Spec.Tasks[1].Resources == nil || job.Spec.Tasks[1].Resources.GPU == nil || job.Spec.Tasks[1].Resources.GPU.Count != 4 {
+		t.Errorf("compute gpu: %+v", job.Spec.Tasks[1].Resources)
+	}
+}
+
 const scriptJob = `
 apiVersion: v1
 kind: ConfigMap
@@ -59,6 +116,113 @@ spec:
       restartPolicy: Never
 `
 
+const jobSetManifest = `
+apiVersion: jobset.x-k8s.io/v1alpha2
+kind: JobSet
+metadata:
+  name: ml-training
+  namespace: research
+  labels:
+    kueue.x-k8s.io/queue-name: gpu-hpc
+    team: nlp
+spec:
+  replicatedJobs:
+    - name: driver
+      replicas: 1
+      template:
+        spec:
+          backoffLimit: 3
+          template:
+            spec:
+              containers:
+                - name: driver
+                  image: ray:2.9
+                  command: ["ray", "start", "--head"]
+                  resources:
+                    requests:
+                      cpu: "4"
+                      memory: 8Gi
+              restartPolicy: Never
+    - name: worker
+      replicas: 1
+      template:
+        spec:
+          parallelism: 3
+          completions: 3
+          backoffLimit: 3
+          template:
+            spec:
+              nodeSelector:
+                gpu: "true"
+              containers:
+                - name: worker
+                  image: ubuntu:22.04
+                  command: ["/bin/bash", "-c"]
+                  args: ["echo hi && sleep 1"]
+                  resources:
+                    requests:
+                      cpu: "8"
+                      memory: 16Gi
+                      nvidia.com/gpu: "4"
+              restartPolicy: Never
+`
+
+func TestParse_JobSet(t *testing.T) {
+	job, err := kueue.Parse([]byte(jobSetManifest))
+	if err != nil {
+		t.Fatalf("Parse: %v", err)
+	}
+
+	if job.Metadata.Name != "ml-training" {
+		t.Errorf("name: got %q", job.Metadata.Name)
+	}
+	if job.Metadata.Annotations["bammm.io/namespace"] != "research" {
+		t.Errorf("namespace: got %q", job.Metadata.Annotations["bammm.io/namespace"])
+	}
+	if job.Spec.Schedule.Queue != "gpu-hpc" {
+		t.Errorf("queue: got %q", job.Spec.Schedule.Queue)
+	}
+	if job.Metadata.Labels["team"] != "nlp" {
+		t.Errorf("label team: got %q", job.Metadata.Labels["team"])
+	}
+	if job.Spec.Lifecycle.MaxRetries != 3 {
+		t.Errorf("maxRetries: got %d want 3", job.Spec.Lifecycle.MaxRetries)
+	}
+
+	if len(job.Spec.Tasks) != 2 {
+		t.Fatalf("tasks: got %d want 2", len(job.Spec.Tasks))
+	}
+
+	driver := job.Spec.Tasks[0]
+	if driver.Name != "driver" || driver.Replicas != 1 {
+		t.Errorf("driver: name=%q replicas=%d", driver.Name, driver.Replicas)
+	}
+	if driver.Execution == nil || driver.Execution.Container == nil || driver.Execution.Container.Image != "ray:2.9" {
+		t.Errorf("driver execution: %+v", driver.Execution)
+	}
+	if driver.Resources == nil || driver.Resources.CPUsPerTask != 4 {
+		t.Errorf("driver cpus: %+v", driver.Resources)
+	}
+
+	worker := job.Spec.Tasks[1]
+	if worker.Name != "worker" || worker.Replicas != 3 {
+		t.Errorf("worker: name=%q replicas=%d want worker/3", worker.Name, worker.Replicas)
+	}
+	// An inlined /bin/bash -c on the placeholder image round-trips to a script.
+	if worker.Execution == nil || worker.Execution.Container != nil {
+		t.Errorf("worker should be a script, got %+v", worker.Execution)
+	}
+	if worker.Execution.Script != "echo hi && sleep 1" {
+		t.Errorf("worker script: got %q", worker.Execution.Script)
+	}
+	if worker.Resources == nil || worker.Resources.GPU == nil || worker.Resources.GPU.Count != 4 {
+		t.Errorf("worker gpu: %+v", worker.Resources)
+	}
+	if worker.Placement == nil || worker.Placement.NodeSelector["gpu"] != "true" {
+		t.Errorf("worker nodeSelector: %+v", worker.Placement)
+	}
+}
+
 func TestParse_ScriptJob(t *testing.T) {
 	job, err := kueue.Parse([]byte(scriptJob))
 	if err != nil {
diff --git a/internal/parser/volcano/parser.go b/internal/parser/volcano/parser.go
index 914ee66..c7d750e 100644
--- a/internal/parser/volcano/parser.go
+++ b/internal/parser/volcano/parser.go
@@ -7,7 +7,6 @@ import (
 	"fmt"
 	"time"
 
-	corev1 "k8s.io/api/core/v1"
 	"sigs.k8s.io/yaml"
 
 	"github.com/InsightSoftmax/BAMMM/internal/k8senc"
@@ -87,7 +86,7 @@ func Parse(data []byte) (*splat.Job, error) {
 	for i := range vc.Spec.Tasks {
 		job.Spec.Tasks = append(job.Spec.Tasks, taskFromVolcano(&vc.Spec.Tasks[i], volumes))
 	}
-	job.Spec.Volumes = collectVolumes(volumes)
+	job.Spec.Volumes = k8senc.SortVolumes(volumes)
 
 	applyExtensions(job, &vc.Spec)
 	return job, nil
@@ -118,64 +117,12 @@ func taskFromVolcano(vt *volcanotypes.Task, volumes map[string]splat.Volume) spl
 	task.Execution = exec
 
 	if len(pod.Tolerations) > 0 {
-		task.Placement = &splat.Placement{Tolerations: tolerations(pod.Tolerations)}
+		task.Placement = &splat.Placement{Tolerations: k8senc.Tolerations(pod.Tolerations)}
 	}
-	recordVolumes(&c, pod.Volumes, volumes)
+	k8senc.VolumesFromPod(&c, &pod, volumes)
 	return task
 }
 
-// recordVolumes maps a container's volume mounts (paired with the pod's volume
-// sources) into shared SPLAT volumes, keyed by name so identical per-task
-// volumes are deduplicated at the job level.
-func recordVolumes(c *corev1.Container, podVolumes []corev1.Volume, out map[string]splat.Volume) {
-	sources := map[string]corev1.Volume{}
-	for _, v := range podVolumes {
-		sources[v.Name] = v
-	}
-	for _, m := range c.VolumeMounts {
-		if _, seen := out[m.Name]; seen {
-			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
-			}
-		}
-		out[m.Name] = v
-	}
-}
-
-func collectVolumes(m map[string]splat.Volume) []splat.Volume {
-	if len(m) == 0 {
-		return nil
-	}
-	names := make([]string, 0, len(m))
-	for n := range m {
-		names = append(names, n)
-	}
-	// Stable order by name.
-	for i := 1; i < len(names); i++ {
-		for j := i; j > 0 && names[j] < names[j-1]; j-- {
-			names[j], names[j-1] = names[j-1], names[j]
-		}
-	}
-	out := make([]splat.Volume, 0, len(names))
-	for _, n := range names {
-		out = append(out, m[n])
-	}
-	return out
-}
-
 func applyExtensions(job *splat.Job, spec *volcanotypes.JobSpec) {
 	ext := map[string]interface{}{}
 	if spec.SchedulerName != "" {
@@ -191,11 +138,3 @@ func applyExtensions(job *splat.Job, spec *volcanotypes.JobSpec) {
 		job.Spec.Extensions.Volcano = ext
 	}
 }
-
-func tolerations(ts []corev1.Toleration) []interface{} {
-	out := make([]interface{}, 0, len(ts))
-	for i := range ts {
-		out = append(out, ts[i])
-	}
-	return out
-}
diff --git a/internal/parser/yunikorn/parser.go b/internal/parser/yunikorn/parser.go
index 56ac9b9..25006fe 100644
--- a/internal/parser/yunikorn/parser.go
+++ b/internal/parser/yunikorn/parser.go
@@ -12,6 +12,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"
@@ -38,13 +39,22 @@ const (
 func Parse(data []byte) (*splat.Job, error) {
 	var k8sJob *batchv1.Job
 	for _, doc := range k8senc.SplitYAMLDocs(data) {
-		if k8senc.DocumentKind(doc) == "Job" {
+		switch k8senc.DocumentKind(doc) {
+		case jobset.Kind:
+			var js jobset.JobSet
+			if err := yaml.Unmarshal(doc, &js); err != nil {
+				return nil, fmt.Errorf("yunikorn: unmarshal JobSet: %w", err)
+			}
+			return jobFromJobSet(&js), nil
+		case "Job":
+			if k8sJob != nil {
+				continue
+			}
 			var j batchv1.Job
 			if err := yaml.Unmarshal(doc, &j); err != nil {
 				return nil, fmt.Errorf("yunikorn: unmarshal Job: %w", err)
 			}
 			k8sJob = &j
-			break
 		}
 	}
 	if k8sJob == nil {
@@ -64,6 +74,33 @@ func Parse(data []byte) (*splat.Job, error) {
 	return job, nil
 }
 
+// jobFromJobSet builds a multi-role SPLAT job from a YuniKorn-scheduled JobSet,
+// recovering gang scheduling from the task-groups annotation the emitter stamps
+// on every pod template.
+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": "yunikorn"}
+	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
+	}
+	for i := range js.Spec.ReplicatedJobs {
+		applyGang(job, js.Spec.ReplicatedJobs[i].Template.Spec.Template.Annotations)
+		if job.Spec.Gang != nil {
+			break
+		}
+	}
+	return job
+}
+
 func applyLabels(job *splat.Job, labels map[string]string) {
 	for k, v := range labels {
 		switch k {
diff --git a/internal/parser/yunikorn/parser_test.go b/internal/parser/yunikorn/parser_test.go
index 2acde25..ba4f97f 100644
--- a/internal/parser/yunikorn/parser_test.go
+++ b/internal/parser/yunikorn/parser_test.go
@@ -6,6 +6,84 @@ import (
 	"github.com/InsightSoftmax/BAMMM/internal/parser/yunikorn"
 )
 
+const jobSetYAML = `
+apiVersion: jobset.x-k8s.io/v1alpha2
+kind: JobSet
+metadata:
+  name: gang-training
+  namespace: default
+  labels:
+    yunikorn.apache.org/app-id: gang-training
+    yunikorn.apache.org/queue: root.gpu
+spec:
+  replicatedJobs:
+    - name: ps
+      replicas: 1
+      template:
+        spec:
+          template:
+            metadata:
+              annotations:
+                yunikorn.apache.org/task-group-name: ps
+                yunikorn.apache.org/task-groups: '[{"name":"ps","minMember":1},{"name":"worker","minMember":2}]'
+            spec:
+              schedulerName: yunikorn
+              containers:
+                - name: ps
+                  image: tf:2.15
+                  resources:
+                    requests:
+                      cpu: "2"
+              restartPolicy: Never
+    - name: worker
+      replicas: 1
+      template:
+        spec:
+          parallelism: 2
+          completions: 2
+          template:
+            metadata:
+              annotations:
+                yunikorn.apache.org/task-group-name: worker
+                yunikorn.apache.org/task-groups: '[{"name":"ps","minMember":1},{"name":"worker","minMember":2}]'
+            spec:
+              schedulerName: yunikorn
+              containers:
+                - name: worker
+                  image: tf:2.15
+                  resources:
+                    requests:
+                      cpu: "8"
+                      nvidia.com/gpu: "1"
+              restartPolicy: Never
+`
+
+func TestParse_JobSet_Gang(t *testing.T) {
+	job, err := yunikorn.Parse([]byte(jobSetYAML))
+	if err != nil {
+		t.Fatalf("Parse: %v", err)
+	}
+	if job.Spec.Schedule.Queue != "gpu" {
+		t.Errorf("queue: got %q want gpu", job.Spec.Schedule.Queue)
+	}
+	if len(job.Spec.Tasks) != 2 {
+		t.Fatalf("tasks: got %d want 2", len(job.Spec.Tasks))
+	}
+	if job.Spec.Tasks[0].Name != "ps" || job.Spec.Tasks[1].Name != "worker" {
+		t.Errorf("task names: got %q, %q", job.Spec.Tasks[0].Name, job.Spec.Tasks[1].Name)
+	}
+	if job.Spec.Tasks[1].Replicas != 2 {
+		t.Errorf("worker replicas: got %d want 2", job.Spec.Tasks[1].Replicas)
+	}
+	if job.Spec.Tasks[1].Resources == nil || job.Spec.Tasks[1].Resources.GPU == nil || job.Spec.Tasks[1].Resources.GPU.Count != 1 {
+		t.Errorf("worker gpu: %+v", job.Spec.Tasks[1].Resources)
+	}
+	// Gang minAvailable is the sum of the task groups' minMembers (1 + 2).
+	if job.Spec.Gang == nil || job.Spec.Gang.MinAvailable != 3 {
+		t.Errorf("gang: %+v want MinAvailable=3", job.Spec.Gang)
+	}
+}
+
 const jobYAML = `
 apiVersion: batch/v1
 kind: Job