Merge branch 'tensorzero:main' into main

Author: hongantran3804

.config/nextest.toml

@@ -11,6 +11,10 @@ slow-timeout = { period = "10s", terminate-after = 3 }
 # On CI, we use the 'ci' profile, which runs all tests.
 default-filter = "not test(no_aws_credentials)"
 
+[[profile.default.overrides]]
+filter = 'test(evaluations)'
+slow-timeout = { period = "20s", terminate-after = 3 }
+
 # Profiles config
 # We use these profiles to define our major test groups.
 # By using `default-filter` to specify our tests, we can further restrict the tests
@@ -36,6 +40,8 @@ default-filter = 'binary(optimization-live)'
 [profile.optimization-mock]
 default-filter = 'binary(optimization-mock)'
 
+[profile.ci-unit.junit]
+path = "junit.xml"
 
 [[profile.optimization.overrides]]
 # Settings for running optimization tests

.github/workflows/batch-test.yml

@@ -41,7 +41,19 @@ jobs:
     steps:
       - uses: actions/checkout@11bd71901bbe5b1630ceea73d27597364c9af683
 
-      - uses: dtolnay/rust-toolchain@stable
+      - name: Install Rust toolchain
+        run: |
+          for attempt in 1 2 3; do
+            if rustup toolchain install stable && rustup default stable; then
+              break
+            fi
+            if [ $attempt -eq 3 ]; then
+              echo "Failed to install Rust toolchain after 3 attempts"
+              exit 1
+            fi
+            sleep $((10 * attempt))
+          done
+        shell: bash
 
       - name: Login to DockerHub
         uses: docker/login-action@74a5d142397b4f367a81961eba4e8cd7edddf772

.github/workflows/build-gateway-container.yml

@@ -5,7 +5,7 @@ on:
 
 jobs:
   build-gateway-container:
-    runs-on: namespace-profile-tensorzero-8x16;overrides.cache-tag=build-gateway-cache
+    runs-on: namespace-profile-tensorzero-2x8;overrides.cache-tag=build-gateway-cache
 
     steps:
       # TODO - investigate why using the Namespace checkout action causes
@@ -36,4 +36,4 @@ jobs:
           path: gateway-container.tar
           retention-days: 1
           if-no-files-found: error
-          overwrite: false
+          overwrite: false

.github/workflows/build-ui-container.yml

@@ -5,7 +5,7 @@ on:
 
 jobs:
   build-ui-container:
-    runs-on: namespace-profile-tensorzero-8x16;overrides.cache-tag=build-ui-cache
+    runs-on: namespace-profile-tensorzero-2x8;overrides.cache-tag=build-ui-cache
 
     steps:
       - uses: namespacelabs/nscloud-checkout-action@953fed31a6113cc2347ca69c9d823743c65bc84b
@@ -33,4 +33,4 @@ jobs:
           path: ui-container.tar
           retention-days: 1
           if-no-files-found: error
-          overwrite: false
+          overwrite: false

.github/workflows/buildkite-test-suite.yml

@@ -0,0 +1,19 @@
+name: Buildkite Test Suite
+run-name: "Buildkite Test Suite: ${{ github.event.pull_request.title || github.ref }}"
+
+on:
+  merge_group:
+  pull_request:
+    branches: ["main"]
+
+jobs:
+  unit-tests:
+    # Run for non-PR events, or PRs where the head repo is the same as this repo (not a fork).
+    if: ${{ github.event_name != 'pull_request' || (github.event.pull_request.head.repo.full_name == github.repository && github.actor != 'dependabot[bot]')}}
+    runs-on: ubuntu-latest
+    steps:
+      - name: Trigger a Buildkite Build on Push using v2.0.0
+        uses: buildkite/[email protected]
+        with:
+          buildkite_api_access_token: ${{ secrets.TRIGGER_BK_BUILD_TOKEN }}
+          pipeline: "tensorzero/pr-tests"

.github/workflows/caching-hack.yml

@@ -25,7 +25,19 @@ jobs:
     runs-on: ${{ matrix.runner }}
     steps:
       - uses: actions/checkout@11bd71901bbe5b1630ceea73d27597364c9af683
