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

[jam-jee]

What this changes

Replaces the CodeBuild integ suite for sagemaker-train on the PR gate with a faster selection that keeps meaningful server-side coverage.

CreateTrainingJob returns a TrainingJobArn only after the request has cleared every synchronous server-side gate: public-model shape validation, SigV4, sagemaker:CreateTrainingJob authorization (including condition keys), iam:PassRole on the execution role, the training backend's synchronous request validators, its role-assuming validators (which make real S3/ECR/FSx calls as the customer), post-validator business logic (training-plan capacity, routing, recipe filtering), and the final conditional write that rejects duplicate job names.

So "the ARN came back" proves the SDK-shaped payload was accepted as sent, and the caller held the permissions needed to submit it — without paying for a training run. The job is stopped immediately.

New: tests/integ/train/shallow (97 tests, 82 on the PR gate)

One file per trainer, mirroring the existing deep-suite layout so the shallow counterpart of any deep test is obvious:

Shallow file	Deep counterpart
test_model_trainer.py	test_model_trainer.py
test_sft_trainer.py	test_sft_trainer_integration.py
test_dpo_trainer.py	test_dpo_trainer_integration.py
test_rlvr_trainer.py	test_rlvr_trainer_integration.py
test_rlaif_trainer.py	test_rlaif_trainer_integration.py
test_cpt_trainer.py	test_cpt_hyperpod.py
test_multi_turn_rl_trainer.py	test_multi_turn_rl_trainer_integration.py
test_tuner.py	test_tuner_distributed.py
test_nova_data_mixing.py	test_sft_trainer_data_mixing_integration.py
test_nova_trainers.py	::test_sft_trainer_nova_workflow, ::test_rlvr_trainer_nova_workflow

harness.py provides submitted() / assert_submitted() / assert_rejected(): forces wait=False, resolves the submitted job across the four attribute names trainers use for it, and stops the job in a finally so a failed assertion still cleans up.

