From 3c2ec9706fb09f8ebf7e4c43bd0f00aa6626d5f9 Mon Sep 17 00:00:00 2001
From: Vladislav Kozlov <vkozlov@nvidia.com>
Date: Wed, 27 Nov 2024 00:02:09 -0800
Subject: [PATCH] Clone airio

---
 .github/container/Dockerfile.t5x | 17 +++++++++++++++--
 .github/container/manifest.yaml  |  4 ++--
 .github/workflows/_ci.yaml       |  1 +
 3 files changed, 18 insertions(+), 4 deletions(-)

diff --git a/.github/container/Dockerfile.t5x b/.github/container/Dockerfile.t5x
index dec39fce8..1d6494f09 100644
--- 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
index 5ccf5fe75..b9c06e2e6 100644
--- 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
index 6cef657db..2f77c07f4 100644
--- 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: