action.yml

name: Setup Linux

description: Set up Docker workspace on EC2

runs:
  using: composite
  steps:
    - name: Display EC2 information
      shell: bash
      run: |
        set -euo pipefail
        function get_ec2_metadata() {
          # Pulled from instance metadata endpoint for EC2
          # see https://docs.aws.amazon.com/AWSEC2/latest/UserGuide/instancedata-data-retrieval.html
          category=$1
          curl -H "X-aws-ec2-metadata-token: $(curl -s -X PUT "http://169.254.169.254/latest/api/token" -H "X-aws-ec2-metadata-token-ttl-seconds: 30")" -fsSL "http://169.254.169.254/latest/meta-data/${category}"
        }
        echo "ami-id: $(get_ec2_metadata ami-id)"
        echo "instance-id: $(get_ec2_metadata instance-id)"
        echo "instance-type: $(get_ec2_metadata instance-type)"
        echo "system info $(uname -a)"

    - name: Check if in a container runner
      shell: bash
      id: check_container_runner
      run: echo "IN_CONTAINER_RUNNER=$(if [ -f /.inarc ] || [ -f /.incontainer ]; then echo true ; else echo false; fi)" >> "$GITHUB_OUTPUT"

    - name: Start docker if docker deamon is not running
      shell: bash
      if: ${{ steps.check_container_runner.outputs.IN_CONTAINER_RUNNER == 'false' }}
      run: |
        if systemctl is-active --quiet docker; then
            echo "Docker daemon is running...";
        else
            echo "Starting docker deamon..." && sudo systemctl start docker;
        fi

    - name: Log in to ECR
      shell: bash
      env:
        AWS_RETRY_MODE: standard
        AWS_MAX_ATTEMPTS: "5"
        AWS_DEFAULT_REGION: us-east-1
      run: |
        AWS_ACCOUNT_ID=$(aws sts get-caller-identity|grep Account|cut -f4 -d\")
        retry () { "$@"  || (sleep 1 && "$@") || (sleep 2 && "$@") }
        retry aws ecr get-login-password --region "$AWS_DEFAULT_REGION" | docker login --username AWS \
            --password-stdin "$AWS_ACCOUNT_ID.dkr.ecr.$AWS_DEFAULT_REGION.amazonaws.com"

    - name: Preserve github env variables for use in docker
      shell: bash
      run: |
        env | grep '^GITHUB' >> "${RUNNER_TEMP}/github_env_${GITHUB_RUN_ID}"
        env | grep '^CI' >> "${RUNNER_TEMP}/github_env_${GITHUB_RUN_ID}"

    - name: Setup useful environment variables
      shell: bash
      working-directory: ${{ inputs.repository }}
      run: |
        RUNNER_ARTIFACT_DIR="${RUNNER_TEMP}/artifacts"
        sudo rm -rf "${RUNNER_ARTIFACT_DIR}"
        mkdir -p "${RUNNER_ARTIFACT_DIR}"
        echo "RUNNER_ARTIFACT_DIR=${RUNNER_ARTIFACT_DIR}" >> "${GITHUB_ENV}"

        RUNNER_TEST_RESULTS_DIR="${RUNNER_TEMP}/test-results"
        sudo rm -rf "${RUNNER_TEST_RESULTS_DIR}"
        mkdir -p "${RUNNER_TEST_RESULTS_DIR}"
        echo "RUNNER_TEST_RESULTS_DIR=${RUNNER_TEST_RESULTS_DIR}" >> "${GITHUB_ENV}"

        RUNNER_DOCS_DIR="${RUNNER_TEMP}/docs"
        sudo rm -rf "${RUNNER_DOCS_DIR}"
        mkdir -p "${RUNNER_DOCS_DIR}"
        echo "RUNNER_DOCS_DIR=${RUNNER_DOCS_DIR}" >> "${GITHUB_ENV}"

    - name: Needs nvidia driver?
      if: ${{ steps.check_container_runner.outputs.IN_CONTAINER_RUNNER == 'false' }}
      shell: bash
      id: needs-nvidia-driver
      run: |
        needs=0
        if lspci -v | grep -e 'controller.*NVIDIA' >/dev/null 2>/dev/null; then
          needs=1
        fi
        echo "does=${needs}" >> $GITHUB_OUTPUT

    - name: Install nvidia driver, nvidia-docker runtime, set GPU_FLAG
      uses: pytorch/test-infra/.github/actions/setup-nvidia@main
      if: ${{ steps.needs-nvidia-driver.outputs.does == 1 && steps.check_container_runner.outputs.IN_CONTAINER_RUNNER == 'false'}}

    - name: Set GPU_FLAG if in a container that has a NVIDIA GPU
      if: ${{ steps.check_container_runner.outputs.IN_CONTAINER_RUNNER == 'true' }}
      shell: bash
      run: |
        if compgen -G "/dev/nvidia[0-9]" ; then
          echo "Found NVIDIA GPU device in /dev/nvidiaX, exporting GPU_FLAG"
          echo "GPU_FLAG=--gpus all -e NVIDIA_DRIVER_CAPABILITIES=all" >> "${GITHUB_ENV}"
        else
          echo "No NVIDIA GPU device found in /dev/nvidiaX, not exporting GPU_FLAG"
        fi

    - name: Kill any existing containers, clean up images
      if: ${{ steps.check_container_runner.outputs.IN_CONTAINER_RUNNER == 'false' }}
      shell: bash
      run: |
        # ignore expansion of "docker ps -q" since it could be empty
        # shellcheck disable=SC2046
        docker stop $(docker ps -q) || true
        # Prune all of the docker images
        docker system prune -af