Use Clang to build MJX

.github/container/Dockerfile.mjx

@@ -4,7 +4,8 @@ ARG BASE_IMAGE=ghcr.io/nvidia/jax:mealkit
 ARG SRC_PATH_MUJOCO=/opt/mujoco
 ARG SRC_PATH_MUJOCO_MPC=/opt/mujoco-mpc
 ARG SRC_PATH_L2R=/opt/language-to-reward-2023
-
+ARG CC_COMPILER=/usr/bin/clang
+ARG CXX_COMPILER=/usr/bin/clang++
 
 ###############################################################################
 ## Download source and add auxiliary scripts
@@ -14,6 +15,8 @@ FROM ${BASE_IMAGE} as mealkit
 ARG SRC_PATH_MUJOCO
 ARG SRC_PATH_MUJOCO_MPC
 ARG SRC_PATH_L2R
+ARG CC_COMPILER
+ARG CXX_COMPILER
 
 # Install system dependencies for Mujuco/MPC
 RUN <<"EOF" bash -ex
@@ -36,7 +39,7 @@ get-source.sh -l mujoco-mpc -m ${MANIFEST_FILE}
 get-source.sh -l language-to-reward-2023 -m ${MANIFEST_FILE}
 echo "-f https://py.mujoco.org/" >> /opt/pip-tools.d/requirements-mjx.in
 echo "-e file://${SRC_PATH_MUJOCO}/mjx" >> /opt/pip-tools.d/requirements-mjx.in
-echo "-e file://${SRC_PATH_MUJOCO_MPC}/python" >> /opt/pip-tools.d/requirements-l2r.in
+echo "-e file://${SRC_PATH_MUJOCO_MPC}/python" >> /opt/pip-tools.d/requirements-mpc.in
 echo "-e file://${SRC_PATH_L2R}" >> /opt/pip-tools.d/requirements-l2r.in
 EOF
 
@@ -47,4 +50,8 @@ EOF
 
 FROM mealkit as final
 
-RUN pip-finalize.sh
+# compile MuJoCo-MPC with Clang separately 
+RUN CC=${CC_COMPILER} CXX=${CXX_COMPILER} pip-finalize.sh requirements-mpc.in
+
+# compiler the rest dependencies
+RUN pip-finalize.sh requirements-mjx.in requirements-l2r.in

.github/container/manifest.yaml

@@ -135,7 +135,7 @@ mujoco:
 mujoco-mpc:
   url: https://github.com/google-deepmind/mujoco_mpc.git
   tracking_ref: main
-  latest_verified_commit: 73633d7da1900c428a7315d2ffe1120c5393a7d8
+  latest_verified_commit: c5c7ead065b7f4034ab265a13023231900dbfaa7
   mode: git-clone
 language-to-reward-2023:
   url: https://github.com/google-deepmind/language_to_reward_2023.git

.github/container/pip-finalize.sh

@@ -4,11 +4,16 @@ set -eoux pipefail
 
 pushd /opt/pip-tools.d
 
+echo "Run with the following argumnets: ${@}"
+
+REQUIREMENTS_LIST=${@:-$(ls requirements-*.in)}
+echo REQUIREMENTS_LIST=${REQUIREMENTS_LIST}
+
 # First pip-compile gathers all reqs, but we are care only about VCS installs
 # It's possible there are 2nd degree transitive dependencies that are VCS, so
 # this is more robust to gather VCS requirements at the cost of pip-compiling
 # twice
-pip-compile -o requirements.pre $(ls requirements-*.in)
+pip-compile -o requirements.pre ${REQUIREMENTS_LIST}
 
 IFS=$'\n'
 for line in $(cat requirements.pre | egrep '^[^#].+ @ git\+' || true); do
@@ -28,7 +33,7 @@ unset IFS
 #
 # JAX_TOOLBOX_VCS_EQUIVALENCY is an environment variable enabling custom logic in pip
 # that treats the above as equivalent and prefers the URI wit the SHA
-JAX_TOOLBOX_VCS_EQUIVALENCY=true pip-compile -o requirements.txt requirements.vcs $(ls requirements-*.in)
+JAX_TOOLBOX_VCS_EQUIVALENCY=true pip-compile -o requirements.txt requirements.vcs ${REQUIREMENTS_LIST}
 
 # If there are unpinned VCS dependencies, error since these should be included in the manifest
 unpinned_vcs_dependencies=$(cat requirements.txt | egrep '^[^#].+ @ git\+' | egrep -v '^[^#].+ @ git\+.+@' || true)

.github/workflows/_sandbox.yaml

@@ -2,40 +2,73 @@ name: "~Sandbox"
 
 on:
   workflow_dispatch:
+    inputs:
+      PUBLISH:
+        type: boolean
+        description: Publish dated images and update the 'latest' tag?
+        default: false
+        required: false
+      BUMP_MANIFEST:
+        type: boolean
+        description: Whether or not to bump software versions in manifest.yaml
+        default: false
+        required: false
 
 jobs:
-  sandbox:
+  metadata:
     runs-on: ubuntu-22.04
+    outputs:
+      BUILD_DATE: ${{ steps.date.outputs.BUILD_DATE }}
+      PUBLISH: ${{ steps.if-publish.outputs.PUBLISH }}
+      BUMP_MANIFEST: ${{ steps.if-bump-manifest.outputs.BUMP_MANIFEST }}
     steps:
