rocm-ci.yml

name: ROCM CI

on:
  # Trigger the workflow on push or pull request,
  # but only for the rocm-main branch
  push:
    branches:
      - rocm-main
  pull_request:
    branches:
      - rocm-main

jobs:
  build-docker: # strategy and matrix come here 
    runs-on:  mi-250
    steps:
      - name: Print Environment Variables
        run: printenv
      - uses: actions/checkout@11bd71901bbe5b1630ceea73d27597364c9af683  # v4.2.2
          with:
            path: ${{ github.sha }}
      - name: Build Docker 
        env: 
          DOCKER_IMG_NAME: rocm_jax_r6_1_3_py3_10_id${{ github.run_id }}
          # XLA_CLONE_DIR: 
        run: |
          ./build/rocm/ci_build.sh --rocm_version 6.1.3 \
          --keep_image --py_version 3.10 
      - name: Archive jax wheels
        uses: actions/upload-artifact@v4
        with:
          name: rocm_jax_r6_1_3_py3_10_id${{ github.run_id }}
          path: ${{ github.workspace }}/wheelhouse/*.whl
      - name: Detect GPUs
        run: |
          docker run \
            --group-add video \
            $RENDER_DEVICES \
            --cap-add=SYS_PTRACE \
            --security-opt seccomp=unconfined \
            --shm-size=64G \
            rocm/rocm-terminal rocm-smi
      - name: Run tests
        env: 
          DOCKER_IMG_NAME: rocm_jax_r6_1_3_py3_10_id${{ github.run_id }}
        run: |
          ./build/rocm/ci_build test ${DOCKER_IMG_NAME}