add 1st iteration of pre-merge pipeline

Signed-off-by: Johnny Yang <johnnyyang@google.com>

Author: jcyang43

.buildkite/pipeline_jax.yml

@@ -107,23 +107,6 @@ steps:
            exit 0
          fi
 
-
-   - label: "JAX unit tests"
-     key: test_7
-     soft_fail: true
-     agents:
-       queue: tpu_v6e_queue
-     commands:
-       - |
-         .buildkite/scripts/run_in_docker.sh \
-           python3 -m pytest -s -v -x /workspace/tpu_inference/tests/ \
-           --ignore=/workspace/tpu_inference/tests/kernels \
-           --ignore=/workspace/tpu_inference/tests/lora \
-           --ignore=/workspace/tpu_inference/tests/e2e \
-           --ignore=/workspace/tpu_inference/tpu_inference/mock \
-           --ignore=/workspace/tpu_inference/tests/layers/vllm/test_compressed_tensors_moe.py \
-           --cov-config=/workspace/tpu_inference/.coveragerc --cov tpu_inference --cov-report term-missing --cov-fail-under=69
-
    - label: "JAX unit tests - kernels"
      key: test_8
      soft_fail: true
@@ -268,7 +251,6 @@ steps:
        - test_4
        - test_5
        - test_6
-       - test_7
        - test_8
        - test_9
        - test_10
@@ -282,4 +264,4 @@ steps:
      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 test_15 test_16
+           "TPU JAX Tests Failed" test_0 test_1 test_2 test_3 test_4 test_5 test_6 test_8 test_9 test_10 test_11 test_12 test_13 test_15 test_16

.buildkite/pipeline_pre_merge.yml

@@ -0,0 +1,39 @@
+steps:
+  - label: "JAX unit tests"
+    soft_fail: true
+    agents:
+      queue: tpu_v6e_queue
+    commands:
+      - |
+        .buildkite/scripts/run_in_docker.sh \
+          python3 -m pytest -s -v -x /workspace/tpu_inference/tests/ \
+          --ignore=/workspace/tpu_inference/tests/kernels \
+          --ignore=/workspace/tpu_inference/tests/lora \
+          --ignore=/workspace/tpu_inference/tests/e2e \
+          --ignore=/workspace/tpu_inference/tpu_inference/mock \
+          --cov-config=/workspace/tpu_inference/.coveragerc --cov tpu_inference --cov-report term-missing --cov-fail-under=69
+
+  - label: "Upload modified feature(s)/model(s) introduced in this PR"
+    soft_fail: true
+    # nightly run automatically uploads all features & models, so the next line is needed to dedup
+    if: build.env("NIGHTLY") != "1"
+    agents:
+      queue: cpu
+    command: |
+      git fetch origin $BUILDKITE_PULL_REQUEST_BASE_BRANCH
+      BASE_REF="$BUILDKITE_PULL_REQUEST_BASE_BRANCH"
+      HEAD_REF="$BUILDKITE_COMMIT"
+
+      # Get newly added or modified yml files in '.buildkite/models/' OR '.buildkite/features/'
+      MODIFIED_YML_PATHS=$(git diff --name-only --diff-filter=AM "$BASE_REF" "$HEAD_REF" | grep -E '^\.buildkite\/(models|features)\/.*\.yml$' | tr '\n' ' ')
+
+      if [ -n "$MODIFIED_YML_PATHS" ]; then
+        echo "Detected new models/features yml files: $MODIFIED_YML_PATHS"
+        for FILE_PATH in $MODIFIED_YML_PATHS; do
+          echo "Processing and uploading pipeline: ${FILE_PATH}"
+          buildkite-agent pipeline upload "${FILE_PATH}"
+          echo "Successfully uploaded ${FILE_PATH}"
+        done
+      else
+        echo "No new models/features yml files detected"
+      fi

.buildkite/scripts/bootstrap.sh

@@ -20,18 +20,29 @@ else
   echo "Code files changed. Proceeding with pipeline upload."
 fi
 
+upload_pre_merge_pipeline() {
+    buildkite-agent pipeline upload .buildkite/pipeline_pre_merge.yml
+}
+
 upload_pipeline() {
-    VLLM_COMMIT_HASH=$(git ls-remote https://github.com/vllm-project/vllm.git HEAD | awk '{ print $1}')
-    buildkite-agent meta-data set "VLLM_COMMIT_HASH" "${VLLM_COMMIT_HASH}"
-    echo "Using vllm commit hash: $(buildkite-agent meta-data get "VLLM_COMMIT_HASH")"
     buildkite-agent pipeline upload .buildkite/pipeline_jax.yml
     # buildkite-agent pipeline upload .buildkite/pipeline_torch.yml
     buildkite-agent pipeline upload .buildkite/main.yml
     buildkite-agent pipeline upload .buildkite/nightly_releases.yml
+    upload_pre_merge_pipeline
 }
 
-echo "--- Starting Buildkite Bootstrap ---"
+fetch_latest_upstream_vllm_commit() {
+    # To help with debugging (when needed), perform setup to:
+    #    1. Use the same upstream vllm commit for all jobs in this CI run for consistency
+    #    2. Record which upstream commit this CI run is using
+    VLLM_COMMIT_HASH=$(git ls-remote https://github.com/vllm-project/vllm.git HEAD | awk '{ print $1}')
+    buildkite-agent meta-data set "VLLM_COMMIT_HASH" "${VLLM_COMMIT_HASH}"
+    echo "Using vllm commit hash: $(buildkite-agent meta-data get "VLLM_COMMIT_HASH")"
+}
 
+echo "--- Starting Buildkite Bootstrap ---"
+fetch_latest_upstream_vllm_commit
 # Check if the current build is a pull request
 if [ "$BUILDKITE_PULL_REQUEST" != "false" ]; then
   echo "This is a Pull Request build."
@@ -42,8 +53,8 @@ if [ "$BUILDKITE_PULL_REQUEST" != "false" ]; then
     echo "Found 'ready' label on PR. Uploading main pipeline..."
     upload_pipeline
   else
-    echo "No 'ready' label found on PR. Skipping main pipeline upload."
-    exit 0 # Exit with 0 to indicate success (no error, just skipped)
+    echo "No 'ready' label found on PR. Uploading fast check pipeline"
+    upload_pre_merge_pipeline
   fi
 else
   # If it's NOT a Pull Request (e.g., branch push, tag, manual build)