-      - name: Login to GitHub Container Registry
-        uses: docker/login-action@v3
-        with:
-          registry: ghcr.io
-          username: ${{ github.repository_owner }}
-          password: ${{ secrets.GITHUB_TOKEN }}
+      - name: Set build date
+        id: date
+        shell: bash -x -e {0}
+        run: |
+          BUILD_DATE=$(TZ='US/Los_Angeles' date '+%Y-%m-%d')
+          echo "BUILD_DATE=${BUILD_DATE}" >> $GITHUB_OUTPUT
+
+      - name: Determine whether results will be 'published'
+        id: if-publish
+        shell: bash -x -e {0}
+        run: |
+          echo "PUBLISH=${{ github.event_name == 'schedule' || inputs.PUBLISH }}" >> $GITHUB_OUTPUT
 
-      - name: Print usage
+      - name: Determine whether results will be 'published'
+        id: if-bump-manifest
+        shell: bash -x -e {0}
         run: |
-          cat << EOF
-          This is an empty workflow file located in the main branch of your
-          repository. It serves as a testing ground for new GitHub Actions on
-          development branches before merging them to the main branch. By
-          defining and overloading this workflow on your development branch,
-          you can test new actions without affecting your main branch, ensuring
-          a smooth integration process once the changes are ready to be merged.
+          echo "BUMP_MANIFEST=${{ github.event_name == 'schedule' || inputs.BUMP_MANIFEST }}" >> $GITHUB_OUTPUT
+
+  amd64-base:
+    needs: [metadata]
+    uses: ./.github/workflows/_build_base.yaml
+    with:
+      ARCHITECTURE: amd64
+      BUILD_DATE: ${{ needs.metadata.outputs.BUILD_DATE }}
+      BUMP_MANIFEST: ${{ needs.metadata.outputs.BUMP_MANIFEST == 'true' }}
+    secrets: inherit
+
+  amd64-jax:
+    needs: [ amd64-base ]
+    uses: ./.github/workflows/_build_jax.yaml
+    with:
+      ARCHITECTURE: amd64
+      BASE_IMAGE: ${{ needs.amd64-base.outputs.DOCKER_TAG }}
+      BUILD_DATE: "2024-02-13"
+    secrets: inherit
+
+  amd64-mjx:
+    needs: [ amd64-jax ]
+    uses: ./.github/workflows/_build.yaml
+    with:
+      ARCHITECTURE: amd64
+      ARTIFACT_NAME: artifact-mjx-build
+      BADGE_FILENAME: badge-mjx-build
+      BASE_IMAGE: ${{ needs.amd64-jax.outputs.DOCKER_TAG_MEALKIT }}
+      BUILD_DATE: "2024-02-13"
+      CONTAINER_NAME: mjx
+      DOCKERFILE: .github/container/Dockerfile.mjx
+    secrets: inherit
 
-          Usage:
-
-          1. In your development branch, modify the sandbox.yml workflow file
-             to include the new actions you want to test. Make sure to commit
-             the changes to the development branch.
-          2. Navigate to the 'Actions' tab in your repository, select the
-             '~Sandbox' workflow, and choose your development branch from the
-             branch dropdown menu. Click on 'Run workflow' to trigger the
-             workflow on your development branch.
-          3. Once you have tested and verified the new actions in the Sandbox
-             workflow, you can incorporate them into your main workflow(s) and
-             merge the development branch into the main branch. Remember to
-             revert the changes to the sandbox.yml file in the main branch to
-             keep it empty for future testing.
-          EOF

.github/workflows/nightly-jax-build.yaml

@@ -8,7 +8,12 @@ on:
     branches: [main]
   workflow_dispatch:
     inputs:
-      BASE_IMAGE:
+      BASE_IMAGE_AMD64:
+        type: string
+        description: 'CUDA base image built by NVIDIA/JAX-Toolbox'
+        default: 'ghcr.io/nvidia/jax-toolbox:base'
+        required: false
+      BASE_IMAGE_ARM64:
         type: string
         description: 'CUDA base image built by NVIDIA/JAX-Toolbox'
         default: 'ghcr.io/nvidia/jax-toolbox:base'
@@ -50,7 +55,7 @@ jobs:
     uses: ./.github/workflows/_build_jax.yaml
     with:
       ARCHITECTURE: amd64
-      BASE_IMAGE: ${{ inputs.BASE_IMAGE || 'ghcr.io/nvidia/jax-toolbox:base' }}
+      BASE_IMAGE: ${{ inputs.BASE_IMAGE_AMD64 || 'ghcr.io/nvidia/jax-toolbox:base' }}
       BUILD_DATE: ${{ needs.metadata.outputs.BUILD_DATE }}
     secrets: inherit
 
@@ -59,7 +64,7 @@ jobs:
     uses: ./.github/workflows/_build_jax.yaml
     with:
       ARCHITECTURE: arm64
-      BASE_IMAGE: ${{ inputs.BASE_IMAGE || 'ghcr.io/nvidia/jax-toolbox:base' }}
+      BASE_IMAGE: ${{ inputs.BASE_IMAGE_ARM64 || 'ghcr.io/nvidia/jax-toolbox:base' }}
       BUILD_DATE: ${{ needs.metadata.outputs.BUILD_DATE }}
     secrets: inherit