Skip to content
New issue

Have a question about this project? Sign up for a free GitHub account to open an issue and contact its maintainers and the community.

Sign up for GitHub

By clicking “Sign up for GitHub”, you agree to our terms of service and privacy statement. We’ll occasionally send you account related emails.

Already on GitHub? Sign in to your account

Jump to bottom

Introduces manifest.yaml that is the "last working world state" #405

Merged
merged 30 commits into from
Dec 18, 2023
Merged

Introduces manifest.yaml that is the "last working world state" #405

Show file tree
Hide file tree
Changes from 27 commits
Commits
Show all changes
30 commits
Select commit Hold shift + click to select a range
fa48558
wip
yhtang Oct 23, 2023
2ccf1a9
parent abb6f9704b8e32ed331ac56bd767338796224c3c
yhtang Oct 23, 2023
404d629
Adds support for building rosetta with local patches and an already
terrykong Oct 31, 2023
4404710
parent 8a43f4a741108ad610fb990a1c10b65dec607200
terrykong Nov 17, 2023
4951c31
fix pip patch
terrykong Dec 1, 2023
7a852f8
fix pip-finalize
terrykong Dec 1, 2023
6206f4c
add praxis TE patch from upstream
terrykong Dec 2, 2023
5688cac
apply xla fix reversion from main
terrykong Dec 2, 2023
83c07fd
style
terrykong Dec 2, 2023
27036f7
short -b switch
terrykong Dec 2, 2023
2ef9313
remove [optional]
terrykong Dec 2, 2023
ecf9a6a
submodule init
terrykong Dec 2, 2023
c5f0b87
.ref -> .latest_verified_commit
terrykong Dec 4, 2023
645d4af
.ref -> .latest_verified_commit (pt2)
terrykong Dec 4, 2023
e780f9b
better description in create-distribution.sh
terrykong Dec 4, 2023
53ee07e
typo
terrykong Dec 4, 2023
f2f13f5
remove quotes
terrykong Dec 5, 2023
48a3a09
EOF & INNEREOF
terrykong Dec 5, 2023
e8b857b
bump.sh documentation
terrykong Dec 5, 2023
68b001a
rm [Optional]
terrykong Dec 5, 2023
7281f59
get-source.sh documentation
terrykong Dec 5, 2023
2c54e8d
trial branch description
terrykong Dec 5, 2023
2e9fccd
revert sandbox
terrykong Dec 5, 2023
a554f01
small comment + set -u flag in pip-finalize
terrykong Dec 7, 2023
723f91d
fix merge conflict but also allow pass thru of base_build to jax build
terrykong Dec 8, 2023
c0b683a
Fix BASE_IMAGE description and upstream t5x/pax builds now allow
terrykong Dec 11, 2023
0828455
cleanup
terrykong Dec 11, 2023
73df444
_build_rosetta.yaml BADGE_FILENAME -> BADGE_FILENAME_PREFIX
terrykong Dec 13, 2023
811c9d2
Remove trial-branch default 'main' which is confusing
terrykong Dec 13, 2023
aa7cea8
Merge branch 'main' into the-manifest
yhtang Dec 18, 2023
File filter

Filter by extension

