feat: manifest schema and generation

README.md

@@ -113,6 +113,36 @@ find the following examples of 64 GPU (8 node) runs:
 - [H100](./mpi-operator/nccl-test-distributed-h100-64-mpijob.yaml)
 - [H100 with SHARP](./mpi-operator/nccl-test-distributed-h100-64-sharp-mpijob.yaml)
 - [GB200 NVL72](./mpi-operator/nccl-test-distributed-gb200-nvl72-mpijob.yaml)
+- [GB200 NVL72 Multi-rack](./mpi-operator/nccl-test-distributed-gb200-128-multirack-mpijob.yaml)
+
+#### Generating Custom Manifests
+
+For custom configurations, use `nccl.nu` to generate manifests dynamically.
+Requires [Nushell](https://www.nushell.sh/).
+
+```bash
+# Generate manifest for 64 H100 GPUs
+./nccl.nu generate --gpu h100 --scale 64
+
+# Generate with SHARP and multi-rack topology
+./nccl.nu generate --gpu h100 --scale 128 --sharp --multirack
+
+# Generate and apply directly to cluster
+./nccl.nu generate --gpu a100 --scale 64 --apply
+
+# Use custom container image
+./nccl.nu generate --gpu h100 --scale 64 --image ghcr.io/custom/image:tag
+```
+
+Available options:
+- `--gpu` - GPU type (a100, h100, a40, gb200, rtxp6000_8x)
+- `--scale` - Total GPU count (must be divisible by GPUs per node)
+- `--multirack` - Enable multi-rack topology spread constraints
+- `--noib` - Disable InfiniBand
+- `--gdrcopy` - Enable GDRCopy
+- `--sharp` - Enable SHARP/COLLNET
+- `--apply` - Apply manifest directly to cluster
+- `--image` - Override container image
 
 #### Running Jobs
 

cue.mod/module.cue

@@ -0,0 +1,9 @@
+module: "github.com/coreweave/nccl-tests"
+language: {
+	version: "v0.15.3"
+}
+deps: {
+	"cue.dev/x/k8s.io@v0": {
+		v: "v0.6.0"
+	}
+}

generate.cue

@@ -0,0 +1,239 @@
+package main
+
+import (
+	"strings"
+
+	corev1 "cue.dev/x/k8s.io/api/core/v1"
+	metav1 "cue.dev/x/k8s.io/apimachinery/pkg/apis/meta/v1"
+	"github.com/coreweave/nccl-tests/gpus"
+)
+
+// Input parameters (set via -t flag)
+gpu:          *"a100" | string @tag(gpu)
+scale:        *64 | int        @tag(scale,type=int)
+name:         *"" | string     @tag(name)
+multirack:    *false | bool    @tag(multirack,type=bool)
+noib:         *false | bool    @tag(noib,type=bool)
+gdrcopy:      *false | bool    @tag(gdrcopy,type=bool)
+sharp:        *false | bool    @tag(sharp,type=bool)
+image:        *"" | string     @tag(image)
+defaultImage: string           @tag(defaultImage)
+
+// Lookup the GPU config
+_gpu: gpus.gpus[gpu]
+
+// Calculate workers (integer division)
+_workers: scale div _gpu.gpusPerNode
+
+// Resolve image (override takes precedence, then GPU-specific, then default)
+_image: string
+if image != "" {
+	_image: image
+}
+if image == "" {
+	_image: defaultImage
+}
+
+// Generate job name (lowercase)
+_name: string
+if name != "" {
+	_name: name
+}
+if name == "" {
+	_name: "nccl-test-\(scale)-\(strings.ToLower(_gpu.name))"
+}
+
+// Build env var flags string from mpiEnv list
+_gpuEnvFlags: strings.Join([for k, v in _gpu.mpiEnv {"-x \(k)=\(v)"}], " \\\n  ")
+
+// Build gdrcopy env flags
+_gdrcopyEnvFlags: string
+if gdrcopy {
+	_gdrcopyEnvFlags: "-x NCCL_GDRCOPY_ENABLE=1 \\\n  -x NCCL_DEBUG=INFO \\\n  -x NCCL_DEBUG_SUBSYS=INIT"
+}
+if !gdrcopy {
+	_gdrcopyEnvFlags: ""
+}
+
+// Build sharp env flags (COLLNET + ALGO for H100)
+_sharpEnvFlags: string
+if sharp && gpu == "h100" {
+	_sharpEnvFlags: "-x NCCL_COLLNET_ENABLE=1 \\\n  -x NCCL_ALGO=COLLNETCHAIN"
+}
+if sharp && gpu != "h100" {
+	_sharpEnvFlags: "-x NCCL_COLLNET_ENABLE=1"
+}
+if !sharp {
+	_sharpEnvFlags: ""
+}
+
+// Combine all extra env flags
+_allExtraEnv: [for s in [_gpuEnvFlags, _gdrcopyEnvFlags, _sharpEnvFlags] if s != "" {s}]
+_extraEnvStr: strings.Join(_allExtraEnv, " \\\n  ")
+_extraEnv:    string
+if _extraEnvStr != "" {
+	_extraEnv: " \\\n  \(_extraEnvStr)"
+}
+if _extraEnvStr == "" {
+	_extraEnv: ""
+}
+
+// Determine IB HCA setting (eth0 when noib, else ibp)
+_ibHca: string
+if noib {
+	_ibHca: "eth0"
+}
+if !noib {
+	_ibHca: "ibp"
+}
+
+// Build UCX_NET_DEVICES part (only when has ibDevices and not noib and not gdrcopy)
+_ucxPart: string
+if _gpu.ibDevices != "" && !noib && !gdrcopy {
+	_ucxPart: " \\\n  -x UCX_NET_DEVICES=\(_gpu.ibDevices)"
+}
+if _gpu.ibDevices == "" || noib || gdrcopy {
+	_ucxPart: ""
+}
+
+// Build bind-to part (only when not gdrcopy)
+_bindPart: string
+if !gdrcopy {
+	_bindPart: " \\\n  -bind-to none"
+}
+if gdrcopy {
+	_bindPart: ""
+}
+
+// Build the mpirun command
+_mpirunArgs: "mpirun \\\n  -np \(scale)\(_bindPart) \\\n  -x LD_LIBRARY_PATH \\\n  -x NCCL_SOCKET_IFNAME=eth0 \\\n  -x NCCL_IB_HCA=\(_ibHca)\(_ucxPart)\(_extraEnv) \\\n  /opt/nccl_tests/build/all_reduce_perf -b 512M -e 8G -f 2 -g 1\n"
+
+// Launcher container spec (typed)
+_launcherContainer: corev1.#Container & {
+	name:  "nccl"
+	image: _image
+	env: [
+		{name: "OMPI_ALLOW_RUN_AS_ROOT", value: "1"},
+		{name: "OMPI_ALLOW_RUN_AS_ROOT_CONFIRM", value: "1"},
+	]
+	command: ["/bin/bash", "-c"]
+	args: [_mpirunArgs]
+	resources: requests: {cpu: "2", memory: "128Mi"}
+}
+
+// Determine if we should request IB resources (GPU supports IB and noib is not set)
+_useIB: _gpu.useIB && !noib
+
+// Worker volume mounts (always dshm, optionally gdrdrv)
+_workerVolumeMounts: [...corev1.#VolumeMount]
+if !gdrcopy {
+	_workerVolumeMounts: [{mountPath: "/dev/shm", name: "dshm"}]
+}
+if gdrcopy {
+	_workerVolumeMounts: [
+		{mountPath: "/dev/shm", name: "dshm"},
+		{mountPath: "/dev/gdrdrv", name: "gdrdrv"},
+	]
+}
+
+// Worker container spec (typed)
+_workerContainer: corev1.#Container & {
+	name:  "nccl"
+	image: _image
+	resources: {
+		requests: {
+			cpu:              "\(_gpu.cpu)"
+			memory:           _gpu.memory
+			"nvidia.com/gpu": "\(_gpu.gpusPerNode)"
+			if _useIB {"rdma/ib": "1"}
+		}
+		limits: {
+			"nvidia.com/gpu": "\(_gpu.gpusPerNode)"
+			if _useIB {"rdma/ib": "1"}
+			memory: _gpu.memory
+			cpu:    "\(_gpu.cpu)"
+		}
+	}
+	volumeMounts: _workerVolumeMounts
+}
+
+// Worker pod affinity (typed)
+_workerAffinity: corev1.#Affinity & {
+	if _gpu.podAffinity {
+		podAffinity: preferredDuringSchedulingIgnoredDuringExecution: [{
+			weight: 100
+			podAffinityTerm: {
+				labelSelector: matchLabels: "metadata.coreweave.cloud/job": "nccl-test"
+				topologyKey: "topology.kubernetes.io/zone"
+			}
+		}]
+	}
+	nodeAffinity: requiredDuringSchedulingIgnoredDuringExecution: nodeSelectorTerms: [{
+		matchExpressions: [{
+			key:      _gpu.nodeKey
+			operator: "In"
+			values: [_gpu.nodeType]
+		}]
+	}]
+}
+
+// Worker volumes (always dshm, optionally gdrdrv)
+_workerVolumes: [...corev1.#Volume]
+if !gdrcopy {
+	_workerVolumes: [{name: "dshm", emptyDir: medium: "Memory"}]
+}
+if gdrcopy {
+	_workerVolumes: [
+		{name: "dshm", emptyDir: medium: "Memory"},
+		{name: "gdrdrv", hostPath: path: "/dev/gdrdrv"},
+	]
+}
+
+// Topology spread constraints (typed)
+_topologySpreadConstraints: [...corev1.#TopologySpreadConstraint]
+if multirack {
+	_topologySpreadConstraints: [{
+		maxSkew:           1
+		minDomains:        2
+		topologyKey:       "topology.kubernetes.io/zone"
+		whenUnsatisfiable: "DoNotSchedule"
+		labelSelector: matchLabels: "metadata.coreweave.cloud/job": "nccl-test"
+	}]
+}
+
+// The MPIJob manifest
+manifest: {
+	apiVersion: "kubeflow.org/v2beta1"
+	kind:       "MPIJob"
+	metadata: metav1.#ObjectMeta & {
+		name: _name
+	}
+	spec: {
+		slotsPerWorker: _gpu.gpusPerNode
+		runPolicy: cleanPodPolicy: "Running"
+		mpiReplicaSpecs: {
+			Launcher: {
+				replicas: 1
+				template: {
+					spec: corev1.#PodSpec & {
+						containers: [_launcherContainer]
+					}
+				}
+			}
+			Worker: {
+				replicas: _workers
+				template: {
+					metadata: metav1.#ObjectMeta & {
+						labels: "metadata.coreweave.cloud/job": "nccl-test"
+					}
+					spec: corev1.#PodSpec & {
+						topologySpreadConstraints: _topologySpreadConstraints
+						containers: [_workerContainer]
+						affinity: _workerAffinity
+						volumes:  _workerVolumes
+					}
+				}
+			}
+		}
+	}
+}

