Skip to content

Commit 496bd26

Browse files
committed
Move test_assets and logits generation script inside tests directory
1 parent 644eda6 commit 496bd26

52 files changed

Lines changed: 69 additions & 61 deletions

File tree

Some content is hidden

Large Commits have some content hidden by default. Use the searchbox below for content that may be hidden.

.github/workflows/build_and_upload_images.sh

Lines changed: 2 additions & 2 deletions
Original file line numberDiff line numberDiff line change
@@ -59,8 +59,8 @@ docker push gcr.io/$PROJECT/${dependency_image_name}:latest
5959
docker tag ${LOCAL_IMAGE_NAME} gcr.io/$PROJECT/${dependency_image_name}:${image_date}
6060
docker push gcr.io/$PROJECT/${dependency_image_name}:${image_date}
6161

62-
# Download other test assets from GCS into "${MAXTEXT_PKG_DIR:-${MAXTEXT_REPO_ROOT:-$PWD}/src/MaxText}"/test_assets
63-
if ! gcloud storage cp gs://maxtext-test-assets/* "${MAXTEXT_TEST_ASSETS_ROOT:-${MAXTEXT_REPO_ROOT:-$PWD}/test_assets}"; then
62+
# Download other test assets from GCS into ${MAXTEXT_TEST_ASSETS_ROOT:-${MAXTEXT_REPO_ROOT:-$PWD}/tests/assets/golden_logits}
63+
if ! gcloud storage cp gs://maxtext-test-assets/* "${MAXTEXT_TEST_ASSETS_ROOT:-${MAXTEXT_REPO_ROOT:-$PWD}/tests/assets/golden_logits}"; then
6464
echo "WARNING: Failed to download test assets from GCS. These files are only used for end-to-end tests; you may not have access to the bucket."
6565
fi
6666

.github/workflows/run_pathways_tests.yml

Lines changed: 2 additions & 2 deletions
Original file line numberDiff line numberDiff line change
@@ -83,7 +83,7 @@ jobs:
8383
python3 --version
8484
python3 -m pip freeze
8585
- name: Copy test assets files
86-
run : gcloud storage cp gs://maxtext-test-assets/* src/MaxText/test_assets
86+
run : gcloud storage cp gs://maxtext-test-assets/* tests/assets
8787
- name: Run Tests
8888
run: |
8989
if [ "${{ inputs.is_scheduled_run }}" = "true" ]; then
@@ -93,7 +93,7 @@ jobs:
9393
fi
9494
export MAXTEXT_REPO_ROOT=$(pwd)
9595
export MAXTEXT_ASSETS_ROOT=$(pwd)/src/MaxText/assets
96-
export MAXTEXT_TEST_ASSETS_ROOT=$(pwd)/src/MaxText/test_assets
96+
export MAXTEXT_TEST_ASSETS_ROOT=$(pwd)/tests/assets
9797
export MAXTEXT_PKG_DIR=$(pwd)/src/MaxText
9898
# TODO(b/454659463): Enable test_default_hlo_match after volume mount is supported.
9999
.venv/bin/python3 -m pytest ${{ inputs.pytest_addopts }} -v -m "${FINAL_PYTEST_MARKER}" -k "not AotHloIdenticalTest and not CompileThenLoad" --durations=0

.github/workflows/run_tests_against_package.yml

Lines changed: 2 additions & 2 deletions
Original file line numberDiff line numberDiff line change
@@ -91,7 +91,7 @@ jobs:
9191
python3 -m pip freeze
9292
uv pip install pytest-cov
9393
- name: Copy test assets files
94-
run : gcloud storage cp gs://maxtext-test-assets/* src/MaxText/test_assets
94+
run : gcloud storage cp gs://maxtext-test-assets/* tests/assets
9595
- name: Run Tests
9696
shell: bash
9797
run: |
@@ -103,7 +103,7 @@ jobs:
103103
# TODO: Use package data for testing and remove the env vars
104104
export MAXTEXT_REPO_ROOT=$(pwd)
105105
export MAXTEXT_ASSETS_ROOT=$(pwd)/src/MaxText/assets
106-
export MAXTEXT_TEST_ASSETS_ROOT=$(pwd)/src/MaxText/test_assets
106+
export MAXTEXT_TEST_ASSETS_ROOT=$(pwd)/tests/assets
107107
export MAXTEXT_PKG_DIR=$(pwd)/src/MaxText
108108
# omit this libtpu init args for gpu tests
109109
if [ "${{ inputs.device_type }}" != "cuda12" ]; then

codecov.yml