recipe_cases.py holds the cases every recipe trainer shares. Each per-trainer class subclasses RecipeTrainerCases and sets TRAINER, so adding a trainer is a two-line file and per-trainer deviations are declared rather than duplicated — EXTRA_KWARGS (RLAIF's reward model), SUPPORTS_SERVERFUL = False (RLAIF takes no compute), SUPPORTS_TRAINING_TYPE = False (CPT has no LoRA/full split).

Negative tests are included deliberately: without them the suite would stay green even if the SDK started sending a permissive-but-wrong payload.

Not every trainer creates a TrainingJob — HyperparameterTuner creates a HyperParameterTuningJob and MultiTurnRLTrainer creates an AgentRFT Job — so assert_submitted takes the expected ARN resource segment.

Coverage of every gpu_intensive test

The rule: a deep test belongs off the PR gate only if this suite covers the same code path. There are 46 gpu_intensive tests in tests/integ/train. All are accounted for:

Covered by this suite — test_model_trainer.py (8: tar source, py/sh entry, MPI, torchrun, HP json/yaml, custom driver), SFT (complete workflow incl. MLflow, validation dataset, sequence length, Nova, serverful SMTJ), DPO (both), RLAIF (complete workflow, reward-prompt ARN, continued fine-tuning), RLVR (complete workflow, all three reward-function forms, recipe+overrides, sequence length, Nova), tuner (sm_drivers channel), MTRL (3, needs prerequisites), CPT HyperPod (needs a cluster), Nova data mixing.

The full test-by-test mapping is in tests/integ/train/shallow/README.md.

MLflow is worth calling out: every *_complete_workflow deep test configures it, so RecipeTrainerCases covers both forms — experiment/run names (always runs) and mlflow_resource_arn (skips when the account has no app).

Not covered, and why:

• 11 evaluator tests (test_benchmark_evaluator.py, test_custom_scorer_evaluator.py, test_mtrl_evaluator_3p_agent.py, test_mtrl_trainer_integration.py) — evaluate() is a different API surface returning pipeline executions rather than jobs, so it needs its own harness support. Already gpu_intensive on master, so this PR loses no coverage. Clearest follow-up.
• 3 HyperPod tests — submit to a pre-provisioned cluster, not CreateTrainingJob. test_cpt_trainer.py is written in the shallow style and activates when SHALLOW_HYPERPOD_CLUSTER is set.

This PR newly marks only 10 tests: the 8 in test_model_trainer.py, test_sft_trainer_lora_with_sequence_length, and test_tuner_includes_sm_drivers_channel. Everything else listed above was already marked on master.

Deliberately not marked, because this suite does not cover them: all evaluator tests, test_notifications.py (EventBridge/SNS side effects), test_local_model_trainer.py (no service call).

The rule is documented in the suite README so a future change cannot silently erode the gate.

tox.ini: widened the gpu_intensive description. Despite the name it gates anything consuming real training capacity, including serverless and CPU-instance jobs.

Fixtures that look up rather than create

mlflow_arn, reward_lambda_arn and reward_evaluator only look their resources up and skip when absent. The deep suite's equivalents create them — IAM roles, Lambdas, MLflow apps, registry entries — which is a durable side effect a fast PR-gate suite should not have.

Workflow change

The PR job now runs the whole tests/integ/train tree with -m "not gpu_intensive and not us_east_1" rather than only shallow/. That keeps the ~170 client-side tests (recipe resolution, data utils, dry-run, log streaming, docker-compose detection) on the gate — they make no service call and were never the expensive part.

Verification

Run against a real account (us-west-2):

81 passed, 1 skipped, 0 failed in 7m04s

The skip is RLAIF's serverful case, which reports its own reason: RLAIFTrainer takes no compute argument. Tests needing us-east-1 credentials or a HyperPod cluster are deselected on the gate and skip cleanly when their prerequisites are absent.

Cost model measured, not assumed. Across 100 training jobs created by these runs, every one ended Stopped and every BillableTimeInSeconds was null — jobs are torn down while still in Starting/Pending, before instances become billable.

Real bugs the AWS runs found in these tests (all wrong assumptions on my part, not service problems), each now fixed with the evidence recorded in a comment:

• the public botocore model has no ServerlessJobConfig.SequenceLength, so those requests failed client-side before reaching the service — fixed with an AWS_DATA_PATH fixture mirroring the one in test_recipe_override_integration.py
• tuning job names cap at 32 characters, not the 63 allowed for training jobs
• the tuner derives its own job name from the training image plus a second-granularity timestamp and ignores base_job_name, so two tuner tests in the same second collided with ResourceInUse
• ModelTrainer.__init__ validates the execution role, so a bad role is rejected at construction and never reaches CreateTrainingJob
• RLAIFTrainer takes no compute; CPTTrainer takes no training_type and requires HyperPod
• for this model the recipe catalogue offers only 4K sequence length and no serverless recipe for TrainingType.FULL
• RLVR reward functions are pre-validated by invoking them over sample records, so they need GSM8k-shaped data — this suite's generic chat-format fixture fails with GSM8k scoring failed
• list_mlflow_apps is not a paginatable operation

What this deliberately does NOT cover

Training behaviour: no model artifacts, no metrics, no container logs, no convergence.

A regression that breaks training itself — a bad entry script, a broken container command, a distributed-launch bug — will pass this gate and be caught by the scheduled suites instead. That is the accepted trade for the runtime and cost reduction, and it is stated plainly in tests/integ/train/shallow/README.md.

Also out of scope for this pattern: HyperPod (submits to a pre-provisioned cluster, not CreateTrainingJob) and local container mode (no service call).

Known remaining gap: evaluator evaluate() submissions. Same pattern applies; it is a distinct API surface, so it is left for a follow-up — and the corresponding deep tests stay on the PR gate until then.

Note for reviewers

Because this workflow triggers on pull_request_target, CI runs the workflow definition from the base branch, not this PR's. So the fast-integ-tests job added here does not appear in this PR's checks, and the old integ-tests (sagemaker-train) job still runs — the workflow change only takes effect once merged. The test selection was verified locally (numbers above); the job definition itself is unproven in CI.

[github-actions]

Reviewing this PR

• Read the full PR diff
• Check workflow trigger & security model
• Verify harness/test correctness
• Post review

Working…

[github-actions]

Claude finished @jam-jee's task in 2m 51s —— View job

---

Review

Thorough, well-documented PR. The rationale (submit-then-stop as a cheap proxy for "the service accepted this payload and the caller was authorized") is sound, the negative tests are the right instinct, and the README/docstrings are unusually good. Comments below, none blocking.

Note: the create_inline_comment tool isn't available to me in this run, so findings are inline-referenced here rather than as line comments.

---

1. Security / isolation change worth flagging to approvers — .github/workflows/pr-checks-master.yml:56

The new fast-integ-tests job runs the PR's code directly on the GitHub-hosted runner with the CI AWS role assumed into the environment, whereas every pre-existing integ/unit job runs PR code inside CodeBuild via source-version-override. Two things execute untrusted PR code with credentials present:

• pip install ./sagemaker-train / ./sagemaker-core → runs the PR's setup.py;
• the test + conftest code itself (which already makes arbitrary boto calls).

This is gated behind collab-check → wait-for-approval (transitively, via needs: [detect-changes]), which is the right mitigation and matches the existing trust model. The thing to make explicit for reviewers: under pull_request_target, approving a fork PR to run checks now grants arbitrary code execution on the runner with the CI_AWS_ROLE_ARN credentials and the workflow's id-token: write, not just a sandboxed CodeBuild run. If the CodeBuild isolation was being relied on as a boundary, moving to the runner removes it. Worth a sentence in the PR description and, ideally, confirming the CI role is scoped to what this suite needs. (I can't edit workflow files, so this is FYI only.)

2. Two negative tests can leak a running job if the request is unexpectedly accepted

assert_rejected and assert_submitted always stop the job in a finally — good. But the two hand-rolled pytest.raises blocks don't:

• test_recipe_trainers_submission.py:2431 (test_unknown_model_is_rejected) — if SFTTrainer(...)/train(wait=False) unexpectedly succeeds, the submitted job is never resolved or stopped.
• test_model_trainer_submission.py:1503 (test_duplicate_job_name_is_rejected) — the second create() leaks if the duplicate is unexpectedly accepted (the first is captured and stopped).

Both are low-probability (they assert failures that should be deterministic), but they're the exact cases where a regression would both fail the test and leave capacity running — which is what this suite is otherwise careful to avoid. Consider resolving + stop_quietly in a finally for these two, matching the discipline everywhere else.

3. Account/region constants are hardcoded — test_recipe_trainers_submission.py:2130, test_recipe_customization_submission.py:1855, harness.py:580

MODEL_PACKAGE_GROUP (...729646638167...), the Nova MODEL_PACKAGE_GROUP, RLAIF_REWARD_MODEL_ID, and CPU_IMAGE (pinned to 763104351884.dkr.ecr.us-west-2...) bake this suite to one account + us-west-2. That's a little at odds with the MTRL prerequisites being deliberately moved to env vars "so this does not bake in another account-specific constant." It's consistent with the existing dry-run suite so not a regression, but since this is now the PR gate, it's worth a note that the gate only runs in an account that has these exact resources. (CPU_IMAGE's region-pinned registry also means a non-us-west-2 run does a cross-region ECR resolve.)

4. AWS_DATA_PATH env mutation outlives the shallow tests on a worker — conftest.py:377

bundled_service_model is autouse, scope="session", and restores AWS_DATA_PATH only at session teardown. It correctly applies only to the shallow/ subtree, but because the gate runs the whole tests/integ/train tree with --dist loadfile, once a worker executes a shallow file the env var stays set for any later non-shallow file that lands on the same worker. Almost certainly benign (the bundled model is a superset), but the restore is coarser than the intent — a narrower scope or explicit unset after the shallow files would be tighter.

5. Minor

• pr-checks-master.yml:61 uses actions/checkout@v3 while the surrounding steps use setup-python@v5 / configure-aws-credentials@v4 / upload-artifact@v4. Consider checkout@v4 for consistency.
• test_other_job_types_submission.py:1731 — mtrl_prerequisites requests account_id and region but doesn't use either; can be dropped from the signature.

---

Overall this is a net improvement to PR-gate cost/latency with coverage that's honestly scoped (the README's "what this does NOT catch" section is exactly right). Nothing here blocks merge.
· shallow-pr-checks-sagemaker-train

[github-actions]

Claude finished @jam-jee's task in 2m 10s —— View job

---

Review

Re-review after the latest push. The suite has been restructured into one file per trainer since the first automated review, so a few of that review's line references (e.g. test_recipe_trainers_submission.py, test_other_job_types_submission.py) no longer exist — I've re-checked against the current tree. This is a well-documented, honestly-scoped PR; the submit-then-stop rationale is sound and the negative tests are the right instinct. Findings below, none blocking.

Review checklist

• Read the full PR diff
• Verified the workflow trigger & trust model against the checked-out pr-checks-master.yml
• Reviewed harness / conftest / test correctness
• Posted review

