JAX-vLLM Offloading k8s (GKE)

.github/actions/gke-xpk/action.yml

@@ -28,7 +28,7 @@ inputs:
     required: false
     type: string
   MAIN_CONTAINER_NAME: 
-    description: 'Name of the main contianer in an XPK JobSet (fixed)'
+    description: 'Name of the main contianer in an XPK JobSet (fixed in xpk)'
     default: gpu-image
     required: false
     type: string
@@ -57,6 +57,11 @@ inputs:
     required: false
     default: 'nvidia-smi; free -h;'
     type: string
+  ENVS:
+    description: 'Environment variables to pass to xpk for setting in JobSet (delimited by ;)'
+    required: false
+    default: ''
+    type: string
   EXIT_COMMAND:
     description: 'Command to set exit code'
     required: false
@@ -124,6 +129,7 @@ runs:
   - name: Set workload commands
     shell: bash -x -e -u {0}
     run: |
+      # install dependencies to enable export artifacts from container to gcs bucket
       PRELUDE="
           apt install -y ripgrep > /dev/null;
           curl -LO https://dl.google.com/dl/cloudsdk/channels/rapid/downloads/google-cloud-cli-linux-x86_64.tar.gz;
@@ -133,6 +139,7 @@ runs:
 
           mkdir -p /usr/share/workload;
           mkdir -p ${{ inputs.CONTAINER_OUTPUT_PATH }};"
+
       # Work around GCP's deployment model that munges together three
       # mostly unrelated things: (1) the host machine's CUDA driver/libs,
       # (2) the version of NCCL installed on the host machine, and (3)
@@ -151,6 +158,7 @@ runs:
           env;
       "
 
+      # gsutil command to export logs from container's /opt/output to bucket
       POSTLUDE="
           ./google-cloud-sdk/bin/gsutil cp -r ${{ inputs.CONTAINER_OUTPUT_PATH }}/ ${GCS_ARTIFACT_PATH}/node-0\$NODE_RANK;
           ${{ inputs.EXIT_COMMAND }}
@@ -163,9 +171,7 @@ runs:
       POSTLUDE=$(echo ${POSTLUDE} | sed 's/\n/\ /g')
       CMD=$(echo ${CMD} | sed 's/\n/\ /g')
 
-      echo "PRELUDE=${PRELUDE}" >> ${GITHUB_ENV}
-      echo "CMD=${CMD}" >> ${GITHUB_ENV}
-      echo "POSTLUDE=${POSTLUDE}" >> ${GITHUB_ENV}
+      echo "CMD=${PRELUDE} ${CMD} ${POSTLUDE}" >> ${GITHUB_ENV}
 
   - name: Create workload on cluster with XPK
     shell: bash -x -e -u {0}
@@ -194,16 +200,21 @@ runs:
       }
 
       if version_greater "${{ inputs.XPK_VERSION }}" "v0.10.0"; then
-          args+=(
-              --docker-image-pull-secret=${{ inputs.IMAGE_PULL_SECRET_NAME }}
-              --env="JAX_COORDINATOR_PORT=3389"
-              --env="JAX_COORDINATOR_ADDRESS=\$(JOBSET_NAME)-\$(REPLICATED_JOB_NAME)-0-0.\$(JOBSET_NAME):3389"
-          )
-      fi
+        args+=(
+            --docker-image-pull-secret=${{ inputs.IMAGE_PULL_SECRET_NAME }}
+        )
 
+        envs_flat=$(echo "${{ inputs.ENVS }}" | tr '\n' ' ')
+        IFS=';' read -ra env_vars <<< "${envs_flat}"
+        for env in "${env_vars[@]}"; do
+            env=$(echo "${env}" | xargs)
+            [[ -n "${env}" ]] && args+=(--env="${env}")
+        done
+      fi
+
       python xpk.py workload create \
-        ${args[@]} \
-        --command="${PRELUDE} ${CMD} ${POSTLUDE}"
+        "${args[@]}" \
+        --command="${CMD}"
 
   - name: Wait for JobSet to unsuspend on cluster
     shell: bash -u {0}

