Copybara import of gpu-recipes:

  - 969ea416145c7a2486135ac16ca137fc49a3d268 Changing base container
  - f2f2d1a35b78252fc05ae2f6c16023c304f5c31e Updating MaxText version to fix grpc error after training...
  - 6232aaf0f3b073f5125c1ed29df34b81f5fa429c Merge "Clean the nemo-launcher-job helm chart" into main

GitOrigin-RevId: 6232aaf0f3b073f5125c1ed29df34b81f5fa429c

Author: Copybara

src/docker/maxtext/README.md

@@ -1,6 +1,5 @@
-# MaxText Nightly
+# MaxText Benchmarks
 
 This is the Dockerfile for building a container image for MaxText/JAX training workloads.
 Using the following versions:
-- BASE_IMAGE: ghcr.io/nvidia/jax:jax-2024-12-04
-- JAV_VERSION: 0.4.36.dev20241202
+- BASE_IMAGE: ghcr.io/nvidia/jax:maxtext-2025-01-10

src/docker/maxtext/cloudbuild.yml

@@ -1,24 +1,11 @@
 steps:
-- name: 'gcr.io/cloud-builders/git'
-  args: ['clone', 'https://github.com/AI-Hypercomputer/maxtext.git']
 - name: 'docker'
-  args: ['build',
-         '--network', 'host',
-         '--build-arg', 'MODE=${_MODE}',
-         '--build-arg', 'JAX_VERSION=${_JAX_VERSION}',
-         '--build-arg', 'DEVICE=${_DEVICE}',
-         '--build-arg', 'BASEIMAGE=${_BASEIMAGE}',
-         '-f', './maxtext_gpu_dependencies.Dockerfile',
-         '-t', '${_ARTIFACT_REGISTRY}/maxtext-nightly',
-         '.'
-        ]
-  dir: 'maxtext' # Set the working directory to 'maxtext'
-options:
-  substitution_option: 'ALLOW_LOOSE'
-substitutions:
-  _MODE: 'nightly'
-  _JAX_VERSION: '0.4.36.dev20241202'
-  _DEVICE: 'gpu'
-  _BASEIMAGE: 'ghcr.io/nvidia/jax:jax-2024-12-04'
+  args:
+    - 'build'
+    - '--tag=${_ARTIFACT_REGISTRY}/maxtext-benchmark'
+    - '--file=maxtext.Dockerfile'
+    - '.'
+  automapSubstitutions: true
+
 images:
-- '${_ARTIFACT_REGISTRY}/maxtext-nightly'
+- '${_ARTIFACT_REGISTRY}/maxtext-benchmark'

src/docker/maxtext/maxtext.Dockerfile

