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)

.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)

.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)

.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)

.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)

.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)

.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)

.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)