Skip to content

TESTING ONLY (TPU test) #20

TESTING ONLY (TPU test)

TESTING ONLY (TPU test) #20

Workflow file for this run

name: TPU Tests
on:
push:
branches: [ master ]
pull_request:
release:
types: [created]
permissions:
contents: read # Only read permission is needed for checkout
packages: write
# Define base environment variables at the workflow level
# These can still be used inside steps, just not for the container image definition
env:
PROJECT_ID: gtech-rmi-dev
GAR_LOCATION: us-central1
GAR_REPO: keras-docker-images
IMAGE_NAME: keras-jax-tpu-amd64
IMAGE_TAG: latest
jobs:
build-and-push:
name: Build and Push to GHCR
runs-on: ubuntu-latest # This job doesn't need the special TPU runner
steps:
- name: Checkout Repository
uses: actions/checkout@v4
- name: Log in to GitHub Container Registry
uses: docker/login-action@v3
with:
registry: ghcr.io
# GITHUB_TOKEN is automatically created by Actions and has permissions to push to your repo's package registry.
username: ${{ github.actor }}
password: ${{ secrets.GHCR_PAT }}
- name: Build and Push Docker Image
uses: docker/build-push-action@v6
with:
context: .
# Push the image to ghcr.io
push: true
# Create a unique tag using the commit SHA for this specific build
tags: ghcr.io/${{ github.repository }}:${{ github.sha }}
test-in-container:
name: Test in Custom Container
# This job must run after the build-and-push job is complete
needs: build-and-push
# Use the required TPU runner
runs-on: linux-x86-ct6e-44-1tpu
# CRITICAL: Use the container image we just pushed in the previous job.
# This satisfies the runner's requirement for a container to be specified.
container:
image: ghcr.io/${{ github.repository }}:${{ github.sha }}
options: --privileged --network host
steps:
- name: Checkout Repository
uses: actions/checkout@v4
# We need the code available inside the container's workspace to run tests.
- name: Run Verification and Tests
run: |
echo "Successfully running inside the custom container from GHCR!"
echo "Current working directory:"
pwd
echo "Contents of current directory:"
ls -la
echo "Verifying JAX installation..."
python3 -c "import jax; print(f'JAX backend: {jax.default_backend()}'); print(f'JAX devices: {jax.devices()}')"