docker-build.yml

name: Docker Build

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

env:
  AWS_REGION: us-west-2
  ECR_REPOSITORY: openpi-flash
  ECR_REGISTRY: ${{ vars.AWS_ECR_REGISTRY }}
  AWS_ROLE_TO_ASSUME: ${{ vars.AWS_ROLE_TO_ASSUME }}

jobs:
  lint:
    runs-on: ubuntu-latest
    steps:
      - name: Checkout hosting
        uses: actions/checkout@v6
        with:
          path: hosting

      - name: Lint Dockerfile
        uses: hadolint/hadolint-action@v3.1.0
        with:
          dockerfile: hosting/Dockerfile
          config: hosting/.hadolint.yaml

  # Fast checks that do not need the full server environment.
  test:
    needs: lint
    runs-on: ubuntu-latest
    steps:
      - uses: actions/checkout@v6
        with:
          path: hosting

      # Only the client package is needed for msgpack_numpy in metadata tests.
      - uses: actions/checkout@v6
        with:
          repository: Hebbian-Robotics/openpi
          path: openpi
          sparse-checkout: packages/openpi-client

      - uses: dtolnay/rust-toolchain@stable
        with:
          components: rustfmt, clippy

      - uses: actions/cache@v4
        with:
          path: |
            ~/.cargo/registry
            ~/.cargo/git
            hosting/flash-transport/target
          key: ${{ runner.os }}-cargo-${{ hashFiles('hosting/flash-transport/Cargo.lock') }}

      - name: Rust lints + tests
        working-directory: hosting/flash-transport
        run: |
          cargo fmt --check
          cargo clippy --all-targets --all-features -- -D warnings
          cargo test
          cargo build   # builds the echo binaries the Python drift tests exec

      - uses: astral-sh/setup-uv@v5

      # Install just enough for tests; avoid syncing torch/jax/modal.
      - name: Python lints + tests
        working-directory: hosting
        run: |
          uv venv --python 3.11
          uv pip install \
            "numpy>=1.26.4,<2.0.0" \
            "msgpack>=1.2.1" \
            "pytest>=9.1.1" \
            "ruff>=0.15.20" \
            -e ../openpi/packages/openpi-client
          uv pip install --no-deps -e packages/openpi-flash-transport -e .
          .venv/bin/ruff check .
          .venv/bin/ruff format --check .
          .venv/bin/pytest tests/

  # ---------------------------------------------------------------------------
  # Parallel: build transport binaries + Docker base image
  # ---------------------------------------------------------------------------

  build-transport-x86:
    needs: lint
    runs-on: ubuntu-latest
    steps:
      - name: Checkout hosting
        uses: actions/checkout@v6
        with:
          path: hosting

      - name: Install Rust toolchain
        uses: dtolnay/rust-toolchain@stable

      - name: Cache Cargo dependencies
        uses: actions/cache@v4
        with:
          path: |
            ~/.cargo/registry
            ~/.cargo/git
            hosting/flash-transport/target
          key: ${{ runner.os }}-x86_64-cargo-${{ hashFiles('hosting/flash-transport/Cargo.lock') }}

      - name: Build transport (x86_64)
        run: cargo build --release --manifest-path hosting/flash-transport/Cargo.toml

      - name: Upload transport artifact
        uses: actions/upload-artifact@v4
        with:
          name: flash-transport-x86_64-linux
          path: hosting/flash-transport/target/release/openpi-flash-transport
          retention-days: 30

  build-transport-arm:
    needs: lint
    runs-on: ubuntu-latest
    steps:
      - name: Checkout hosting
        uses: actions/checkout@v6
        with:
          path: hosting

      - name: Install Rust toolchain
        uses: dtolnay/rust-toolchain@stable

      - name: Cache Cargo dependencies
        uses: actions/cache@v4
        with:
          path: |
            ~/.cargo/registry
            ~/.cargo/git
            hosting/flash-transport/target
          key: ${{ runner.os }}-aarch64-cargo-${{ hashFiles('hosting/flash-transport/Cargo.lock') }}

      - name: Install cross
        run: cargo install cross

      - name: Build transport (aarch64)
        run: cross build --release --target aarch64-unknown-linux-gnu --manifest-path hosting/flash-transport/Cargo.toml

      - name: Upload transport artifact
        uses: actions/upload-artifact@v4
        with:
          name: flash-transport-aarch64-linux
          path: hosting/flash-transport/target/aarch64-unknown-linux-gnu/release/openpi-flash-transport
          retention-days: 30

  build-base-image:
    needs: lint
    runs-on: ubuntu-latest
    permissions:
      id-token: write
      contents: read
    steps:
      - name: Checkout hosting
        uses: actions/checkout@v6
        with:
          path: hosting

      - name: Checkout openpi
        uses: actions/checkout@v6
        with:
          repository: Hebbian-Robotics/openpi
          path: openpi

      - name: Configure AWS credentials
        uses: aws-actions/configure-aws-credentials@v6
        with:
          role-to-assume: ${{ env.AWS_ROLE_TO_ASSUME }}
          aws-region: ${{ env.AWS_REGION }}

      - name: Login to Amazon ECR
        uses: aws-actions/amazon-ecr-login@v2

      - name: Set up Docker Buildx
        uses: docker/setup-buildx-action@v4

      - name: Build and push base image
        uses: docker/build-push-action@v7
        with:
          context: .
          file: hosting/Dockerfile
          target: base
          push: true
          tags: ${{ env.ECR_REGISTRY }}/${{ env.ECR_REPOSITORY }}:base-${{ github.sha }}
          cache-from: type=gha,scope=base
          cache-to: type=gha,mode=max,scope=base

  # ---------------------------------------------------------------------------
  # Combine: layer transport binary onto base image
  # ---------------------------------------------------------------------------

  combine:
    needs: [build-transport-x86, build-transport-arm, build-base-image, test]
    runs-on: ubuntu-latest
    permissions:
      id-token: write
      contents: read
    steps:
      - name: Checkout hosting
        uses: actions/checkout@v6
        with:
          path: hosting

      - name: Download x86_64 transport artifact
        uses: actions/download-artifact@v4
        with:
          name: flash-transport-x86_64-linux
          path: ./transport

      - name: Mark transport binary as executable
        run: chmod +x ./transport/openpi-flash-transport

      - name: Configure AWS credentials
        uses: aws-actions/configure-aws-credentials@v6
        with:
          role-to-assume: ${{ env.AWS_ROLE_TO_ASSUME }}
          aws-region: ${{ env.AWS_REGION }}

      - name: Login to Amazon ECR
        uses: aws-actions/amazon-ecr-login@v2

      - name: Set up Docker Buildx
        uses: docker/setup-buildx-action@v4

      - name: Compute branch tag
        id: branch-tag
        run: |
          # Use PR branch name for pull_requests, ref_name for pushes.
          # Replace slashes with dashes for Docker tag compatibility.
          RAW="${{ github.head_ref || github.ref_name }}"
          echo "tag=${RAW//\//-}" >> "$GITHUB_OUTPUT"

      - name: Build and push final image
        uses: docker/build-push-action@v7
        with:
          context: ./transport
          file: hosting/Dockerfile.final
          push: true
          build-args: |
            BASE_IMAGE=${{ env.ECR_REGISTRY }}/${{ env.ECR_REPOSITORY }}:base-${{ github.sha }}
          tags: |
            ${{ env.ECR_REGISTRY }}/${{ env.ECR_REPOSITORY }}:${{ github.sha }}
            ${{ env.ECR_REGISTRY }}/${{ env.ECR_REPOSITORY }}:${{ steps.branch-tag.outputs.tag }}
            ${{ github.ref == 'refs/heads/main' && format('{0}/{1}:latest', env.ECR_REGISTRY, env.ECR_REPOSITORY) || '' }}