Lines changed: 0 additions & 1 deletion
Original file line numberDiff line numberDiff line change
@@ -38,7 +38,6 @@ ignore:
3838
- "src/MaxText/inference"
3939
- "src/MaxText/inference_mlperf"
4040
- "src/MaxText/scratch_code"
41-
- "src/MaxText/test_assets"
4241

4342

4443
flags:

dependencies/dockerfiles/maxtext_gpu_dependencies.Dockerfile

Lines changed: 2 additions & 2 deletions
Original file line numberDiff line numberDiff line change
@@ -39,7 +39,7 @@ ARG DEVICE
3939
ENV ENV_DEVICE=$DEVICE
4040

4141
ENV MAXTEXT_ASSETS_ROOT=/deps/src/MaxText/assets
42-
ENV MAXTEXT_TEST_ASSETS_ROOT=/deps/src/MaxText/test_assets
42+
ENV MAXTEXT_TEST_ASSETS_ROOT=/deps/tests/assets
4343
ENV MAXTEXT_PKG_DIR=/deps/src/MaxText
4444
ENV MAXTEXT_REPO_ROOT=/deps
4545

@@ -62,7 +62,7 @@ COPY . .
6262
ARG INCLUDE_TEST_ASSETS=false
6363
RUN if [ "$INCLUDE_TEST_ASSETS" = "true" ]; then \
6464
echo "Downloading test assets from GCS..."; \
65-
if ! gcloud storage cp -r gs://maxtext-test-assets/* "${MAXTEXT_TEST_ASSETS_ROOT}"; then \
65+
if ! gcloud storage cp -r gs://maxtext-test-assets/* "${MAXTEXT_TEST_ASSETS_ROOT}/golden_logits"; then \
6666
echo "WARNING: Failed to download test assets from GCS. These files are only used for end-to-end tests; you may not have access to the bucket."; \
6767
fi; \
6868
fi

dependencies/dockerfiles/maxtext_runner.Dockerfile

Lines changed: 2 additions & 2 deletions
Original file line numberDiff line numberDiff line change
@@ -6,7 +6,7 @@ FROM $BASEIMAGE
66
#FROM maxtext_base_image
77

88
ENV MAXTEXT_ASSETS_ROOT=/deps/src/MaxText/assets
9-
ENV MAXTEXT_TEST_ASSETS_ROOT=/deps/src/MaxText/test_assets
9+
ENV MAXTEXT_TEST_ASSETS_ROOT=/deps/tests/assets
1010
ENV MAXTEXT_PKG_DIR=/deps/src/MaxText
1111
ENV MAXTEXT_REPO_ROOT=/deps
1212

@@ -15,7 +15,7 @@ WORKDIR /deps
1515

1616
# Copy assets separately
1717
COPY src/MaxText/assets/ "${MAXTEXT_ASSETS_ROOT}"
18-
COPY src/MaxText/test_assets/ "${MAXTEXT_TEST_ASSETS_ROOT}"
18+
COPY tests/assets/ "${MAXTEXT_TEST_ASSETS_ROOT}"
1919

2020
# Copy all files except assets from local workspace into docker container
2121
COPY --exclude="${MAXTEXT_ASSETS_ROOT}" --exclude="${MAXTEXT_TEST_ASSETS_ROOT}" . .

dependencies/dockerfiles/maxtext_tpu_dependencies.Dockerfile

Lines changed: 2 additions & 2 deletions
Original file line numberDiff line numberDiff line change
@@ -33,7 +33,7 @@ ARG DEVICE
3333
ENV ENV_DEVICE=$DEVICE
3434

3535
ENV MAXTEXT_ASSETS_ROOT=/deps/src/MaxText/assets
36-
ENV MAXTEXT_TEST_ASSETS_ROOT=/deps/src/MaxText/test_assets
36+
ENV MAXTEXT_TEST_ASSETS_ROOT=/deps/tests/assets
3737
ENV MAXTEXT_PKG_DIR=/deps/src/MaxText
3838
ENV MAXTEXT_REPO_ROOT=/deps
3939

@@ -59,7 +59,7 @@ COPY . .
5959
ARG INCLUDE_TEST_ASSETS=false
6060
RUN if [ "$INCLUDE_TEST_ASSETS" = "true" ]; then \
6161
echo "Downloading test assets from GCS..."; \
62-
if ! gcloud storage cp -r gs://maxtext-test-assets/* "${MAXTEXT_TEST_ASSETS_ROOT}"; then \
62+
if ! gcloud storage cp -r gs://maxtext-test-assets/* "${MAXTEXT_TEST_ASSETS_ROOT}/golden_logits"; then \
6363
echo "WARNING: Failed to download test assets from GCS. These files are only used for end-to-end tests; you may not have access to the bucket."; \
6464
fi; \
6565
fi