@@ -0,0 +1,32 @@
+# Copyright 2024 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.
+
+FROM ghcr.io/nvidia/jax:maxtext-2025-01-10
+
+# GCSfuse components (used to provide shared storage, not intended for high performance)
+RUN apt-get update && apt-get install --yes --no-install-recommends \
+    ca-certificates \
+    curl \
+    gnupg \
+  && echo "deb https://packages.cloud.google.com/apt gcsfuse-buster main" \
+    | tee /etc/apt/sources.list.d/gcsfuse.list \
+  && echo "deb https://packages.cloud.google.com/apt cloud-sdk main" \
+    | tee -a /etc/apt/sources.list.d/google-cloud-sdk.list \
+  && curl https://packages.cloud.google.com/apt/doc/apt-key.gpg | apt-key add - \
+  && apt-get update \
+  && apt-get install --yes gcsfuse \
+  && apt-get install --yes google-cloud-cli \
+  && apt-get install --yes dnsutils \
+  && apt-get clean && rm -rf /var/lib/apt/lists/* /tmp/* /var/tmp/* \
+  && mkdir /gcs

src/helm-charts/a3ultra/maxtext-training/templates/maxtext-launcher-job.yaml

@@ -49,11 +49,32 @@ spec:
         gke-gcsfuse/memory-limit: "0"
         gke-gcsfuse/ephemeral-storage-limit: "0"
         {{- end}}
+        {{- if not $root.Values.network.hostNetwork }}
+        networking.gke.io/default-interface: "eth0"
+        networking.gke.io/interfaces: |
+        {{- if $root.Values.network.subnetworks }}
+          [
+            {{- range $i, $subnetwork := $root.Values.network.subnetworks }}
+            {"interfaceName":"eth{{ $i }}","network":"{{ $subnetwork }}"}{{ eq $i 9 | ternary "" ","}}
+            {{- end }}
+          ]
+        {{- else }}
+          [
+            {"interfaceName":"eth0","network":"default"},
+            {"interfaceName":"eth1","network":"{{ $root.Values.clusterName }}-sub-1"},
+            {{- range  $i := until 8 }}
+            {"interfaceName":"eth{{ add 2 $i }}","network":"{{ $root.Values.clusterName }}-rdma-sub-{{ $i }}"}{{ eq $i 7 | ternary "" ","}}
+            {{- end }}
+          ]
+        {{- end }}
+        {{- end }}
     spec:
       schedulingGates:
       - name: "gke.io/topology-aware-auto-scheduling"
+      {{- if $root.Values.network.hostNetwork }}
       hostNetwork: true
       dnsPolicy: ClusterFirstWithHostNet
+      {{- end }}
       subdomain: "{{.Release.Name}}"
       restartPolicy: Never
       {{ if $root.Values.targetNodes }}
@@ -84,8 +105,6 @@ spec:
       - name: workload-configuration
         configMap:
           name: "{{.Release.Name}}"
-      - name: workload-terminated-volume
-        emptyDir: {}
       - name: local-ssd
         hostPath:
           path: /mnt/stateful_partition/kube-ephemeral-ssd
@@ -143,8 +162,10 @@ spec:
       - name: workload
         image: "{{ $root.Values.workload.image }}"
         imagePullPolicy: Always
+        {{- if $root.Values.network.hostNetwork }}
         securityContext:
           privileged: true
+        {{- end }}
         env:
         - name: JOB_IDENTIFIER
           value: "{{ .Release.Name }}-{{ $timestamp }}-{{ $jobSuffix }}"
@@ -160,7 +181,7 @@ spec:
         - name: GCS_BUCKET
           value: {{ .bucketName }}
         {{- end }}
-        
+
         # JAX-specific environment variables
         - name: JAX_COORDINATOR_ADDRESS
           value: "{{.Release.Name}}-0.{{.Release.Name}}.default.svc.cluster.local"
@@ -211,19 +232,13 @@ spec:
         - bash
         - -c
         - |
-          function on_script_completion {
-            touch /semaphore/workload_terminated
-          }
-          trap on_script_completion EXIT
           echo "Pod on $(hostname --fqdn) is running"
           echo "Pod is assigned job index of $JOB_COMPLETION_INDEX"
           echo "Job ID is $JOB_IDENTIFIER"
 
           echo "Running nvidia-smi"
           nvidia-smi
 
-
-          echo "Warning: Set LD_LIBRARY_PATH=$LD_LIBRARY_PATH to override the NCCL library"
           ldconfig $LD_LIBRARY_PATH
           echo "Added ${LD_LIBRARY_PATH} to ldconfig:"
           ldconfig -p | grep libcuda | sed 's/^/  /'
@@ -247,19 +262,16 @@ spec:
           echo "{{ . }}"
           {{- end }}
 
-          sleep 10 # <- Hack to allow some time for service to boot
-
           export NODE_RANK=$JOB_COMPLETION_INDEX
           echo "Launching MaxText as node rank $NODE_RANK out of $NNODES nodes"
 
-          if [ "$NODE_RANK" -eq "1" ]; then
-            echo "Launching nvidia-smi in daemon mode with (20 sec delay)"
-            nvidia-smi dmon -d 20 -s pum &
-          fi
-
           echo "XLA Flags: $XLA_FLAGS"
 
+          sleep 10 # <- Allow some time for service to boot
+
+          echo "Setting JAX_COORDINATOR_IP"
           export JAX_COORDINATOR_IP=$(nslookup "$JAX_COORDINATOR_ADDRESS" 2>/dev/null | awk '/^Address: / { print $2 }' | head -n 1)
+          echo $JAX_COORDINATOR_IP
 
           # Parsing Configuration
           while IFS= read -r line || [[ -n "$line" ]]; \
@@ -284,8 +296,6 @@ spec:
             mountPath: /usr/local/nvidia
           - name: gib
             mountPath: /usr/local/gib
-          - name: workload-terminated-volume
-            mountPath: /semaphore
           - name: workload-configuration
             mountPath: /etc/workload-configuration
           - name: shared-memory