Updated container path in build-and-test-on-tpu job #7
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
| name: Tests | ||
| # TODO: Consider enabling all tests (pytest, applications, etc.) with NNX in the future | ||
| # Currently only basic flow tests run with NNX enabled | ||
| on: | ||
| push: | ||
| branches: [ master ] | ||
| pull_request: | ||
| release: | ||
| types: [created] | ||
| permissions: | ||
| contents: read | ||
| id-token: write | ||
| env: | ||
| PYTHON: ${{ matrix.python-version }} | ||
| KERAS_HOME: .github/workflows/config/${{ matrix.backend }} | ||
| KERAS_BACKEND: jax | ||
| PROJECT_ID: gtech-rmi-dev # Replace with your GCP project ID | ||
| GAR_LOCATION: us-central1 # Replace with your Artifact Registry location (e.g., us-central1) | ||
| IMAGE_REPO: keras-docker-images | ||
| IMAGE_NAME: keras-jax-tpu-amd64:latest # Name of your Docker image | ||
| TPU_VM_NAME: kharshith-jax-tpu # Replace with your TPU VM instance name | ||
| TPU_VM_ZONE: us-central1-b # Replace with your TPU VM zone | ||
| jobs: | ||
| build-and-test-on-tpu: | ||
| strategy: | ||
| fail-fast: false | ||
| matrix: | ||
| python-version: ['3.10'] | ||
| backend: [jax] | ||
| name: Run TPU tests | ||
| runs-on: | ||
| # - keras-jax-tpu-runner | ||
| # - linux-x86-ct5lp-112-4tpu | ||
| # - linux-x86-ct5lp-112-4tpu-fvn6n-runner-6kb8n | ||
| - linux-x86-ct6e-44-1tpu | ||
| # - linux-x86-ct6e-44-1tpu-4khbn-runner-x4st4 | ||
| # - linux-x86-ct6e-44-1tpu-4khbn-runner-45nmc | ||
| container: us-central1-docker.pkg.dev/gtech-rmi-dev/keras-docker-images/keras-jax-tpu-amd64:latest | ||
| # container: | ||
| # image: docker:latest # Provides the Docker CLI within the job container | ||
| # volumes: | ||
| # - /var/run/docker.sock:/var/run/docker.sock # Mounts host's Docker socket for control | ||
| # options: --privileged | ||
| # steps: | ||
| # - name: Checkout Repository | ||
| # uses: actions/checkout@v4 | ||
| # - name: Set up Docker BuildX | ||
| # uses: docker/setup-buildx-action@v3 | ||
| # - name: Authenticate to Google Cloud (Workload Identity Federation) | ||
| # id: 'auth' | ||
| # uses: 'google-github-actions/auth@v2' | ||
| # with: | ||
| # # Replace with your Workload Identity Federation provider details. | ||
| # # This service account needs 'Artifact Registry Writer' role. | ||
| # workload_identity_provider: 'projects/YOUR_PROJECT_NUMBER/locations/global/workloadIdentityPools/YOUR_POOL_ID/providers/YOUR_PROVIDER_ID' | ||
| # service_account: 'your-github-actions-sa@${{ env.PROJECT_ID }}.iam.gserviceaccount.com' | ||
| # - name: Configure Docker to use Google Artifact Registry | ||
| # run: gcloud auth configure-docker ${{ env.GAR_LOCATION }}-docker.pkg.dev | ||
| # - name: Build Docker Image | ||
| # run: | | ||
| # IMAGE_TAG="${{ env.GAR_LOCATION }}-docker.pkg.dev/${{ env.PROJECT_ID }}/${{ env.IMAGE_REPO }}/${{ env.IMAGE_NAME_BASE }}:${{ github.sha }}" | ||
| # echo "Building Docker image: $IMAGE_TAG" | ||
| # docker build \ | ||
| # --platform=linux/amd64 \ | ||
| # -f .github/workflows/tpu/Dockerfile \ | ||
| # -t "$IMAGE_TAG" \ | ||
| # -t "${{ env.GAR_LOCATION }}-docker.pkg.dev/${{ env.PROJECT_ID }}/${{ env.IMAGE_REPO }}/${{ env.IMAGE_NAME_BASE }}:latest" \ | ||
| # . | ||
| # echo "Built Docker image: $IMAGE_TAG" | ||
| # echo "LOCAL_TEST_IMAGE_TAG=$IMAGE_TAG" >> $GITHUB_ENV # Store for immediate use in run step | ||
| # - name: Push Docker Image to Artifact Registry | ||
| # run: | | ||
| # echo "Pushing Docker image to Artifact Registry: ${{ env.LOCAL_TEST_IMAGE_TAG }}" | ||
| # docker push "${{ env.LOCAL_TEST_IMAGE_TAG }}" | ||
| # docker push "${{ env.GAR_LOCATION }}-docker.pkg.dev/${{ env.PROJECT_ID }}/${{ env.IMAGE_REPO }}/${{ env.IMAGE_NAME_BASE }}:latest" | ||
| # echo "Pushed Docker image." | ||
| # - name: Run Docker container and execute tests on TPU | ||
| # run: | | ||
| # echo "Running Docker container with TPU access and executing tests..." | ||
| # docker run --rm \ | ||
| # --privileged \ | ||
| # --network host \ | ||
| # -e PYTHON=3.10 \ # Use a specific version or derive from matrix | ||
| # -e KERAS_HOME=.github/workflows/config/jax \ | ||
| # -e KERAS_BACKEND=jax \ | ||
| # ${{ env.LOCAL_TEST_IMAGE_TAG }} \ | ||
| # /bin/bash -c ' \ | ||
| # echo "Verifying JAX TPU backend inside container..." && \ | ||
| # python3 -c "import jax; print(\"JAX Version:\", jax.__version__); print(\"Default Backend:\", jax.default_backend()); assert jax.default_backend().lower() == \"tpu\", \"TPU backend not found or not default\"; print(\"TPU verification successful!\")" \ | ||
| # # Add your actual pytest command here. Ensure pytest is installed inside your Docker image. | ||
| # # && pytest keras --ignore keras/src/applications --ignore keras/src/layers/merging/merging_test.py --cov=keras --cov-config=pyproject.toml | ||
| # ' | ||
| # echo "Docker container finished running tests." | ||
| build: | ||
| strategy: | ||
| fail-fast: false | ||
| matrix: | ||
| python-version: ['3.10'] | ||
| backend: [tensorflow, jax, torch, numpy, openvino] | ||
| nnx_enabled: [false] | ||
| include: | ||
| - python-version: '3.10' | ||
| backend: jax | ||
| nnx_enabled: true | ||
| name: ${{ matrix.backend == 'jax' && format('Run tests ({0}, {1}, nnx_enabled = {2})', matrix.python-version, matrix.backend, matrix.nnx_enabled) || format('Run tests ({0}, {1})', matrix.python-version, matrix.backend) }} | ||
| runs-on: ubuntu-latest | ||
| env: | ||
| PYTHON: ${{ matrix.python-version }} | ||
| KERAS_HOME: .github/workflows/config/${{ matrix.backend }} | ||
| steps: | ||
| - uses: actions/checkout@v4 | ||
| - name: Check for changes in keras/src/applications | ||
| uses: dorny/paths-filter@v3 | ||
| id: filter | ||
| with: | ||
| filters: | | ||
| applications: | ||
| - 'keras/src/applications/**' | ||
| - name: Set up Python | ||
| uses: actions/setup-python@v5 | ||
| with: | ||
| python-version: ${{ matrix.python-version }} | ||
| - name: Get pip cache dir | ||
| id: pip-cache | ||
| run: | | ||
| python -m pip install --upgrade pip setuptools | ||
| echo "dir=$(pip cache dir)" >> $GITHUB_OUTPUT | ||
| - name: pip cache | ||
| uses: actions/cache@v4 | ||
| with: | ||
| path: ${{ steps.pip-cache.outputs.dir }} | ||
| key: ${{ runner.os }}-pip-${{ hashFiles('pyproject.toml') }}-${{ hashFiles('requirements.txt') }} | ||
| - name: Install dependencies | ||
| run: | | ||
| pip install -r requirements.txt --progress-bar off --upgrade | ||
| if [ "${{ matrix.nnx_enabled }}" == "true" ]; then | ||
| pip install --upgrade git+https://github.com/google/flax.git | ||
| fi | ||
| pip uninstall -y keras keras-nightly | ||
| pip install -e "." --progress-bar off --upgrade | ||
| - name: Test applications with pytest | ||
| if: ${{ steps.filter.outputs.applications == 'true' && matrix.nnx_enabled == false }} | ||
| run: | | ||
| pytest keras/src/applications --cov=keras/src/applications --cov-config=pyproject.toml | ||
| coverage xml --include='keras/src/applications/*' -o apps-coverage.xml | ||
| - name: Codecov keras.applications | ||
| if: ${{ steps.filter.outputs.applications == 'true' && matrix.nnx_enabled == false }} | ||
| uses: codecov/codecov-action@v5 | ||
| with: | ||
| env_vars: PYTHON,KERAS_HOME | ||
| flags: keras.applications,keras.applications-${{ matrix.backend }} | ||
| files: apps-coverage.xml | ||
| token: ${{ secrets.CODECOV_TOKEN }} | ||
| fail_ci_if_error: false | ||
| - name: Test integrations | ||
| if: ${{ matrix.backend != 'numpy' && matrix.nnx_enabled == false }} | ||
| run: | | ||
| python integration_tests/import_test.py | ||
| python integration_tests/numerical_test.py | ||
| - name: Test JAX-specific integrations | ||
| if: ${{ matrix.backend == 'jax' && matrix.nnx_enabled == false }} | ||
| run: | | ||
| python integration_tests/jax_custom_fit_test.py | ||
| - name: Test basic flow with NNX | ||
| if: ${{ matrix.nnx_enabled == true }} | ||
| env: | ||
| KERAS_NNX_ENABLED: true | ||
| run: | | ||
| python integration_tests/import_test.py | ||
| python integration_tests/basic_full_flow.py | ||
| - name: Test TF-specific integrations | ||
| if: ${{ matrix.backend == 'tensorflow'}} | ||
| run: | | ||
| python integration_tests/tf_distribute_training_test.py | ||
| python integration_tests/tf_custom_fit_test.py | ||
| - name: Test Torch-specific integrations | ||
| if: ${{ matrix.backend == 'torch'}} | ||
| run: | | ||
| pytest integration_tests/torch_workflow_test.py | ||
| python integration_tests/torch_custom_fit_test.py | ||
| - name: Test with pytest | ||
| if: ${{ matrix.nnx_enabled == false }} | ||
| run: | | ||
| if [ "${{ matrix.backend }}" == "openvino" ]; then | ||
| IGNORE_FILE="keras/src/backend/openvino/excluded_tests.txt" | ||
| IGNORE_ARGS=$(awk '{print "--ignore=" $0}' "$IGNORE_FILE") | ||
| else | ||
| IGNORE_ARGS="" | ||
| fi | ||
| pytest keras --ignore keras/src/applications --cov=keras --cov-config=pyproject.toml $IGNORE_ARGS | ||
| coverage xml --omit='keras/src/applications/*,keras/api' -o core-coverage.xml | ||
| - name: Codecov keras | ||
| if: ${{ matrix.nnx_enabled == false }} | ||
| uses: codecov/codecov-action@v5 | ||
| with: | ||
| env_vars: PYTHON,KERAS_HOME,KERAS_NNX_ENABLED | ||
| flags: keras,keras-${{ matrix.backend }}${{ matrix.nnx_enabled == 'true' && '-nnx' || '' }} | ||
| files: core-coverage.xml | ||
| token: ${{ secrets.CODECOV_TOKEN }} | ||
| fail_ci_if_error: false | ||
| format: | ||
| name: Check the code format | ||
| runs-on: ubuntu-latest | ||
| steps: | ||
| - uses: actions/checkout@v4 | ||
| - name: Set up Python 3.10 | ||
| uses: actions/setup-python@v5 | ||
| with: | ||
| python-version: '3.10' | ||
| - name: Get pip cache dir | ||
| id: pip-cache | ||
| run: | | ||
| python -m pip install --upgrade pip setuptools | ||
| echo "dir=$(pip cache dir)" >> $GITHUB_OUTPUT | ||
| - name: pip cache | ||
| uses: actions/cache@v4 | ||
| with: | ||
| path: ${{ steps.pip-cache.outputs.dir }} | ||
| key: ${{ runner.os }}-pip-${{ hashFiles('pyproject.toml') }}-${{ hashFiles('requirements.txt') }} | ||
| - name: Install dependencies | ||
| run: | | ||
| pip install -r requirements.txt --progress-bar off --upgrade | ||
| pip uninstall -y keras keras-nightly | ||
| pip install -e "." --progress-bar off --upgrade | ||
| - name: Run pre-commit | ||
| run: pre-commit run --all-files --hook-stage manual | ||
