Clone airio

NVIDIA · Nov 27, 2024 · 3c2ec97 · 3c2ec97
1 parent beb4f82
commit 3c2ec97
Show file tree

Hide file tree

Showing 3 changed files with 18 additions and 4 deletions.
diff --git a/.github/container/Dockerfile.t5x b/.github/container/Dockerfile.t5x
@@ -5,8 +5,11 @@
 ARG BASE_IMAGE=ghcr.io/nvidia/jax-mealkit:jax
 ARG URLREF_TFTEXT=https://github.com/tensorflow/text.git#master
 ARG URLREF_T5X=https://github.com/google-research/t5x.git#main
+ARG URLREF_AIRIO=https://github.com/google/airio.git#main
 ARG SRC_PATH_TFTEXT=/opt/tensorflow-text
 ARG SRC_PATH_T5X=/opt/t5x
+ARG SRC_PATH_AIRIO=/opt/airio
+
 
 ###############################################################################
 ## build several packages which do not have working arm64 pip wheels
@@ -23,8 +26,8 @@ ARG URLREF_TFTEXT
 ARG SRC_PATH_TFTEXT
 
 RUN pip install tensorflow_datasets==4.9.2 auditwheel tensorflow==2.18.0
-RUN git-clone.sh ${URLREF_TFTEXT} ${SRC_PATH_TFTEXT}
 RUN <<"EOF" bash -exu -o pipefail
+git-clone.sh ${URLREF_TFTEXT} ${SRC_PATH_TFTEXT}
 cd ${SRC_PATH_TFTEXT}
 
 # The tftext build script queries GitHub, but these requests are sometimes
@@ -45,14 +48,16 @@ EOF
 ## T5X
 ###############################################################################
 
+ARG BASE_IMAGE
 FROM ${BASE_IMAGE} AS mealkit
 ARG URLREF_T5X
+ARG URLREF_AIRIO
 ARG SRC_PATH_TFTEXT
 ARG SRC_PATH_T5X
+ARG SRC_PATH_AIRIO
 
 # Preserve version information of tensorflow-text
 COPY --from=tftext-builder /opt/manifest.d/git-clone.yaml /opt/manifest.d/git-clone.yaml
-
 COPY --from=tftext-builder ${SRC_PATH_TFTEXT}/tensorflow_text*.whl /opt/
 RUN echo "tensorflow-text @ file://$(ls /opt/tensorflow_text*.whl)" >> /opt/pip-tools.d/requirements-t5x.in
 
@@ -78,6 +83,14 @@ if [[ "$(dpkg --print-architecture)" == "arm64" ]]; then
 
     sed -i "s/'protobuf/#'protobuf/" setup.py
     sed -i "s/'numpy/#'numpy/" setup.py
+
+
+    # airio pins grain==0.2.0, but the later does not have arm64 wheel.
+    # Need to bump grain to 0.2.2 to resolve the issue (https://github.com/google/airio/issues/257)
+    git-clone.sh ${URLREF_AIRIO} ${SRC_PATH_AIRIO}
+    sed -i "s/grain==0.2.0/grain/g" ${SRC_PATH_AIRIO}/setup.py
+    sed -i "s/'airio/#'airio/g" setup.py
+    echo "-e file://${SRC_PATH_AIRIO}" >> /opt/pip-tools.d/requirements-t5x.in
 fi
 if git diff --quiet; then
     echo "URL specs no longer present in select dependencies of t5x"

diff --git a/.github/container/manifest.yaml b/.github/container/manifest.yaml
@@ -78,8 +78,8 @@ fiddle:
 airio:
   url: https://github.com/google/airio.git
   tracking_ref: main
-  latest_verified_commit: cfca4a10de1491d76d2d00fcbd7142079837ca99
-  mode: pip-vcs
+  latest_verified_commit: 37109ff0d1059f885b9b11ef9058eca5d219d7cb
+  mode: git-clone
 clu:
   url: https://github.com/google/CommonLoopUtils.git
   tracking_ref: main

diff --git a/.github/workflows/_ci.yaml b/.github/workflows/_ci.yaml
@@ -144,6 +144,7 @@ jobs:
       EXTRA_BUILD_ARGS: |
         URLREF_T5X=${{ fromJson(inputs.SOURCE_URLREFS).T5X }}
         URLREF_TFTEXT=${{ fromJson(inputs.SOURCE_URLREFS).TENSORFLOW_TEXT }}
+        URLREF_AIRIO=${{ fromJson(inputs.SOURCE_URLREFS).AIRIO }}
     secrets: inherit
 
   build-upstream-pax: