Run Tests Against Tunix Package workflow (google/tunix)

 
# Copyright 2025 Google LLC
 
# Licensed under the Apache License, Version 2.0 (the "License");
# you may not use this file except in compliance with the License.
# You may obtain a copy of the License at
 
#     https://www.apache.org/licenses/LICENSE-2.0
 
# Unless required by applicable law or agreed to in writing, software
# distributed under the License is distributed on an "AS IS" BASIS,
# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
# See the License for the specific language governing permissions and
# limitations under the License.
 
# This file defines a module for running tests against the built tunix package.
 
name: Run Tests Against Tunix Package
 
on:
  workflow_call:
 
permissions:
  contents: read
jobs:
  run:
    timeout-minutes: 30
    runs-on: latchkey-small
    steps:
      - uses: actions/checkout@v4
      - uses: actions/setup-python@v4
        with:
          cache: 'pip'
          python-version: '3.11'
 
      - name: Download the tunix wheel
        uses: actions/download-artifact@634f93cb2916e3fdff6788551b99b062d0335ce0 # v5.0.0
        with:
          name: tunix-wheel
 
      - name: Install the tunix wheel
        run: |
          python -m pip install --upgrade pip
          WHEEL=$(ls -1 google_tunix-*-py3-none-any.whl | head -n1)
          python -m pip install "${WHEEL}[prod,test]"
 
      - name: Verify Tunix imports from installed package
        run: |
          python3 -c "
          import tunix
          import tunix.models
          import tunix.generate
          import tunix.sft
          import tunix.distillation
          import tunix.rl
 
          assert tunix.__version__ != '0.0.0.dev0', 'Tunix version not set correctly'
          print('All tunix modules imported successfully and version is', tunix.__version__)
          "
      - name: Run agentic RL tests
        run: |
          python -m pytest tests/rl/agentic/ -v --tb=short
 
      - name: Run Cli utils tests
        run: |
          python -m pytest tests/cli/utils/ -v --tb=short
 
      - name: Run shared mesh and topology tests
        run: |
          python -m pytest tests/utils/mesh_test.py tests/utils/topology_test.py -v --tb=short
 
      - name: Run perf tests
        run: |
          python -m pytest tests/perf/ -v --tb=short
 
      - name: Run model registry tests
        run: |
          python -m pytest tests/models/registry_test.py -v --tb=short
 
      - name: Run CPU-only flash attention cache regression test
        run: |
          python -m pytest tests/models/flash_attention_cache_test.py -v --tb=short