gpus/a100.cue

@@ -0,0 +1,10 @@
+package gpus
+
+gpus: a100: #GPU & {
+	name:        "A100"
+	gpusPerNode: 8
+	nodeType:    "gd-8xa100-i128"
+	cpu:         112
+	memory:      "1000Gi"
+	ibDevices:   "ibp0:1,ibp1:1,ibp2:1,ibp3:1"
+}

gpus/a40.cue

@@ -0,0 +1,11 @@
+package gpus
+
+gpus: a40: #GPU & {
+	name:        "A40"
+	gpusPerNode: 8
+	nodeType:    "A40"
+	nodeKey:     "gpu.nvidia.com/model"
+	cpu:         90
+	memory:      "400Gi"
+	ibDevices:   "ibp0:1,ibp1:1,ibp2:1,ibp3:1"
+}

gpus/gb200.cue

@@ -0,0 +1,21 @@
+package gpus
+
+gpus: gb200: #GPU & {
+	name:        "GB200"
+	gpusPerNode: 4
+	nodeType:    "gb200-4x"
+	cpu:         140
+	memory:      "900Gi"
+	useIB:       false
+	podAffinity: true
+	mpiEnv: {
+		NVIDIA_IMEX_CHANNELS:       "0"
+		NCCL_NET_GDR_C2C:           "1"
+		NCCL_MNNVL_ENABLE:          "1"
+		NCCL_CUMEM_ENABLE:          "1"
+		NCCL_SHM_DISABLE:           "0"
+		UCX_TLS:                    "tcp"
+		UCX_NET_DEVICES:            "eth0"
+		OMPI_MCA_coll_hcoll_enable: "0"
+	}
+}

