Repository navigation
Support InsertOption::kInsertAndPin and InsertOption::kAdmitStrict in KVCacheStoreBackend::Insert. #1854
Workflow file for this run
This file contains hidden or bidirectional Unicode text that may be interpreted or compiled differently than what appears below. To review, open the file in an editor that reveals hidden Unicode characters.
Learn more about bidirectional Unicode characters
| # CPU unit tests and E2E JAX build verification for every pull request and | |
| # Copybara presubmit CL. | |
| # | |
| # What runs (see ci/run_unit_tests.sh): | |
| # 1. Bazel CPU unit test suite: every cc_test and py_test under | |
| # //tpu_sync/core, //tpu_sync/kv_cache, //tpu_sync/weight_sync, and | |
| # //tpu_sync/rpc that is not tagged no_oss, notap, manual, or | |
| # requires-<resource>. | |
| # 2. E2E JAX build (./build.sh jax): compiles _tpu_raiden_jax.so, C++ | |
| # control-plane service binaries (global_registry_server, | |
| # kv_cache_host_store_node_main), and Python protobuf modules. | |
| # 3. Dynamic module binding verification: verifies that _tpu_raiden_jax.so | |
| # and generated protobufs import cleanly in Python 3.12 (with JAX mocked | |
| # on CPU) and runs framework-neutral Python CPU unit tests. | |
| # | |
| # Builds run in the ml-build container used by the wheel workflows. Runs from | |
| # a branch of this repo (including Copybara presubmit branches) authenticate | |
| # to the remote bazel cache (--config=ci); fork PRs have no credentials for it | |
| # and build cold, which is slower but needs no secrets exposed to fork code. | |
| name: Unit Tests | |
| on: | |
| pull_request: | |
| push: | |
| branches: | |
| - main | |
| workflow_dispatch: | |
| inputs: | |
| test_scope: | |
| description: "Bazel package patterns to search for tests" | |
| type: string | |
| default: "//tpu_sync/core:all //tpu_sync/kv_cache:all //tpu_sync/weight_sync:all //tpu_sync/rpc:all" | |
| permissions: | |
| contents: read | |
| concurrency: | |
| group: ${{ github.workflow }}-${{ github.event.pull_request.number || github.ref }} | |
| cancel-in-progress: true | |
| jobs: | |
| test: | |
| name: "Bazel CPU unit tests" | |
| runs-on: linux-x86-n2-32 | |
| container: | |
| image: us-docker.pkg.dev/ml-oss-artifacts-published/ml-public-container/ml-build:latest@sha256:9644db6e0e969079b138547dd99818c02e6ab51d892c4675cf2f4edf02a1e71a | |
| options: --shm-size=16g | |
| timeout-minutes: 300 | |
| env: | |
| RAIDEN_TEST_SCOPE: ${{ inputs.test_scope || '//tpu_sync/core:all //tpu_sync/kv_cache:all //tpu_sync/weight_sync:all //tpu_sync/rpc:all' }} | |
| # Fork PRs run without the ambient GCP service account, so the | |
| # remote cache is only requested when the head branch lives in this | |
| # repository. | |
| EXTRA_BAZEL_FLAGS: ${{ (github.event_name != 'pull_request' || github.event.pull_request.head.repo.full_name == github.repository) && '--config=ci' || '' }} | |
| steps: | |
| - name: Checkout tpu-sync | |
| uses: actions/checkout@fbc6f3992d24b796d5a048ff273f7fcc4a7b6c09 # v5.1.0 | |
| with: | |
| persist-credentials: false | |
| path: tpu-sync | |
| # Git security patch (CVE-2022-24765) blocks root from operating on | |
| # runner-owned directories. Tell Git explicitly to trust the workspace | |
| # inside this ephemeral container. | |
| - name: Fix Git Workspace Ownership | |
| run: git config --global --add safe.directory "$GITHUB_WORKSPACE/tpu-sync" | |
| # The golden-commit tooling runs on a stock python3 outside Bazel (see | |
| # golden_commit.yml), so its tests run the same way, before the build. | |
| - name: Test the golden-commit tooling | |
| working-directory: tpu-sync | |
| run: python3 -m unittest ci/golden/golden_commit_test.py | |
| - name: Run CPU unit tests, E2E JAX build, and binding verification | |
| working-directory: tpu-sync | |
| run: | | |
| export BAZEL_CACHE_DIR="${RUNNER_TEMP}/bazel_cache" | |
| export BAZEL_OUTPUT_BASE="${BAZEL_CACHE_DIR}/output_base" | |
| bash ci/run_unit_tests.sh |