-
Notifications
You must be signed in to change notification settings - Fork 275
/
run_tests.sh
executable file
·67 lines (54 loc) · 1.89 KB
/
run_tests.sh
1
2
3
4
5
6
7
8
9
10
11
12
13
14
15
16
17
18
19
20
21
22
23
24
25
26
27
28
29
30
31
32
33
34
35
36
37
38
39
40
41
42
43
44
45
46
47
48
49
50
51
52
53
54
55
56
57
58
59
60
61
62
63
64
65
66
67
#!/usr/bin/env bash
set -e -x
# Install the package (necessary for CLI tests).
# Requirements should already be cached in the docker image.
pip install -e .
# Log installed versions
echo "PIP FREEZE:"
pip freeze
exit_if_error() {
local exit_code=$1
shift
printf 'ERROR: %s\n' "$@" >&2
exit "$exit_code"
}
download_assets() {
set -e -x
mkdir -p axlearn/data/tokenizers/sentencepiece
mkdir -p axlearn/data/tokenizers/bpe
curl https://huggingface.co/t5-base/resolve/main/spiece.model -o axlearn/data/tokenizers/sentencepiece/t5-base
curl https://huggingface.co/FacebookAI/roberta-base/raw/main/merges.txt -o axlearn/data/tokenizers/bpe/roberta-base-merges.txt
curl https://huggingface.co/FacebookAI/roberta-base/raw/main/vocab.json -o axlearn/data/tokenizers/bpe/roberta-base-vocab.json
}
precommit_checks() {
set -e -x
pre-commit install
pre-commit run --all-files || exit_if_error $? "pre-commit failed."
# Run pytype separately to utilize all cpus and for better output.
pytype -j auto . || exit_if_error $? "pytype failed."
}
# Collect all background PIDs explicitly.
TEST_PIDS=()
download_assets
if [[ "${1:-x}" = "--skip-pre-commit" ]] ; then
SKIP_PRECOMMIT=true
shift
fi
# Skip pre-commit on parallel CI because it is run as a separate job.
if [[ "${SKIP_PRECOMMIT:-false}" = "false" ]] ; then
precommit_checks &
TEST_PIDS[$!]=1
fi
UNQUOTED_PYTEST_FILES=$(echo $1 | tr -d "'")
pytest --durations=100 -v -n auto \
-m "not (gs_login or tpu or high_cpu or fp64)" ${UNQUOTED_PYTEST_FILES} \
--dist worksteal &
TEST_PIDS[$!]=1
JAX_ENABLE_X64=1 pytest --durations=100 -v -n auto -v -m "fp64" --dist worksteal &
TEST_PIDS[$!]=1
# Use Bash 5.1's new wait -p feature to quit immediately if any subprocess fails to make error
# finding a bit easier.
while [ ${#TEST_PIDS[@]} -ne 0 ]; do
wait -n -p PID ${!TEST_PIDS[@]} || exit_if_error $? "Test failed."
unset TEST_PIDS[$PID]
done