.github/gke-workflow/jax-vllm-offloading/deploy-transfer.sh

@@ -0,0 +1,7 @@
+kubectl apply -f transfer/deployment/gateway-pod.yml
+kubectl apply -f transfer/deployment/gateway-svc.yml
+
+kubectl apply -f huggingface-secret.yml
+
+kubectl apply -f transfer/deployment/rollout.yml
+kubectl apply -f transfer/deployment/trainer.yml

.github/gke-workflow/jax-vllm-offloading/grpo/jobset.yaml

@@ -0,0 +1,280 @@
+apiVersion: jobset.x-k8s.io/v1alpha2
+kind: JobSet
+metadata:
+  annotations:
+  name: jax-vllm-grpo
+  namespace: default
+spec:
+  network:
+    enableDNSHostnames: true
+    publishNotReadyAddresses: true
+  replicatedJobs:
+  - name: slice-job
+    replicas: 1
+    template:
+      metadata: {}
+      spec:
+        backoffLimit: 0
+        completionMode: Indexed
+        completions: 2
+        parallelism: 2
+        template:
+          metadata:
+            annotations:
+              devices.gke.io/container.tcpxo-daemon: |
+                - path: /dev/nvidia0
+                - path: /dev/nvidia1
+                - path: /dev/nvidia2
+                - path: /dev/nvidia3
+                - path: /dev/nvidia4
+                - path: /dev/nvidia5
+                - path: /dev/nvidia6
+                - path: /dev/nvidia7
+                - path: /dev/nvidiactl
+                - path: /dev/nvidia-uvm
+                - path: /dev/dmabuf_import_helper
+              networking.gke.io/default-interface: eth0
+              networking.gke.io/interfaces: |-
+                [
+                    {"interfaceName":"eth0","network":"default"},
+                    {"interfaceName":"eth1","network":"jtb-2025-10-07-gpunet-0-subnet"},
+                    {"interfaceName":"eth2","network":"jtb-2025-10-07-gpunet-1-subnet"},
+                    {"interfaceName":"eth3","network":"jtb-2025-10-07-gpunet-2-subnet"},
+                    {"interfaceName":"eth4","network":"jtb-2025-10-07-gpunet-3-subnet"},
+                    {"interfaceName":"eth5","network":"jtb-2025-10-07-gpunet-4-subnet"},
+                    {"interfaceName":"eth6","network":"jtb-2025-10-07-gpunet-5-subnet"},
+                    {"interfaceName":"eth7","network":"jtb-2025-10-07-gpunet-6-subnet"},
+                    {"interfaceName":"eth8","network":"jtb-2025-10-07-gpunet-7-subnet"}
+                ]
+          spec:
+            imagePullSecrets:
+            - name: jax-toolbox-ghcr
+            containers:
+            - name: gpu-image
+              image: ghcr.io/nvidia/jax-toolbox-internal:19751502075-jio-amd64
+              imagePullPolicy: Always
+              command:
+              - bash
+              - -c
+              - |
+                pip install jax[k8s]
+                python -c "
+                import jax
+                jax.distributed.initialize()
+                print(jax.devices())
+                print(jax.local_devices())
+                assert jax.process_count() > 1
+                assert len(jax.devices()) > len(jax.local_devices())"
+
+                PIDS=()
+                # hard-code split of vLLM-JAX on 1x node each on 2x slice jobset
+                if [ ${NODE_RANK} = "0" ]; then
+                  echo "Starting gateway"
+                  cd /opt/jtbx/jax-inference-offloading
+                  python jax_inference_offloading/controller/gateway.py 2>&1 | tee -a gateway.log &
+                  PIDS+=($!)
+
+                  echo "Starting rollout"
+                  cd /opt/jtbx/jax-inference-offloading/examples
+                  python rollout.py 2>&1 | tee -a rollout.log &
+                  PIDS+=($!)
+                else
+                  echo "Starting trainer"
+                  export MODEL_PATH=$(python "download_model.py" --hub=hf --model=${MODEL_NAME} --ignore="*.pth")
+                  python trainer_grpo.py 2>&1 | tee -a trainer_grpo.log &
+                  PIDS+=($!)
+                fi
+
+                wait "${PIDS[@]}"
+                echo "All done"
+              env:
+              # jobset
+              - name: REPLICATED_JOB_NAME
+                valueFrom:
+                  fieldRef:
+                    fieldPath: metadata.annotations['jobset.sigs.k8s.io/replicatedjob-name']
+              - name: JOBSET_NAME
+                valueFrom:
+                  fieldRef:
+                    fieldPath: metadata.annotations['jobset.sigs.k8s.io/jobset-name']
+              - name: NODE_RANK
+                valueFrom:
+                  fieldRef:
+                    fieldPath: metadata.annotations['batch.kubernetes.io/job-completion-index']
+              - name: USE_GPUDIRECT
+                value: tcpxo
+              - name: GPUS_PER_NODE
+                value: "8"
+
+              - name: LD_LIBRARY_PATH
+                value: "/usr/lib/x86_64-linux-gnu:/usr/local/cuda-12.9/compat/lib.real:/usr/local/nvidia/lib64"
+
+              # huggingface
+              - name: HF_TOKEN
+                valueFrom:
+                  secretKeyRef:
+                    name: hf-token-secret
+                    key: token
+              - name: MODEL_NAME
+                value: "meta-llama/Llama-3.1-8B-Instruct"
+              - name: SCRATCHDIR
+                value: "/opt/scratch"
+
+              # gateway
+              - name: GATEWAY_PORT
+                value: "50051"
+              - name: GATEWAY_URL
+                value: "$(JOBSET_NAME):$(GATEWAY_PORT)"
+
+              # JAX
+              - name: JAX_COORDINATOR_PORT
+                value: "3389"
+              - name: JAX_COORDINATOR_ADDRESS
+                value: $(JOBSET_NAME)-$(REPLICATED_JOB_NAME)-0-0.$(JOBSET_NAME):3389
+
+              # CUDA
+              - name: CUDA_VISIBLE_DEVICES
+                value: "0,1,2,3,4,5,6,7"
+              - name: CUDA_DEVICE_ORDER
+                value: "PCI_BUS_ID"
+              - name: CUDA_DEVICE_MAX_CONNECTIONS
+                value: "16"
+
+              # vLLM
+              - name: VLLM_ENFORCE_EAGER
+                value: "1"
+              - name: VLLM_GPU_MEMORY_UTILIZATION
+                value: "0.7"
+              - name: VLLM_TENSOR_PARALLEL_SIZE
+                value: "8"
+              - name: VLLM_DISTRIBUTED_BACKEND
+                value: "mp"
+              - name: VLLM_ATTENTION_BACKEND
+                value: "TRITON_ATTN"
+              - name: VLLM_LOAD_FORMAT
+                value: "dummy"
+
+              # NCCL
+              - name: NCCL_NET_PLUGIN
+                value: "/opt/hpcx/nccl_rdma_sharp_plugin/lib/libnccl-net.so"
+              - name: NCCL_TUNER_PLUGIN
+                value: "none"
+              - name: NCCL_FASTRAK_LLCM_DEVICE_DIRECTORY
+                value: /dev/aperture_devices
+              - name: NCCL_CUMEM_ENABLE
+                value: "0"  # https://docs.vllm.ai/en/v0.9.1/usage/troubleshooting.html#known-issues
+              - name: NCCL_BUFFSIZE
+                value: "16777216"
+
+              # XLA
+              - name: XLA_PYTHON_CLIENT_MEM_FRACTION
+                value: "0.95"
+              - name: XLA_FLAGS
+                value: "--xla_gpu_enable_latency_hiding_scheduler=true
+                        --xla_gpu_enable_command_buffer=FUSION,CUBLAS,CUDNN,CUSTOM_CALL
+                        --xla_gpu_collective_permute_combine_threshold_bytes=8589934592
+                        --xla_gpu_reduce_scatter_combine_threshold_bytes=8589934592
+                        --xla_gpu_all_gather_combine_threshold_bytes=8589934592
+                        --xla_gpu_all_reduce_combine_threshold_bytes=8589934592"
+
+              # trainer
+              - name: TRANSFER_MODE
+                value: "grouped"
+              - name: USE_POLYMORPHIC_MESH
+                value: "0"
+              - name: JAX_COMPILATION_CACHE_DIR
+                value: /opt/jax-compilation
+              - name: JAX_PERSISTENT_CACHE_MIN_COMPILE_TIME_SECS
+                value: "0.1"
+              - name: RUN_MODE
+                value: "timing"
+              - name: ROLLOUT_ENGINE
+                value: "vllm_gpu"
+              - name: GRPO_TRAIN_MICRO_BATCH_SIZE
+                value: "2"
+
+
+              ports:
+              - containerPort: 50051
+                protocol: TCP
+              - containerPort: 3389
+                protocol: TCP
+              resources:
+                limits:
+                  nvidia.com/gpu: "8"
+              securityContext:
+                privileged: true
+              volumeMounts:
+              - mountPath: /dev/aperture_devices
+                name: aperture-devices
+              - mountPath: /usr/local/nvidia
+                name: libraries
+              - mountPath: /dev/shm
+                name: dshm
+              - mountPath: /opt/scratch
+                name: scratch
+            dnsPolicy: ClusterFirstWithHostNet
+            initContainers:
+            - args:
+              - |-
+                set -ex
+                chmod 755 /fts/entrypoint_rxdm_container.sh
+                /fts/entrypoint_rxdm_container.sh --num_hops=2 --num_nics=8 --uid= --alsologtostderr
+              command:
+              - /bin/sh
+              - -c
+              env:
+              - name: LD_LIBRARY_PATH
+                value: /usr/local/nvidia/lib64
+              image: us-docker.pkg.dev/gce-ai-infra/gpudirect-tcpxo/tcpgpudmarxd-dev:v1.0.12
+              imagePullPolicy: Always
+              name: tcpxo-daemon
+              resources: {}
+              restartPolicy: Always
+              securityContext:
+                capabilities:
+                  add:
+                  - NET_ADMIN
+                  - NET_BIND_SERVICE
+              volumeMounts:
+              - mountPath: /usr/local/nvidia
+                name: libraries
+              - mountPath: /hostsysfs
+                name: sys
+              - mountPath: /hostprocsysfs
+                name: proc-sys
+            nodeSelector:
+              cloud.google.com/gke-accelerator: nvidia-h100-mega-80gb
+            priorityClassName: high
+            terminationGracePeriodSeconds: 30
+            tolerations:
+            - key: nvidia.com/gpu
+              operator: Exists
+            - effect: NoSchedule
+              key: user-workload
+              operator: Equal
+              value: "true"
+            volumes:
+            - hostPath:
+                path: /home/kubernetes/bin/nvidia
+              name: libraries
+            - hostPath:
+                path: /sys
+              name: sys
+            - hostPath:
+                path: /proc/sys
+              name: proc-sys
+            - hostPath:
+                path: /dev/aperture_devices
+              name: aperture-devices
+            - emptyDir:
+                medium: Memory
+              name: dshm
+            - emptyDir:
+                sizeLimit: 2Gi
+              name: scratch
+  startupPolicy:
+    startupPolicyOrder: AnyOrder
+  successPolicy:
+    operator: All
+  ttlSecondsAfterFinished: 100000

.github/gke-workflow/jax-vllm-offloading/huggingface-secret.yml

@@ -0,0 +1,8 @@
+apiVersion: v1
+kind: Secret
+metadata:
+  name: hf-token-secret
+  namespace: default
+type: Opaque
+stringData:
+  token: {{ HF_TOKEN}}