Add lora unit tests to the CI

Signed-off-by: Xiongfei Wei <isaacwxf23@gmail.com>

Author: vanbasten23

.buildkite/pipeline_jax.yml

@@ -151,28 +151,20 @@ steps:
            exit 0
          fi
 
-   - label: "lora tests for JAX + vLLM models single chip"
+   - label: "lora e2e tests for JAX + vLLM models single chip"
      key: test_10
      soft_fail: true
      agents:
        queue: tpu_v6e_queue
      commands:
        - |
-<<<<<<< HEAD
          if [[ "$$NIGHTLY" == "1" ]]; then
            .buildkite/scripts/run_in_docker.sh \
-             bash -c 'MODEL_IMPL_TYPE=vllm TPU_BACKEND_TYPE=jax python3 -m pytest -s -v -x /workspace/tpu_inference/tests/lora/test_lora.py && \
-             python3 -m pytest -s -v -x /workspace/tpu_inference/tests/lora/test_bgmv.py'
+             bash -c 'MODEL_IMPL_TYPE=vllm TPU_BACKEND_TYPE=jax python3 -m pytest -s -v -x /workspace/tpu_inference/tests/lora/test_lora.py'
          else
            echo "Skipping: NIGHTLY environment variable not set"
            exit 0
          fi
-=======
-         .buildkite/scripts/run_in_docker.sh \
-           bash -c 'MODEL_IMPL_TYPE=vllm TPU_BACKEND_TYPE=jax python3 -m pytest -s -v -x /workspace/tpu_inference/tests/lora/test_lora.py && \
-           python3 -m pytest -s -v -x /workspace/tpu_inference/tests/lora/test_bgmv.py && \
-           python3 -m pytest -s -v -x /workspace/tpu_inference/tests/lora/test_layers.py'
->>>>>>> c17bacea (add multi-chip test case)
 
    - label: "E2E MLPerf tests for JAX + vLLM models on multiple chips"
      key: test_11
@@ -210,7 +202,7 @@ steps:
            exit 0
          fi
 
-   - label: "lora tests for JAX + vLLM models multi chips"
+   - label: "lora e2e tests for JAX + vLLM models multi chips"
      key: test_13
      soft_fail: true
      env:
@@ -220,19 +212,13 @@ steps:
        queue: tpu_v6e_8_queue
      commands:
        - |
-<<<<<<< HEAD
          if [[ "$$NIGHTLY" == "1" ]]; then
            .buildkite/scripts/run_in_docker.sh \
              bash -c 'MODEL_IMPL_TYPE=vllm TPU_BACKEND_TYPE=jax python3 -m pytest -s -v -x /workspace/tpu_inference/tests/lora/test_lora.py'
          else
            echo "Skipping: NIGHTLY environment variable not set"
            exit 0
          fi
-=======
-         .buildkite/scripts/run_in_docker.sh \
-           bash -c 'MODEL_IMPL_TYPE=vllm TPU_BACKEND_TYPE=jax python3 -m pytest -s -v -x /workspace/tpu_inference/tests/lora/test_lora.py && \
-                    python3 -m pytest -s -v -x /workspace/tpu_inference/tests/lora/test_layers.py'
->>>>>>> c17bacea (add multi-chip test case)
 
    - label: "E2E data parallelism test"
      key: test_14
@@ -246,6 +232,29 @@ steps:
          .buildkite/scripts/run_in_docker.sh \
            bash -c 'python3 -m pytest -s -v -x /workspace/tpu_inference/tests/e2e/test_data_parallel.py'
 
+   - label: "lora unit tests on single chip"
+     key: test_15
+     soft_fail: true
+     agents:
+       queue: tpu_v6e_queue
+     commands:
+       - |
+         .buildkite/scripts/run_in_docker.sh \
+           bash -c ' python3 -m pytest -s -v -x /workspace/tpu_inference/tests/lora/test_bgmv.py && \
+           python3 -m pytest -s -v -x /workspace/tpu_inference/tests/lora/test_layers.py'
+
+   - label: "lora unit tests on multi chips"
+     key: test_16
+     soft_fail: true
+     env:
+       USE_V6E8_QUEUE: "True"
+       VLLM_LOG_LEVEL: "INFO"
+     agents:
+       queue: tpu_v6e_8_queue
+     commands:
+       - |
+         .buildkite/scripts/run_in_docker.sh \
+           bash -c 'python3 -m pytest -s -v -x /workspace/tpu_inference/tests/lora/test_layers.py'
   # -----------------------------------------------------------------
   # NOTIFICATION STEP
   # -----------------------------------------------------------------
@@ -266,9 +275,11 @@ steps:
        - test_12
        - test_13
        - test_14
+       - test_15
+       - test_16
      agents:
        queue: cpu
      commands:
        - |
          .buildkite/scripts/check_results.sh \
-           "TPU JAX Tests Failed" test_0 test_1 test_2 test_3 test_4 test_5 test_6 test_7 test_8 test_9 test_10 test_11 test_12 test_13
+           "TPU JAX Tests Failed" test_0 test_1 test_2 test_3 test_4 test_5 test_6 test_7 test_8 test_9 test_10 test_11 test_12 test_13 test_14 test_15 test_16