updating the readme and adding the sample workload files

Author: Ishita-Chail

examples/dws-sample-workloads/dws-flex-start.yaml

@@ -0,0 +1,53 @@
+# Copyright 2026 "Google LLC"
+#
+# Licensed under the Apache License, Version 2.0 (the "License");
+# you may not use this file except in compliance with the License.
+# You may obtain a copy of the License at
+#
+#     http://www.apache.org/licenses/LICENSE-2.0
+#
+# Unless required by applicable law or agreed to in writing, software
+# distributed under the License is distributed on an "AS IS" BASIS,
+# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+# See the License for the specific language governing permissions and
+# limitations under the License.
+
+apiVersion: batch/v1
+kind: Job
+metadata:
+  name: gpu-job-1
+spec:
+  template:
+    spec:
+      nodeSelector:
+        cloud.google.com/gke-flex-start: "true"
+      containers:
+      - name: gpu-container-1
+        image: gcr.io/k8s-staging-perf-tests/sleep:latest
+        args: ["10s"] # Sleep for 10 seconds
+        resources:
+          requests:
+            nvidia.com/gpu: 1
+          limits:
+            nvidia.com/gpu: 1
+      restartPolicy: OnFailure
+---
+apiVersion: batch/v1
+kind: Job
+metadata:
+  name: gpu-job-2
+spec:
+  template:
+    spec:
+      nodeSelector:
+        cloud.google.com/gke-flex-start: "true"
+      containers:
+      - name: gpu-container-2
+        image: gcr.io/k8s-staging-perf-tests/sleep:latest
+        args: ["10s"] # Sleep for 10 seconds
+        resources:
+          requests:
+            nvidia.com/gpu: 1
+          limits:
+            nvidia.com/gpu: 1
+      restartPolicy: OnFailure

examples/dws-sample-workloads/nccl-jobset-dws-queue.yaml

