UnitTests.yml

# Copyright 2024 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 workflow will install Python dependencies, run tests and lint with a variety of Python versions
# For more information see: https://docs.github.com/en/actions/automating-builds-and-tests/building-and-testing-python

name: Unit Test

on:
  pull_request:
  push:
    branches: [ "main" ]
  workflow_dispatch:

jobs:
  maxtext_workload:
    name: "Run MaxText Workload"
    # IMPORTANT: Replace with the label for your runner (e.g., v5p-8)
    runs-on: ["linux-x86-a3-megagpu-h100-8gpu"]
    container:
      image: us-docker.pkg.dev/tpu-prod-env-multipod/jax-stable-stack/maxtext-gpu-custom:latest
    steps:
      - name: Checkout MaxText Repo
        uses: actions/checkout@v4
        with:
          repository: AI-Hypercomputer/maxtext
          path: maxtext
      
      - name: Print dependencies
        run: |
          pip freeze

      - name: Run MaxText Training
        run: |
          # This command is adapted from your DAG for a single-slice configuration.
          cd maxtext && \

          export XLA_PYTHON_CLIENT_MEM_FRACTION=0.65
          export TF_FORCE_GPU_ALLOW_GROWTH=true
          export NVTE_FUSED_ATTN=1

          python3 -m MaxText.train MaxText/configs/base.yml \
            steps=5 \
            enable_checkpointing=false \
            attention=cudnn_flash_te \
            dataset_type=synthetic \
            run_name=rbierneni-test-maxtext-gpu \
            base_output_directory=gs://rbierneni-multipod-dev/maxtext/${{ github.run_id }}