wip

Author: wenxindongwork

.buildkite/pipeline_jax.yml

@@ -229,8 +229,13 @@ steps:
        queue: tpu_v6e_8_queue
      commands:
        - |
-         .buildkite/scripts/run_in_docker.sh \
-           bash -c 'python3 -m pytest -s -v -x /workspace/tpu_inference/tests/e2e/test_data_parallel.py'
+         if [[ "$$NIGHTLY" == "1" ]]; then
+           .buildkite/scripts/run_in_docker.sh \
+             bash -c 'python3 -m pytest -s -v -x /workspace/tpu_inference/tests/e2e/test_data_parallel.py'
+         else
+           echo "Skipping: NIGHTLY environment variable not set"
+           exit 0
+         fi
 
    - label: "lora unit tests on single chip"
      key: test_15