Filter by extension
Conversations
Failed to load comments.
Loading
Jump to
Jump to file
Failed to load files.
Loading
Diff view
Diff view
14 changes: 13 additions & 1 deletion .github/container/Dockerfile.base
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -38,12 +38,24 @@ RUN <<"EOF" bash -ex
git config --global user.name "${GIT_USER_NAME}"
git config --global user.email "${GIT_USER_EMAIL}"
EOF
RUN pip install --upgrade --no-cache-dir pip pip-tools && rm -rf ~/.cache/*
RUN mkdir -p /opt/pip-tools.d
ADD --chmod=777 \
get-source.sh \
pip-finalize.sh \
/usr/local/bin/
RUN wget https://github.com/mikefarah/yq/releases/latest/download/yq_linux_$(dpkg --print-architecture) -O /usr/local/bin/yq && \
Copy link
Contributor

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

[question] Will it be better to install it thru apt install:

$> add-apt-repository ppa:rmescandon/yq
$> apt update
$> apt install yq -y

?

Copy link
Contributor Author

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

I guess it would be better to have it be managed by apt, then the package can be inspected by dpkg. FWIW, I had to install apt-get install software-properties-common to get add-apt-respository which is 68.4MB according to apt. Is that okay with you?

Copy link
Collaborator

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

We don't have a guideline yet regarding third-party PPAs. Let's stick with this simple wget for the time being.

chmod 777 /usr/local/bin/yq
RUN git clone -b 23.3.1 https://github.com/pypa/pip.git /opt/pip
# Patch is specific to 23.3.1
# Generated via: "git diff > pip-vcs-equivalency.patch"
ADD pip-vcs-equivalency.patch /opt/pip/
terrykong marked this conversation as resolved.
Show resolved Hide resolved
RUN <<EOF bash -e -x
cd /opt/pip
git apply </opt/pip/pip-vcs-equivalency.patch
git add -u
git commit -m 'Adds JAX_TOOLBOX_VCS_EQUIVALENCY as a trigger to treat all github VCS installs for a package as equivalent. The spec of the last encountered version will be used'
EOF
RUN pip install --upgrade --no-cache-dir -e /opt/pip pip-tools && rm -rf ~/.cache/*

###############################################################################
## Install cuDNN
Expand Down
66 changes: 37 additions & 29 deletions .github/container/Dockerfile.jax
View file
Open in desktop
Original file line number Diff line number Diff line change
@@ -1,16 +1,10 @@
ARG BASE_IMAGE=ghcr.io/nvidia/jax-toolbox:base
ARG REPO_JAX="https://github.com/google/jax.git"
ARG REPO_XLA="https://github.com/openxla/xla.git"
ARG REPO_FLAX="https://github.com/google/flax.git"
ARG REPO_TE="https://github.com/NVIDIA/TransformerEngine.git"
ARG REF_JAX=main
ARG REF_XLA=main
ARG REF_FLAX=main
ARG REF_TE=main
ARG SRC_MANIFEST_FILE=manifest.yaml
ARG DEST_MANIFEST_DIR=/opt/manifest.d
ARG SRC_PATH_JAX=/opt/jax
ARG SRC_PATH_XLA=/opt/xla-source
ARG SRC_PATH_XLA=/opt/xla
ARG SRC_PATH_FLAX=/opt/flax
ARG SRC_PATH_TE=/opt/transformer-engine-source
ARG SRC_PATH_TE=/opt/transformer-engine
ARG GIT_USER_NAME="JAX Toolbox"
ARG [email protected]

Expand All @@ -22,22 +16,28 @@ ARG BUILD_DATE
###############################################################################

FROM ${BASE_IMAGE} as builder
ARG REPO_JAX
ARG REPO_XLA
ARG REF_JAX
ARG REF_XLA
ARG SRC_MANIFEST_FILE
ARG DEST_MANIFEST_DIR
ARG SRC_PATH_JAX
ARG SRC_PATH_XLA
ARG BAZEL_CACHE
ARG GIT_USER_NAME
ARG GIT_USER_EMAIL

RUN git clone "${REPO_JAX}" "${SRC_PATH_JAX}" && cd "${SRC_PATH_JAX}" && git checkout ${REF_JAX}
# Set the manifest env vars
ENV MANIFEST_FILE=${DEST_MANIFEST_DIR}/${SRC_MANIFEST_FILE}

ADD --chmod=777 create-distribution.sh ${DEST_MANIFEST_DIR}/
COPY ${SRC_MANIFEST_FILE} ${DEST_MANIFEST_DIR}/${SRC_MANIFEST_FILE}
terrykong marked this conversation as resolved.
Show resolved Hide resolved
COPY patches/ ${DEST_MANIFEST_DIR}/patches/

RUN get-source.sh -l jax -m ${MANIFEST_FILE}
RUN --mount=type=ssh \
--mount=type=secret,id=SSH_KNOWN_HOSTS,target=/root/.ssh/known_hosts \
git clone "${REPO_XLA}" "${SRC_PATH_XLA}" && cd "${SRC_PATH_XLA}" && git checkout ${REF_XLA}
get-source.sh -l xla -m ${MANIFEST_FILE}

ADD build-jax.sh local_cuda_arch test-jax.sh /usr/local/bin/
# TODO: move this patch into the manifest
ADD xla-arm64-neon.patch /opt
RUN build-jax.sh \
--bazel-cache ${BAZEL_CACHE} \
Expand All @@ -53,46 +53,53 @@ RUN build-jax.sh \

ARG BASE_IMAGE
FROM ${BASE_IMAGE} as mealkit
ARG SRC_MANIFEST_FILE
ARG DEST_MANIFEST_DIR
ARG SRC_PATH_JAX
ARG SRC_PATH_XLA
ARG SRC_PATH_TE
ARG BUILD_DATE

ENV BUILD_DATE=${BUILD_DATE}
# The following environment variables tune performance
ENV XLA_FLAGS="--xla_gpu_enable_latency_hiding_scheduler=true --xla_gpu_enable_async_all_gather=true --xla_gpu_enable_async_reduce_scatter=true --xla_gpu_enable_triton_gemm=false"
ENV CUDA_DEVICE_MAX_CONNECTIONS=1
ENV NCCL_IB_SL=1
ENV NCCL_NVLS_ENABLE=0
ENV CUDA_MODULE_LOADING=EAGER
# Set the manifest env vars
ENV MANIFEST_FILE=${DEST_MANIFEST_DIR}/${SRC_MANIFEST_FILE}

ADD --chmod=777 create-distribution.sh ${DEST_MANIFEST_DIR}/
COPY ${SRC_MANIFEST_FILE} ${DEST_MANIFEST_DIR}/${SRC_MANIFEST_FILE}
terrykong marked this conversation as resolved.
Show resolved Hide resolved
COPY patches/ ${DEST_MANIFEST_DIR}/patches/

COPY --from=builder ${SRC_PATH_JAX} ${SRC_PATH_JAX}
COPY --from=builder ${SRC_PATH_XLA} ${SRC_PATH_XLA}
ADD build-jax.sh local_cuda_arch test-jax.sh /usr/local/bin/

COPY --from=builder ${SRC_PATH_JAX} ${SRC_PATH_JAX}
COPY --from=builder ${SRC_PATH_XLA} ${SRC_PATH_XLA}
ADD build-jax.sh local_cuda_arch test-jax.sh /usr/local/bin/

RUN mkdir -p /opt/pip-tools.d
RUN <<"EOF" bash -ex
echo "-e file://${SRC_PATH_JAX}" >> /opt/pip-tools.d/manifest.jax
echo "jaxlib @ file://$(ls ${SRC_PATH_JAX}/dist/*.whl)" >> /opt/pip-tools.d/manifest.jax
echo "-e file://${SRC_PATH_JAX}" >> /opt/pip-tools.d/requirements-jax.in
echo "jaxlib @ file://$(ls ${SRC_PATH_JAX}/dist/*.whl)" >> /opt/pip-tools.d/requirements-jax.in
EOF

## Flax
ARG REPO_FLAX
ARG REF_FLAX
ARG SRC_PATH_FLAX
RUN get-source.sh -f ${REPO_FLAX} -r ${REF_FLAX} -d ${SRC_PATH_FLAX} -m /opt/pip-tools.d/manifest.flax
RUN get-source.sh -l flax -m ${MANIFEST_FILE} -o /opt/pip-tools.d/requirements-flax.in

## Transformer engine: check out source and build wheel
ARG REPO_TE
ARG REF_TE
ARG SRC_PATH_TE
ENV NVTE_FRAMEWORK=jax
ENV SRC_PATH_TE=${SRC_PATH_TE}
RUN <<"EOF" bash -ex
set -o pipefail
RUN <<"EOF" bash -ex -o pipefail
pip install ninja && rm -rf ~/.cache/pip
get-source.sh -f ${REPO_TE} -r ${REF_TE} -d ${SRC_PATH_TE}
get-source.sh -l transformer-engine -m ${MANIFEST_FILE}
pushd ${SRC_PATH_TE}
python setup.py bdist_wheel && rm -rf build
echo "transformer-engine @ file://$(ls ${SRC_PATH_TE}/dist/*.whl)" >> /opt/pip-tools.d/manifest.te
echo "transformer-engine @ file://$(ls ${SRC_PATH_TE}/dist/*.whl)" >> /opt/pip-tools.d/requirements-te.in
EOF

# TODO: properly configure entrypoint
Expand All @@ -104,3 +111,4 @@ EOF
FROM mealkit as final

RUN pip-finalize.sh

18 changes: 5 additions & 13 deletions .github/container/Dockerfile.pax.amd64
View file
Open in desktop
Original file line number Diff line number Diff line change
@@ -1,10 +1,6 @@
# syntax=docker/dockerfile:1-labs

ARG BASE_IMAGE=ghcr.io/nvidia/jax:mealkit
ARG REPO_PAXML=https://github.com/google/paxml.git
ARG REPO_PRAXIS=https://github.com/google/praxis.git
ARG REF_PAXML=main
ARG REF_PRAXIS=main
ARG SRC_PATH_PAXML=/opt/paxml
ARG SRC_PATH_PRAXIS=/opt/praxis

Expand All @@ -13,21 +9,17 @@ ARG SRC_PATH_PRAXIS=/opt/praxis
###############################################################################

FROM ${BASE_IMAGE} as mealkit
ARG REPO_PAXML
ARG REPO_PRAXIS
ARG REF_PAXML
ARG REF_PRAXIS
ARG SRC_PATH_PAXML
ARG SRC_PATH_PRAXIS

# update TE manifest file to install the [test] extras
Copy link
Contributor

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

Why here but not in Dockerfile.jax? TE is a part of JAX image now, so I would assume someone wants to test TE in JAX container.

Copy link
Contributor Author

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

I don't mind. Any reason you kept it separate @yhtang ?

Copy link
Contributor Author

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

So I took a look at it, and it appears we could merge the two without any issue since there's just a sed in Dockerfile.pax.amd64 that replaces the line, so I think it'd be fine to do that in requirements-jax.in

Any objections? @DwarKapex @yhtang

Actually, another issue I found looking at this is that Dockerfile.pax.arm64 doesn't include this sed. Is that intentional @yhtang ?

Copy link
Collaborator

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

Yes. As I mentioned during the CI sync, IIRC the [test] extra could not be installed in ARM64.

Copy link
Collaborator

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

We can address this as part of #338.

RUN sed -i "s/transformer-engine @/transformer-engine[test] @/g" /opt/pip-tools.d/manifest.te
RUN sed -i "s/transformer-engine @/transformer-engine[test] @/g" /opt/pip-tools.d/requirements-te.in

RUN <<"EOF" bash -ex
get-source.sh -f ${REPO_PAXML} -r ${REF_PAXML} -d ${SRC_PATH_PAXML}
get-source.sh -f ${REPO_PRAXIS} -r ${REF_PRAXIS} -d ${SRC_PATH_PRAXIS}
echo "-e file://${SRC_PATH_PAXML}[gpu]" >> /opt/pip-tools.d/manifest.pax
echo "-e file://${SRC_PATH_PRAXIS}" >> /opt/pip-tools.d/manifest.pax
get-source.sh -l paxml -m ${MANIFEST_FILE}
get-source.sh -l praxis -m ${MANIFEST_FILE}
echo "-e file://${SRC_PATH_PAXML}[gpu]" >> /opt/pip-tools.d/requirements-paxml.in
echo "-e file://${SRC_PATH_PRAXIS}" >> /opt/pip-tools.d/requirements-paxml.in

for src in ${SRC_PATH_PAXML} ${SRC_PATH_PRAXIS}; do
pushd ${src}
Expand Down
67 changes: 29 additions & 38 deletions .github/container/Dockerfile.pax.arm64
View file
Open in desktop
Original file line number Diff line number Diff line change
@@ -1,12 +1,10 @@
# syntax=docker/dockerfile:1-labs

ARG BASE_IMAGE=ghcr.io/nvidia/jax:mealkit
ARG REPO_PAXML=https://github.com/google/paxml.git
ARG REPO_PRAXIS=https://github.com/google/praxis.git
ARG REF_PAXML=main
ARG REF_PRAXIS=main
ARG SRC_PATH_PAXML=/opt/paxml
ARG SRC_PATH_PRAXIS=/opt/praxis
ARG SRC_PATH_TFTEXT=/opt/tensorflow-text
ARG SRC_PATH_LINGVO=/opt/lingvo

###############################################################################
## build tensorflow-text and lingvo, which do not have working arm64 pip wheels
Expand All @@ -24,39 +22,35 @@ RUN wget https://github.com/bazelbuild/bazelisk/releases/download/v1.17.0/bazeli
#------------------------------------------------------------------------------

FROM wheel-builder as tftext-builder

RUN <<"EOT" bash -exu
set -o pipefail
ARG SRC_PATH_TFTEXT
RUN <<"EOF" bash -exu -o pipefail
pip install tensorflow_datasets==4.9.2 auditwheel tensorflow==2.13.0
git clone http://github.com/tensorflow/text.git /opt/tensorflow-text
cd /opt/tensorflow-text
git checkout v2.13.0
get-source.sh -l tensorflow-text -m ${MANIFEST_FILE}
cd ${SRC_PATH_TFTEXT}
./oss_scripts/run_build.sh
EOT
EOF

#------------------------------------------------------------------------------
# build lingvo
#------------------------------------------------------------------------------

FROM wheel-builder as lingvo-builder
ARG REPO_LINGVO=https://github.com/tensorflow/lingvo.git
ARG REF_LINGVO=master
ARG SRC_PATH_LINGVO=/opt/lingvo
ARG SRC_PATH_TFTEXT
ARG SRC_PATH_LINGVO

COPY --from=tftext-builder /opt/tensorflow-text/tensorflow_text*.whl /opt/
COPY --from=tftext-builder ${SRC_PATH_TFTEXT}/tensorflow_text*.whl /opt/

RUN get-source.sh -f ${REPO_LINGVO} -r ${REF_LINGVO} -d ${SRC_PATH_LINGVO}
RUN get-source.sh -l lingvo -m ${MANIFEST_FILE}

# build lingvo
RUN <<"EOT" bash -exu
set -o pipefail
RUN <<"EOF" bash -exu -o pipefail
Copy link
Collaborator

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

I used EOT in consistency with Docker's official guide. This way, you don't have to rename the inner EOF to EOFINNER.

Copy link
Contributor Author

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

I mean we don't have to rename, but I feel that it's more clear since a question already came up about its significance. This way the identifier is also self-documenting


pushd ${SRC_PATH_LINGVO}
git fetch origin pull/329/head:pr329
git cherry-pick --allow-empty pr329

# Disable 2 flaky tests here
terrykong marked this conversation as resolved.
Show resolved Hide resolved
patch -p1 <<"EOF"
patch -p1 <<"EOFINNER"
diff --git a/pip_package/build.sh b/pip_package/build.sh
index ef62c432e..659e78956 100755
--- a/pip_package/build.sh
Expand All @@ -70,7 +64,7 @@ index ef62c432e..659e78956 100755
fi

DST_DIR="/tmp/lingvo/dist"
EOF
EOFINNER

pip install tensorflow_datasets==4.9.2 auditwheel tensorflow==2.13.0 /opt/tensorflow_text*.whl
sed -i 's/tensorflow=/#tensorflow=/' docker/dev.requirements.txt
Expand All @@ -82,38 +76,35 @@ pip install -r docker/dev.requirements.txt
# running the tests entirely by uncommentin the following line.
# SKIP_TEST=1
PYTHON_MINOR_VERSION=$(python --version | cut -d ' ' -f 2 | cut -d '.' -f 2) pip_package/build.sh
EOT
EOF

###############################################################################
## Pax for AArch64
###############################################################################

ARG BASE_IMAGE
FROM ${BASE_IMAGE} as mealkit
ARG REPO_PAXML
ARG REPO_PRAXIS
ARG REF_PAXML
ARG REF_PRAXIS
ARG SRC_PATH_PAXML
ARG SRC_PATH_PRAXIS
ARG SRC_PATH_TFTEXT

COPY --from=lingvo-builder /tmp/lingvo/dist/lingvo*linux_aarch64.whl /opt/
RUN echo "lingvo @ file://$(ls /opt/lingvo*.whl)" >> /opt/pip-tools.d/manifest.pax
RUN echo "lingvo @ file://$(ls /opt/lingvo*.whl)" >> /opt/pip-tools.d/requirements-paxml.in

COPY --from=tftext-builder /opt/tensorflow-text/tensorflow_text*.whl /opt/
RUN echo "tensorflow-text @ file://$(ls /opt/tensorflow_text*.whl)" >> /opt/pip-tools.d/manifest.pax
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-paxml.in

# paxml + praxis
RUN <<"EOT" bash -ex
echo "tensorflow==2.13.0" >> /opt/pip-tools.d/manifest.pax
echo "tensorflow_datasets==4.9.2" >> /opt/pip-tools.d/manifest.pax
echo "chex==0.1.7" >> /opt/pip-tools.d/manifest.pax
echo "auditwheel" >> /opt/pip-tools.d/manifest.pax
RUN <<"EOF" bash -ex
echo "tensorflow==2.13.0" >> /opt/pip-tools.d/requirements-paxml.in
echo "tensorflow_datasets==4.9.2" >> /opt/pip-tools.d/requirements-paxml.in
echo "chex==0.1.7" >> /opt/pip-tools.d/requirements-paxml.in
Copy link
Contributor

@DwarKapex DwarKapex Dec 5, 2023 •
edited
Loading

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

Why do I see chex==0.1.7 here, but not in line 27?
I wonder, can we do something like this:

26: ENV EXTRA_PACKAGES="tensorflow==2.13.0 tensorflow_datasets==4.9.2 chex==0.1.7 auditwheel"
27: RUN <<"EOT" bash -exu -o pipefail
28: pip install ${EXTRA_PACKAGES}
29: get-source.sh -l tensorflow-text -m ${MANIFEST_FILE}
30: cd ${SRC_PATH_TFTEXT}
31: ./oss_scripts/run_build.sh
32: EOT
<blah-blah-blah>
98:  RUN <<EOT bash -ex
99:  for spec_pkg in $(tr " " "\n"); do
100:    echo ${spec_pkg} >> /opt/pip-tools.d/requirements-paxml.in
101: done

?
Btw, what is the difference between EOT vs EOF?

Copy link
Contributor Author

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

I didn't author this dockerfile so I defer to @joker-eph regarding chex

W.r.t EOT vs EOF, either works, but one section of the dockerfile has two nested here-docs, so you need to choose a different char sequence, otherwise bash will end the heredoc early. There's one above it using EOT, but not nested; I can change that to EOF, which is more common

Copy link
Contributor Author

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

I can rename the nested one to EOF and EOFINNER to be more clear

Copy link
Member

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

If I wrote it, that was likely because it was required at the time to make the container work on ARM64.
We have lot of issues with incompleteness of the ARM64 pip packages, and the order in which we install things, their incompatible requirements, etc.

Ultimately line 27 and here are trying to get to a different state (earlier is just the requirement for building PAXML I believe?).
It is also very possible that the situation evolved since (that was over the summer if I remember correctly) and remove this from here would work.

Copy link
Collaborator

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

Regarding EOT vs EOF, EOT is what Docker's official documentation used for demoing heredoc RUN. But ultimately this choice is arbitrary and up to us to reach a consensus.

echo "auditwheel" >> /opt/pip-tools.d/requirements-paxml.in

get-source.sh -f ${REPO_PAXML} -r ${REF_PAXML} -d ${SRC_PATH_PAXML}
get-source.sh -f ${REPO_PRAXIS} -r ${REF_PRAXIS} -d ${SRC_PATH_PRAXIS}
echo "-e file://${SRC_PATH_PAXML}[gpu]" >> /opt/pip-tools.d/manifest.pax
echo "-e file://${SRC_PATH_PRAXIS}" >> /opt/pip-tools.d/manifest.pax
get-source.sh -l paxml -m ${MANIFEST_FILE}
get-source.sh -l praxis -m ${MANIFEST_FILE}
echo "-e file://${SRC_PATH_PAXML}[gpu]" >> /opt/pip-tools.d/requirements-paxml.in
echo "-e file://${SRC_PATH_PRAXIS}" >> /opt/pip-tools.d/requirements-paxml.in

for src in ${SRC_PATH_PAXML} ${SRC_PATH_PRAXIS}; do
pushd ${src}
Expand All @@ -139,7 +130,7 @@ for src in ${SRC_PATH_PAXML} ${SRC_PATH_PRAXIS}; do
fi
popd
done
EOT
EOF

ADD test-pax.sh /usr/local/bin

Expand Down
8 changes: 2 additions & 6 deletions .github/container/Dockerfile.t5x
View file
Open in desktop
Original file line number Diff line number Diff line change
@@ -1,8 +1,6 @@
# syntax=docker/dockerfile:1-labs

ARG BASE_IMAGE=ghcr.io/nvidia/jax:mealkit
ARG REPO_T5X=https://github.com/google-research/t5x.git
ARG REF_T5X=main
ARG SRC_PATH_T5X=/opt/t5x

###############################################################################
Expand All @@ -11,13 +9,11 @@ ARG SRC_PATH_T5X=/opt/t5x

FROM ${BASE_IMAGE} as mealkit

ARG REPO_T5X
ARG REF_T5X
ARG SRC_PATH_T5X

RUN <<"EOF" bash -ex
get-source.sh -f ${REPO_T5X} -r ${REF_T5X} -d ${SRC_PATH_T5X}
echo "-e file://${SRC_PATH_T5X}[gpu]" >> /opt/pip-tools.d/manifest.t5x
get-source.sh -l t5x -m ${MANIFEST_FILE}
echo "-e file://${SRC_PATH_T5X}[gpu]" >> /opt/pip-tools.d/requirements-t5x.in

# remove head-of-tree specs from select dependencies
pushd ${SRC_PATH_T5X}
Expand Down
Loading
Loading