Disable build-jax for now

Author: aybchan

.github/workflows/_ci.yaml

@@ -38,33 +38,33 @@ permissions:
 
 jobs:
 
-  build-base:
-    uses: ./.github/workflows/_build_base.yaml
-    with:
-      ARCHITECTURE: ${{ inputs.ARCHITECTURE }}
-      BASE_IMAGE: ${{ inputs.CUDA_IMAGE }}
-      BUILD_DATE: ${{ inputs.BUILD_DATE }}
-      MANIFEST_ARTIFACT_NAME: ${{ inputs.MANIFEST_ARTIFACT_NAME }}
-    secrets: inherit
-
-  build-jax:
-    needs: build-base
-    uses: ./.github/workflows/_build.yaml
-    with:
-      ARCHITECTURE: ${{ inputs.ARCHITECTURE }}
-      ARTIFACT_NAME: artifact-jax-build
-      BADGE_FILENAME: badge-jax-build
-      BUILD_DATE: ${{ inputs.BUILD_DATE }}
-      BASE_IMAGE: ${{ needs.build-base.outputs.DOCKER_TAG }}
-      CONTAINER_NAME: jax
-      DOCKERFILE: .github/container/Dockerfile.jax
-      RUNNER_SIZE: large
-      EXTRA_BUILD_ARGS: |
-        URLREF_JAX=${{ fromJson(inputs.SOURCE_URLREFS).JAX }}
-        URLREF_XLA=${{ fromJson(inputs.SOURCE_URLREFS).XLA }}
-        URLREF_FLAX=${{ fromJson(inputs.SOURCE_URLREFS).FLAX }}
-        URLREF_TRANSFORMER_ENGINE=${{ fromJson(inputs.SOURCE_URLREFS).TRANSFORMER_ENGINE }}
-    secrets: inherit
+#  build-base:
+#    uses: ./.github/workflows/_build_base.yaml
+#    with:
+#      ARCHITECTURE: ${{ inputs.ARCHITECTURE }}
+#      BASE_IMAGE: ${{ inputs.CUDA_IMAGE }}
+#      BUILD_DATE: ${{ inputs.BUILD_DATE }}
+#      MANIFEST_ARTIFACT_NAME: ${{ inputs.MANIFEST_ARTIFACT_NAME }}
+#    secrets: inherit
+#
+#  build-jax:
+#    needs: build-base
+#    uses: ./.github/workflows/_build.yaml
+#    with:
+#      ARCHITECTURE: ${{ inputs.ARCHITECTURE }}
+#      ARTIFACT_NAME: artifact-jax-build
+#      BADGE_FILENAME: badge-jax-build
+#      BUILD_DATE: ${{ inputs.BUILD_DATE }}
+#      BASE_IMAGE: ${{ needs.build-base.outputs.DOCKER_TAG }}
+#      CONTAINER_NAME: jax
+#      DOCKERFILE: .github/container/Dockerfile.jax
+#      RUNNER_SIZE: large
+#      EXTRA_BUILD_ARGS: |
+#        URLREF_JAX=${{ fromJson(inputs.SOURCE_URLREFS).JAX }}
+#        URLREF_XLA=${{ fromJson(inputs.SOURCE_URLREFS).XLA }}
+#        URLREF_FLAX=${{ fromJson(inputs.SOURCE_URLREFS).FLAX }}
+#        URLREF_TRANSFORMER_ENGINE=${{ fromJson(inputs.SOURCE_URLREFS).TRANSFORMER_ENGINE }}
+#    secrets: inherit
 
 #  build-triton:
 #    needs: build-jax
@@ -489,11 +489,12 @@ jobs:
   #       test-equinox.log
   #   secrets: inherit
   test-transformerengine-eks:
-    needs: build-jax
+#    needs: build-jax
     if: inputs.ARCHITECTURE == 'amd64'
     runs-on: eks
     env:
-      JAX_DOCKER_IMAGE: ${{  needs.build-jax.outputs.DOCKER_TAG_FINAL }}
+#      JAX_DOCKER_IMAGE: ${{  needs.build-jax.outputs.DOCKER_TAG_FINAL }}
+      JAX_DOCKER_IMAGE: ghcr.io/nvidia/jax-toolbox-internal:13654704499-base-amd64
       JOB_NAME: transformerengine-${{ github.run_id }}
     steps:
     - name: Check out the repository