---

1. Trust-boundary change worth calling out to approvers — .github/workflows/pr-checks-master.yml:279

I confirmed the gating: fast-integ-tests → needs: [detect-changes] → wait-for-approval → collab-check, so it is behind the same approval gate as every other job. Good. The thing to make explicit: unlike every pre-existing job (which runs PR code inside CodeBuild via source-version-override), this new job checks out and runs the PR's code directly on the GitHub-hosted runner with CI_AWS_ROLE_ARN assumed and id-token: write present. Both pip install ./sagemaker-train (runs the PR's setup.py) and the test/conftest code execute in that context.

Two consequences to be aware of, not necessarily to fix:

• For a collaborator author, collab-check returns auto-approve, so this runs with no human in the loop — arbitrary PR code on the runner with those credentials.
• The CodeBuild isolation other jobs rely on as a boundary is removed for this path.

The approval gate is the right mitigation and matches the existing model; just worth a sentence in the description and confirming CI_AWS_ROLE_ARN is scoped to what this suite actually needs.

2. test_duplicate_job_name_is_rejected can leak the second job — sagemaker-train/tests/integ/train/shallow/test_model_trainer.py:1868

first = None
try:
    first = create()
    ...
    with pytest.raises(Exception):
        create()          # <- return value discarded
    ...
finally:
    stop_quietly(first)   # only stops the first job

If the duplicate is unexpectedly accepted, the second create() returns a real, running TrainingJob whose handle is thrown away and never stopped. Low-probability (it asserts a deterministic ResourceInUse), but it's exactly the regression case this suite is otherwise careful to guard against. Capture the second job and stop_quietly it in the finally too. (The harness.assert_rejected/assert_submitted paths already handle this correctly — this is the one hand-rolled pytest.raises block that doesn't.)

3. AWS_DATA_PATH mutation outlives the shallow tests on a worker — sagemaker-train/tests/integ/train/shallow/conftest.py:397

bundled_service_model is autouse, scope="session", and only restores AWS_DATA_PATH at session teardown. It auto-applies only to shallow/, but the env var it sets is process-global — and the gate runs the whole tests/integ/train tree with --dist loadfile, so once a worker runs a shallow file, any later non-shallow file on the same worker inherits the bundled (superset) model. Almost certainly benign, and behavior even depends on file ordering on the worker. A narrower scope, or an explicit unset after the shallow files, would match the intent more tightly.

4. Account/region constants pin the gate to one account + us-west-2

MODEL_PACKAGE_GROUP (...729646638167..., recipe_cases.py:903), CPU_IMAGE (763104351884.dkr.ecr.us-west-2..., harness.py:600), and the RLAIF reward-model id bake the suite to a specific account and region. Consistent with the existing dry-run suite, so not a regression — but since this is now the PR gate, it's worth noting the gate only passes in an account holding those exact resources, and a non-us-west-2 run does a cross-region ECR resolve for CPU_IMAGE. Slightly at odds with the MTRL prerequisites being deliberately moved to env vars for the same reason.

5. Minor

• pr-checks-master.yml:284 uses actions/checkout@v3 while the surrounding steps use setup-python@v5 / configure-aws-credentials@v4 / upload-artifact@v4. checkout@v4 would be consistent (and detect-changes at line 59 also still uses @v3).
• test_sft_trainer_integration.py:180 now has both the stale # @pytest.mark.gpu_intensive comment and the real @pytest.mark.gpu_intensive decorator — drop the comment.
• test_multi_turn_rl_trainer.py:1936 — mtrl_prerequisites takes account_id and region fixtures but uses neither; can be dropped from the signature.
• test_tuner.py:2429 — the metric regex r"eval_loss: ([0-9\\.]+)" produces a literal-backslash char class (\\ in a raw string). Harmless here since the job is stopped before metrics are parsed, but not what's intended.

---

Net: a real improvement to PR-gate cost/latency with coverage that's scoped honestly (the README's "what this does NOT catch" section is exactly right). Nothing here blocks merge; #1 and #2 are the two worth a second look.
· shallow-pr-checks-sagemaker-train

[github-actions]

Claude finished @jam-jee's task in 1m 31s —— View job

---

Review

