change(train): gate deep integ tests behind gpu_intensive, add shallow submit-then-stop suite

.github/workflows/pr-checks-master.yml

@@ -243,6 +243,118 @@ jobs:
           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]

sagemaker-train/pyproject.toml

@@ -84,6 +84,15 @@ addopts = ["-vv"]
 testpaths = ["tests"]
 markers = [
     "serial: marks tests that must run serially (not in parallel)",
+    # gpu_intensive and us_east_1 are declared in tox.ini too, but pytest reads
+    # its config from this file (it is the first of the candidates present), so
+    # markers listed only there are unregistered at runtime and raise
+    # PytestUnknownMarkWarning. Registering them here matters because the PR gate
+    # selects with -m "not gpu_intensive and not us_east_1": a typo'd marker name
+    # would otherwise silently put an expensive deep test back on the gate instead
+    # of warning.
+    "gpu_intensive: marks a test that consumes real training capacity (scheduled CI, not PR checks); see tests/integ/train/shallow",
+    "us_east_1: marks a test that must run in us-east-1 (Nova); runs in the us-east-1 integ job",
 ]
 
 [tool.black]

sagemaker-train/tests/integ/conftest.py

@@ -31,10 +31,25 @@
   its own). ``adaptive`` mode adds client-side rate limiting so bursts of
   ``SimulatePrincipalPolicy`` calls ride out transient throttling.
 
-Throttling that still exhausts the adaptive retry budget is deliberately left to
-fail the test loudly (rather than being converted to a skip), so a persistent
-rate-limit regression stays visible instead of silently disappearing from the
-results.
+* ``_memoize_role_validation`` (autouse) — retries alone were not enough. A PR-gate
+  run failed four tests with ``(Throttling) ... SimulatePrincipalPolicy (reached
+  max retries: 9)``: the adaptive budget was exhausted, not merely stressed. The
+  cause is volume, not burstiness — ~190 tests each construct a trainer, every
+  construction calls ``get_role``, and each of those runs a *paginated*
+  ``SimulatePrincipalPolicy`` over ~20 action names. Under ``-n auto`` on a large
+  CodeBuild container that is thousands of calls against a low, account-wide TPS
+  limit, so raising the retry budget only trades failures for a slower build.
+
+  Since the arguments repeat, the result does too: this memoizes
+  ``resolve_and_validate_role`` per worker, collapsing those calls to one per
+  distinct ``(provided_role, role_type, region)``. Validation still happens — once,
+  and its outcome (including a raised ``RoleValidationError``) is what gets reused,
+  so a genuinely bad role still fails every test that uses it.
+
+Throttling that still exhausts the retry budget after memoization is deliberately
+left to fail the test loudly (rather than being converted to a skip), so a
+persistent rate-limit regression stays visible instead of silently disappearing
+from the results.
 """
 from __future__ import absolute_import
 
@@ -66,3 +81,82 @@ def _configure_boto_adaptive_retries():
             os.environ.pop(key, None)
         else:
             os.environ[key] = value
+
+
+# Modules that did `from ...iam_role_resolver import resolve_and_validate_role`
+# hold their own reference to the original function, so patching only the defining
+# module would leave those bindings calling IAM directly. Each importer is patched
+# too. Kept as a list of (module path, attribute) so adding a caller is one line.
+_ROLE_RESOLVER_CALLERS = (
+    ("sagemaker.core.helper.iam_role_resolver", "resolve_and_validate_role"),
+    ("sagemaker.train.defaults", "resolve_and_validate_role"),
+    ("sagemaker.train.evaluate.base_evaluator", "resolve_and_validate_role"),
+)
+
+
+@pytest.fixture(autouse=True, scope="session")
+def _memoize_role_validation():
+    """Validate each distinct role once per xdist worker instead of once per test.
+
+    See this module's docstring for why retries alone were insufficient. Caches on
+    ``(provided_role, role_type, region)`` -- region is part of the key because the
+    Nova tests validate the same role against us-east-1, and a role's resolution is
+    region-scoped. Exceptions are cached alongside successes so a bad role keeps
+    failing rather than being silently retried per test.
+    """
+    import importlib
+
+    patched = []
+    cache = {}
+
+    try:
+        source = importlib.import_module(_ROLE_RESOLVER_CALLERS[0][0])
+    except ImportError:  # pragma: no cover - SDK layout changed
+        yield
+        return
+
+    original = source.resolve_and_validate_role
+
+    def memoized(provided_role=None, role_type=None, sagemaker_session=None, **kwargs):
+        region = None
+        if sagemaker_session is not None:
+            region = getattr(sagemaker_session, "boto_region_name", None)
+        key = (provided_role, role_type, region)
+
+        if key not in cache:
+            try:
+                cache[key] = (
+                    original(
+                        provided_role=provided_role,
+                        role_type=role_type,
+                        sagemaker_session=sagemaker_session,
+                        **kwargs,
+                    ),
+                    None,
+                )
+            except Exception as exc:  # cache the verdict, not just the happy path
+                cache[key] = (None, exc)
+
+        result, error = cache[key]
+        if error is not None:
+            raise error
+        return result
+
+    for module_path, attribute in _ROLE_RESOLVER_CALLERS:
+        try:
+            module = importlib.import_module(module_path)
+        except ImportError:
+            continue  # optional/renamed caller; the others still get patched
+        if getattr(module, attribute, None) is original:
+            setattr(module, attribute, memoized)
+        patched.append((module, attribute))
+
+    yield
+
+    # Restore by checking for `memoized` rather than only undoing what was patched
+    # above: a caller imported *after* the source module was patched binds the
+    # memoized function at its own import time, so it needs restoring too even
+    # though this fixture never set it.
+    for module, attribute in patched:
+        if getattr(module, attribute, None) is memoized:
+            setattr(module, attribute, original)