dependencies/scripts/docker_upload_runner.sh

Lines changed: 2 additions & 2 deletions
Original file line numberDiff line numberDiff line change
@@ -76,8 +76,8 @@ if [ -n "$ABSOLUTE_LINKS" ]; then
7676
exit 1
7777
fi
7878

79-
# Download other test assets from GCS into "${MAXTEXT_PKG_DIR:-${MAXTEXT_REPO_ROOT:-$PWD}/src/MaxText}"/test_assets
80-
# if ! gcloud storage cp gs://maxtext-test-assets/* "${MAXTEXT_TEST_ASSETS_ROOT:-${MAXTEXT_REPO_ROOT:-$PWD}/test_assets}"; then
79+
# Download other test assets from GCS into ${MAXTEXT_TEST_ASSETS_ROOT:-${MAXTEXT_REPO_ROOT:-$PWD}}/tests/assets/golden_logits
80+
# if ! gcloud storage cp gs://maxtext-test-assets/* "${MAXTEXT_TEST_ASSETS_ROOT:-${MAXTEXT_REPO_ROOT:-$PWD}/tests/assets/golden_logits}"; then
8181
# echo "WARNING: Failed to download test assets from GCS. These files are only used for end-to-end tests; you may not have access to the bucket."
8282
# fi
8383

docs/tutorials/posttraining/multimodal.md

Lines changed: 1 addition & 1 deletion
Original file line numberDiff line numberDiff line change
@@ -84,7 +84,7 @@ python -m MaxText.decode \
8484
scan_layers=false \
8585
use_multimodal=true \
8686
prompt='Describe image <start_of_image>' \
87-
image_path='src/MaxText/test_assets/test_image.jpg' \
87+
image_path='tests/assets/test_image.jpg' \
8888
attention='dot_product'
8989
```
9090

end_to_end/tpu/deepseek/v2-16b/test_deepseek.sh

Lines changed: 3 additions & 3 deletions
Original file line numberDiff line numberDiff line change
@@ -10,7 +10,7 @@
1010
# Example Usage: export HF_TOKEN=<huggingface_access_token>; export BASE_OUTPUT_PATH=<GCS_bucket_path>; bash test_deepseek.sh
1111

1212
# The golden logit can be generated by:
13-
# python3 -m MaxText.scratch_code.generate_hf_golden_logits --model-id=deepseek-ai/DeepSeek-V2-Lite --output-path=golden_data_deepseek2-16b.jsonl --prompts='I love to;Today is a;What is the' --hf-model-path=$local_bf16_path --trust-remote-code=False
13+
# python3 -m tests.assets.logits_generation.generate_hf_golden_logits --model-id=deepseek-ai/DeepSeek-V2-Lite --output-path=golden_data_deepseek2-16b.jsonl --prompts='I love to;Today is a;What is the' --hf-model-path=$local_bf16_path --trust-remote-code=False
1414

1515
set -ex
1616

@@ -57,8 +57,8 @@ export UNSCANNED_CKPT_PATH=${BASE_OUTPUT_PATH}/unscanned/0/items
5757
export DATASET_PATH=gs://maxtext-dataset
5858

5959
# Test whether the forward pass logits match the golden logits
60-
# default golden_logits_path=/deps/src/MaxText/test_assets/golden_data_{MODEL_NAME}.jsonl, copied from gs://maxtext-test-assets/golden_data_${MODEL_NAME}.jsonl
61-
GOLDEN_LOGITS_DISK_LOCATION="/deps/src/MaxText/test_assets/golden_data_${MODEL_NAME}.jsonl"
60+
# default golden_logits_path=/deps/tests/assets/golden_logits/golden_data_{MODEL_NAME}.jsonl, copied from gs://maxtext-test-assets/golden_data_${MODEL_NAME}.jsonl
61+
GOLDEN_LOGITS_DISK_LOCATION="/deps/tests/assets/golden_logits/golden_data_${MODEL_NAME}.jsonl"
6262
if [ ! -f "${GOLDEN_LOGITS_DISK_LOCATION}" ]; then
6363
GOLDEN_LOGITS_PATH="gs://maxtext-test-assets/golden_data_${MODEL_NAME}.jsonl"
6464
GOLDEN_LOGITS_DISK_LOCATION=/tmp/golden_data.jsonl

0 commit comments

Comments
 (0)