# Copyright 2025 The HuggingFace Inc. team. All rights reserved.
#
# 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
#
#     http://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 handles Docker image publishing & testing.
name: Docker Publish & Test

on:
  # Allows running this workflow manually from the Actions tab
  workflow_dispatch:

  # Runs at 02:00
  schedule:
    - cron: "0 2 * * *"

# Sets up the environment variables
env:
  UV_VERSION: "0.11.30"
  PYTHON_VERSION: "3.12"
  DOCKER_IMAGE_NAME_CPU: huggingface/lerobot-cpu:latest
  DOCKER_IMAGE_NAME_GPU: huggingface/lerobot-gpu:latest

# Ensures that only the latest commit is built, canceling older runs.
concurrency:
  group: ${{ github.workflow }}-${{ github.head_ref || github.run_id }}
  cancel-in-progress: true

permissions: {}

jobs:
  # This job builds a CPU image for testing & distribution
  build-docker-cpu:
    permissions:
      contents: read
    name: Build CPU Docker
    runs-on:
      group: aws-general-8-plus
    if: github.repository == 'huggingface/lerobot'
    outputs:
      image_tag: ${{ env.DOCKER_IMAGE_NAME_CPU }}
    steps:
      - name: Install Git LFS
        run: |
          sudo apt-get update
          sudo apt-get install git-lfs
          git lfs install
      - uses: actions/checkout@3d3c42e5aac5ba805825da76410c181273ba90b1 # v7.0.1
        with:
          lfs: true
          persist-credentials: false
      - name: Set up Docker Buildx
        uses: docker/setup-buildx-action@f87e5991a6d7451dcb8d9637bfbc97413f497069 # v4.4.1
        with:
          cache-binary: false
      - name: Login to Docker Hub
        uses: docker/login-action@af1e73f918a031802d376d3c8bbc3fe56130a9b0 # v4.4.0
        with:
          username: ${{ secrets.DOCKERHUB_LEROBOT_USERNAME }}
          password: ${{ secrets.DOCKERHUB_LEROBOT_PASSWORD }}
      - name: Build and push Docker image CPU
        uses: docker/build-push-action@c3c9e263c25d99ce0380d002d59b67737d91b0dc # v7.4.0
        with:
          context: .
          file: ./docker/Dockerfile.user
          push: true
          tags: ${{ env.DOCKER_IMAGE_NAME_CPU }}

  # This job builds a GPU image for testing & distribution
  build-docker-gpu:
    permissions:
      contents: read
    name: Build GPU Docker
    runs-on:
      group: aws-general-8-plus
    if: github.repository == 'huggingface/lerobot'
    outputs:
      image_tag: ${{ env.DOCKER_IMAGE_NAME_GPU }}
    steps:
      - name: Install Git LFS
        run: |
          sudo apt-get update
          sudo apt-get install git-lfs
          git lfs install
      - uses: actions/checkout@3d3c42e5aac5ba805825da76410c181273ba90b1 # v7.0.1
        with:
          lfs: true
          persist-credentials: false
      - name: Set up Docker Buildx
        uses: docker/setup-buildx-action@f87e5991a6d7451dcb8d9637bfbc97413f497069 # v4.4.1
        with:
          cache-binary: false
      - name: Login to Docker Hub
        uses: docker/login-action@af1e73f918a031802d376d3c8bbc3fe56130a9b0 # v4.4.0
        with:
          username: ${{ secrets.DOCKERHUB_LEROBOT_USERNAME }}
          password: ${{ secrets.DOCKERHUB_LEROBOT_PASSWORD }}
      - name: Build and push Docker image GPU
        uses: docker/build-push-action@c3c9e263c25d99ce0380d002d59b67737d91b0dc # v7.4.0
        with:
          context: .
          file: ./docker/Dockerfile.internal
          push: true
          tags: ${{ env.DOCKER_IMAGE_NAME_GPU }}

  # This job runs the E2E tests + pytest with all extras in the CPU image
  cpu-tests:
    permissions: {}
    name: CPU Tests
    needs: [build-docker-cpu]
    runs-on:
      group: aws-g6-4xlarge-plus
    env:
      HF_HOME: /home/user_lerobot/.cache/huggingface
      HF_LEROBOT_HOME: /home/user_lerobot/.cache/huggingface/lerobot
      TORCH_HOME: /home/user_lerobot/.cache/torch
      TRITON_CACHE_DIR: /home/user_lerobot/.cache/triton
      HF_USER_TOKEN: ${{ secrets.LEROBOT_HF_USER }}
    container:
      image: ${{ needs.build-docker-cpu.outputs.image_tag }} # zizmor: ignore[unpinned-images]
      options: --shm-size "16gb"
      credentials:
        username: ${{ secrets.DOCKERHUB_LEROBOT_USERNAME }}
        password: ${{ secrets.DOCKERHUB_LEROBOT_PASSWORD }}
    defaults:
      run:
        shell: bash
        working-directory: /lerobot
    steps:
      - name: Login to Hugging Face
        if: env.HF_USER_TOKEN != ''
        run: |
          hf auth login --token "$HF_USER_TOKEN" --add-to-git-credential
          hf auth whoami
      - name: Run pytest on CPU
        run: pytest tests -vv --maxfail=10
      - name: Run end-to-end tests
        run: make test-end-to-end

  # This job runs the E2E tests + pytest with all extras in the GPU image
  gpu-tests:
    permissions: {}
    name: GPU Tests
    needs: [build-docker-gpu]
    runs-on:
      group: aws-g6-4xlarge-plus
    env:
      HF_HOME: /home/user_lerobot/.cache/huggingface
      HF_LEROBOT_HOME: /home/user_lerobot/.cache/huggingface/lerobot
      TORCH_HOME: /home/user_lerobot/.cache/torch
      TRITON_CACHE_DIR: /home/user_lerobot/.cache/triton
      HF_USER_TOKEN: ${{ secrets.LEROBOT_HF_USER }}
    container:
      image: ${{ needs.build-docker-gpu.outputs.image_tag }} # zizmor: ignore[unpinned-images]
      options: --gpus all --shm-size "16gb"
      credentials:
        username: ${{ secrets.DOCKERHUB_LEROBOT_USERNAME }}
        password: ${{ secrets.DOCKERHUB_LEROBOT_PASSWORD }}
    defaults:
      run:
        shell: bash
        working-directory: /lerobot
    steps:
      - name: Login to Hugging Face
        if: env.HF_USER_TOKEN != ''
        run: |
          hf auth login --token "$HF_USER_TOKEN" --add-to-git-credential
          hf auth whoami
      - name: Run pytest on GPU
        run: pytest tests -vv --maxfail=10
      - name: Run end-to-end tests
        run: make test-end-to-end

  # This job runs the multi-GPU training tests (the suite needs 2+ GPUs)
  multi-gpu-tests:
    permissions: {}
    name: Multi-GPU Tests
    needs: [build-docker-gpu]
    runs-on:
      group: aws-g4dn-12xlarge  # Instance with 4 GPUs
    env:
      HF_HOME: /home/user_lerobot/.cache/huggingface
      HF_LEROBOT_HOME: /home/user_lerobot/.cache/huggingface/lerobot
      TORCH_HOME: /home/user_lerobot/.cache/torch
      TRITON_CACHE_DIR: /home/user_lerobot/.cache/triton
      HF_USER_TOKEN: ${{ secrets.LEROBOT_HF_USER }}
    container:
      image: ${{ needs.build-docker-gpu.outputs.image_tag }} # zizmor: ignore[unpinned-images]
      options: --gpus all --shm-size "16gb"
      credentials:
        username: ${{ secrets.DOCKERHUB_LEROBOT_USERNAME }}
        password: ${{ secrets.DOCKERHUB_LEROBOT_PASSWORD }}
    defaults:
      run:
        shell: bash
        working-directory: /lerobot
    steps:
      - name: Login to Hugging Face
        if: env.HF_USER_TOKEN != ''
        run: |
          hf auth login --token "$HF_USER_TOKEN" --add-to-git-credential
          hf auth whoami
      - name: Verify GPU availability
        run: |
          nvidia-smi
          python -c "import torch; print(f'PyTorch CUDA available: {torch.cuda.is_available()}'); print(f'Number of GPUs: {torch.cuda.device_count()}')"

      - name: Run multi-GPU training tests
        run: pytest -vv tests/training/
