name: DeepMD-Kit Test

on:
  workflow_call:
    inputs:
      docker_build_image:
        type: string
        required: true
      can-skip:
        type: string
        required: false
        default: "false"
      MANUALLY_PR_ID:
        type: string
        required: false
      MANUALLY_COMMIT_ID:
        type: string
        required: false
      build-can-skip:
        type: string
        required: false
        default: 'false'

env:
  PR_ID: ${{ github.event.pull_request.number || '0' }}
  COMMIT_ID: ${{ github.event.pull_request.head.sha || github.sha }}
  work_dir: /paddle
  PADDLE_ROOT: /paddle
  TASK: paddle-CI-${{ github.event.pull_request.number }}-deepmd
  ci_scripts: /paddle/ci
  BRANCH: ${{ github.event.pull_request.base.ref || github.ref_name }}
  CI_name: deepmd
  no_proxy: "bcebos.com,apiin.im.baidu.com,gitee.com,aliyun.com,.baidu.com,.tuna.tsinghua.edu.cn"

defaults:
  run:
    shell: bash

jobs:
  check-bypass:
    name: Check bypass
    if: ${{ github.repository_owner == 'PaddlePaddle' && inputs.build-can-skip != 'true' }}
    uses: ./.github/workflows/check-bypass.yml
    with:
      workflow-name: "deepmd"
    secrets:
      github-token: ${{ secrets.GITHUB_TOKEN }}

  deepmd-test:
    name: DeepMD-Kit Integration Test
    needs: [check-bypass]
    if: ${{ needs.check-bypass.outputs.can-skip != 'true' && inputs.can-skip != 'true' }}
    runs-on:
      group: BD_BJ-V100
    timeout-minutes: 20
    steps:
      - name: Check docker image and run container
        env:
          wheel_link: https://paddle-github-action.bj.bcebos.com/PR/build/${{ env.PR_ID }}/${{ env.COMMIT_ID }}/paddlepaddle_gpu-0.0.0-cp310-cp310-linux_x86_64.whl
          RUN_ID: ${{ github.run_id }}
        run: |
          container_name="deepmd_test_ci_${RUN_ID}"
          echo "container_name=${container_name}" >> $GITHUB_ENV
          docker_image=${{ inputs.docker_build_image }}
          docker run -d -t --gpus all --name ${container_name} --shm-size=128g \
            -v "/home/data/cfs:/home/data/cfs" \
            -v ${{ github.workspace }}/../../..:${{ github.workspace }}/../../.. \
            -v ${{ github.workspace }}:/paddle \
            -e BRANCH \
            -e PR_ID \
            -e COMMIT_ID \
            -e RUN_ID \
            -e wheel_link \
            -e work_dir \
            -e PADDLE_ROOT \
            -e ci_scripts \
            -e no_proxy \
            -e CI_name \
            -w /paddle --network host ${docker_image}

      - name: Print PR information
        run: |
          docker exec -t ${{ env.container_name }} /bin/bash -c '
          set -e
          echo "Testing PR #${{ env.PR_ID }}"
          echo "Commit: ${{ env.COMMIT_ID }}"
          echo "Branch: ${{ env.BRANCH }}"
          '

      - name: Install PaddlePaddle
        run: |
          docker exec -t ${{ env.container_name }} /bin/bash -c '
          set -e
          mkdir -p ${{ github.workspace }}/../../../.cache/pip
          source ${{ github.workspace }}/../../../proxy
          python3.10 -m pip config set global.cache-dir ${{ github.workspace }}/../../../.cache/pip

          if [[ "${{ inputs.MANUALLY_PR_ID }}" == "" ]]; then
            echo "Installing PaddlePaddle from: $wheel_link"
            python3.10 -m pip install $wheel_link
          else
            echo "Installing PaddlePaddle from manual PR: ${{ inputs.MANUALLY_PR_ID }}"
            python3.10 -m pip install https://paddle-github-action.bj.bcebos.com/PR/build/${{ inputs.MANUALLY_PR_ID }}/${{ inputs.MANUALLY_COMMIT_ID }}/paddlepaddle_gpu-0.0.0-cp310-cp310-linux_x86_64.whl
          fi

          python3.10 -c "import paddle; print(f\"PaddlePaddle version: {paddle.__version__}\")"
          '

      - name: Clone and setup deepmd-kit
        run: |
          docker exec -t ${{ env.container_name }} /bin/bash -c '
          set -e
          cd /paddle
          rm -rf deepmd-kit
          echo "Cloning deepmd-kit repository..."
          source ${{ github.workspace }}/../../../proxy
          git clone https://github.com/deepmodeling/deepmd-kit.git -b v3.1.2
          ls -lah /paddle/deepmd-kit
          cd deepmd-kit
          '

      - name: Install deepmd-kit dependencies
        run: |
          docker exec -t ${{ env.container_name }} /bin/bash -c '
          set -e
          cd /paddle/deepmd-kit
          echo "Installing deepmd-kit..."
          source ${{ github.workspace }}/../../../proxy
          DP_ENABLE_TENSORFLOW=0 DP_ENABLE_PYTORCH=0 python3.10 -m pip install -v -e . -i https://pypi.tuna.tsinghua.edu.cn/simple
          '

      - name: Install test dependencies
        run: |
          docker exec -t ${{ env.container_name }} /bin/bash -c '
          set -e
          echo "Installing pytest..."
          python3.10 -m pip install \
            pytest pytest-xdist pytest-json-report pytest-metadata pytest-mock pytest-rich pytest-sugar pytest-timeout \
            -i https://pypi.tuna.tsinghua.edu.cn/simple
          '

      - name: Run deepmd-kit pytest
        id: pytest_run
        run: |
          docker exec -t ${{ env.container_name }} /bin/bash -c '
          set -e
          cd /paddle/deepmd-kit
          echo "=========================================="
          echo "Running DeepMD-Kit Integration Tests"
          echo "=========================================="

          python3.10 -m pytest -vvv --tb=long --color=yes \
            source/tests/pd/test_auto_batch_size.py \
            source/tests/pd/test_change_bias.py \
            source/tests/pd/test_decomp.py \
            source/tests/pd/test_dp_show.py \
            source/tests/pd/test_finetune.py \
            source/tests/pd/test_multitask.py \
            source/tests/pd/test_training.py \
            source/tests/pd/test_update_sel.py \
            source/tests/pd/test_utils.py \
            source/tests/pd/model/test_atomic_model_atomic_stat.py \
            source/tests/pd/model/test_atomic_model_global_stat.py \
            source/tests/pd/model/test_autodiff.py \
            source/tests/pd/model/test_deeppot.py \
            source/tests/pd/model/test_descriptor_dpa1.py \
            source/tests/pd/model/test_descriptor_dpa2.py \
            source/tests/pd/model/test_dp_atomic_model.py \
            source/tests/pd/model/test_dp_model.py \
            source/tests/pd/model/test_dpa1.py \
            source/tests/pd/model/test_dpa2.py \
            source/tests/pd/model/test_dpa3.py \
            source/tests/pd/model/test_dynamic_sel.py \
            source/tests/pd/model/test_ener_fitting.py \
            source/tests/pd/model/test_env_mat.py \
            source/tests/pd/model/test_exclusion_mask.py \
            source/tests/pd/model/test_force_grad.py \
            source/tests/pd/model/test_forward_lower.py \
            source/tests/pd/model/test_get_model.py \
            source/tests/pd/model/test_jit.py \
            source/tests/pd/model/test_mlp.py \
            source/tests/pd/model/test_nlist.py \
            source/tests/pd/model/test_null_input.py \
            source/tests/pd/model/test_permutation_denoise.py \
            source/tests/pd/model/test_permutation.py \
            source/tests/pd/model/test_region.py \
            source/tests/pd/model/test_rot_denoise.py \
            source/tests/pd/model/test_rot.py \
            source/tests/pd/model/test_rotation.py \
            source/tests/pd/model/test_se_e2_a.py \
            source/tests/pd/model/test_smooth.py \
            source/tests/pd/model/test_trans_denoise.py \
            source/tests/pd/model/test_trans.py \
            source/tests/pd/model/test_unused_params.py 2>&1 | tee /paddle/pytest_output.log
          TEST_RESULT=${PIPESTATUS[0]}

          echo "DeepMD-Kit Test Finished with status: $TEST_RESULT"
          exit $TEST_RESULT
          '

      - name: Terminate and delete the container
        if: always()
        run: |
          set +e
          docker exec -t ${{ env.container_name }} /bin/bash -c 'cd /paddle && rm -rf deepmd-kit'
          docker rm -f ${{ env.container_name }}