Re-reviewed the current tree after the latest push (whole diff read from /tmp/pr.diff). This is a well-scoped, unusually well-documented change — the submit-then-stop rationale is sound, the negative tests are the right instinct, and the README's "what this does NOT catch" section is honest. Two prior automated reviews already covered the big items; below I confirm which still stand in the current one-file-per-trainer layout and add a couple of small things. Nothing blocks merge.

Checklist

• Read the full PR diff
• Verified workflow trigger & trust model
• Checked harness / conftest / test correctness
• Posted review

---

1. test_duplicate_job_name_is_rejected still leaks the second job — shallow/test_model_trainer.py (create() inside pytest.raises)

first = None
try:
    first = create()
    assert_submitted(first, expected_name=job_name)
    with pytest.raises(Exception):
        create()          # <- return value discarded
    ...
finally:
    stop_quietly(first)   # only stops the first job

If the duplicate is unexpectedly accepted, the second create() returns a real, running TrainingJob whose handle is thrown away and never stopped. Low-probability (it asserts a deterministic ResourceInUse), but it is exactly the regression case this suite is otherwise careful to guard — and the only hand-rolled pytest.raises block that doesn't route cleanup through the harness. Capture the second job and stop_quietly it in the finally too:

second = None
try:
    first = create()
    ...
    with pytest.raises(Exception) as excinfo:
        second = create()
    ...
finally:
    stop_quietly(first)
    stop_quietly(second)

(The assert_rejected / submitted paths already handle this correctly.)

2. Tuner metric regex has a literal-backslash char class — shallow/test_tuner.py (_tuner, metric_definitions)

metric_definitions=[{"Name": "eval_loss", "Regex": r"eval_loss: ([0-9\\.]+)"}],

In a raw string \\ is two characters, so the class is [0-9\.] — it matches a literal backslash in addition to digits/dot, not what's intended. Harmless here since the job is stopped before any metric is parsed, but it should be r"eval_loss: ([0-9.]+)" (inside a class the dot is already literal).

3. AWS_DATA_PATH mutation outlives the shallow tests on a worker — shallow/conftest.py (bundled_service_model)

autouse, scope="session", restores AWS_DATA_PATH only at session teardown. It auto-applies only under shallow/, but the env var is process-global and the gate runs the whole tests/integ/train tree with --dist loadfile — so once a worker runs a shallow file, any later non-shallow file landing on that worker inherits the bundled (superset) model. Almost certainly benign, and outcome depends on file ordering. A narrower scope, or an explicit unset after the shallow files, would match the stated intent more tightly.

4. Trust-boundary change — worth one sentence in the description (FYI, can't edit workflows)

Confirmed the gating is correct: fast-integ-tests → needs: [detect-changes] → the same approval chain as every other job. The thing to make explicit for approvers: unlike the pre-existing jobs (PR code inside CodeBuild via source-version-override), this job checks out and runs the PR's code directly on the runner with CI_AWS_ROLE_ARN assumed and id-token: write present — both pip install ./sagemaker-train (runs the PR's setup.py) and the test/conftest code. Approving a fork PR now grants arbitrary code execution on the runner with those credentials, not a sandboxed CodeBuild run. The approval gate is the right mitigation; just worth stating, and confirming CI_AWS_ROLE_ARN is scoped to what this suite needs.

5. Minor

• test_sft_trainer_integration.py:180 now carries both the stale # @pytest.mark.gpu_intensive comment and the real decorator — drop the comment.
• shallow/test_multi_turn_rl_trainer.py — mtrl_prerequisites takes account_id and region fixtures but uses neither; drop them from the signature.
• shallow/test_model_trainer.py imports MAX_RUNTIME_IN_SECONDS / stop_quietly etc. — all used; no dead imports spotted. Compute/InputData/Networking/SourceCode all used. Good.
• .github/workflows/pr-checks-master.yml uses actions/checkout@v3 for the new job while surrounding steps use @v4/@v5; checkout@v4 would be consistent (pre-existing detect-changes also still @v3).

I verified HyperParameterTuningJob.refresh() returns self (so the sm_drivers read-back in test_distributed_tuning_job_is_accepted is safe) and that unique_name fits within MAX_TUNING_JOB_NAME=32 — both fine.

Net: a real improvement to PR-gate cost/latency with honestly-scoped coverage. #1 is the one worth fixing before merge (it defeats the suite's own no-leaked-capacity discipline in the exact case it's meant to catch); the rest are polish.
· shallow-pr-checks-sagemaker-train