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

[jam-jee]

What this changes

Adds a fast server-side acceptance suite for sagemaker-train to the PR gate, and marks the deep tests that submit a job and wait for it so the existing gate stops running them.

The existing integ-tests CodeBuild job is unchanged and still runs. This PR adds a second job next to it rather than replacing it.

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 (100 tests, 84 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, test_sft_trainer_serverful_smtj.py

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 there. Clearest follow-up, and per "What bounded the integ-tests job" below it is also what would let the six newly marked evaluator tests be covered rather than merely deselected.
• 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 16 tests, in two groups:

• 10 with shallow counterparts — the 8 in test_model_trainer.py, test_sft_trainer_lora_with_sequence_length, and test_tuner_includes_sm_drivers_channel. These satisfy the rule above outright.
• 6 pipeline-waiting evaluator tests — see "What bounded the integ-tests job" below. These do not have a shallow counterpart yet; the reasoning for marking them anyway, and what coverage that costs, is spelled out there and in the suite README rather than glossed over.

Everything else listed above was already marked on master.

Deliberately not marked: test_notifications.py (EventBridge/SNS side effects) and test_local_model_trainer.py (no service call, but it does run real containers).

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

Two things, and the first is smaller than it sounds:

1. integ-tests is untouched. The job definition is byte-identical to master — sagemaker-train is still in its matrix and still runs the full CodeBuild suite. What changes is what that suite selects: its buildspec already filters -m "not gpu_intensive and not us_east_1", so the marks this PR adds are what take the expensive tests off the gate. No workflow edit was needed for that, and no client-side coverage is lost — the CodeBuild job runs the whole tests/integ tree, so the ~170 tests that make no service call keep running exactly as before.

2. New fast-integ-tests job, scoped to tests/integ/train/shallow only, running directly on the runner:

python -m pytest tests/integ/train/shallow -m "not gpu_intensive and not us_east_1" -n 8 --dist loadfile

84 of the suite's 100 tests; the 16 deselected are the 5 us_east_1 Nova cases (this job holds us-west-2 credentials only — they run in integ-tests-us-east-1) and the 11 gpu_intensive CPT/MTRL cases, which are written in the shallow style but need a HyperPod cluster and an agent runtime respectively.

Deliberately not widened to the whole tests/integ/train tree: the CodeBuild job already covers the client-side tests, so widening would duplicate them and double the training jobs this suite creates. Separate job rather than folded into the buildspec because the buildspec is CDK-managed outside this repo, while this selection is reviewable in the PR that changes it — and because a shallow failure then reports as its own check, distinguishable at a glance from a deep-suite failure.

Verification

Run against a real account. The shallow suite alone (us-west-2):

83 passed, 1 skipped, 0 failed in 3m33s

(5m20s before the role-validation memoization described below; the saved time is SimulatePrincipalPolicy round-trips.)

The skip is RLAIF's serverful case, which reports its own reason: RLAIFTrainer takes no compute argument.

The five us_east_1 tests, run separately with us-east-1 credentials:

5 passed in 47s

Those five had never executed before this PR — see the bug list below.

Also run: the whole tests/integ/train tree under the same selection the CodeBuild job uses (-m "not gpu_intensive and not us_east_1", -n 8 --dist loadfile) — 198 passed, 9 skipped, 0 failed. Skips all report their own prerequisite: Docker absent locally (5), RLAIF's serverful case, and three pre-existing skip/prerequisite cases. That run predates the marker changes below, so it includes the five us-west-2 evaluator tests now deselected; it is the "before" measurement, and the runtime caveat in the next section is what came out of it.

Marker selection after this PR, verified by collection: 266 of 342 collected, 76 deselected, none of the six newly marked evaluator tests among them, and no PytestUnknownMarkWarning remaining.

What bounded the integ-tests job (pre-existing, now fixed here)

Not the shallow suite, and not the new job — fast-integ-tests is the 5-minute one. This is about the CodeBuild job, and it is worth stating plainly because I got it wrong earlier in this PR's own description.

CI proved it on this branch. From the gate's own run, the serial pass durations:

2783.83s  test_llm_as_judge_base_model_fix.py::...::test_base_model_evaluation_uses_correct_weights
2504.59s  test_llm_as_judge_base_model_fix.py::...::test_base_model_false_still_works
  91.44s  the next-slowest test in that pass

88 minutes for two tests, against the project's 180-minute build timeout, and they were the entire tail — everything else in that pass finished in under 92s. Six tests block on execution.wait(..., timeout=14400), a 4-hour ceiling each, and none of them was marked:

Test	Was	Now
test_llm_as_judge_base_model_fix.py::test_base_model_evaluation_uses_correct_weights	serial	gpu_intensive (class)
test_llm_as_judge_base_model_fix.py::test_base_model_false_still_works	serial	gpu_intensive (class)
test_benchmark_evaluator.py::test_benchmark_evaluation_full_flow	none	gpu_intensive
test_custom_scorer_evaluator.py::test_custom_scorer_evaluation_full_flow	xdist_group	gpu_intensive
test_llm_as_judge_evaluator.py::test_llm_as_judge_evaluation_full_flow	none	gpu_intensive
test_llmaj_custom_model.py::TestLLMAJCustomModelIntegration	@pytest.mark.slow — unregistered, silent no-op	gpu_intensive

The last row was a genuine mismarking: the registered name is slow_test, so @pytest.mark.slow did nothing and only raised a PytestUnknownMarkWarning. us_east_1 already kept it off the us-west-2 gate, so marking it changes nothing there; it now stays off the us-east-1 job too.

An earlier revision of this description said I deliberately would not mark these, on the grounds that they have no shallow counterpart. The 88-minute measurement changed my mind, and it is a real trade rather than a free win, so here is what it costs. Three of the files are marked per-test and keep their constructor/validation tests on the gate — that is where SDK-side regressions surface. The two class-level ones leave nothing behind, and what the gate stops checking there is that a submitted evaluation pipeline is accepted and succeeds. Their already-marked siblings (test_benchmark_evaluation_base_model_only, test_custom_scorer_base_model_only) show this was already the established call for pipeline-waiting tests — these six were unmarked by omission, not by decision. Shallow evaluate() coverage is the follow-up that closes the gap properly.

fast-integ-tests existing as its own check is the other half of the answer: it returns in minutes regardless of what the deep suite is doing.

Also fixed: SimulatePrincipalPolicy throttling

The same CI run failed four shallow tests, and they were not test defects:

FAILED tests/integ/train/shallow/test_model_trainer.py::TestSourceCodePackaging::test_shell_entry_script
FAILED tests/integ/train/shallow/test_rlaif_trainer.py::TestRLAIFTrainerSubmission::test_mlflow_resource_arn
FAILED tests/integ/train/shallow/test_rlvr_trainer.py::TestRLVRTrainerSubmission::test_with_validation_dataset
FAILED tests/integ/train/shallow/test_rlvr_trainer.py::TestRLVRTrainerSubmission::test_dataset_passed_to_train_overrides_constructor
ClientError: (Throttling) ... SimulatePrincipalPolicy (reached max retries: 9): Rate exceeded

All four pass locally in isolation and under a local -n 36 run. Every ModelTrainer construction calls TrainDefaults.get_role → resolve_and_validate_role, which paginates SimulatePrincipalPolicy over ~20 action names against a low, account-wide TPS limit. CodeBuild runs the whole tree under -n auto (~36 workers on a 2XLARGE) — 188 trainer constructions, enough to exhaust even the adaptive 10-attempt budget the existing retry fixture grants. The cause is volume, not burstiness, so more retries would not have helped.

Fixed with a _memoize_role_validation autouse session fixture: each distinct (role, role_type, region) is validated once per xdist worker instead of once per test. Measured with an instrumented botocore _make_api_call — 3 trainers → 3 calls unpatched, 10 trainers → 1 call memoized. Exceptions are cached alongside successes so a bad role still fails (test_unassumable_role_is_rejected still passes), and teardown restores any caller now holding the memoized function rather than only the ones the fixture explicitly patched, since a module imported after the patch binds it at its own import time.

Cost model measured, not assumed. Across ~110 training jobs created by these runs, in both us-west-2 and us-east-1, 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
• the five us_east_1 tests had never actually run. They referenced a bucket hardcoded to one test account (s3://sagemaker-us-east-1-784379639078/...), which other accounts cannot read — AccessDenied on ListObjectsV2 from 729646638167. They now derive every path from default_bucket() and resolve the reward function from the caller's own hub, the way test_sft_trainer_serverful_smtj.py::training_resources already does. All five pass.
• a model package group must be in the job's region — passing the us-west-2 ARN the other trainers use is rejected with Model package group ARN region 'us-west-2' does not match expected region 'us-east-1', so the Nova files use a bare name that resolves per-session
• gpu_intensive and us_east_1 were declared only in tox.ini, but pytest reads its config from pyproject.toml here, so both markers were unregistered at runtime. That matters when the gate selects with -m "not gpu_intensive and not us_east_1": a typo'd marker name would silently put an expensive deep test back on the gate instead of warning. Now registered where pytest reads them.

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 (test_local_model_trainer.py — no service call, but it runs real containers, and it stays on the CodeBuild job which already has Docker).

Known remaining gap: evaluator evaluate() submissions. Same pattern applies — assert the pipeline execution ARN comes back without waiting for the pipeline to finish — but it is a distinct API surface, so it is left for a follow-up. Until then the gate checks that evaluators construct and validate correctly, but not that a submitted pipeline is accepted.

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 own checks; it starts running once merged. The test selection was verified locally (numbers above); the job definition itself is unproven in CI.

integ-tests (sagemaker-train) does run on this PR, as it will after merge, and it runs this branch's test code — which is how the 88-minute evaluator tail and the SimulatePrincipalPolicy throttling were both found. What it runs against is master's workflow definition, so the new job's absence from these checks is expected; the marker changes, being in the test files themselves, are exercised on the next run of this PR.

[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

[github-actions]

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

---

Review

Re-reviewed the current tree after the latest push (full 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. Prior automated reviews covered the big items; below I confirm what still stands in the current one-file-per-trainer layout. Nothing blocks merge.

Note: the inline-comment tool isn't available in this run, so findings are referenced by file/anchor here rather than as line comments.

---

1. test_duplicate_job_name_is_rejected still leaks the second job — shallow/test_model_trainer.py

first = None
try:
    first = create()
    assert_submitted(first, expected_name=job_name)
    with pytest.raises(Exception) as excinfo:
        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 — 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:

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

(assert_rejected / submitted already handle this correctly — this is the lone exception.)

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

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

In a raw string \\ is two characters, so the class is [0-9\.] plus a literal backslash — not the intent. Harmless here (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, with the outcome depending on file ordering. Almost certainly benign; 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)

The gating is correct: fast-integ-tests → needs: [detect-changes] → the same approval chain as the other jobs. 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

• Account/region constants (MODEL_PACKAGE_GROUP at recipe_cases.py, CPU_IMAGE at harness.py, the RLAIF/RLVR reward-model ARNs) pin the gate to one account + us-west-2; consistent with the existing dry-run suite, so not a regression, but the gate only passes where those exact resources exist, and a non-us-west-2 run does a cross-region ECR resolve for CPU_IMAGE.
• 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.
• .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.

---

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

[github-actions]

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

---

Review

Re-reviewed the current tree from the full 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. Prior automated reviews flagged the substantive items; below I confirm which still stand after the latest push. Nothing blocks merge; #1 is the one worth fixing first.

Note: the inline-comment tool isn't available in this run, so findings are referenced by file/anchor rather than as line comments.

---

1. test_duplicate_job_name_is_rejected still leaks the second job — shallow/test_model_trainer.py (test_duplicate_job_name_is_rejected)

first = None
try:
    first = create()
    assert_submitted(first, expected_name=job_name)
    with pytest.raises(Exception) as excinfo:
        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 — and the only hand-rolled pytest.raises block that doesn't route cleanup through the harness (assert_rejected/submitted both handle this correctly). Capture the second job and stop it in the finally too:

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

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

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

In a raw string \\ is two characters, so the class is [0-9\.] plus a literal backslash — not the intent. Harmless here (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", and restores AWS_DATA_PATH only at session teardown. It auto-applies only under shallow/, but the env var is process-global and the gate can run 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, with the outcome depending on file ordering. Almost certainly benign (the bundled model is a superset), but a narrower scope or an explicit unset after the shallow files would match the stated intent more tightly. Note: fast-integ-tests scopes pytest to tests/integ/train/shallow only, so this matters mainly for a local full-tree run, not the new job.

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

The gating is correct: fast-integ-tests → needs: [detect-changes] → the same approval chain as the other jobs. 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 (pr-checks-master.yml, fast-integ-tests) 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. 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

• tests/integ/train/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.
• Account/region constants (MODEL_PACKAGE_GROUP in recipe_cases.py, CPU_IMAGE in harness.py, the RLVR/RLAIF reward ARNs) pin the gate to one account + us-west-2. Consistent with the existing dry-run suite, so not a regression — but since this is now the PR gate, the gate only passes where those exact resources exist, 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.
• .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 is also still @v3).

I also confirmed HyperParameterTuningJob.refresh() returning self makes the sm_drivers read-back in test_distributed_tuning_job_is_accepted safe, and that unique_name(..., max_length=MAX_TUNING_JOB_NAME) stays within the 32-char cap.

---

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

[github-actions]

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

---

Review

Re-reviewed the current tree from the full diff. This remains 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. Five prior automated reviews covered the standing items; I confirm which still stand and add one new finding (#2) they missed. Nothing blocks merge.

Note: the inline-comment tool isn't available in this run, so findings are referenced by file/anchor rather than as line comments.

Checklist

• Read the full PR diff (/tmp/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

first = None
try:
    first = create()
    assert_submitted(first, expected_name=job_name)
    with pytest.raises(Exception) as excinfo:
        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 — and the only hand-rolled pytest.raises block that doesn't route cleanup through the harness (assert_rejected/submitted both handle this). Capture and stop it too:

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

This is the one worth fixing before merge — it defeats the suite's own no-leaked-capacity discipline in the exact case it exists to catch.

2. Dead imports (and a stale docstring) in shallow/test_tuner.py — NEW

The file imports pytest (line ~3176) and MultiTurnRLTrainer (from sagemaker.train.multi_turn_rl_trainer import MultiTurnRLTrainer, line ~3182), but neither is used anywhere in the file — TestTuningJobSubmission has no marks/pytest.* calls, and the only MTRL tests live in test_multi_turn_rl_trainer.py. The module docstring likewise still describes MultiTurnRLTrainer.train() at length, a leftover from when the two shared a file. If flake8 runs over the integ tree (F401), this fails lint. Drop both imports and trim the MTRL paragraph from the docstring.

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

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

In a raw string \\ is two characters, so the class is [0-9\.] plus a literal backslash — not the intent. Harmless here (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).

4. 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; a local full-tree run with --dist loadfile means once a worker runs a shallow file, any later non-shallow file on that worker inherits the bundled (superset) model, with the outcome depending on file ordering. Almost certainly benign (superset), and the fast-integ-tests job scopes pytest to shallow/ only — so this matters mainly for a local full-tree run. A narrower scope or explicit unset after the shallow files would match the stated intent more tightly.

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

Gating is correct: fast-integ-tests → needs: [detect-changes] → the same approval chain as the other jobs. 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 execute in that context. 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.

6. Minor

• Temp-file leak — shallow/test_rlvr_trainer.py (test_explicit_recipe_file, test_recipe_and_overrides_together) create NamedTemporaryFile(..., delete=False) and never os.unlink them, so each run leaves a .yaml in the temp dir. Cheap to clean up in a finally or via tmp_path.
• shallow/test_multi_turn_rl_trainer.py — mtrl_prerequisites takes account_id and region fixtures but uses neither; drop them from the signature.
• tests/integ/train/test_sft_trainer_integration.py:180 now carries both the stale # @pytest.mark.gpu_intensive comment and the real decorator — drop the comment.
• Account/region constants (MODEL_PACKAGE_GROUP in recipe_cases.py, CPU_IMAGE in harness.py, the RLVR/RLAIF reward ARNs) pin the gate to one account + us-west-2. Consistent with the existing dry-run suite, so not a regression — but since this is now the PR gate, it only passes where those exact resources exist, and a non-us-west-2 run does a cross-region ECR resolve for CPU_IMAGE.
• .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 is also still @v3).

I also re-confirmed the memoize fixture caches exceptions so test_unassumable_role_is_rejected still fails a bad role, that assert_rejected/submitted clean up on unexpected acceptance, and that unique_name(..., max_length=MAX_TUNING_JOB_NAME) stays within the 32-char cap.

---

Net: a real improvement to PR-gate cost/latency with honestly-scoped coverage. #1 (leaked second job) and #2 (dead imports → possible lint failure) are the two worth fixing before merge; the rest are polish.
· shallow-pr-checks-sagemaker-train