gpus/gpus.cue

@@ -0,0 +1,13 @@
+package gpus
+
+import "github.com/coreweave/nccl-tests/schema"
+
+// Default image passed via tag from nccl.nu (which reads workflow YAML and git SHA)
+defaultImage: string @tag(defaultImage)
+
+// Alias for schema.#GPU with our default image injected
+#GPU: schema.#GPU & {
+	image: string | *defaultImage
+}
+
+gpus: [string]: #GPU

gpus/h100.cue

@@ -0,0 +1,11 @@
+package gpus
+
+gpus: h100: #GPU & {
+	name:        "H100"
+	gpusPerNode: 8
+	nodeType:    "gd-8xh100ib-i128"
+	cpu:         110
+	memory:      "960Gi"
+	ibDevices:   "ibp0:1,ibp1:1,ibp2:1,ibp3:1,ibp4:1,ibp5:1,ibp6:1,ibp7:1"
+	mpiEnv: SHARP_COLL_ENABLE_PCI_RELAXED_ORDERING: "1"
+}

gpus/rtxp6000_8x.cue

@@ -0,0 +1,10 @@
+package gpus
+
+gpus: rtxp6000_8x: #GPU & {
+	name:        "RTXP6000-8X"
+	gpusPerNode: 8
+	nodeType:    "rtxp6000-8x"
+	cpu:         112
+	memory:      "960Gi"
+	useIB:       false
+}