version: 2.1

#examples:
#https://github.com/facebookresearch/ParlAI/blob/master/.circleci/config.yml
#https://github.com/facebookresearch/hydra/blob/master/.circleci/config.yml
#https://github.com/facebookresearch/habitat-api/blob/master/.circleci/config.yml

#drive tests with nox or tox or pytest?

# -------------------------------------------------------------------------------------
# environments where we run our jobs
# -------------------------------------------------------------------------------------


setupcuda: &setupcuda
  run:
    name: Setup CUDA
    working_directory: ~/
    command: |
      # download and install nvidia drivers, cuda, etc
      wget --no-verbose --no-clobber -P ~/nvidia-downloads https://developer.download.nvidia.com/compute/cuda/11.3.1/local_installers/cuda_11.3.1_465.19.01_linux.run
      sudo sh ~/nvidia-downloads/cuda_11.3.1_465.19.01_linux.run --silent
      echo "Done installing CUDA."
      pyenv versions
      nvidia-smi
      pyenv global 3.9.1

binary_common: &binary_common
  parameters:
    # Edit these defaults to do a release`
    build_version:
      description: "version number of release binary; by default, build a nightly"
      type: string
      default: ""
    pytorch_version:
      description: "PyTorch version to build against; by default, use a nightly"
      type: string
      default: ""
    # Don't edit these
    python_version:
      description: "Python version to build against (e.g., 3.7)"
      type: string
    cu_version:
      description: "CUDA version to build against, in CU format (e.g., cpu or cu100)"
      type: string
    wheel_docker_image:
      description: "Wheel only: what docker image to use"
      type: string
      default: "pytorch/manylinux-cuda101"
    conda_docker_image:
      description: "what docker image to use for docker"
      type: string
      default: "pytorch/conda-cuda"
  environment:
    PYTHON_VERSION: << parameters.python_version >>
    BUILD_VERSION: << parameters.build_version >>
    PYTORCH_VERSION: << parameters.pytorch_version >>
    CU_VERSION: << parameters.cu_version >>
    TESTRUN_DOCKER_IMAGE: << parameters.conda_docker_image >>

jobs:
  main:
    environment:
      CUDA_VERSION: "11.3"
    resource_class: gpu.nvidia.small.multi
    machine:
      image: linux-cuda-11:default
    steps:
      - checkout
      - <<: *setupcuda
      - run: pip3 install --progress-bar off imageio wheel matplotlib 'pillow<7'
      - run: pip3 install --progress-bar off torch==1.10.0+cu113 torchvision==0.11.1+cu113 -f https://download.pytorch.org/whl/cu113/torch_stable.html
      # - run: conda create -p ~/conda_env python=3.7 numpy
      # - run: conda activate ~/conda_env
      # - run: conda install -c pytorch pytorch torchvision

      - run: pip3 install --progress-bar off 'git+https://github.com/facebookresearch/fvcore'
      - run: pip3 install --progress-bar off 'git+https://github.com/facebookresearch/iopath'
      - run:
          name: build
          command: |
            export LD_LIBRARY_PATH=$LD_LIBARY_PATH:/usr/local/cuda-11.3/lib64
            python3 setup.py build_ext --inplace
      - run: LD_LIBRARY_PATH=$LD_LIBARY_PATH:/usr/local/cuda-11.3/lib64 python -m unittest discover -v -s tests -t .
      - run: python3 setup.py bdist_wheel

  binary_linux_wheel:
    <<: *binary_common
    docker:
      - image: << parameters.wheel_docker_image >>
        auth:
          username: $DOCKERHUB_USERNAME
          password: $DOCKERHUB_TOKEN
    resource_class: 2xlarge+
    steps:
      - checkout
      - run: MAX_JOBS=15 packaging/build_wheel.sh
      - store_artifacts:
          path: dist
      - persist_to_workspace:
          root: dist
          paths:
            - "*"

  binary_linux_conda:
    <<: *binary_common
    docker:
      - image: "<< parameters.conda_docker_image >>"
        auth:
          username: $DOCKERHUB_USERNAME
          password: $DOCKERHUB_TOKEN
    resource_class: 2xlarge+
    steps:
      - checkout
      # This is building with cuda but no gpu present,
      # so we aren't running the tests.
      - run:
          name: build
          no_output_timeout: 40m
          command: MAX_JOBS=15 TEST_FLAG=--no-test python3 packaging/build_conda.py
      - store_artifacts:
          path: /opt/conda/conda-bld/linux-64
      - persist_to_workspace:
          root: /opt/conda/conda-bld/linux-64
          paths:
            - "*"

  binary_linux_conda_cuda:
    <<: *binary_common
    machine:
      image: linux-cuda-11:default
    resource_class: gpu.nvidia.small.multi
    steps:
    - checkout

    - run:
        name: Pull docker image
        command: |
          nvidia-smi
          set -e

          { docker login -u="$DOCKERHUB_USERNAME" -p="$DOCKERHUB_TOKEN" ; } 2> /dev/null

          echo Pulling docker image $TESTRUN_DOCKER_IMAGE
          docker pull $TESTRUN_DOCKER_IMAGE
    - run:
        name: Build and run tests
        no_output_timeout: 40m
        command: |
          set -e

          cd ${HOME}/project/

          export JUST_TESTRUN=1
          VARS_TO_PASS="-e PYTHON_VERSION -e BUILD_VERSION -e PYTORCH_VERSION -e CU_VERSION -e JUST_TESTRUN"

          docker run --gpus all  --ipc=host -v $(pwd):/remote -w /remote ${VARS_TO_PASS} ${TESTRUN_DOCKER_IMAGE} python3 ./packaging/build_conda.py

workflows:
  version: 2
  build_and_test:
    jobs:
      # - main:
      #     context: DOCKERHUB_TOKEN
      {{workflows()}}
      - binary_linux_conda_cuda:
          name: testrun_conda_cuda_py310_cu117_pyt201
          context: DOCKERHUB_TOKEN
          python_version: "3.10"
          pytorch_version: '2.0.1'
          cu_version: "cu117"