-      - uses: dtolnay/rust-toolchain@stable
+      - name: Install Rust toolchain
+        run: |
+          for attempt in 1 2 3; do
+            if rustup toolchain install stable && rustup default stable; then
+              break
+            fi
+            if [ $attempt -eq 3 ]; then
+              echo "Failed to install Rust toolchain after 3 attempts"
+              exit 1
+            fi
+            sleep $((10 * attempt))
+          done
+        shell: bash
       - uses: Swatinem/rust-cache@9d47c6ad4b02e050fd481d890b2ea34778fd09d6
         with:
           cache-provider: "buildjet"

.github/workflows/cancel-stale-merge-queue.yml

@@ -0,0 +1,40 @@
+name: Cancel stale merge queue workflows
+on:
+  merge_group:
+    types:
+      - destroyed
+
+permissions:
+  actions: write
+  contents: read
+
+jobs:
+  cancel-workflows:
+    name: Cancel Workflow Runs
+    runs-on: ubuntu-latest
+    if: github.event.reason != 'merged'
+    steps:
+      - name: Get Merge Queue Commit SHA
+        id: get-sha
+        run: |
+          echo "Repository: ${{ github.repository }}"
+          echo "Head SHA: ${{ github.sha }}"
+          echo "Merge Group Reason: ${{ github.event.reason }}"
+
+      - name: Cancel Workflow Runs by SHA
+        run: |
+          # Get all workflow runs for the specific SHA
+          workflow_runs=$(curl -s -H "Authorization: Bearer ${{ github.token }}" \
+            "https://api.github.com/repos/${{ github.repository }}/actions/runs?head_sha=${{ github.sha }}")
+          
+          # Extract run IDs and cancel them (except current workflow)
+          current_run_id="${{ github.run_id }}"
+          echo "Current Run ID: $current_run_id"
+          echo "$workflow_runs" | jq -r '.workflow_runs[] | select(.status != "completed") | .id' | while read -r run_id; do
+            echo "Run ID: $run_id"
+            if [ -n "$run_id" ] && [ "$run_id" != "$current_run_id" ]; then
+              echo "Cancelling workflow run $run_id"
+              curl -X POST -H "Authorization: Bearer ${{ github.token }}" \
+                "https://api.github.com/repos/${{ github.repository }}/actions/runs/$run_id/cancel"
+            fi
+          done

.github/workflows/general.yml

@@ -3,6 +3,8 @@ run-name: "General Checks for: ${{ github.event.pull_request.title || github.ref
 
 on:
   merge_group:
+    paths-ignore:
+      - "docs/**"
   pull_request:
     branches: ["main"]
 
@@ -109,7 +111,18 @@ jobs:
           node-version: "22.9.0"
 
       - name: Install pnpm
-        uses: pnpm/action-setup@a7487c7e89a18df4991f7f222e4898a00d66ddda
+        run: |
+          for attempt in 1 2 3; do
+            if npm install -g pnpm@latest; then
+              break
+            fi
+            if [ $attempt -eq 3 ]; then
+              echo "Failed to install pnpm after 3 attempts"
+              exit 1
+            fi
+            sleep $((10 * attempt))
+          done
+        shell: bash
 
       - name: Install dependencies
         working-directory: internal/tensorzero-node
@@ -138,7 +151,19 @@ jobs:
     runs-on: windows-latest
     steps:
       - uses: actions/checkout@11bd71901bbe5b1630ceea73d27597364c9af683
-      - uses: dtolnay/rust-toolchain@stable
+      - name: Install Rust toolchain
+        run: |
+          for attempt in 1 2 3; do
+            if rustup toolchain install stable && rustup default stable; then
+              break
+            fi
+            if [ $attempt -eq 3 ]; then
+              echo "Failed to install Rust toolchain after 3 attempts"
+              exit 1
+            fi
+            sleep $((10 * attempt))
+          done
+        shell: bash
       - uses: Swatinem/rust-cache@9d47c6ad4b02e050fd481d890b2ea34778fd09d6
         with:
           cache-provider: "buildjet"
@@ -147,29 +172,27 @@ jobs:
         run: cargo build --workspace
 
   validate:
-    runs-on: namespace-profile-tensorzero-8x16
+    runs-on: namespace-profile-tensorzero-32x64
 
     timeout-minutes: 30
 
     steps:
       - uses: actions/checkout@11bd71901bbe5b1630ceea73d27597364c9af683
 
-      # We allow the namespace builder setup to fail on Dependabot PRs and PRs from forks
-      # (where the oidc token is not available)
-
-      - name: Install Namespace CLI
-        uses: namespacelabs/nscloud-setup@d1c625762f7c926a54bd39252efff0705fd11c64
-        continue-on-error: ${{ github.event.pull_request.head.repo.full_name != github.repository || github.actor == 'dependabot[bot]' }}
-
-      - name: Configure Namespace-powered Buildx
-        uses: namespacelabs/nscloud-setup-buildx-action@84ca8c58fdf372d6a4750476cd09b7b96ee778ca
-        continue-on-error: ${{ github.event.pull_request.head.repo.full_name != github.repository || github.actor == 'dependabot[bot]' }}
-
       # We deliberately install our MSRV here (rather than 'stable') to ensure that everything compiles with that version
       - name: Install Rust 1.86.0
         run: |
-          rustup install 1.86.0 --component clippy,rustfmt
-          rustup default 1.86.0
+          for attempt in 1 2 3; do
+            if rustup install 1.86.0 --component clippy,rustfmt && rustup default 1.86.0; then
+              break
+            fi
+            if [ $attempt -eq 3 ]; then
+              echo "Failed to install Rust 1.86.0 after 3 attempts"
+              exit 1
+            fi
+            sleep $((10 * attempt))
+          done
+        shell: bash
 
       - name: Print Rust version
         run: rustc --version
@@ -180,10 +203,21 @@ jobs:
           node-version: "22.9.0"
 
       - name: Install pnpm
-        uses: pnpm/action-setup@a7487c7e89a18df4991f7f222e4898a00d66ddda
+        run: |
+          for attempt in 1 2 3; do
+            if npm install -g pnpm@latest; then
+              break
+            fi
+            if [ $attempt -eq 3 ]; then
+              echo "Failed to install pnpm after 3 attempts"
+              exit 1
+            fi
+            sleep $((10 * attempt))
+          done
+        shell: bash
 
       - name: Install uv
-        run: curl -LsSf https://astral.sh/uv/0.6.17/install.sh | sh
+        run: curl -LsSf --retry 2 --retry-delay 10 --retry-max-time 60 https://astral.sh/uv/0.6.17/install.sh | sh
 
       - name: Configure Namespace cache for Rust, Python (pip), and pnpm
         uses: namespacelabs/nscloud-cache-action@2f50e7d0f70475e6f59a55ba0f05eec9108e77cc
@@ -359,23 +393,29 @@ jobs:
     # We don't run many tests here, so use a normal runner with Github Actions caching
     # to avoid unnecessarily using Namespace credits (it should still always finish before
     # the main 'validate' job)
-    runs-on: ${{ matrix.replicated && 'namespace-profile-tensorzero-16x32' || 'ubuntu-latest' }}
+    runs-on: ${{ matrix.replicated && 'namespace-profile-tensorzero-32x64' || 'namespace-profile-tensorzero-16x32' }}
     continue-on-error: ${{ matrix.clickhouse_version.allow_failure }}
     strategy:
       matrix:
         # Only include replicated: true when running in merge queue
         replicated: ${{ github.event_name == 'merge_group' && fromJSON('[true, false]') || fromJSON('[false]') }}
-        clickhouse_version:
-          - tag: "24.12-alpine"
-            prefix: "24.12"
-            allow_failure: false
-          - tag: "latest-alpine"
-            prefix: ""
-            allow_failure: false
+        clickhouse_version: ${{ github.event_name == 'merge_group' && fromJSON('[{"tag":"24.12-alpine","prefix":"24.12","allow_failure":false},{"tag":"25.7-alpine","prefix":"25.7","allow_failure":false}]') || fromJSON('[{"tag":"25.7-alpine","prefix":"25.7","allow_failure":false}]') }}
 
     steps:
       - uses: actions/checkout@11bd71901bbe5b1630ceea73d27597364c9af683
