EMO/.github/workflows/docker.yml at main · allenai/EMO · GitHub

1
2
3
4
5
6
7
8
9
10
11
12
13
14
15
16
17
18
19
20
21
22
23
24
25
26
27
28
29
30
31
32
33
34
35
36
37
38
39
40
41
42
43
44
45
46
47
48
49
50
51
52
53
54
55
56
57
58
59
60
61
62
63
64
65
66
67
68
69
70
71
72
73
74
75
76
77
78
79
80
81
82
83
84
85
86
87
88
89
90
91
92
93
94
95
96
97
98
99
100
101
name: Docker

concurrency:
  group: ${{ github.workflow }}-${{ github.ref }}
  cancel-in-progress: true

on:
  workflow_dispatch:
  # TODO: disabled for now because it OOMs GHA runners in CI. Can re-enable by
  # migrating to Beaker when it has the capability to build docker images from
  # within batch jobs (currently Beaker only supports building images inside of
  # interactive sessions).
  # pull_request:
  #   branches:
  #     - main
  #   paths:
  #     - 'Makefile'
  #     - 'pyproject.toml'
  #     - 'src/olmo_core/version.py'
  #     - 'src/Dockerfile'
  #     - '.github/workflows/docker.yml'
  # push:
  #   # branches:
  #   #   - main
  #   tags:
  #     - 'v*.*.*'

jobs:
  docker:
    name: CUDA ${{ matrix.cuda }} Torch ${{ matrix.torch.version }}
    runs-on: 8-Core-XL-Runner-Ubuntu-Latest  # 8 cores, 32GB memory
    timeout-minutes: 60
    env:
      BEAKER_TOKEN: ${{ secrets.BEAKER_TOKEN }}
      DOCKER_HUB_TOKEN: ${{ secrets.DOCKER_HUB_TOKEN }}
      DOCKER_HUB_USER: ${{ secrets.DOCKER_HUB_USER }}
    strategy:
      fail-fast: false
      matrix:
        cuda: ["12.8.1", "12.9.1"]  # NOTE: check tags for options: https://hub.docker.com/r/nvidia/cuda/tags
        torch:
          - version: 2.7.1
            channel: whl/test
          - version: 2.9.0
            channel: whl/test
        exclude:
          - cuda: "12.9.1"
            torch:
              version: 2.7.1
              channel: whl/test
    steps:
      - uses: actions/checkout@v4

      - name: Setup Python
        uses: actions/setup-python@v5
        with:
          python-version: '3.12'

      - name: Set env vars
        run: |
          echo "BEAKER_WORKSPACE=$(make get-beaker-workspace)" >> $GITHUB_ENV

      - name: Authenticate with Beaker
        uses: allenai/setup-beaker@v2
        if: env.BEAKER_TOKEN != ''
        with:
          token: ${{ env.BEAKER_TOKEN }}
          workspace: ${{ env.BEAKER_WORKSPACE }}

      - name: Authenticate with Docker Hub
        if: env.DOCKER_HUB_TOKEN != ''
        run: |
          echo ${{ env.DOCKER_HUB_TOKEN }} | docker login -u ${{ env.DOCKER_HUB_USER }} --password-stdin

      - name: Build image
        run: |
          #rm -rf /opt/hostedtoolcache  # clear up some disk space
          make docker-image \
            CUDA_VERSION=${{ matrix.cuda }} \
            TORCH_VERSION=${{ matrix.torch.version }} \
            INSTALL_CHANNEL=${{ matrix.torch.channel }} \
            FA3_MAX_JOBS=4

      - name: Push to GHCR
        if: startsWith(github.ref, 'refs/tags/')
        run: |
          echo ${{ secrets.GITHUB_TOKEN }} | docker login ghcr.io -u ${{ github.actor }} --password-stdin
          make ghcr-image \
            CUDA_VERSION=${{ matrix.cuda }} \
            TORCH_VERSION=${{ matrix.torch.version }} \
            INSTALL_CHANNEL=${{ matrix.torch.channel }} \
            FA3_MAX_JOBS=4

      - name: Push to Beaker
        if: env.BEAKER_TOKEN != '' && startsWith(github.ref, 'refs/tags/')
        run: |
          make beaker-image \
            CUDA_VERSION=${{ matrix.cuda }} \
            TORCH_VERSION=${{ matrix.torch.version }} \
            INSTALL_CHANNEL=${{ matrix.torch.channel }} \
            FA3_MAX_JOBS=4