Removes flax from rosetta-pax build and unit rosetta unit tests from …

…rosetta-pax since those are t5x specific (#253)

1. Previously flax was installed as a distribution for both t5x and pax,
but only t5x uses the flax distribution, so removing for now b/c pax
build does not need to fail if flax distribution cannot be built
2. Unit tests in rosetta were mainly for ViT, and since that requires
t5x, I am removing the unit tests for pax for now. They can be
reintroduced once we introduce a mechanism for testing only t5x stuff
vs. only pax stuff.

CC @maanug-nv

.github/workflows/nightly-rosetta-pax-build.yaml

@@ -83,14 +83,6 @@ jobs:
         echo "MESSAGE='${MSG}'" >> $GITHUB_OUTPUT
         echo "COLOR='${BADGE_COLOR}'" >> $GITHUB_OUTPUT
 
-  test-unit:
-    if: (github.event_name == 'workflow_run' && github.event.workflow_run.conclusion == 'success') || github.event_name == 'workflow_dispatch'
-    needs: build
-    uses: ./.github/workflows/_test_rosetta.yaml
-    with:
-      ROSETTA_IMAGE: ${{ needs.build.outputs.DOCKER_TAGS }}
-    secrets: inherit
-
   test-pax:
     needs: build
     uses: ./.github/workflows/_test_pax.yaml
@@ -101,32 +93,25 @@ jobs:
     secrets: inherit
 
   publish-test:
-    needs: [metadata, build, test-unit, test-pax]
+    needs: [metadata, build, test-pax]
     uses: ./.github/workflows/_publish_badge.yaml
     if: ( always() )
     secrets: inherit
     with:
       ENDPOINT_FILENAME: 'rosetta-pax-overall-test-status.json'
       PUBLISH: ${{ github.event_name == 'workflow_run' || needs.metadata.outputs.PUBLISH == 'true' }}
       SCRIPT: |
-        UNIT_STATUS=${{ needs.test-unit.outputs.TEST_STATUS }}
         PAX_STATUS=${{ needs.test-pax.outputs.TEST_STATUS }}
 
         echo "LABEL='Tests'" >> $GITHUB_OUTPUT
 
         if [[ ${{ needs.build.result }} == "success" ]]; then
-          if [[ $UNIT_STATUS == "success" ]] && [[ $PAX_STATUS == "success" ]]; then
+          if [[ $PAX_STATUS == "success" ]]; then
             COLOR=brightgreen
-            MESSAGE="Unit passed / MGMN passed"
-          elif [[ $UNIT_STATUS == "success" ]]; then 
-            COLOR=yellow
-            MESSAGE="Unit passed / MGMN failed"
-          elif [[ $PAX_STATUS == "success" ]]; then
-            COLOR=yellow
-            MESSAGE="Unit failed / MGMN passed"
+            MESSAGE="MGMN passed"
           else
             COLOR=red
-            MESSAGE="Unit failed / MGMN failed"
+            MESSAGE="MGMN failed"
           fi
         else
           MESSAGE="n/a"
@@ -137,8 +122,8 @@ jobs:
         echo "COLOR='${COLOR}'" >> $GITHUB_OUTPUT
 
   publish-latest-container:
-    needs: [metadata, build, test-unit, test-pax]
-    if: ( ${{ needs.test-unit.outputs.TEST_STATUS == 'success' }} && ${{ needs.test-pax.outputs.TEST_STATUS == 'success' }} ) && ((github.event_name == 'workflow_run' && github.event.workflow_run.conclusion == 'success') || (github.event_name == 'workflow_dispatch' && inputs.PUBLISH))
+    needs: [metadata, build, test-pax]
+    if: ( ${{ needs.test-pax.outputs.TEST_STATUS == 'success' }} ) && ((github.event_name == 'workflow_run' && github.event.workflow_run.conclusion == 'success') || (github.event_name == 'workflow_dispatch' && inputs.PUBLISH))
     uses: ./.github/workflows/_publish_container.yaml
     secrets: inherit
     with:

rosetta/Dockerfile.pax

@@ -12,9 +12,6 @@ ADD --keep-git-dir=true https://github.com/google/paxml.git#main /
 FROM scratch as praxis-mirror-source
 ADD --keep-git-dir=true https://github.com/google/praxis.git#main /
 
-FROM scratch as flax-mirror-source
-ADD --keep-git-dir=true https://github.com/google/flax.git#main /
-
 FROM ${BASE_IMAGE} AS rosetta
 
 ARG GIT_USER_EMAIL
@@ -27,8 +24,7 @@ EOF
 COPY --from=rosetta-source / /opt/rosetta
 WORKDIR /opt/rosetta
 RUN --mount=target=/opt/pax-mirror,from=pax-mirror-source,readwrite \
-    --mount=target=/opt/praxis-mirror,from=praxis-mirror-source,readwrite \
-    --mount=target=/opt/flax-mirror,from=flax-mirror-source,readwrite <<"EOF" bash -e
+    --mount=target=/opt/praxis-mirror,from=praxis-mirror-source,readwrite <<"EOF" bash -e
 bash create-distribution.sh \
   -p patchlist-paxml.txt \
   -u https://github.com/google/paxml.git \
@@ -39,11 +35,6 @@ bash create-distribution.sh \
   -u https://github.com/google/praxis.git \
   -d $(dirname $(python -c "import praxis; print(*praxis.__path__)")) \
   -e /opt/praxis-mirror
-bash create-distribution.sh \
-  -p patchlist-flax.txt \
-  -u https://github.com/google/flax.git \
-  -d $(dirname $(python -c "import flax; print(*flax.__path__)")) \
-  -e /opt/flax-mirror
 rm -rf $(find /opt -name "__pycache__")
 EOF