Merge branch 'main' into alechan/upgrade-xpk-v0.13.0

Author: aybchan

.github/container/Dockerfile.axlearn

@@ -1,19 +1,35 @@
 # syntax=docker/dockerfile:1-labs
 ARG BASE_IMAGE=ghcr.io/nvidia/jax-mealkit:jax
-ARG URLREF_AXLEARN=https://github.com/Steboss/axlearn.git#main
+ARG URLREF_AXLEARN=https://github.com/apple/axlearn.git
 ARG SRC_PATH_AXLEARN=/opt/axlearn
+ARG DEST_MANIFEST_DIR=/opt/manifest.d
+ARG GIT_USER_NAME=JAX Toolbox
+ARG [email protected]
 
 ###############################################################################
 ## Download source and configure dependencies
 ###############################################################################
 FROM ${BASE_IMAGE} AS mealkit
+ARG DEST_MANIFEST_DIR
+ARG SRC_PATH_AXLEARN
+ARG GIT_USER_NAME
+ARG GIT_USER_EMAIL
 ARG URLREF_AXLEARN
 ARG SRC_PATH_AXLEARN
 
-RUN git-clone.sh "${URLREF_AXLEARN}" "${SRC_PATH_AXLEARN}"
+# Run the patch with cloning
+RUN <<"EOF" bash -exu
+  git config --global user.email "${GIT_USER_EMAIL}"
+  git config --global user.name "${GIT_USER_NAME}"
+  git-clone.sh "${URLREF_AXLEARN}" "${SRC_PATH_AXLEARN}"
+  ${DEST_MANIFEST_DIR}/create-distribution.sh \
+    --manifest ${DEST_MANIFEST_DIR}/manifest.yaml \
+    --package axlearn
+EOF
 
 # these packages are needed to run axlearn tests
 # https://github.com/apple/axlearn/blob/main/pyproject.toml as reference
+WORKDIR /opt/axlearn
 RUN <<"EOF" bash -ex
   echo "-e ${SRC_PATH_AXLEARN}" > /opt/pip-tools.d/requirements-axlearn.in
   cat <<REQUIREMENTS >> /opt/pip-tools.d/requirements-axlearn.in

.github/container/fuji-train-perf.py

@@ -22,6 +22,8 @@
     HybridMeshShape,
 )
 
+import ast
+
 FLAGS = flags.FLAGS
 
 MODELS_INFO = {
@@ -129,10 +131,9 @@
 
 parser.add_argument(
     "--trace_steps",
-    nargs="+",
-    type=int,
+    type=ast.literal_eval,
     default=None,
-    help="Steps to trace (e.g. [1, 20, 50]). To profile the training give `--jax_profiler_port 9999` to the script.",
+    help="Steps where we want JAX tracing. Example: [1,20,30]. AXLearn will pick up the 3 consecutive steps. Please, remember to add also `--jax_profiler_port 9999` to the input command",
 )
 
 parser.add_argument("--world_size", type=int, help="Total number of GPUs")
@@ -253,8 +254,8 @@ def main(parsed_args):
     save_checkpoint_steps = parsed_args.save_checkpoint_steps
     write_summary_steps = parsed_args.write_summary_steps
     output_log_file = parsed_args.output_log_file
-    trace_steps = parsed_args.trace_steps
     world_size = parsed_args.world_size
+    trace_steps = parsed_args.trace_steps
 
     print(
         f"=== Parameter Check ===\n"

.github/container/git-clone.sh

@@ -41,7 +41,7 @@ while [ : ]; do
         ;;
     --)
         shift;
-        break 
+        break
         ;;
   esac
 done
@@ -80,7 +80,6 @@ git submodule update --init --recursive
 popd
 
 ## update the manifest file
-
 mkdir -p $(dirname ${MANIFEST})
 touch ${MANIFEST}
 PACKAGE=$(basename "${DESTINATION}")

.github/container/manifest.yaml

@@ -96,9 +96,12 @@ pathwaysutils:
   latest_verified_commit: 359776d454940ffaa337c36d1df16308d44a95a9
   mode: pip-vcs
 axlearn:
-  url: https://github.com/Steboss/axlearn.git
+  url: https://github.com/apple/axlearn.git
+  mirror_url: https://github.com/nvjax-svc-0/axlearn.git
   tracking_ref: main
   mode: git-clone
+  patches:
+    pull/1339/head: file://patches/axlearn/PR-1339.patch
 qwix:
   url: https://github.com/google/qwix.git
   tracking_ref: main