name: Build default models

on:
  workflow_dispatch:
    inputs:
      target:
        description: 'Model target to build'
        required: true
        type: choice
        options:
          - small
          - big
          - dm
  workflow_call:
    inputs:
      target:
        description: 'Model target to build (small, big, or dm)'
        required: true
        type: string

concurrency:
  group: build-default-models-${{ inputs.target }}
  cancel-in-progress: false

env:
  HF_REPO: sunnypilot/sunnypilot_models_v1

jobs:
  resolve:
    runs-on: ubuntu-24.04
    outputs:
      model_name: ${{ steps.resolve.outputs.model_name }}
      onnx_ref: ${{ steps.resolve.outputs.onnx_ref }}
      onnx_path: ${{ steps.resolve.outputs.onnx_path }}
      hf_defaults_path: ${{ steps.resolve.outputs.hf_defaults_path }}
      tinygrad_ref: ${{ steps.resolve.outputs.tinygrad_ref }}
    steps:
      - uses: actions/checkout@v4
        with:
          submodules: recursive

      - id: resolve
        run: |
          export PYTHONPATH=${{ github.workspace }}

          if [ "${{ inputs.target }}" = "big" ]; then
            NAME=$(python3 -c "from openpilot.sunnypilot.models.model_name import DEFAULT_BIG_MODEL; print(DEFAULT_BIG_MODEL)")
            ONNX_PATH="openpilot/selfdrive/modeld/models/big_driving_supercombo.onnx"
            HF_DEFAULTS_PATH="models/defaults/big"
          elif [ "${{ inputs.target }}" = "dm" ]; then
            ONNX_PATH="openpilot/selfdrive/modeld/models/dmonitoring_model.onnx"
            HF_DEFAULTS_PATH="models/defaults/dm"
            NAME="dmonitoring_model ($(git log -1 --format=%cd --date=format:'%B %d, %Y' -- "$ONNX_PATH"))"
          else
            NAME=$(python3 -c "from openpilot.sunnypilot.models.model_name import DEFAULT_MODEL; print(DEFAULT_MODEL)")
            ONNX_PATH="openpilot/selfdrive/modeld/models/driving_supercombo.onnx"
            HF_DEFAULTS_PATH="models/defaults/small"
          fi

          ONNX_REF=$(git log -1 --format='%H' -- "$ONNX_PATH")
          TINYGRAD_REF=$(python3 openpilot/sunnypilot/models/tinygrad_ref.py)
          if [ -z "$TINYGRAD_REF" ]; then
            echo "::error::Failed to resolve tinygrad ref"
            exit 1
          fi

          echo "model_name=${NAME}" >> $GITHUB_OUTPUT
          echo "onnx_ref=${ONNX_REF}" >> $GITHUB_OUTPUT
          echo "onnx_path=${ONNX_PATH}" >> $GITHUB_OUTPUT
          echo "hf_defaults_path=${HF_DEFAULTS_PATH}" >> $GITHUB_OUTPUT
          echo "tinygrad_ref=${TINYGRAD_REF}" >> $GITHUB_OUTPUT

  build_small_model:
    needs: resolve
    if: ${{ inputs.target == 'small' }}
    runs-on: [self-hosted, tici]
    env:
      SMALL_ONNX: openpilot/selfdrive/modeld/models/driving_supercombo.onnx
      SMALL_PKL: openpilot/selfdrive/modeld/models/driving_tinygrad.pkl
    steps:
      - uses: actions/checkout@v4
        with:
          submodules: recursive

      - name: Pull ONNX via LFS
        run: git lfs pull -I "${{ env.SMALL_ONNX }}"

      - name: Set environment variables
        run: |
          source /etc/profile
          export UV_PROJECT_ENVIRONMENT=${HOME}/venv
          export UV_PYTHON_PREFERENCE=managed
          export UV_PYTHON_INSTALL_DIR=${HOME}/uv/python
          export VIRTUAL_ENV=$UV_PROJECT_ENVIRONMENT
          uv sync --frozen
          printenv >> $GITHUB_ENV

      - name: Disable powersave
        run: |
          source ${UV_PROJECT_ENVIRONMENT}/bin/activate
          PYTHONPATH=$PYTHONPATH:${{ github.workspace }}/ ${{ github.workspace }}/scripts/manage-powersave.py --disable

      - name: Compile small model with stock compiler
        run: |
          source ${UV_PROJECT_ENVIRONMENT}/bin/activate
          export PYTHONPATH="${PYTHONPATH}:${{ github.workspace }}/tinygrad_repo:${{ github.workspace }}"

          MODEL_SIZE=$(python3 -c "from openpilot.common.transformations.model import MEDMODEL_INPUT_SIZE as s; print(f'{s[0]}x{s[1]}')")
          CAMERA_RES=$(python3 -c "from openpilot.common.transformations.camera import _ar_ox_fisheye as a, _os_fisheye as o; print(f'{a.width}x{a.height} {o.width}x{o.height}')")
          FRAME_SKIP=$(python3 -c "from openpilot.selfdrive.modeld.constants import ModelConstants as MC; print(MC.MODEL_RUN_FREQ // MC.MODEL_CONTEXT_FREQ)")

          TG_FLAGS="DEV=QCOM IMAGE=1 FLOAT16=1 NOLOCALS=1 JIT_BATCH_SIZE=0 OPENPILOT_HACKS=1"

          env ${TG_FLAGS} python3 \
            ${{ github.workspace }}/openpilot/selfdrive/modeld/compile_modeld.py \
            --onnx ${{ github.workspace }}/${{ env.SMALL_ONNX }} \
            --model-size $MODEL_SIZE \
            --camera-resolutions $CAMERA_RES \
            --frame-skip $FRAME_SKIP \
            --output ${{ github.workspace }}/${{ env.SMALL_PKL }}

      - name: Chunk small pkl
        run: |
          source ${UV_PROJECT_ENVIRONMENT}/bin/activate
          export PYTHONPATH=${{ github.workspace }}
          python3 -c "
          from openpilot.common.file_chunker import chunk_file, get_chunk_targets
          import os
          pkl = '${{ github.workspace }}/${{ env.SMALL_PKL }}'
          size = os.path.getsize(pkl)
          targets = get_chunk_targets(pkl, size)
          chunk_file(pkl, targets)
          print(f'Chunked into {len(targets)} files')
          "

      - name: Prepare output
        env:
          MODEL_NAME: ${{ needs.resolve.outputs.model_name }}
        run: |
          source ${UV_PROJECT_ENVIRONMENT}/bin/activate
          export PYTHONPATH=${{ github.workspace }}
          MODELS_DIR="${{ github.workspace }}/openpilot/selfdrive/modeld/models"
          OUTPUT_DIR="${{ github.workspace }}/small_output"
          PKL_BASE="driving_tinygrad.pkl"
          mkdir -p "$OUTPUT_DIR"

          cp "$MODELS_DIR/${PKL_BASE}".chunk* "$OUTPUT_DIR/"
          cp "$MODELS_DIR/${PKL_BASE}.chunkmanifest" "$OUTPUT_DIR/"

          python3 "${{ github.workspace }}/release/ci/model_generator.py" \
            --model-dir "$MODELS_DIR" \
            --output-dir "$OUTPUT_DIR" \
            --custom-name "$MODEL_NAME" \
            --upstream-branch "${{ needs.resolve.outputs.onnx_ref }}"

          echo "model-${MODEL_NAME}-${{ github.run_number }}" > "$OUTPUT_DIR/artifact_name.txt"

      - name: Upload small model artifact
        uses: actions/upload-artifact@v4
        with:
          name: model-${{ needs.resolve.outputs.model_name }}-${{ github.run_number }}
          path: ${{ github.workspace }}/small_output/

      - name: Upload artifact name file
        uses: actions/upload-artifact@v4
        with:
          name: artifact-name-${{ needs.resolve.outputs.model_name }}
          path: ${{ github.workspace }}/small_output/artifact_name.txt

      - name: Re-enable powersave
        if: always()
        run: |
          source ${UV_PROJECT_ENVIRONMENT}/bin/activate
          PYTHONPATH=$PYTHONPATH:${{ github.workspace }}/ ${{ github.workspace }}/scripts/manage-powersave.py --enable

  build_big_model:
    needs: resolve
    if: ${{ inputs.target == 'big' }}
    runs-on: [self-hosted, chestnut]
    env:
      BIG_ONNX: openpilot/selfdrive/modeld/models/big_driving_supercombo.onnx
      BIG_PKL: openpilot/selfdrive/modeld/models/big_driving_tinygrad.pkl
    steps:
      - uses: actions/checkout@v4
        with:
          submodules: recursive

      - name: Pull big ONNX via LFS
        run: git lfs pull -I "${{ env.BIG_ONNX }}"

      - name: Set environment variables
        run: |
          source /etc/profile
          export UV_PROJECT_ENVIRONMENT=${HOME}/venv
          export UV_PYTHON_PREFERENCE=managed
          export UV_PYTHON_INSTALL_DIR=${HOME}/uv/python
          export VIRTUAL_ENV=$UV_PROJECT_ENVIRONMENT
          uv sync --frozen
          printenv >> $GITHUB_ENV

      - name: Disable powersave
        run: |
          source ${UV_PROJECT_ENVIRONMENT}/bin/activate
          PYTHONPATH=$PYTHONPATH:${{ github.workspace }}/ ${{ github.workspace }}/scripts/manage-powersave.py --disable

      - name: Wait for chestnut PCIe link
        run: |
          source ${UV_PROJECT_ENVIRONMENT}/bin/activate
          export PYTHONPATH="${PYTHONPATH}:${{ github.workspace }}/tinygrad_repo:${{ github.workspace }}"
          python3 -c "
          import time
          from openpilot.system.hardware.chestnut.flash import link_up
          for i in range(10):
              if link_up():
                  print(f'PCIe link up after {i+1} attempt(s)')
                  break
              time.sleep(1)
          else:
              raise RuntimeError('Chestnut PCIe link not ready after 10 attempts')
          "

      - name: Compile big model with stock compiler
        run: |
          source ${UV_PROJECT_ENVIRONMENT}/bin/activate
          export PYTHONPATH="${PYTHONPATH}:${{ github.workspace }}/tinygrad_repo:${{ github.workspace }}"

          MODEL_SIZE=$(python3 -c "from openpilot.common.transformations.model import MEDMODEL_INPUT_SIZE as s; print(f'{s[0]}x{s[1]}')")
          CAMERA_RES=$(python3 -c "from openpilot.common.transformations.camera import _ar_ox_fisheye as a, _os_fisheye as o; print(f'{a.width}x{a.height} {o.width}x{o.height}')")
          FRAME_SKIP=$(python3 -c "from openpilot.selfdrive.modeld.constants import ModelConstants as MC; print(MC.MODEL_RUN_FREQ // MC.MODEL_CONTEXT_FREQ)")

          TG_FLAGS="DEBUG=2 DEV=USB+AMD:LLVM WARP_DEV=QCOM FLOAT16=1 JIT_BATCH_SIZE=0 GMMU=0 TC_OPT=2"

          env ${TG_FLAGS} python3 \
            ${{ github.workspace }}/openpilot/selfdrive/modeld/compile_modeld.py \
            --onnx ${{ github.workspace }}/${{ env.BIG_ONNX }} \
            --model-size $MODEL_SIZE \
            --camera-resolutions $CAMERA_RES \
            --frame-skip $FRAME_SKIP \
            --output ${{ github.workspace }}/${{ env.BIG_PKL }}

      - name: Chunk big pkl
        run: |
          source ${UV_PROJECT_ENVIRONMENT}/bin/activate
          export PYTHONPATH=${{ github.workspace }}
          python3 -c "
          from openpilot.common.file_chunker import chunk_file, get_chunk_targets
          import os
          pkl = '${{ github.workspace }}/${{ env.BIG_PKL }}'
          size = os.path.getsize(pkl)
          targets = get_chunk_targets(pkl, size)
          chunk_file(pkl, targets)
          print(f'Chunked into {len(targets)} files')
          "

      - name: Prepare output
        env:
          MODEL_NAME: ${{ needs.resolve.outputs.model_name }}
        run: |
          source ${UV_PROJECT_ENVIRONMENT}/bin/activate
          export PYTHONPATH=${{ github.workspace }}
          MODELS_DIR="${{ github.workspace }}/openpilot/selfdrive/modeld/models"
          OUTPUT_DIR="${{ github.workspace }}/big_output"
          PKL_BASE="big_driving_tinygrad.pkl"
          mkdir -p "$OUTPUT_DIR"

          cp "$MODELS_DIR/${PKL_BASE}".chunk* "$OUTPUT_DIR/"
          cp "$MODELS_DIR/${PKL_BASE}.chunkmanifest" "$OUTPUT_DIR/"

          python3 "${{ github.workspace }}/release/ci/model_generator.py" \
            --model-dir "$MODELS_DIR" \
            --output-dir "$OUTPUT_DIR" \
            --custom-name "$MODEL_NAME" \
            --upstream-branch "${{ needs.resolve.outputs.onnx_ref }}"

          echo "model-${MODEL_NAME}-${{ github.run_number }}" > "$OUTPUT_DIR/artifact_name.txt"

      - name: Upload big model artifact
        uses: actions/upload-artifact@v4
        with:
          name: model-${{ needs.resolve.outputs.model_name }}-${{ github.run_number }}
          path: ${{ github.workspace }}/big_output/

      - name: Upload artifact name file
        uses: actions/upload-artifact@v4
        with:
          name: artifact-name-${{ needs.resolve.outputs.model_name }}
          path: ${{ github.workspace }}/big_output/artifact_name.txt

      - name: Re-enable powersave
        if: always()
        run: |
          source ${UV_PROJECT_ENVIRONMENT}/bin/activate
          PYTHONPATH=$PYTHONPATH:${{ github.workspace }}/ ${{ github.workspace }}/scripts/manage-powersave.py --enable

  upload_defaults:
    needs: [ resolve, build_small_model, build_big_model, build_dm_model ]
    if: |
      ${{
        !cancelled() &&
        (inputs.target == 'big' && needs.build_big_model.result == 'success' ||
         inputs.target == 'small' && needs.build_small_model.result == 'success' ||
         inputs.target == 'dm' && needs.build_dm_model.result == 'success')
      }}
    runs-on: ubuntu-24.04
    permissions:
      id-token: write
      contents: write
    steps:
      - uses: actions/checkout@v4

      - name: Pull ONNX via LFS
        run: git lfs pull -I "${{ needs.resolve.outputs.onnx_path }}"

      - name: Install huggingface_hub
        run: pip install --upgrade "huggingface_hub>=0.22.0"

      - name: Download artifact name
        if: ${{ inputs.target == 'small' || inputs.target == 'big' }}
        uses: actions/download-artifact@v4
        with:
          name: artifact-name-${{ needs.resolve.outputs.model_name }}
          path: artifact_name

      - name: Read artifact name
        if: ${{ inputs.target == 'small' || inputs.target == 'big' }}
        id: artifact
        run: |
          ARTIFACT_NAME=$(cat artifact_name/artifact_name.txt)
          echo "artifact_name=$ARTIFACT_NAME" >> $GITHUB_OUTPUT

      - name: Download model artifact
        if: ${{ inputs.target == 'small' || inputs.target == 'big' }}
        uses: actions/download-artifact@v4
        with:
          name: ${{ steps.artifact.outputs.artifact_name }}
          path: output

      - name: Upload model to HF
        if: ${{ inputs.target == 'small' || inputs.target == 'big' }}
        env:
          HF_OIDC_RESOURCE: datasets/${{ env.HF_REPO }}
          ARTIFACT_NAME: ${{ steps.artifact.outputs.artifact_name }}
        run: |
          rm -f output/artifact_name.txt
          export PYTHONPATH=$(pwd)
          python3 release/ci/upload_default_model.py \
            --hf-repo "${{ env.HF_REPO }}" \
            --hf-defaults-path "${{ needs.resolve.outputs.hf_defaults_path }}" \
            --artifact-name "$ARTIFACT_NAME" \
            --model-dir output \
            --onnx-path "${{ needs.resolve.outputs.onnx_path }}" \
            --onnx-ref "${{ needs.resolve.outputs.onnx_ref }}" \
            --model-name "${{ needs.resolve.outputs.model_name }}" \
            --tinygrad-ref "${{ needs.resolve.outputs.tinygrad_ref }}" \
            --run-number "${{ github.run_number }}"

      - name: Download DM artifact
        if: ${{ inputs.target == 'dm' }}
        uses: actions/download-artifact@v4
        with:
          name: dm-model-${{ github.run_number }}
          path: dm_output

      - name: Generate DM metadata and upload to HF
        if: ${{ inputs.target == 'dm' }}
        env:
          HF_OIDC_RESOURCE: datasets/${{ env.HF_REPO }}
        run: |
          export PYTHONPATH=$(pwd)
          python3 -c "
          import json, hashlib
          from pathlib import Path
          from datetime import datetime, UTC

          dm_dir = Path('dm_output')
          manifest = list(dm_dir.glob('*.chunkmanifest'))
          assert manifest, 'No chunkmanifest found'
          pkl_name = manifest[0].name.removesuffix('.chunkmanifest')
          num_chunks = int(manifest[0].read_text().strip())

          chunks = []
          for i in range(num_chunks):
              chunk = dm_dir / f'{pkl_name}.chunk{i+1:02d}of{num_chunks:02d}'
              chunks.append({
                  'file_name': chunk.name,
                  'sha256': hashlib.sha256(chunk.read_bytes()).hexdigest()
              })

          digest = hashlib.sha256()
          for c in chunks:
              with open(dm_dir / c['file_name'], 'rb') as f:
                  while block := f.read(1024*1024):
                      digest.update(block)

          metadata = {
              'bundles': [{
                  'short_name': 'DMMODEL',
                  'display_name': '${{ needs.resolve.outputs.model_name }}',
                  'ref': '${{ needs.resolve.outputs.onnx_ref }}',
                  'runner': 'tinygrad',
                  'build_time': datetime.now(UTC).strftime('%Y-%m-%dT%H:%M:%SZ'),
                  'models': [{
                      'type': 'chunked',
                      'artifact': {
                          'file_name': pkl_name,
                          'download_uri': {'url': '', 'sha256': digest.hexdigest()},
                          'chunks': chunks
                      }
                  }]
              }]
          }
          with open(dm_dir / 'metadata.json', 'w') as f:
              json.dump(metadata, f, indent=2)
          print('Generated DM metadata.json')
          "

          python3 release/ci/upload_default_model.py \
            --hf-repo "${{ env.HF_REPO }}" \
            --hf-defaults-path "${{ needs.resolve.outputs.hf_defaults_path }}" \
            --artifact-name "dm-model-${{ github.run_number }}" \
            --model-dir dm_output \
            --onnx-path "${{ needs.resolve.outputs.onnx_path }}" \
            --onnx-ref "${{ needs.resolve.outputs.onnx_ref }}" \
            --model-name "${{ needs.resolve.outputs.model_name }}" \
            --tinygrad-ref "${{ needs.resolve.outputs.tinygrad_ref }}" \
            --run-number "${{ github.run_number }}"

  build_dm_model:
    needs: resolve
    if: ${{ inputs.target == 'dm' }}
    runs-on: [self-hosted, tici]
    env:
      DM_ONNX: openpilot/selfdrive/modeld/models/dmonitoring_model.onnx
      DM_PKL: openpilot/selfdrive/modeld/models/dmonitoring_model_tinygrad.pkl
    steps:
      - uses: actions/checkout@v4
        with:
          submodules: recursive

      - name: Pull DM ONNX via LFS
        run: git lfs pull -I "${{ env.DM_ONNX }}"

      - name: Set environment variables
        run: |
          source /etc/profile
          export UV_PROJECT_ENVIRONMENT=${HOME}/venv
          export UV_PYTHON_PREFERENCE=managed
          export UV_PYTHON_INSTALL_DIR=${HOME}/uv/python
          export VIRTUAL_ENV=$UV_PROJECT_ENVIRONMENT
          uv sync --frozen
          printenv >> $GITHUB_ENV

      - name: Disable powersave
        run: |
          source ${UV_PROJECT_ENVIRONMENT}/bin/activate
          PYTHONPATH=$PYTHONPATH:${{ github.workspace }}/ ${{ github.workspace }}/scripts/manage-powersave.py --disable

      - name: Compile DM model
        run: |
          source ${UV_PROJECT_ENVIRONMENT}/bin/activate
          export PYTHONPATH="${PYTHONPATH}:${{ github.workspace }}/tinygrad_repo:${{ github.workspace }}"

          TG_FLAGS="DEV=QCOM IMAGE=1 FLOAT16=1 NOLOCALS=1 JIT_BATCH_SIZE=0 OPENPILOT_HACKS=1"

          taskset -c 7 env ${TG_FLAGS} python3 \
            ${{ github.workspace }}/tinygrad_repo/examples/openpilot/compile3.py \
            ${{ github.workspace }}/${{ env.DM_ONNX }} \
            ${{ github.workspace }}/${{ env.DM_PKL }}

      - name: Chunk DM pkl
        run: |
          source ${UV_PROJECT_ENVIRONMENT}/bin/activate
          export PYTHONPATH=${{ github.workspace }}
          python3 -c "
          from openpilot.common.file_chunker import chunk_file, get_chunk_targets
          import os
          pkl = '${{ github.workspace }}/${{ env.DM_PKL }}'
          size = os.path.getsize(pkl)
          targets = get_chunk_targets(pkl, size)
          chunk_file(pkl, targets)
          print(f'Chunked {pkl} into {len(targets)} chunks')
          "

      - name: Prepare DM output
        run: |
          mkdir -p dm_output
          cp ${{ github.workspace }}/${{ env.DM_PKL }}.chunk* dm_output/
          cp ${{ github.workspace }}/${{ env.DM_PKL }}.chunkmanifest dm_output/

      - name: Upload DM artifact
        uses: actions/upload-artifact@v4
        with:
          name: dm-model-${{ github.run_number }}
          path: dm_output/

      - name: Re-enable powersave
        if: always()
        run: |
          source ${UV_PROJECT_ENVIRONMENT}/bin/activate
          PYTHONPATH=$PYTHONPATH:${{ github.workspace }}/ ${{ github.workspace }}/scripts/manage-powersave.py --enable