name: TheRock Test Packages multi-node for rccl

on:
  workflow_call:
    inputs:
      amdgpu_families:
        type: string
      artifact_group:
        type: string
      test_runs_on:
        type: string
      artifact_run_id:
        type: string
  workflow_dispatch:
    inputs:
      amdgpu_families:
        type: string
      artifact_group:
        type: string
      test_runs_on:
        type: string
      artifact_run_id:
        type: string

concurrency:
  group: ${{ inputs.amdgpu_families == 'gfx94X-dcgpu' && 'gfx94x-multinode-exclusive' || github.run_id }}
  cancel-in-progress: false

permissions:
  contents: read
  id-token: write

jobs:
  test_rccl_multi_node:
    name: 'Test multi-node'
    runs-on: ${{ inputs.test_runs_on }}
    defaults:
      run:
        shell: bash
    permissions:
      contents: read
      id-token: write
    env:
      VENV_DIR: ${{ github.workspace }}/.venv
      ARTIFACT_RUN_ID: "${{ inputs.artifact_run_id }}"
      THEROCK_BIN_DIR: "./build/bin"
    steps:
      - name: Checkout Repository
        uses: actions/checkout@11bd71901bbe5b1630ceea73d27597364c9af683 # v4.2.2
        with:
          repository: "ROCm/TheRock"
          ref: release/therock-10.0

      # Centralized arch-specific configuration
      - name: Set arch-specific paths
        run: |
          case "${{ inputs.amdgpu_families }}" in
            gfx950-dcgpu)
              echo "OUTPUT_ARTIFACTS_DIR=/apps/cvs_tests/dist_new/dist/rocm" >> $GITHUB_ENV
              echo "AWS_SHARED_CREDENTIALS_FILE=/apps/cvs_tests/awsconfig/credentials.ini" >> $GITHUB_ENV
              echo "REPORT_PATH=/apps/cvs_tests/test_reports" >> $GITHUB_ENV
              echo "LOG_DEST=/logs/gfx950-dcgpu" >> $GITHUB_ENV
              ;;

            gfx94X-dcgpu)
              echo "OUTPUT_ARTIFACTS_DIR=/home/arravikum/dist_new/dist/rocm" >> $GITHUB_ENV
              echo "AWS_SHARED_CREDENTIALS_FILE=/home/arravikum/awsconfig/credentials.ini" >> $GITHUB_ENV
              echo "REPORT_PATH=/home/arravikum/cvs_8nic/cvs/test_reports" >> $GITHUB_ENV
              echo "LOG_DEST=/logs/gfx94X-dcgpu" >> $GITHUB_ENV
              ;;

            *)
              echo "Unsupported arch: ${{ inputs.amdgpu_families }}"
              exit 1
              ;;
          esac

      # Only needed for gfx94X, which runs CVS tests directly from this workflow
      # and can leave orphan MPI/RCCL processes after failures or cancellations.
      # gfx950 uses the Slurm-backed ci-entrypoint.sh flow, so Slurm handles cleanup.
      - name: Cleanup orphan RCCL/MPI processes
        if: ${{ inputs.amdgpu_families == 'gfx94X-dcgpu' }}
        run: |
          echo "Cleaning up orphan processes..."
          pkill -9 -f prterun || true
          pkill -9 -f mpirun || true
          pkill -9 -f "_perf" || true

          sleep 10

          echo "Processes still holding artifact directory:"
          fuser -vm "${OUTPUT_ARTIFACTS_DIR}" || true

      - name: Run setup test environment workflow
        uses: './.github/actions/setup_test_environment'
        with:
          ARTIFACT_RUN_ID: ${{ env.ARTIFACT_RUN_ID }}
          AMDGPU_FAMILIES: ${{ inputs.amdgpu_families }}
          ARTIFACT_GROUP: ${{ inputs.artifact_group }}
          OUTPUT_ARTIFACTS_DIR: ${{ env.OUTPUT_ARTIFACTS_DIR }}
          VENV_DIR: ${{ env.VENV_DIR }}
          FETCH_ARTIFACT_ARGS: "--rccl --tests"
          IS_PR_FROM_FORK: ${{ github.event.pull_request.head.repo.fork }}

      # The following step leverages slurm to run multi node rccl tests on the slurm mi350x cluster.
      # salloc will hold 4 nodes while the commands inside the block run. After the block completes, salloc automatically releases the nodes.
      # sbatch script runs rccl_heatmap_cvs script which validates and generates a bandwidth heatmap file for different rccl collectives
      - name: Test gfx950
        if: ${{ inputs.amdgpu_families == 'gfx950-dcgpu' }}
        run: |
          chmod +x /apps/cvs_tests/ci-entrypoint.sh
          /apps/cvs_tests/ci-entrypoint.sh

      - name: Test gfx94X
        if: ${{ inputs.amdgpu_families == 'gfx94X-dcgpu' }}
        run: |
          source /home/arravikum/TheRock/.venv/bin/activate
          cd /home/arravikum/cvs_8nic/cvs/cvs
          pytest -vvv -s ./tests/rccl/rccl_heatmap_cvs.py   --cluster_file ./input/cluster_file/cluster.json \
            --config_file ./input/config_file/rccl/rccl_config.json   --log-file=/tmp/rccl_log.log \
            --html=/home/arravikum/cvs_8nic/cvs/test_reports/ci_test_report.html   --capture=tee-sys \
            --self-contained-html

      - name: Configure AWS Credentials for non-forked repos
        if: ${{ always() && !github.event.pull_request.head.repo.fork }}
        uses: aws-actions/configure-aws-credentials@7474bc4690e29a8392af63c5b98e7449536d5c3a # v4.3.1
        with:
          aws-region: us-east-2
          role-to-assume: arn:aws:iam::692859939525:role/therock-artifacts-external

      - name: Post test report upload
        if: always()
        working-directory: ${{ github.workspace }}
        run: |
          export PYTHONPATH="${PYTHONPATH}:${{ github.workspace }}/build_tools"
          python3 build_tools/github_actions/upload_test_report_script.py \
            --run-id "${{ github.run_id }}" \
            --amdgpu-family "${{ inputs.amdgpu_families }}" \
            --report-path "${REPORT_PATH}" \
            --log-destination "${LOG_DEST}" \
            --index-file-name "index_rccl_test_report.html"