Skip to content

Commit 2970cdc

Browse files
py4Pooya Moradi
andauthored
[Misc] Change buildkite pipeline to run all steps but skip some through command (#1005)
Signed-off-by: Pooya Moradi <pooyam@google.com> Co-authored-by: Pooya Moradi <pooyam@google.com>
1 parent 2bba577 commit 2970cdc

File tree

1 file changed

+65
-24
lines changed

1 file changed

+65
-24
lines changed

.buildkite/pipeline_jax.yml

Lines changed: 65 additions & 24 deletions
Original file line numberDiff line numberDiff line change
@@ -13,24 +13,34 @@ steps:
1313
- label: "E2E MLPerf tests for JAX models with quantization"
1414
key: test_1
1515
soft_fail: true
16-
if: build.env("NIGHTLY") == "1"
1716
env:
1817
QUANTIZATION: "True"
1918
agents:
2019
queue: tpu_v6e_queue
2120
commands:
22-
- .buildkite/scripts/run_in_docker.sh bash /workspace/tpu_inference/tests/e2e/benchmarking/mlperf.sh
21+
- |
22+
if [[ "$$NIGHTLY" == "1" ]]; then
23+
.buildkite/scripts/run_in_docker.sh bash /workspace/tpu_inference/tests/e2e/benchmarking/mlperf.sh
24+
else
25+
echo "Skipping: NIGHTLY environment variable not set"
26+
exit 0
27+
fi
2328
2429
- label: "E2E MLPerf tests for JAX new models"
2530
key: test_2
2631
soft_fail: true
27-
if: build.env("NIGHTLY") == "1"
2832
env:
2933
NEW_MODEL_DESIGN: "True"
3034
agents:
3135
queue: tpu_v6e_queue
3236
commands:
33-
- .buildkite/scripts/run_in_docker.sh bash /workspace/tpu_inference/tests/e2e/benchmarking/mlperf.sh
37+
- |
38+
if [[ "$$NIGHTLY" == "1" ]]; then
39+
.buildkite/scripts/run_in_docker.sh bash /workspace/tpu_inference/tests/e2e/benchmarking/mlperf.sh
40+
else
41+
echo "Skipping: NIGHTLY environment variable not set"
42+
exit 0
43+
fi
3444
3545
- label: "E2E MLPerf tests for JAX + vLLM models on single chip"
3646
key: test_3
@@ -45,38 +55,51 @@ steps:
4555
- label: "E2E MLperf tests for Llama4 models"
4656
key: test_4
4757
soft_fail: true
48-
if: build.env("NIGHTLY") == "1"
4958
env:
5059
NEW_MODEL_DESIGN: "True"
5160
USE_V6E8_QUEUE: "True"
5261
agents:
5362
queue: tpu_v6e_8_queue
5463
commands:
55-
- .buildkite/scripts/run_in_docker.sh bash /workspace/tpu_inference/tests/e2e/benchmarking/mlperf.sh
64+
- |
65+
if [[ "$$NIGHTLY" == "1" ]]; then
66+
.buildkite/scripts/run_in_docker.sh bash /workspace/tpu_inference/tests/e2e/benchmarking/mlperf.sh
67+
else
68+
echo "Skipping: NIGHTLY environment variable not set"
69+
exit 0
70+
fi
5671
5772
5873
- label: "E2E multi modality test"
5974
key: test_5
6075
soft_fail: true
61-
if: build.env("NIGHTLY") == "1"
6276
agents:
6377
queue: tpu_v6e_queue
6478
commands:
6579
- |
66-
.buildkite/scripts/run_in_docker.sh \
67-
bash -c 'python3 -m pytest -s -v -x /workspace/tpu_inference/tests/e2e/test_multi_modal_inference.py && \
68-
bash /workspace/tpu_inference/tests/e2e/benchmarking/mm_bench.sh'
80+
if [[ "$$NIGHTLY" == "1" ]]; then
81+
.buildkite/scripts/run_in_docker.sh \
82+
bash -c 'python3 -m pytest -s -v -x /workspace/tpu_inference/tests/e2e/test_multi_modal_inference.py && \
83+
bash /workspace/tpu_inference/tests/e2e/benchmarking/mm_bench.sh'
84+
else
85+
echo "Skipping: NIGHTLY environment variable not set"
86+
exit 0
87+
fi
6988
7089
- label: "E2E speculative decoding test"
7190
key: test_6
7291
soft_fail: true
73-
if: build.env("NIGHTLY") == "1"
7492
agents:
7593
queue: tpu_v6e_queue
7694
commands:
7795
- |
78-
.buildkite/scripts/run_in_docker.sh \
79-
bash -c 'python3 -m pytest -s -v -x /workspace/tpu_inference/tests/e2e/test_speculative_decoding.py'
96+
if [[ "$$NIGHTLY" == "1" ]]; then
97+
.buildkite/scripts/run_in_docker.sh \
98+
bash -c 'python3 -m pytest -s -v -x /workspace/tpu_inference/tests/e2e/test_speculative_decoding.py'
99+
else
100+
echo "Skipping: NIGHTLY environment variable not set"
101+
exit 0
102+
fi
80103
81104
- label: "JAX unit tests"
82105
key: test_7
@@ -129,30 +152,38 @@ steps:
129152
- label: "lora tests for JAX + vLLM models single chip"
130153
key: test_10
131154
soft_fail: true
132-
if: build.env("NIGHTLY") == "1"
133155
agents:
134156
queue: tpu_v6e_queue
135157
commands:
136158
- |
137-
.buildkite/scripts/run_in_docker.sh \
138-
bash -c 'MODEL_IMPL_TYPE=vllm TPU_BACKEND_TYPE=jax python3 -m pytest -s -v -x /workspace/tpu_inference/tests/lora/test_lora.py && \
139-
python3 -m pytest -s -v -x /workspace/tpu_inference/tests/lora/test_bgmv.py'
159+
if [[ "$$NIGHTLY" == "1" ]]; then
160+
.buildkite/scripts/run_in_docker.sh \
161+
bash -c 'MODEL_IMPL_TYPE=vllm TPU_BACKEND_TYPE=jax python3 -m pytest -s -v -x /workspace/tpu_inference/tests/lora/test_lora.py && \
162+
python3 -m pytest -s -v -x /workspace/tpu_inference/tests/lora/test_bgmv.py'
163+
else
164+
echo "Skipping: NIGHTLY environment variable not set"
165+
exit 0
166+
fi
140167
141168
- label: "E2E MLPerf tests for JAX + vLLM models on multiple chips"
142169
key: test_11
143170
soft_fail: true
144-
if: build.env("NIGHTLY") == "1"
145171
env:
146172
MODEL_IMPL_TYPE: "vllm"
147173
agents:
148174
queue: tpu_v6e_8_queue
149175
commands:
150-
- .buildkite/scripts/run_in_docker.sh bash /workspace/tpu_inference/tests/e2e/benchmarking/mlperf.sh
176+
- |
177+
if [[ "$$NIGHTLY" == "1" ]]; then
178+
.buildkite/scripts/run_in_docker.sh bash /workspace/tpu_inference/tests/e2e/benchmarking/mlperf.sh
179+
else
180+
echo "Skipping: NIGHTLY environment variable not set"
181+
exit 0
182+
fi
151183
152184
- label: "E2E MLperf tests for DeepSeek-R1 (no accuracy, 12-decoder layers only)"
153185
key: test_12
154186
soft_fail: true
155-
if: build.env("NIGHTLY") == "1"
156187
env:
157188
NEW_MODEL_DESIGN: "True"
158189
USE_V6E8_QUEUE: "True"
@@ -162,21 +193,31 @@ steps:
162193
agents:
163194
queue: tpu_v6e_8_queue
164195
commands:
165-
- .buildkite/scripts/run_in_docker.sh bash /workspace/tpu_inference/tests/e2e/benchmarking/mlperf.sh -m deepseek-ai/DeepSeek-R1-0528
196+
- |
197+
if [[ "$$NIGHTLY" == "1" ]]; then
198+
.buildkite/scripts/run_in_docker.sh bash /workspace/tpu_inference/tests/e2e/benchmarking/mlperf.sh -m deepseek-ai/DeepSeek-R1-0528
199+
else
200+
echo "Skipping: NIGHTLY environment variable not set"
201+
exit 0
202+
fi
166203
167204
- label: "lora tests for JAX + vLLM models multi chips"
168205
key: test_13
169206
soft_fail: true
170-
if: build.env("NIGHTLY") == "1"
171207
env:
172208
USE_V6E8_QUEUE: "True"
173209
VLLM_LOG_LEVEL: "INFO"
174210
agents:
175211
queue: tpu_v6e_8_queue
176212
commands:
177213
- |
178-
.buildkite/scripts/run_in_docker.sh \
179-
bash -c 'MODEL_IMPL_TYPE=vllm TPU_BACKEND_TYPE=jax python3 -m pytest -s -v -x /workspace/tpu_inference/tests/lora/test_lora.py'
214+
if [[ "$$NIGHTLY" == "1" ]]; then
215+
.buildkite/scripts/run_in_docker.sh \
216+
bash -c 'MODEL_IMPL_TYPE=vllm TPU_BACKEND_TYPE=jax python3 -m pytest -s -v -x /workspace/tpu_inference/tests/lora/test_lora.py'
217+
else
218+
echo "Skipping: NIGHTLY environment variable not set"
219+
exit 0
220+
fi
180221
181222
182223
# -----------------------------------------------------------------

0 commit comments

Comments
 (0)