@@ -0,0 +1,205 @@
+# Copyright 2026 Google Inc. All rights reserved.
+#
+# Licensed under the Apache License, Version 2.0 (the "License");
+# you may not use this file except in compliance with the License.
+# You may obtain a copy of the License at
+#
+#     http://www.apache.org/licenses/LICENSE-2.0
+#
+# Unless required by applicable law or agreed to in writing, software
+# distributed under the License is distributed on an "AS IS" BASIS,
+# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+# See the License for the specific language governing permissions and
+# limitations under the License.
+
+apiVersion: jobset.x-k8s.io/v1alpha2
+kind: JobSet
+metadata:
+  generateName: ag-2-
+  namespace: default
+  labels:
+    kueue.x-k8s.io/queue-name: dws-local-queue
+  annotations:
+    provreq.kueue.x-k8s.io/maxRunDurationSeconds: "600"
+spec:
+  ttlSecondsAfterFinished: 1200
+  suspend: False
+  network:
+    enableDNSHostnames: true
+  replicatedJobs:
+  - name: w
+    template:
+      spec:
+        parallelism: 2
+        completions: 2
+
+        template:
+          metadata:
+            annotations:
+              networking.gke.io/default-interface: 'eth0'
+              networking.gke.io/interfaces: |
+                [
+                  {"interfaceName":"eth0","network":"default"},
+                  {"interfaceName":"eth1","network":"vpc1"},
+                  {"interfaceName":"eth2","network":"vpc2"},
+                  {"interfaceName":"eth3","network":"vpc3"},
+                  {"interfaceName":"eth4","network":"vpc4"},
+                  {"interfaceName":"eth5","network":"vpc5"},
+                  {"interfaceName":"eth6","network":"vpc6"},
+                  {"interfaceName":"eth7","network":"vpc7"},
+                  {"interfaceName":"eth8","network":"vpc8"}
+                ]
+          spec:
+            activeDeadlineSeconds: 3600
+            restartPolicy: Never
+            nodeSelector:
+              cloud.google.com/gke-nodepool: a3-megagpu-8g-a3-megagpu-pool
+              cloud.google.com/gke-flex-start: "true"
+            tolerations:
+            - key: cloud.google.com/gke-queued
+              effect: NoSchedule
+              value: "true"
+
+            - key: "nvidia.com/gpu"
+              operator: "Exists"
+              effect: "NoSchedule"
+
+            setHostnameAsFQDN: true
+            volumes:
+            - name: nvidia
+              hostPath:
+                path: /home/kubernetes/bin/nvidia
+            - name: lib64
+              hostPath:
+                path: /lib64
+            - name: shared-memory
+              emptyDir:
+                medium: "Memory"
+                sizeLimit: 250Gi
+            - name: sys
+              hostPath:
+                path: /sys
+            - name: proc-sys
+              hostPath:
+                path: /proc/sys
+
+            initContainers:
+            - name: gpu-healthcheck
+              image: alpine:latest
+              command: ["/bin/sh", "-c"]
+              args:
+              - |
+                apk add --no-cache bash  # Install bash
+                /bin/bash -c "set -ex
+                NUM_GPUS=$(/usr/local/nvidia/bin/nvidia-smi --query-gpu=driver_version --format=csv,noheader,nounits | wc -l)
+                if [ \${NUM_GPUS} -lt 8 ]; then
+                  echo \"Error: Only \${NUM_GPUS} GPUs and expected 8\"
+                  exit 1
+                fi
+                gpu_errors=(\$(/usr/local/nvidia/bin/nvidia-smi --query-gpu=ecc.errors.uncorrected.volatile.total --format=csv,noheader,nounits))
+                for gpu_index in \${!gpu_errors[@]}; do
+                    if [ \${gpu_errors[\$gpu_index]} == '[N/A]' ]; then
+                        echo 'Error: ERR detected in GPU index '\$gpu_index
+                        exit 1
+                    elif [ \${gpu_errors[\$gpu_index]} -gt 0 ]; then
+                        echo 'Error: Unrecoverable ECC errors detected in GPU index '\$gpu_index
+                        exit 1
+                    fi
+                done
+                echo \${NUM_GPUS} GPUs found with no ERR or Unrecoverable ECC errors"
+
+              volumeMounts:
+              - name: nvidia
+                mountPath: /usr/local/nvidia
+              - name: lib64
+                mountPath: /lib64
+              securityContext:
+                privileged: true
+              env:
+              - name: LD_LIBRARY_PATH
+                value: /usr/local/nvidia/lib64
+
+            containers:
+            - name: nccl
+              stdin: true
+              tty: true
+              image: us-docker.pkg.dev/gce-ai-infra/gpudirect-tcpxo/nccl-plugin-gpudirecttcpx-dev:v1.0.15
+              securityContext:
+                privileged: true
+              env:
+              - name: MY_NODE_NAME
+                valueFrom:
+                  fieldRef:
+                    fieldPath: spec.nodeName
+              - name: OMPI_ALLOW_RUN_AS_ROOT
+                value: "1"
+              - name: OMPI_ALLOW_RUN_AS_ROOT_CONFIRM
+                value: "1"
+              command:
+              - bash
+              - -c
+              - |
+                set -x
+                export N_NODES=2
+                echo "Starting workload container on ${MY_NODE_NAME} for $N_NODES benchmark"
+
+                # Load all the cuda libs
+                /sbin/ldconfig
+
+                # Install ping
+                apt update -y
+                apt install -y iputils-ping
+
+                # Start sshd
+                /scripts/container_entry.sh daemon &
+
+                # Get helper variables to form all hostnames
+                export POSTFIX=$(hostname | cut -d . -f 2-)
+                export WORKERS_BASENAME=$(hostname | cut -d . -f 1 | rev | cut -d - -f 2- | rev )
+                export NODE_RANK=$JOB_COMPLETION_INDEX
+
+                # For every worker, wait till online and add to hostfile
+                for i in `seq 0 $(($N_NODES-1))`; do
+                  OTHER=${WORKERS_BASENAME}-${i}.${POSTFIX}
+                  until ssh -p 222 -o StrictHostKeyChecking=no $OTHER hostname; do
+                    echo Waiting for ${OTHER}...
+                    sleep 10
+                  done
+                  echo ${OTHER} port=222 slots=8 | tee -a /tmp/hostfile;
+                done
+
+                cat /tmp/hostfile
+
+                # Launch from head node
+                if [[ "${NODE_RANK}" -eq "0" ]]; then
+                    export NCCL_TESTS_SPLIT_MASK="0x0";
+                    ENV_VARS=$(echo ${!NCCL*} ${!OMPI*} LD_LIBRARY_PATH PATH | sed 's/ / -x /g')
+
+                    mpirun --hostfile /tmp/hostfile \
+                      -x $ENV_VARS  \
+                      -mca plm_rsh_no_tree_spawn 1 \
+                      --mca orte_keep_fqdn_hostnames 1 \
+                      --mca btl self,tcp \
+                      --mca btl_tcp_if_include eth0 \
+                      --bind-to none \
+                      --mca plm_rsh_agent "ssh -q -o LogLevel=ERROR -o StrictHostKeyChecking=no -p 222" \
+                      /third_party/nccl-tests/build/all_gather_perf -b 1K -e 8G -f 2 -g 1 -w 5 --iters 100 -c 1
+
+                else
+                    while ping -c 1 ${WORKERS_BASENAME}-0.${POSTFIX}; do
+                    sleep 5
+                done
+                fi
+
+                exit 0
+
+              volumeMounts:
+              - name: nvidia
+                mountPath: /usr/local/nvidia
+              - name: shared-memory
+                mountPath: /dev/shm
+              resources:
+                limits:
+                  nvidia.com/gpu: 8
+                requests:
+                  nvidia.com/gpu: 8