pr-checks-master.yml

name: Sagemaker PR Checks (Master)
on:
  pull_request_target:
    branches:
      - "master"
      - "master-mtrl-trainer"
      - "master-mtrl-release"
      - "master-nova-reconcillation"
    paths:
      - 'sagemaker-train/**'
      - 'sagemaker-serve/**'
      - 'sagemaker-mlops/**'
      - 'sagemaker-core/**'

concurrency:
  group: ${{ github.workflow }}-${{ github.event.pull_request.number || github.head_ref }}
  cancel-in-progress: true

permissions:
  id-token: write

jobs:
  collab-check:
    runs-on: ubuntu-latest
    outputs:
      approval-env: ${{ steps.collab-check.outputs.result }}
    steps:
      - name: Collaborator Check
        uses: actions/github-script@v7
        id: collab-check
        with:
          github-token: ${{ secrets.COLLAB_CHECK_TOKEN }}
          result-encoding: string
          script: | 
            try {
              const res = await github.rest.repos.checkCollaborator({
                owner: context.repo.owner,
                repo: context.repo.repo,
                username: "${{ github.event.pull_request.user.login }}",
              });
              console.log("Verifed ${{ github.event.pull_request.user.login }} is a repo collaborator. Auto Approving PR Checks.")
              return res.status == "204" ? "auto-approve" : "manual-approval"
            } catch (error) {
              console.log("${{ github.event.pull_request.user.login }} is not a collaborator. Requiring Manual Approval to run PR Checks.")
              return "manual-approval"
            }
  wait-for-approval:
    runs-on: ubuntu-latest
    needs: [ collab-check ]
    environment: ${{ needs.collab-check.outputs.approval-env }}
    steps:
      - run: echo "Workflow Approved! Starting PR Checks."
  detect-changes:
    runs-on: ubuntu-latest
    needs: [wait-for-approval]
    outputs:
      submodules: ${{ steps.check-changes.outputs.submodules }}
    steps:
      - uses: actions/checkout@v3
        with:
          fetch-depth: 0
          token: ${{ secrets.GH_PAT }}
          ref: ${{ github.event.pull_request.base.ref }}
      - name: Detect Changes
        id: check-changes
        run: |
          set -e
          
          echo "Target Branch: ${{ github.event.pull_request.base.ref }}"
          echo "Current Target SHA: $(git rev-parse HEAD)"
          echo "PR Number: ${{ github.event.pull_request.number }}"
          echo "PR Latest SHA: ${{ github.event.pull_request.head.sha }}"
          
          git fetch origin pull/${{ github.event.pull_request.number }}/head
          CHANGES=$(git diff --name-only HEAD FETCH_HEAD)
          
          echo "Changed files:"
          echo "$CHANGES"
          
          # Function to extract dependencies from pyproject.toml
          get_dependencies() {
            local module=$1
            grep "sagemaker-" "$module/pyproject.toml" | grep -o 'sagemaker-[a-z]*' | sort -u
          }
          
          # Function to find all modules that depend on a given module (recursively)
          find_dependents() {
            local target=$1
            local all_modules=("sagemaker-core" "sagemaker-train" "sagemaker-serve" "sagemaker-mlops")
            local dependents=()
            
            for module in "${all_modules[@]}"; do
              if [ "$module" != "$target" ]; then
                if get_dependencies "$module" | grep -q "^$target$"; then
                  dependents+=("$module")
                fi
              fi
            done
            
            echo "${dependents[@]}"
          }
          
          # Initialize set of submodules to test (using associative array)
          declare -A SUBMODULES_SET
          
          # Function to recursively add module and all its dependents
          add_module_and_dependents() {
            local module=$1
            
            if [ -z "${SUBMODULES_SET[$module]}" ]; then
              SUBMODULES_SET["$module"]=1
              echo "Adding $module to test set"
              
              # Find all modules that depend on this one and add them recursively
              local dependents=$(find_dependents "$module")
              for dependent in $dependents; do
                add_module_and_dependents "$dependent"
              done
            fi
          }
          
          # Determine whether a module has any non-test changes. A change counts
          # as a source change if it touches anything under the module other than
          # its tests/ directory (e.g. src/, pyproject.toml, tox.ini, VERSION).
          # This is intentionally conservative: only changes confined entirely to
          # tests/ are treated as test-only.
          is_source_changed() {
            local module=$1
            echo "$CHANGES" | grep "^$module/" | grep -qv "^$module/tests/"
          }
          
          all_modules=("sagemaker-core" "sagemaker-train" "sagemaker-serve" "sagemaker-mlops")
          
          # Pass 1: modules with source changes pull in themselves plus every
          # module that (transitively) depends on them, since a source change can
          # affect downstream behaviour. This preserves the original logic.
          for module in "${all_modules[@]}"; do
            if is_source_changed "$module"; then
              echo "$module has source changes - adding it and all dependents"
              add_module_and_dependents "$module"
            fi
          done
          
          # Pass 2: modules with test-only changes add only themselves and skip
          # dependency propagation, since changing a module's tests cannot affect
          # other modules. Run after Pass 1 so source-change propagation is never
          # short-circuited by a test-only module already being in the set.
          for module in "${all_modules[@]}"; do
            if echo "$CHANGES" | grep -q "^$module/" && ! is_source_changed "$module"; then
              if [ -z "${SUBMODULES_SET[$module]}" ]; then
                echo "$module has test-only changes - adding only $module"
                SUBMODULES_SET["$module"]=1
              fi
            fi
          done
          
          # Convert associative array to JSON array
          SUBMODULES='[]'
          for submodule in "${!SUBMODULES_SET[@]}"; do
            if [ "$SUBMODULES" = '[]' ]; then
              SUBMODULES="[\"$submodule\"]"
            else
              SUBMODULES=$(echo $SUBMODULES | sed "s/\]$/,\"$submodule\"\]/")
            fi
          done
          
          echo "Final SUBMODULES: $SUBMODULES"
          echo "submodules=$SUBMODULES" >> $GITHUB_OUTPUT

  codestyle-doc-tests:
    runs-on: ubuntu-latest
    needs: [detect-changes]
    if: needs.detect-changes.outputs.submodules != '[]'
    strategy:
      fail-fast: false
      matrix:
        submodule: ${{ fromJson(needs.detect-changes.outputs.submodules) }}
    steps:
      - name: Configure AWS Credentials
        uses: aws-actions/configure-aws-credentials@v4
        with:
          role-to-assume: ${{ secrets.CI_AWS_ROLE_ARN }}
          aws-region: us-west-2
          role-duration-seconds: 10800

      - name: Run CodeBuild for ${{ matrix.submodule }}
        uses: aws-actions/aws-codebuild-run-build@v1
        with:
          project-name: ${{ github.event.repository.name }}-ci-${{ matrix.submodule }}-codestyle-doc-tests
          source-version-override: 'refs/pull/${{ github.event.pull_request.number }}/head^{${{ github.event.pull_request.head.sha }}}'

  unit-tests:
    runs-on: ubuntu-latest
    needs: [detect-changes]
    if: needs.detect-changes.outputs.submodules != '[]'
    strategy:
      fail-fast: false
      matrix:
        submodule: ${{ fromJson(needs.detect-changes.outputs.submodules) }}
    steps:
      - name: Configure AWS Credentials
        uses: aws-actions/configure-aws-credentials@v4
        with:
          role-to-assume: ${{ secrets.CI_AWS_ROLE_ARN }}
          aws-region: us-west-2
          role-duration-seconds: 10800

      - name: Run Unit Tests for ${{ matrix.submodule }}
        uses: aws-actions/aws-codebuild-run-build@v1
        with:
          # Use the single CDK-managed V3 unit-test project (driven by the SUBMODULE
          # env var), the same project the CI-health workflow uses. The previous
          # per-submodule projects (sagemaker-python-sdk-ci-<submodule>-unit-tests)
          # were created manually, are not CDK/pipeline-managed, and had drifted
          # stale (e.g. still running `--cov=.` instead of the deployed
          # `--cov=sagemaker`), so PR coverage never reflected buildspec fixes.
          project-name: ${{ github.event.repository.name }}-ci-health-unit-test-v3
          source-version-override: 'refs/pull/${{ github.event.pull_request.number }}/head^{${{ github.event.pull_request.head.sha }}}'
          env-vars-for-codebuild: |
            SUBMODULE
        env:
          SUBMODULE: ${{ matrix.submodule }}

  integ-tests:
    runs-on: ubuntu-latest
    needs: [detect-changes]
    if: needs.detect-changes.outputs.submodules != '[]'
    strategy:
      fail-fast: false
      matrix:
        submodule: ${{ fromJson(needs.detect-changes.outputs.submodules) }}
    steps:
      - name: Configure AWS Credentials
        uses: aws-actions/configure-aws-credentials@v4
        with:
          role-to-assume: ${{ secrets.CI_AWS_ROLE_ARN }}
          aws-region: us-west-2
          role-duration-seconds: 10800

      - name: Run Integ Tests for ${{ matrix.submodule }}
        uses: aws-actions/aws-codebuild-run-build@v1
        with:
          project-name: ${{ github.event.repository.name }}-ci-${{ matrix.submodule }}-integ-tests
          source-version-override: 'refs/pull/${{ github.event.pull_request.number }}/head^{${{ github.event.pull_request.head.sha }}}'

  # Additive: runs the shallow (submit-then-stop) suite for sagemaker-train
  # alongside the existing integ-tests job above, which is unchanged.
  #
  # Why a separate job rather than folding this into the CodeBuild suite: this
  # job's selection is reviewable in the PR that changes it, whereas the
  # sagemaker-train CodeBuild buildspec is CDK-managed outside this repo. It also
  # reports as its own check, so a shallow failure is distinguishable at a glance
  # from a deep-suite failure, and it finishes in minutes -- fast feedback that
  # does not wait on the 2XLARGE CodeBuild container.
  #
  # What runs here: only tests/integ/train/shallow. The client-side tests
  # (recipe resolution, data utils, dry-run, log streaming) are deliberately NOT
  # repeated -- the CodeBuild suite already runs the whole tests/integ tree, so
  # widening this job's scope would duplicate them and double the job creation
  # the shallow suite performs.
  #
  # Why submit-then-stop is worth gating on: CreateTrainingJob returns a
  # TrainingJobArn only after the request has cleared public-model validation,
  # SigV4, sagemaker:CreateTrainingJob authorization, iam:PassRole, the training
  # backend's request validators (including the role-assuming ones that resolve
  # S3 and ECR as the customer) and the final duplicate-name write. So a returned
  # ARN proves the payload and the caller's permissions are both good -- without
  # paying for a training run. The job is stopped immediately.
  #
  # It asserts nothing about training *behaviour* (artifacts, metrics,
  # convergence); that remains the deep suites' job.
  fast-integ-tests:
    runs-on: ubuntu-latest
    needs: [detect-changes]
    if: contains(fromJson(needs.detect-changes.outputs.submodules), 'sagemaker-train')
    steps:
      - uses: actions/checkout@v3
        with:
          # pull_request_target checks out the base ref by default; these tests
          # must run against the PR's code.
          ref: 'refs/pull/${{ github.event.pull_request.number }}/head'

      - name: Set up Python
        uses: actions/setup-python@v5
        with:
          python-version: '3.12'

      - name: Configure AWS Credentials
        uses: aws-actions/configure-aws-credentials@v4
        with:
          role-to-assume: ${{ secrets.CI_AWS_ROLE_ARN }}
          aws-region: us-west-2
          role-duration-seconds: 10800

      - name: Install sagemaker-train and test dependencies
        run: |
          python -m pip install --upgrade pip
          pip install ./sagemaker-core
          pip install ./sagemaker-train
          pip install -r requirements/extras/test_requirements.txt

      - name: Run shallow sagemaker-train integ tests
        working-directory: sagemaker-train
        env:
          AWS_DEFAULT_REGION: us-west-2
          # Role resolution goes through iam:SimulatePrincipalPolicy, which is
          # low-TPS; adaptive retries keep parallel workers from throttling each
          # other.
          AWS_RETRY_MODE: adaptive
          AWS_MAX_ATTEMPTS: '10'
          # Cap the training jobs the *service* counts against the concurrency
          # quota, across all xdist workers, so the suite stays inside the
          # "concurrent model customization serverless jobs per Region" quota
          # (20) with room for the deep integ-tests suite running the same
          # account concurrently. The harness holds each slot until the job is
          # terminal, not until stop() returns -- the service counts a job for
          # ~1-3 min after the stop -- so 10 means "at most 10 jobs counted at
          # once", the batches-of-10 behaviour, not "10 stops in flight".
          #
          # Not redundant with -n 8. -n caps worker processes; this caps what
          # the service counts, and with the slot held to terminal those diverge
          # sharply (each drain outlives the worker's stop() by minutes). It is
          # also what keeps the ceiling stable if -n is raised. A serverful job
          # counts one slot per instance.
          SHALLOW_MAX_CONCURRENT_JOBS: '10'
        run: |
          # Scoped to shallow/ only -- see the comment above this job for why the
          # rest of tests/integ/train is not repeated here.
          #
          # 84 of the suite's 100 tests run; the 16 deselected are:
          #   us_east_1 (5)     -- Nova cases; this job holds us-west-2
          #                        credentials only, so they run in the
          #                        integ-tests-us-east-1 job instead.
          #   gpu_intensive (11) -- the CPT and MTRL classes. These are written in
          #                        the shallow style but cannot be made
          #                        self-contained: CPT submits only via HyperPod
          #                        (a pre-provisioned cluster, not
          #                        CreateTrainingJob) and MTRL needs an agent
          #                        runtime plus an MLflow app. Both become
          #                        gate-eligible by dropping one marker once those
          #                        prerequisites exist in the PR account.
          python -m pytest tests/integ/train/shallow \
            -m "not gpu_intensive and not us_east_1" \
            -n 8 \
            --dist loadfile \
            -v \
            --durations=15 \
            --junitxml=shallow-integ-results.xml

      - name: Upload test results
        if: always()
        uses: actions/upload-artifact@v4
        with:
          name: shallow-integ-test-results
          path: sagemaker-train/shallow-integ-results.xml
          if-no-files-found: warn

  integ-tests-us-east-1:
    runs-on: ubuntu-latest
    needs: [detect-changes]
    if: needs.detect-changes.outputs.submodules != '[]'
    steps:
      - name: Configure AWS Credentials (us-east-1)
        uses: aws-actions/configure-aws-credentials@v4
        with:
          role-to-assume: ${{ secrets.CI_AWS_ROLE_US_EAST_1_ARN }}
          aws-region: us-east-1
          role-duration-seconds: 10800

      - name: Run us-east-1 Integ Tests (all submodules)
        uses: aws-actions/aws-codebuild-run-build@v1
        with:
          project-name: ${{ github.event.repository.name }}-ci-integ-tests-us-east-1
          source-version-override: 'refs/pull/${{ github.event.pull_request.number }}/head^{${{ github.event.pull_request.head.sha }}}'