name: Create Virtual Environment

runs:
  using: "composite"
  steps:
    - id: update-env
      run: |
        sudo apt-get update
        sudo apt-get install -y libaio-dev
        python -m pip install --user --upgrade pip
        python -m pip install --user --upgrade virtualenv
      shell: bash
    - id: create-venv
      run: |
        rm -rf ./unit-test-venv
        python -m venv unit-test-venv
        source ./unit-test-venv/bin/activate
        python -m pip install --upgrade pip
        pip install wheel # required after pip>=23.1
        echo PATH=$PATH >> $GITHUB_ENV # Make it so venv is inherited for other steps
      shell: bash
    - id: set-env-vars
      run: |
        echo TEST_DATA_DIR=/blob/ >> $GITHUB_ENV
        echo TRANSFORMERS_CACHE=/blob/transformers_cache/ >> $GITHUB_ENV
        echo TORCH_EXTENSIONS_DIR=./torch-extensions/ >> $GITHUB_ENV
        echo TORCH_CACHE=/blob/torch_cache/ >> $GITHUB_ENV
        echo HF_DATASETS_CACHE=/blob/datasets_cache/ >> $GITHUB_ENV
        echo MEGATRON_CKPT_DIR=/blob/megatron_ckpt/ >> $GITHUB_ENV
        echo CRITIC_CKPT_DIR=/blob/step2_opt_125m_ckpt/ >> $GITHUB_ENV
        echo PYTEST_OPTS="--color=yes --durations=0 --verbose -rF" >> $GITHUB_ENV
      shell: bash
    - id: print-env
      run: |
        which python
        python --version
        if [[ -z "${AISC_NODE_INSTANCE_ID}" ]]; then
          echo "Not on self-hosted node"
        else
          echo "JobID: ${AISC_NODE_INSTANCE_ID}"
        fi
        if command -v nvidia-smi; then
          nvidia-smi
          which nvcc
          nvcc --version
        elif command -v rocm-smi; then
          rocm-smi --showhw
          which hipcc
          hipcc --version
        fi
      shell: bash