From 1d434b7a9b68e4beaa77a33566bc5b8a0b19282a Mon Sep 17 00:00:00 2001 From: Vladislav Date: Mon, 12 Feb 2024 13:27:58 -0700 Subject: [PATCH 1/7] Pin mjc-mpc --- .github/container/manifest.yaml | 2 +- 1 file changed, 1 insertion(+), 1 deletion(-) diff --git a/.github/container/manifest.yaml b/.github/container/manifest.yaml index 9eccb9768..33a99889c 100644 --- a/.github/container/manifest.yaml +++ b/.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 From ba1d0d3ca9281ed6a795b64b81813c1022d156e2 Mon Sep 17 00:00:00 2001 From: Vladislav Date: Mon, 12 Feb 2024 14:51:41 -0700 Subject: [PATCH 2/7] Use Clang as a compiler --- .github/container/Dockerfile.mjx | 2 ++ 1 file changed, 2 insertions(+) diff --git a/.github/container/Dockerfile.mjx b/.github/container/Dockerfile.mjx index fc7ada514..d1d8f746f 100644 --- a/.github/container/Dockerfile.mjx +++ b/.github/container/Dockerfile.mjx @@ -31,6 +31,8 @@ EOF # Specify installation targets RUN <<"EOF" bash -ex +export CC=/usr/bin/clang +export CXX=/usr/bin/clang++ get-source.sh -l mujoco -m ${MANIFEST_FILE} -b $(dirname ${SRC_PATH_MUJOCO}) get-source.sh -l mujoco-mpc -m ${MANIFEST_FILE} get-source.sh -l language-to-reward-2023 -m ${MANIFEST_FILE} From 5dd8ed3b79924dcda8f755569f93726e636125a4 Mon Sep 17 00:00:00 2001 From: Vladislav Date: Mon, 12 Feb 2024 15:32:43 -0700 Subject: [PATCH 3/7] Use Clang as a compiler --- .github/container/Dockerfile.mjx | 10 +++++++--- 1 file changed, 7 insertions(+), 3 deletions(-) diff --git a/.github/container/Dockerfile.mjx b/.github/container/Dockerfile.mjx index d1d8f746f..eeec40adf 100644 --- a/.github/container/Dockerfile.mjx +++ b/.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=/usr/bin/clang +ARG CXX=/usr/bin/clang++ ############################################################################### ## Download source and add auxiliary scripts @@ -14,6 +15,11 @@ FROM ${BASE_IMAGE} as mealkit ARG SRC_PATH_MUJOCO ARG SRC_PATH_MUJOCO_MPC ARG SRC_PATH_L2R +ARG CC +ARG CXX + +ENV CC=${CC} +ENV CXX=${CXX} # Install system dependencies for Mujuco/MPC RUN <<"EOF" bash -ex @@ -31,8 +37,6 @@ EOF # Specify installation targets RUN <<"EOF" bash -ex -export CC=/usr/bin/clang -export CXX=/usr/bin/clang++ get-source.sh -l mujoco -m ${MANIFEST_FILE} -b $(dirname ${SRC_PATH_MUJOCO}) get-source.sh -l mujoco-mpc -m ${MANIFEST_FILE} get-source.sh -l language-to-reward-2023 -m ${MANIFEST_FILE} From cf821d9f6f641398f47a3b17d17afe391b8123ae Mon Sep 17 00:00:00 2001 From: Vladislav Date: Tue, 13 Feb 2024 13:01:58 -0700 Subject: [PATCH 4/7] Separate compilation --- .github/container/Dockerfile.mjx | 19 ++++++++++--------- .github/container/pip-finalize.sh | 6 ++++-- 2 files changed, 14 insertions(+), 11 deletions(-) diff --git a/.github/container/Dockerfile.mjx b/.github/container/Dockerfile.mjx index eeec40adf..0df96a353 100644 --- a/.github/container/Dockerfile.mjx +++ b/.github/container/Dockerfile.mjx @@ -4,8 +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=/usr/bin/clang -ARG CXX=/usr/bin/clang++ +ARG CC_COMPILER=/usr/bin/clang +ARG CXX_COMPILER=/usr/bin/clang++ ############################################################################### ## Download source and add auxiliary scripts @@ -15,11 +15,8 @@ FROM ${BASE_IMAGE} as mealkit ARG SRC_PATH_MUJOCO ARG SRC_PATH_MUJOCO_MPC ARG SRC_PATH_L2R -ARG CC -ARG CXX - -ENV CC=${CC} -ENV CXX=${CXX} +ARG CC_COMPILER +ARG CXX_COMPILER # Install system dependencies for Mujuco/MPC RUN <<"EOF" bash -ex @@ -42,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 @@ -53,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 /opt/pip-tools.d/requirements-mpc.in + +# compiler the rest dependencies +RUN pip-finalize.sh /opt/pip-tools.d/requirements-mjx.in /opt/pip-tools.d/requirements-l2r.in diff --git a/.github/container/pip-finalize.sh b/.github/container/pip-finalize.sh index 371764f6a..9ba7ef573 100755 --- a/.github/container/pip-finalize.sh +++ b/.github/container/pip-finalize.sh @@ -4,11 +4,13 @@ set -eoux pipefail pushd /opt/pip-tools.d +REQUIREMENTS_LIST=${@:-$(ls requirements-*.in)} + # 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 +30,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) From 27abf3ea7750b88415af4abadf77817535381caa Mon Sep 17 00:00:00 2001 From: Vladislav Date: Tue, 13 Feb 2024 14:04:08 -0700 Subject: [PATCH 5/7] Separate compilation: debug --- .github/container/Dockerfile.mjx | 4 ++-- .github/container/pip-finalize.sh | 3 +++ 2 files changed, 5 insertions(+), 2 deletions(-) diff --git a/.github/container/Dockerfile.mjx b/.github/container/Dockerfile.mjx index 0df96a353..bf77bb845 100644 --- a/.github/container/Dockerfile.mjx +++ b/.github/container/Dockerfile.mjx @@ -51,7 +51,7 @@ EOF FROM mealkit as final # compile MuJoCo-MPC with Clang separately -RUN CC=${CC_COMPILER} CXX=${CXX_COMPILER} pip-finalize.sh /opt/pip-tools.d/requirements-mpc.in +RUN CC=${CC_COMPILER} CXX=${CXX_COMPILER} pip-finalize.sh requirements-mpc.in # compiler the rest dependencies -RUN pip-finalize.sh /opt/pip-tools.d/requirements-mjx.in /opt/pip-tools.d/requirements-l2r.in +RUN pip-finalize.sh requirements-mjx.in requirements-l2r.in diff --git a/.github/container/pip-finalize.sh b/.github/container/pip-finalize.sh index 9ba7ef573..6190c4167 100755 --- a/.github/container/pip-finalize.sh +++ b/.github/container/pip-finalize.sh @@ -4,7 +4,10 @@ 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 From 8fb94458a1eb0901a2c25cdeedff648401b76f3f Mon Sep 17 00:00:00 2001 From: Vladislav Date: Tue, 13 Feb 2024 21:08:05 -0700 Subject: [PATCH 6/7] Build nightly jax from concrete base --- .github/workflows/nightly-jax-build.yaml | 11 ++++++++--- 1 file changed, 8 insertions(+), 3 deletions(-) diff --git a/.github/workflows/nightly-jax-build.yaml b/.github/workflows/nightly-jax-build.yaml index 2f40e3366..5c4c658c4 100644 --- a/.github/workflows/nightly-jax-build.yaml +++ b/.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 From 94f0f25e560cd1257bdf84aeb7dbb3cd90a5f1b9 Mon Sep 17 00:00:00 2001 From: Vladislav Date: Tue, 13 Feb 2024 22:25:54 -0700 Subject: [PATCH 7/7] Build via sandbox --- .github/workflows/_sandbox.yaml | 93 ++++++++++++++++++++++----------- 1 file changed, 63 insertions(+), 30 deletions(-) diff --git a/.github/workflows/_sandbox.yaml b/.github/workflows/_sandbox.yaml index 7b90b72ca..b12d9fa91 100644 --- a/.github/workflows/_sandbox.yaml +++ b/.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