container-physicsnemo.yml

# SPDX-FileCopyrightText: 2026 Samudra Authors
#
# SPDX-License-Identifier: Apache-2.0

# Build the PhysicsNeMo-based container on PRs; on main/manual runs, also publish
# x86_64 and arm64 images, run x86_64 CPU/GPU container tests, and run arm64
# CPU-only container tests on a non-GPU ARM runner.
name: Container PhysicsNeMo 26.05

on:
  push:
    branches:
      - main
  pull_request:
    branches:
      - "**"
  workflow_dispatch:

permissions:
  id-token: write
  contents: read
  packages: write

concurrency:
  group: ${{ github.workflow }}-${{ github.event_name }}-${{ github.ref }}
  cancel-in-progress: ${{ github.event_name == 'pull_request' }}

jobs:
  ec2-build:
    if: github.event_name != 'pull_request' || github.event.pull_request.head.repo.full_name == github.repository
    uses: Open-Athena/ec2-gha/.github/workflows/runner.yml@v2
    with:
      aws_region: us-east-1
      ec2_instance_type: m7i.xlarge
      # Ubuntu 24.04 LTS amd64 hvm ebs-gp3 in us-east-1, Canonical owner 099720109477.
      # To refresh, run:
      # aws ec2 describe-images --region us-east-1 --owners 099720109477 --filters \
      #   'Name=name,Values=ubuntu/images/hvm-ssd-gp3/ubuntu-noble-24.04-amd64-server-*' \
      #   'Name=architecture,Values=x86_64' 'Name=virtualization-type,Values=hvm' \
      #   'Name=root-device-type,Values=ebs' 'Name=state,Values=available' \
      #   --query 'reverse(sort_by(Images, &CreationDate))[:1].ImageId' --output text
      ec2_image_id: ami-04b70fa74e45c3917
      ec2_root_device_size: 200 # (GBs) local x86_64 image build needs substantial headroom
      aws_tags: '[{"Key": "project", "Value": "fomo"}]'
    secrets:
      GH_SA_TOKEN: ${{ secrets.GH_SA_TOKEN }}

  build-and-smoke:
    needs: ec2-build
    runs-on: ${{ needs.ec2-build.outputs.id }}
    env:
      LOCAL_IMAGE_TAG: ocean-emulator:physicsnemo-26.05-${{ github.sha }}
    steps:
      - name: Check out code
        uses: actions/checkout@v6

      - name: Confirm x86_64 runner
        run: |
          arch="$(uname -m)"
          echo "runner arch=${arch}"
          test "${arch}" = "x86_64"

      - name: Install Docker
        run: |
          set -euo pipefail
          if ! command -v docker >/dev/null 2>&1 || ! docker buildx version >/dev/null 2>&1; then
            sudo apt-get update
            sudo DEBIAN_FRONTEND=noninteractive apt-get install -y docker.io docker-buildx
          fi
          sudo systemctl enable --now docker || sudo service docker start
          sudo chmod 666 /var/run/docker.sock
          docker version
          docker buildx version

      - name: Build container
        run: |
          BUILDKIT_PROGRESS=plain \
          IMAGE_TAG="${LOCAL_IMAGE_TAG}" \
          DOCKERFILE=containers/Dockerfile.physicsnemo-26.05 \
          BUILD_APPTAINER=0 \
          bash scripts/container/build_physicsnemo_26_05.sh

      - name: Log in to GHCR
        if: github.event_name == 'workflow_dispatch' || (github.event_name == 'push' && github.ref == 'refs/heads/main')
        uses: docker/login-action@v3
        with:
          registry: ghcr.io
          username: ${{ github.actor }}
          password: ${{ secrets.GITHUB_TOKEN }}

      - name: Publish x86_64 image
        if: github.event_name == 'workflow_dispatch' || (github.event_name == 'push' && github.ref == 'refs/heads/main')
        run: |
          set -euo pipefail

          owner_lower="$(echo "${GITHUB_REPOSITORY_OWNER}" | tr '[:upper:]' '[:lower:]')"
          ghcr_image="ghcr.io/${owner_lower}/ocean-emulator-physicsnemo"

          sha_tag="${ghcr_image}:26.05-${GITHUB_SHA}"
          arch_sha_tag="${ghcr_image}:26.05-x86_64-${GITHUB_SHA}"
          docker tag "${LOCAL_IMAGE_TAG}" "${sha_tag}"
          docker tag "${LOCAL_IMAGE_TAG}" "${arch_sha_tag}"
          docker push "${sha_tag}"
          docker push "${arch_sha_tag}"

          if [ "${GITHUB_EVENT_NAME}" = "workflow_dispatch" ]; then
            ref_tag="${ghcr_image}:26.05-manual-${GITHUB_REF_NAME//\//-}"
            arch_ref_tag="${ghcr_image}:26.05-x86_64-manual-${GITHUB_REF_NAME//\//-}"
            docker tag "${LOCAL_IMAGE_TAG}" "${ref_tag}"
            docker tag "${LOCAL_IMAGE_TAG}" "${arch_ref_tag}"
            docker push "${ref_tag}"
            docker push "${arch_ref_tag}"
          fi

          if [ "${GITHUB_EVENT_NAME}" = "push" ] && [ "${GITHUB_REF}" = "refs/heads/main" ]; then
            latest_tag="${ghcr_image}:26.05-latest"
            arch_latest_tag="${ghcr_image}:26.05-x86_64-latest"
            docker tag "${LOCAL_IMAGE_TAG}" "${latest_tag}"
            docker tag "${LOCAL_IMAGE_TAG}" "${arch_latest_tag}"
            docker push "${latest_tag}"
            docker push "${arch_latest_tag}"
          fi

  ec2:
    if: github.event_name == 'workflow_dispatch' || (github.event_name == 'push' && github.ref == 'refs/heads/main')
    needs: build-and-smoke
    uses: Open-Athena/ec2-gha/.github/workflows/runner.yml@v2
    with:
      ec2_instance_type: g6.xlarge
      ec2_image_id: ami-0365bff494b18bf93 # Deep Learning OSS Nvidia Driver AMI GPU PyTorch 2.7 (Ubuntu 22.04) in oa-ci-dev
      ec2_root_device_size: 200 # large image pull + test data/cache need substantial headroom
      aws_tags: '[{"Key": "project", "Value": "fomo"}]'
    secrets:
      GH_SA_TOKEN: ${{ secrets.GH_SA_TOKEN }}

  ec2-arm:
    if: github.event_name == 'workflow_dispatch' || (github.event_name == 'push' && github.ref == 'refs/heads/main')
    uses: Open-Athena/ec2-gha/.github/workflows/runner.yml@v2
    with:
      aws_region: us-east-1
      ec2_instance_type: m7g.xlarge
      # Ubuntu 24.04 LTS arm64 hvm ebs-gp3 in us-east-1, Canonical owner 099720109477.
      # To refresh, run:
      # aws ec2 describe-images --region us-east-1 --owners 099720109477 --filters \
      #   'Name=name,Values=ubuntu/images/hvm-ssd-gp3/ubuntu-noble-24.04-arm64-server-*' \
      #   'Name=architecture,Values=arm64' 'Name=virtualization-type,Values=hvm' \
      #   'Name=root-device-type,Values=ebs' 'Name=state,Values=available' \
      #   --query 'reverse(sort_by(Images, &CreationDate))[:1].ImageId' --output text
      ec2_image_id: ami-02c57ce16e9a8e906
      ec2_root_device_size: 200 # local arm64 image build plus test data/cache need substantial headroom
      aws_tags: '[{"Key": "project", "Value": "fomo"}]'
    secrets:
      GH_SA_TOKEN: ${{ secrets.GH_SA_TOKEN }}

  test-container-cpu-gpu:
    if: github.event_name == 'workflow_dispatch' || (github.event_name == 'push' && github.ref == 'refs/heads/main')
    needs: [build-and-smoke, ec2]
    runs-on: ${{ needs.ec2.outputs.id }}
    steps:
      - name: Check out code
        uses: actions/checkout@v6
        with:
          ref: ${{ github.sha }}

      - name: Cache test data
        uses: actions/cache@v3
        with:
          path: .data_cache/
          key: ${{ runner.os }}-test-data-${{ hashFiles('tests/conftest.py', 'configs/test/*.yaml') }}

      - name: Log in to GHCR
        uses: docker/login-action@v3
        with:
          registry: ghcr.io
          username: ${{ github.actor }}
          password: ${{ secrets.GITHUB_TOKEN }}

      - name: Show disk before pull
        run: |
          df -h
          docker system df || true

      - name: Pull published image
        run: |
          set -euo pipefail
          owner_lower="$(echo "${GITHUB_REPOSITORY_OWNER}" | tr '[:upper:]' '[:lower:]')"
          image="ghcr.io/${owner_lower}/ocean-emulator-physicsnemo:26.05-${GITHUB_SHA}"
          docker pull "${image}"
          echo "TEST_IMAGE=${image}" >> "${GITHUB_ENV}"

      - name: Run CPU tests in container
        timeout-minutes: 90
        run: |
          docker run --rm \
            --ipc=host \
            -v "${GITHUB_WORKSPACE}/.data_cache:/workspace/.data_cache" \
            --workdir /workspace \
            --entrypoint bash \
            "${TEST_IMAGE}" -lc '
              set -euo pipefail
              . .venv/bin/activate
              python -m pytest -p no:zarr -n 2 -v -m "not manual and not cuda" --durations=20 --durations-min=1.0
            '

      - name: Run GPU tests in container
        timeout-minutes: 90
        run: |
          docker run --rm --gpus all \
            --ipc=host \
            -v "${GITHUB_WORKSPACE}/.data_cache:/workspace/.data_cache" \
            --workdir /workspace \
            --entrypoint bash \
            "${TEST_IMAGE}" -lc '
              set -euo pipefail
              . .venv/bin/activate
              python -m pytest -p no:zarr -v -m "not manual and cuda" -x --durations=20 --durations-min=1.0
            '

  build-test-publish-arm64:
    if: github.event_name == 'workflow_dispatch' || (github.event_name == 'push' && github.ref == 'refs/heads/main')
    needs: ec2-arm
    runs-on: ${{ needs.ec2-arm.outputs.id }}
    env:
      LOCAL_ARM64_IMAGE_TAG: ocean-emulator:physicsnemo-26.05-arm64-${{ github.sha }}
    steps:
      - name: Check out code
        uses: actions/checkout@v6
        with:
          ref: ${{ github.sha }}

      - name: Confirm arm64 runner
        run: |
          arch="$(uname -m)"
          echo "runner arch=${arch}"
          test "${arch}" = "aarch64"

      - name: Install Docker
        run: |
          set -euo pipefail
          if ! command -v docker >/dev/null 2>&1 || ! docker buildx version >/dev/null 2>&1; then
            sudo apt-get update
            sudo DEBIAN_FRONTEND=noninteractive apt-get install -y docker.io docker-buildx
          fi
          sudo systemctl enable --now docker || sudo service docker start
          sudo chmod 666 /var/run/docker.sock
          docker version
          docker buildx version

      - name: Cache test data
        uses: actions/cache@v3
        with:
          path: .data_cache/
          key: ${{ runner.os }}-${{ runner.arch }}-test-data-${{ hashFiles('tests/conftest.py', 'configs/test/*.yaml') }}

      - name: Show disk before build
        run: |
          df -h
          docker system df || true

      - name: Build arm64 container
        run: |
          BUILDKIT_PROGRESS=plain \
          IMAGE_TAG="${LOCAL_ARM64_IMAGE_TAG}" \
          DOCKERFILE=containers/Dockerfile.physicsnemo-26.05 \
          BUILD_APPTAINER=0 \
          bash scripts/container/build_physicsnemo_26_05.sh

      - name: Log in to GHCR
        uses: docker/login-action@v3
        with:
          registry: ghcr.io
          username: ${{ github.actor }}
          password: ${{ secrets.GITHUB_TOKEN }}

      - name: Publish arm64 image
        run: |
          set -euo pipefail

          owner_lower="$(echo "${GITHUB_REPOSITORY_OWNER}" | tr '[:upper:]' '[:lower:]')"
          ghcr_image="ghcr.io/${owner_lower}/ocean-emulator-physicsnemo"

          sha_tag="${ghcr_image}:26.05-arm64-${GITHUB_SHA}"
          docker tag "${LOCAL_ARM64_IMAGE_TAG}" "${sha_tag}"
          docker push "${sha_tag}"
          echo "TEST_ARM64_IMAGE=${sha_tag}" >> "${GITHUB_ENV}"

          if [ "${GITHUB_EVENT_NAME}" = "workflow_dispatch" ]; then
            ref_tag="${ghcr_image}:26.05-arm64-manual-${GITHUB_REF_NAME//\//-}"
            docker tag "${LOCAL_ARM64_IMAGE_TAG}" "${ref_tag}"
            docker push "${ref_tag}"
          fi

          if [ "${GITHUB_EVENT_NAME}" = "push" ] && [ "${GITHUB_REF}" = "refs/heads/main" ]; then
            latest_tag="${ghcr_image}:26.05-arm64-latest"
            docker tag "${LOCAL_ARM64_IMAGE_TAG}" "${latest_tag}"
            docker push "${latest_tag}"
          fi

      - name: Pull published arm64 image
        run: |
          set -euo pipefail
          docker pull "${TEST_ARM64_IMAGE}"

      - name: Run CPU tests in arm64 container
        timeout-minutes: 90
        run: |
          docker run --rm \
            --ipc=host \
            -v "${GITHUB_WORKSPACE}/.data_cache:/workspace/.data_cache" \
            --workdir /workspace \
            --entrypoint bash \
            "${TEST_ARM64_IMAGE}" -lc '
              set -euo pipefail
              . .venv/bin/activate
              # The smaller arm64 runner is more prone to losing the GitHub runner
              # process under xdist CPU-test load after the local image build.
              python -m pytest -p no:zarr -v -m "not manual and not cuda" --durations=20 --durations-min=1.0
            '

  report-container-test-status:
    if: always()
    needs:
      - ec2-build
      - build-and-smoke
      - ec2
      - ec2-arm
      - test-container-cpu-gpu
      - build-test-publish-arm64
    runs-on: ubuntu-latest
    steps:
      - name: Check job results
        if: contains(needs.*.result, 'failure') || contains(needs.*.result, 'cancelled')
        run: exit 1

      - name: Fail if container build did not run
        if: needs['build-and-smoke'].result == 'skipped'
        run: |
          echo "AWS-backed container tests did not run automatically."
          echo "Run this workflow manually from the GitHub Actions UI once the code has been reviewed."
          exit 1