Move benchmarks and src/MaxText to docker image

Author: SurbhiJainUSC

.github/workflows/build_and_push_docker_image.yml

@@ -123,7 +123,6 @@ jobs:
             MODE=${{ inputs.build_mode }}
             WORKFLOW=${{ inputs.workflow }}
             PACKAGE_DIR=./src
-            TESTS_DIR=./tests
             JAX_VERSION=NONE
             LIBTPU_VERSION=NONE
             INCLUDE_TEST_ASSETS=true

src/dependencies/dockerfiles/maxtext_gpu_dependencies.Dockerfile

@@ -41,9 +41,6 @@ ENV ENV_DEVICE=$DEVICE
 ARG PACKAGE_DIR
 ENV PACKAGE_DIR=$PACKAGE_DIR
 
-ARG TESTS_DIR
-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
@@ -66,7 +63,9 @@ RUN --mount=type=cache,target=/root/.cache/uv \
 
 # Now copy the remaining code (source files that may change frequently)
 COPY ${PACKAGE_DIR}/maxtext/ src/maxtext/
-COPY ${TESTS_DIR}*/ tests/
+COPY ${PACKAGE_DIR}/MaxText/ src/MaxText/
+COPY tests*/ tests/
+COPY benchmarks*/ benchmarks/
 
 # Download test assets from GCS if building image with test assets
 ARG INCLUDE_TEST_ASSETS=false

src/dependencies/dockerfiles/maxtext_tpu_dependencies.Dockerfile

@@ -38,9 +38,6 @@ ENV ENV_DEVICE=$DEVICE
 ARG PACKAGE_DIR
 ENV PACKAGE_DIR=$PACKAGE_DIR
 
-ARG TESTS_DIR
-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
@@ -66,7 +63,9 @@ RUN --mount=type=cache,target=/root/.cache/uv \
 
 # Now copy the remaining code (source files that may change frequently)
 COPY ${PACKAGE_DIR}/maxtext/ src/maxtext/
-COPY ${TESTS_DIR}*/ tests/
+COPY ${PACKAGE_DIR}/MaxText/ src/MaxText/
+COPY tests*/ tests/
+COPY benchmarks*/ benchmarks/
 
 # Download test assets from GCS if building image with test assets
 ARG INCLUDE_TEST_ASSETS=false

src/dependencies/scripts/docker_build_dependency_image.sh

@@ -22,8 +22,6 @@
 
 PACKAGE_DIR="${PACKAGE_DIR:-src}"
 echo "PACKAGE_DIR: $PACKAGE_DIR"
-TESTS_DIR="${TESTS_DIR:-tests}"
-echo "TESTS_DIR: $TESTS_DIR"
 
 # Enable "exit immediately if any command fails" option
 set -e
@@ -73,7 +71,6 @@ docker_build_args=(
   "MODE=${MODE}"
   "JAX_VERSION=${JAX_VERSION}"
   "PACKAGE_DIR=${PACKAGE_DIR}"
-  "TESTS_DIR=${TESTS_DIR}"
 )
 
 run_docker_build() {