Modify the MaxText image and the default paths for server_image and proxy_image

Author: DannyLiCom

benchmarks/recipes/user_configs.py

@@ -111,8 +111,8 @@ def __post_init__(self):
     device_type="v6e-256",
     benchmark_steps=20,
     num_slices_list=[2],
-    server_image="us-docker.pkg.dev/cloud-tpu-v2-images/pathways/proxy_server",
-    proxy_image="us-docker.pkg.dev/cloud-tpu-v2-images/pathways/server",
+    server_image="us-docker.pkg.dev/cloud-tpu-v2-images/pathways/server",
+    proxy_image="us-docker.pkg.dev/cloud-tpu-v2-images/pathways/proxy_server",
     runner="us-docker.pkg.dev/path/to/maxtext_runner",
     selected_model_framework=["pathways"],
     selected_model_names=["llama3_1_8b_8192"],

maxtext_runner.Dockerfile

@@ -14,8 +14,8 @@ ENV MAXTEXT_REPO_ROOT=/deps
 WORKDIR /deps
 
 # Copy assets separately
-COPY "${MAXTEXT_ASSETS_ROOT}" src/MaxText/assets/
-COPY "${MAXTEXT_TEST_ASSETS_ROOT}" src/MaxText/test_assets/
+COPY src/MaxText/assets/ "${MAXTEXT_ASSETS_ROOT}"
+COPY src/MaxText/test_assets/ "${MAXTEXT_TEST_ASSETS_ROOT}" 
 
 # Copy all files except assets from local workspace into docker container
 COPY --exclude="${MAXTEXT_ASSETS_ROOT}" --exclude="${MAXTEXT_TEST_ASSETS_ROOT}" . .