-      - uses: dtolnay/rust-toolchain@stable
+      - name: Install Rust toolchain
+        run: |
+          for attempt in 1 2 3; do
+            if rustup toolchain install stable && rustup default stable; then
+              break
+            fi
+            if [ $attempt -eq 3 ]; then
+              echo "Failed to install Rust toolchain after 3 attempts"
+              exit 1
+            fi
+            sleep $((10 * attempt))
+          done
+        shell: bash
       - uses: Swatinem/rust-cache@9d47c6ad4b02e050fd481d890b2ea34778fd09d6
         with:
           cache-provider: "buildjet"
@@ -397,7 +437,7 @@ jobs:
         run: echo "TENSORZERO_CLICKHOUSE_CLUSTER_NAME=tensorzero_e2e_tests_cluster" >> $GITHUB_ENV
 
       - name: Install uv
-        run: curl -LsSf https://astral.sh/uv/0.6.17/install.sh | sh
+        run: curl -LsSf --retry 2 --retry-delay 10 --retry-max-time 60 https://astral.sh/uv/0.6.17/install.sh | sh
 
       - name: Download ClickHouse fixtures
         run: uv run ./ui/fixtures/download-fixtures.py
@@ -431,7 +471,7 @@ jobs:
         run: |
           cargo run-e2e > e2e_logs.txt 2>&1 &
             count=0
-            max_attempts=10
+            max_attempts=20
             while ! curl -s -f http://localhost:3000/health >/dev/null 2>&1; do
               echo "Waiting for gateway to be healthy..."
               sleep 1
@@ -449,7 +489,7 @@ jobs:
           cargo run-e2e --run-migrations-only &&
           cargo run-e2e > e2e_logs.txt 2>&1 &
             count=0
-            max_attempts=10
+            max_attempts=40
             while ! curl -s -f http://localhost:3000/health >/dev/null 2>&1; do
               echo "Waiting for gateway to be healthy..."
               sleep 1
@@ -523,7 +563,19 @@ jobs:
 
     steps:
       - uses: actions/checkout@11bd71901bbe5b1630ceea73d27597364c9af683
-      - uses: dtolnay/rust-toolchain@stable
+      - name: Install Rust toolchain
+        run: |
+          for attempt in 1 2 3; do
+            if rustup toolchain install stable && rustup default stable; then
+              break
+            fi
+            if [ $attempt -eq 3 ]; then
+              echo "Failed to install Rust toolchain after 3 attempts"
+              exit 1
+            fi
+            sleep $((10 * attempt))
+          done
+        shell: bash
       - uses: Swatinem/rust-cache@9d47c6ad4b02e050fd481d890b2ea34778fd09d6
         with:
           cache-provider: "buildjet"
@@ -535,7 +587,7 @@ jobs:
           tool: cargo-nextest
 
       - name: Install uv
-        run: curl -LsSf https://astral.sh/uv/0.6.17/install.sh | sh
+        run: curl -LsSf --retry 2 --retry-delay 10 --retry-max-time 60 https://astral.sh/uv/0.6.17/install.sh | sh
 
       - name: Download ClickHouse fixtures
         run: uv run ./ui/fixtures/download-fixtures.py
@@ -594,6 +646,8 @@ jobs:
         check-python-client-build,
         check-node-bindings,
         build-windows,
+        build-ui-container,
+        build-gateway-container,
         validate,
         clickhouse-tests,
         ui-tests,
@@ -602,5 +656,7 @@ jobs:
       ]
     runs-on: ubuntu-latest
     steps:
-      - if: ${{ contains(needs.*.result, 'failure') || contains(needs.*.result, 'cancelled') }}
+      # When running in the merge queue, jobs should never be skipped.
+      # In PR CI, some jobs may be intentionally skipped (e.g. due to running from a fork, or to save money)
+      - if: ${{ contains(needs.*.result, 'failure') || contains(needs.*.result, 'cancelled' || (github.event_name == 'merge_group' && contains(needs.*.result, 'skipped'))) }}
         run: exit 1