diff --git a/src/dependencies/dockerfiles/maxtext_gpu_dependencies.Dockerfile b/src/dependencies/dockerfiles/maxtext_gpu_dependencies.Dockerfile index b9d7cdf80a..4d6a81a881 100644 --- a/src/dependencies/dockerfiles/maxtext_gpu_dependencies.Dockerfile +++ b/src/dependencies/dockerfiles/maxtext_gpu_dependencies.Dockerfile @@ -46,7 +46,7 @@ ENV TESTS_DIR=$TESTS_DIR ENV MAXTEXT_ASSETS_ROOT=/deps/src/maxtext/assets ENV MAXTEXT_TEST_ASSETS_ROOT=/deps/tests/assets -ENV MAXTEXT_PKG_DIR=/deps/src/MaxText +ENV MAXTEXT_PKG_DIR=/deps/src/maxtext ENV MAXTEXT_REPO_ROOT=/deps # Set the working directory in the container @@ -56,7 +56,7 @@ WORKDIR /deps COPY ${PACKAGE_DIR}/dependencies/github_deps/ src/dependencies/github_deps/ COPY ${PACKAGE_DIR}/dependencies/requirements/ src/dependencies/requirements/ COPY ${PACKAGE_DIR}/dependencies/scripts/ src/dependencies/scripts/ -COPY ${PACKAGE_DIR}/maxtext/integration/vllm/ src/MaxText/integration/vllm/ +COPY ${PACKAGE_DIR}/maxtext/integration/vllm/ src/maxtext/integration/vllm/ # Install dependencies - these steps are cached unless the copied files change RUN echo "Running command: bash setup.sh MODE=$ENV_MODE JAX_VERSION=$ENV_JAX_VERSION DEVICE=${ENV_DEVICE}" @@ -65,7 +65,7 @@ RUN --mount=type=cache,target=/root/.cache/uv \ bash /deps/src/dependencies/scripts/setup.sh MODE=${ENV_MODE} JAX_VERSION=${ENV_JAX_VERSION} DEVICE=${ENV_DEVICE} # Now copy the remaining code (source files that may change frequently) -COPY ${PACKAGE_DIR}/maxtext/ src/MaxText/ +COPY ${PACKAGE_DIR}/maxtext/ src/maxtext/ COPY ${TESTS_DIR}*/ tests/ # Download test assets from GCS if building image with test assets diff --git a/src/dependencies/dockerfiles/maxtext_runner.Dockerfile b/src/dependencies/dockerfiles/maxtext_runner.Dockerfile index f8c721a29a..d85511c848 100644 --- a/src/dependencies/dockerfiles/maxtext_runner.Dockerfile +++ b/src/dependencies/dockerfiles/maxtext_runner.Dockerfile @@ -8,7 +8,7 @@ ENV PACKAGE_DIR=$PACKAGE_DIR ENV MAXTEXT_ASSETS_ROOT=/deps/src/maxtext/assets ENV MAXTEXT_TEST_ASSETS_ROOT=/deps/tests/assets -ENV MAXTEXT_PKG_DIR=/deps/src/MaxText +ENV MAXTEXT_PKG_DIR=/deps/src/maxtext ENV MAXTEXT_REPO_ROOT=/deps # Set the working directory in the container @@ -18,4 +18,4 @@ WORKDIR /deps COPY ${PACKAGE_DIR}/maxtext/assets/ "${MAXTEXT_ASSETS_ROOT}" # Copy all files except assets from local workspace into docker container -COPY --exclude=${PACKAGE_DIR}/maxtext/assets/ ${PACKAGE_DIR}/maxtext/ src/MaxText/ +COPY --exclude=${PACKAGE_DIR}/maxtext/assets/ ${PACKAGE_DIR}/maxtext/ src/maxtext/