Skip to content

Commit 7198091

Browse files
authored
change(train): gate deep integ tests behind gpu_intensive, add shallow submit-then-stop suite (#6176)
* change(train): gate deep integ tests behind gpu_intensive, add shallow submit-then-stop suite 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. Adds tests/integ/train/shallow (70 tests) built on that: submit, assert the ARN, stop immediately. Covers ModelTrainer (payload shaping, source-code packaging, input channels, compute, networking, checkpointing/spot), the recipe trainers (SFT/DPO/RLVR/RLAIF, serverless and serverful), recipe customization (overrides, explicit recipe files, sequence_length, DataMixingConfig), and the non-training job types (HyperParameterTuningJob, AgentRFT Job). Includes negative tests so the suite cannot pass merely because some ARN came back. Marks the 19 previously-unmarked tests that submit a job and wait for it with gpu_intensive, so they continue running on the scheduled CI-health workflows instead of the PR gate. Widens that marker's description: despite the name it gates anything consuming real training capacity, including serverless and CPU-instance jobs. 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/, which keeps the ~170 client-side tests (recipe resolution, data utils, dry-run, log streaming) on the gate -- they make no service call and were never the expensive part. Net: 191 of 251 tests on the PR gate, none of which waits for a training job. This is a deliberate scope reduction: training *behaviour* (artifacts, metrics, convergence) is no longer asserted on the PR gate. A regression that breaks training itself -- a bad entry script, a broken container command -- will pass here and be caught by the scheduled suites. * change(train): fix CPTTrainer construction and role-rejection test after first real AWS run Verified against AWS in account 729646638167 (us-west-2): * test_unassumable_role_is_rejected: ModelTrainer.__init__ validates the role via iam:SimulatePrincipalPolicy, so a bad role raises RoleValidationError at construction and never reaches CreateTrainingJob. Assert around the constructor instead of around train(). * test_cpt_trainer_is_accepted: CPTTrainer takes no training_type, and its compute is HyperPodCompute-only, so it cannot use the shared _trainer helper. WIP: 2 further real failures still to fix (RLAIF compute, tuner job-name collision). See SHALLOW_TEST_RUN_STATE.md. * change(train): fix shallow suite against real AWS; 62/62 passing Ran the suite against account 729646638167 (us-west-2) with PYTHONPATH pointed at this clone, and fixed every failure it surfaced. All were wrong assumptions in the tests, not service problems: * conftest: add a session-scoped bundled_service_model fixture setting AWS_DATA_PATH to sagemaker-core/sample. The public botocore model has no ServerlessJobConfig.SequenceLength, so sequence_length requests were rejected client-side before reaching the service. Mirrors the existing setup_aws_data_path fixture in test_recipe_override_integration.py. * harness: unique_name() now takes max_length. Tuning job names are capped at 32 characters, not the 63 allowed for training jobs, and the service enforces it: Value '...' at 'hyperParameterTuningJobName' failed to satisfy constraint: Member must have length less than or equal to 32 * tuner tests: submit under an explicit job_name via a _tuning() context manager. The tuner derives its default 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. * RLAIF: excluded from TestServerfulSubmission. RLAIFTrainer has no compute parameter, so it has no serverful path. Still covered by every serverless case. * CPT: marked gpu_intensive and skipped unless SHALLOW_HYPERPOD_CLUSTER is set. CPT refuses to submit without HyperPod compute, and HyperPod targets a pre-provisioned cluster rather than CreateTrainingJob. * sequence_length / training_type: narrowed to the values the recipe catalogue actually offers for this model ('4K' only; no serverless recipe for FULL). Both left parametrized so more values can be added against a model that supports them, rather than dropping the distinction. Result: 62 passed, 0 failed, 5m18s serial (~5s/test). Cost model confirmed empirically rather than assumed: across 100 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. * change(train): one shallow file per trainer; only mark deep tests that have shallow coverage Addresses two review points. 1. Only mark deep tests that this suite actually replaces. Reverts gpu_intensive from 9 tests that had no shallow counterpart, so the PR gate no longer loses coverage with nothing replacing it: * all 8 evaluator tests (benchmark, custom scorer, inspect_ai, llm_as_judge x2, llmaj_custom_model) -- evaluate() is a different API surface returning pipeline executions, and this suite has no coverage for it * test_notifications.py -- asserts EventBridge/SNS side effects, not submission 10 marks remain, each with a named shallow equivalent documented in the suite README. The rule is written down there: do not mark a deep test unless a shallow test covers the same path. 2. One file per trainer, matching the existing deep-suite layout. test_recipe_trainers_submission.py -> test_{sft,dpo,rlvr,rlaif,cpt}_trainer.py test_recipe_customization_submission.py (recipe cases folded into rlvr/sft; Nova data mixing to its own file) test_other_job_types_submission.py -> test_tuner.py, test_multi_turn_rl_trainer.py test_model_trainer_submission.py -> test_model_trainer.py The "recipe_*" names described how the SDK groups these internally rather than what a reader looks for; the shallow counterpart of a given deep test is now obvious from the filename. recipe_cases.py holds the cases every recipe trainer shares. Each per-trainer class subclasses RecipeTrainerCases and sets TRAINER, so a new 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). Not named test_* so pytest does not collect the base class. Inheriting the shared cases also widened coverage: DPO and RLAIF now get the full set (output path, dataset override, both negative cases) rather than only the three they had as parametrized entries. 80 tests total, 69 on the PR gate. Verified against AWS (account 729646638167, us-west-2): 68 passed, 1 skipped, 0 failed in 6m59s. The skip is RLAIF's serverful case, reporting "RLAIFTrainer takes no compute argument". * change(train): add shallow coverage for every gpu_intensive test that has an equivalent Previous commits only audited the marks this PR added. This audits all 46 gpu_intensive tests in tests/integ/train -- including those already marked on master -- and adds the missing shallow counterparts. Added (were gaps): * MLflow, in RecipeTrainerCases so all four recipe trainers get it. Every *_complete_workflow deep test configures MLflow, so without this their shallow counterparts missed that half of the payload. Two forms: experiment/run names (always runs) and mlflow_resource_arn (skips if the account has no app). * RLVR reward functions, all three forms the deep suite covers: hub-content ARN, Lambda ARN (auto-creates an Evaluator), and a pre-created Evaluator object. * RLAIF reward_prompt as a hub-content ARN rather than a Builtin.* name, and continued fine-tuning from a model-package ARN. * Nova SFT and Nova RLVR, in test_nova_trainers.py. Nova needs a different recipe family, region and account, so it cannot share RecipeTrainerCases; marked us_east_1. Two real constraints the AWS run surfaced, both now recorded in comments: * The reward-function tests cannot use this suite's generic chat-format fixture. Before submitting, the SDK *invokes* the reward function over sample records and fails if they do not score ("GSM8k scoring failed"). They now use the same dataset as the deep RLVR suite, via a dedicated reward_scored_data_uri fixture. * list_mlflow_apps is not a paginatable operation, so the fixture calls it directly instead of via get_paginator. Also fixed a ScopeMismatch: the three new lookup fixtures were session-scoped but depend on the parent conftest's module-scoped sagemaker_session. All three new fixtures (mlflow_arn, reward_lambda_arn, reward_evaluator) only look 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. Still uncovered, documented in the suite README with the reason: the 11 evaluator tests (evaluate() is a different API surface returning pipeline executions) and the 3 HyperPod tests (submit to a pre-provisioned cluster, not CreateTrainingJob). Neither is newly marked by this PR, so no coverage is lost; the evaluator gap is the clearest follow-up. 97 tests total, 82 on the PR gate. Verified against AWS (729646638167, us-west-2): 81 passed, 1 skipped, 0 failed in 7m04s. The skip is RLAIF's serverful case, which reports its own reason. * change(train): add shallow coverage for recipe overrides, GRPO hyperparameters, Nova serverful Three remaining gpu_intensive tests had no shallow counterpart: * test_sft_trainer_serverful_smtj.py (override half) -> SFT test_recipe_overrides_are_accepted. Asserts both halves: the merge reached the rendered recipe (client-side, exact) and the resulting payload is still accepted (recipe filtering runs after the request validators, so a bad merge only surfaces at submission). Verified against AWS: overrides are written flat under training_config but land nested under training_args, and the recipe default for this model is 5 -- so asserting 1 proves the override applied rather than coinciding with the default. * test_rlvr_trainer_nemotron_with_kl_and_recipe -> RLVR test_kl_and_clipping_hyperparameters. These are separate recipe fields rather than one flag, so the existing max_epochs-only test did not prove they serialize. * test_sft_trainer_serverful_smtj.py (Nova half) -> Nova TestNovaServerfulSubmission. Distinct from the shared serverful case: Nova model, Nova recipe family, Nova-only instance type, us-east-1. Accepts the override under either trainer.max_epochs or training_args.max_epochs, since recipe families nest epoch control differently -- so the test fails on a lost override rather than on a recipe-layout difference. Verified against a real account (us-west-2): 83 passed, 1 skipped, 0 failed in 5m20s. The skip reports its own reason (RLAIFTrainer takes no compute argument). * fix(train): make the shallow Nova tests runnable in any account The five us_east_1 shallow tests referenced resources hardcoded to one test account and had therefore never actually executed. Verified: from 729646638167, `aws s3 ls s3://sagemaker-us-east-1-784379639078/input_data/sft-nova/` returns AccessDenied. Derive everything from the calling account instead, the way test_sft_trainer_serverful_smtj.py::training_resources already does: * nova_sft_data_uri -- uploads the Nova-shaped sample data the deep suite already ships (tests/data/train/sft_smtj_sample_data.jsonl) to the caller's own bucket. Cannot reuse nova_train_data_uri: Nova SFT records carry a schemaVersion the generic chat-format fixture lacks. * nova_rlvr_data_uri -- copies the GSM8k-shaped dataset the us-west-2 RLVR tests use into the us-east-1 bucket. A copy rather than a reference because an S3 input must be in the job's region. * nova_output_path -- default_bucket() rather than a named bucket. * nova_reward_function_arn -- resolves the hub content in the caller's own account, look-up-and-skip like the other reward fixtures. Two service-verified region constraints drove this: * the model package group must be in the job's region -- passing the us-west-2 MODEL_PACKAGE_GROUP ARN is rejected with "Model package group ARN region 'us-west-2' does not match expected region 'us-east-1'". Added NOVA_MODEL_PACKAGE_GROUP (a bare name) alongside it in recipe_cases so the two Nova files cannot drift. * likewise for S3 inputs, hence the RLVR copy above. The Nova RLVR case sets skip_reward_validation=True. The SDK invokes the reward function over sample records before submitting; the function registered under that name in this account returns a shape the verifier rejects ("Each output must include 'id', 'aggregate_reward_score'"), so the test would assert per-account hub contents rather than this payload. The verifier is already covered against a known-compatible function by the three us-west-2 reward-function cases; what is unique here is the Nova recipe family and region. Also register gpu_intensive and us_east_1 in pyproject.toml. They were declared only in tox.ini, but pytest reads its config from pyproject.toml, so both were unregistered at runtime. That matters here: the PR gate selects with -m "not gpu_intensive and not us_east_1", so a typo'd marker name would silently put an expensive deep test back on the gate instead of warning. Verified against a real account: 5 passed in 47s, all five for the first time. Every job ended Stopped with BillableTimeInSeconds null, so the cost model holds in us-east-1 as well. * docs(train): record what actually bounds the PR gate's runtime A full gate run showed the shallow suite is not what makes this job slow. Measured (us-west-2, -n 8 --dist loadfile): 201 of 204 tests finished in ~7 minutes, then three evaluator tests held the run open for another 40+ before being killed. Five evaluator tests are not marked gpu_intensive and each blocks on execution.wait(..., timeout=14400) -- a 4-hour ceiling, ~33 minutes per execution in practice: test_benchmark_evaluator.py::test_benchmark_evaluation_full_flow (no marks) test_custom_scorer_evaluator.py::test_custom_scorer_evaluation_full_flow (xdist_group) test_llm_as_judge_evaluator.py::test_llm_as_judge_evaluation_full_flow (no marks) test_llm_as_judge_base_model_fix.py::test_base_model_evaluation_uses_correct_weights (serial) test_llm_as_judge_base_model_fix.py::test_base_model_false_still_works (serial) They run on master's gate too, so this PR does not add them -- but it does not fix them either, and they now dominate the job's wall clock. Deliberately NOT marking them here: unlike every other gpu_intensive test they have no shallow counterpart, so marking would remove coverage, which is what the rule this PR establishes forbids. Correct order is to add evaluator support to the harness first, then mark. Documented in the suite README so the next person does not have to rediscover it by watching a run stall at 95%. Also flags test_local_model_trainer.py in the workflow: it runs real containers, so it needs Docker and pulls pytorch-training:2.0.0-cpu-py310 (2.3 GB compressed, verified via ECR). That is fine on GitHub-hosted Ubuntu runners, which preinstall Docker, and the ECR read is already covered by the role the shallow tests use -- but it is the slowest non-evaluator thing on the gate and the only step with a disk-space floor, so the note says what to deselect first if the job ever goes flaky on runner capacity. * change(ci): keep the sagemaker-train integ job, add shallow tests alongside it Restores integ-tests to its master definition -- sagemaker-train is back in the matrix, byte-identical to master -- and makes fast-integ-tests additive rather than a replacement. The deep tests still come off the gate, just not by removing the job. The CodeBuild project's buildspec already selects -m "not gpu_intensive and not us_east_1" (verified by reading the live project), so the marks added earlier in this PR are what deselect them. No workflow edit was needed for that. Keeping the CodeBuild job also keeps things the shallow job cannot cover: * the whole tests/integ tree, so the ~170 client-side tests (recipe resolution, data utils, dry-run, log streaming) run without this job repeating them; * test_local_model_trainer.py, which needs a Docker daemon. CodeBuild runs start-dockerd with privilegedMode, which is a better home for it than a GitHub runner pulling a 2.3 GB image -- so the reviewer caveat about that is dropped as moot; * the serial/parallel split the buildspec does for rate-limited tests. fast-integ-tests is therefore scoped to tests/integ/train/shallow only. Widening it would duplicate the client-side tests and double the training jobs this suite creates. It stays a separate job rather than folding into the buildspec because the buildspec is CDK-managed outside this repo, and because a shallow failure then reports as its own check. Corrects a claim in the previous comment: the shallow suite does carry gpu_intensive tests -- 11 of them, the CPT and MTRL classes, which need a HyperPod cluster and an agent runtime. With us_east_1 that is 16 deselected, so 84 of 100 run here. The comment now lists both groups and why. Verified against a real account: 83 passed, 1 skipped, 0 failed in 3m15s (the skip self-reports: RLAIFTrainer takes no compute argument). Faster than the 5m20s measured with the client-side tests bundled in. Every job ended Stopped with BillableTimeInSeconds null; no leaked jobs. * fix: memoize role validation and mark six pipeline-waiting evaluator tests Two problems the PR gate surfaced on its own run of this branch. 1. SimulatePrincipalPolicy throttling (4 shallow tests failed) 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 botocore.exceptions.ClientError: An error occurred (Throttling) when calling the SimulatePrincipalPolicy operation (reached max retries: 9): Rate exceeded Not a test defect: all four pass locally in isolation and in 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. The CodeBuild job runs the whole tests/integ tree under -n auto (~36 workers on a 2XLARGE), which is 188 trainer constructions -- enough to exhaust even the adaptive 10-attempt budget the existing _configure_boto_adaptive_retries fixture grants. The cause is volume, not burstiness, so more retries would not have fixed it. 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. Two details worth keeping: * exceptions are cached alongside successes, so a bad role still fails -- test_unassumable_role_is_rejected still passes; * teardown restores any caller now holding the memoized function, not just the ones this fixture explicitly patched. A module imported after the source module was patched binds the memoized function at its own import time, so restoring only what was patched here would leak across the session. Verified: 83 passed, 1 skipped, 0 failed in 3m33s with memoization active. 2. Six evaluator tests wait on a full evaluation pipeline From the same build's 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 a 180-minute build timeout, and they were the entire tail. Each blocks on execution.wait(..., timeout=14400) -- a 4-hour ceiling per test. Marked gpu_intensive, along with the three *_full_flow tests that wait the same way in test_benchmark_evaluator.py, test_custom_scorer_evaluator.py and test_llm_as_judge_evaluator.py. test_llmaj_custom_model.py was a genuine mismarking: it carried @pytest.mark.slow, but the registered marker name is slow_test, so the mark silently did nothing (PytestUnknownMarkWarning). us_east_1 already kept it off the us-west-2 gate, so this changes nothing there; it now also stays off the us-east-1 job. This is a small, real coverage reduction, and the README says so rather than claiming otherwise. Three of the files are marked per-test and keep their constructor/validation tests on the gate; the two class-level ones leave nothing behind, and what the gate stops checking is that a submitted pipeline is accepted and succeeds. Already-marked siblings in the same files (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. Shallow evaluate() coverage is the follow-up that closes the gap. Verified: 266/342 collected on the gate's selection (76 deselected), none of the six selected, and no PytestUnknownMark warnings remain. * change(train): cap concurrent training jobs in the shallow suite The shallow suite creates a training job per test. That puts it against two different quotas in two different units: * serverless (the default recipe-trainer path, no explicit compute) is bounded by "Maximum number of concurrent model customization serverless jobs per Region" -- a count of jobs, currently 20; * serverful (an explicit Compute/TrainingJobCompute: the ModelTrainer tests, the tuner, test_explicit_compute_is_accepted) is bounded by the per-instance-type quota, e.g. "ml.m5.large for training job usage" -- a count of instances. A slot is one concurrent job; a serverful job also takes one per instance, so a single cap holds the suite inside both quotas without the harness needing to know which kind of job a given test produces. Hold the slot until the job is terminal, not until stop() returns This is the subtle part, and the first cut got it wrong. The service counts a job against the concurrency quota from CreateTrainingJob until the job reaches Completed/Failed/Stopped -- NOT until StopTrainingJob returns. Measured against the service, stop() returns in a few seconds but the job takes ~1-3 min to actually drain (the reservation is torn down without ever becoming billable). Releasing the slot at stop() therefore bounded nothing: with the cap at 10 and 8 workers, each slot recycled ~20x inside a single job's counted lifetime, the suite peaked at ~37 concurrent jobs, and it tripped ResourceLimitExceeded at a utilization of 21 against the limit of 20. _wait_until_terminal closes that gap by holding the slot across the drain, so the cap bounds what the service actually counts. With the fix, live counted concurrency stayed at 4-5 against a cap of 10 for the whole run. The cost is runtime: holding to terminal makes the suite's floor roughly (#jobs * drain) / cap. At ~83 jobs, a ~75s median drain and cap 10 that is ~8-13 min, versus ~2 min if slots released early -- but that fast run is the one that breaches the quota. This is the batches-of-10 behaviour: at most 10 jobs counted at once. Mechanism job_slots() in harness.py, held by submitted() and assert_rejected() until the job is terminal. Slots are O_EXCL-created files under a run-keyed temp directory; xdist workers are separate processes, so an in-process semaphore would bound nothing. Keyed on PYTEST_XDIST_TESTRUNUID (falling back to the parent pid) so two concurrent local runs get separate budgets rather than deadlocking, and a stale directory from a killed run is never mistaken for live slots. Details that matter: * both waits proceed with a warning rather than failing -- acquiring a slot waits up to 900s, _wait_until_terminal up to 300s -- since the cap is a courtesy to the quota, not an assertion about the SDK, and a leaked slot or stuck drain should mean a slower run rather than a red build; * status is read per job type (training_job_status / job_status / hyper_parameter_tuning_job_status), since the SDK is not consistent, and a job that exposes no status releases its slot immediately rather than hanging; * a request larger than the cap is clamped, so a single test cannot deadlock against itself; * enforced in the harness rather than per test, so a new test is capped by default instead of by remembering to opt in. Default 10, overridable via SHALLOW_MAX_CONCURRENT_JOBS; 0 disables gating for a single-worker debugging run. Set explicitly in the workflow so the ceiling is visible at the call site rather than only in a Python default. Verified * Slot mechanism holds under contention: 12 processes x 4 iterations against cap=3, observed peak exactly 3, never 4; slots released on the happy path, on exception, and with correct multi-slot accounting; cap=0 takes none; an oversized request clamps without deadlocking. * Terminal-hold bounds what the service counts: a multi-process simulation where each job stays "counted" past stop() peaked at exactly the cap (3) with 10 workers, versus the pre-fix design that would have peaked far higher. * _wait_until_terminal waits through non-terminal states, releases on terminal, honours each job type's status attribute, and returns rather than hanging on None / a read error / a timeout. * Full suite green with the fix: 83 passed, 1 skipped in 810s (13:30), zero ResourceLimitExceeded, live counted concurrency 4-5 throughout, account fully drained afterward. * docs(train): drop account IDs from the shallow suite's comments This is a public repo, so the comments should not name internal test accounts. Every reference was explanatory -- "the deep test hardcodes a bucket in account X, which other accounts cannot read" -- and the point it makes is that the bucket belongs to *one specific account*, not which account that is. Reworded to say that instead, keeping each rationale (and the verified AccessDenied finding) intact. Comments and docs only; no functional change. Test resource ARNs still name the account they actually live in, since resolving them is what the tests do, and that already matches the convention in the surrounding suite. * change(train): use a bare model package group name in both regions Review feedback: recipe_cases.py pinned MODEL_PACKAGE_GROUP to a full ARN while the Nova path used a bare NOVA_MODEL_PACKAGE_GROUP, for the same group. The bare name is the better form on both paths, so the two constants collapse into one. The SDK accepts either -- _resolve_model_package_group_arn() returns an ARN unchanged and otherwise resolves a name via ModelPackageGroup.get() against the *session's* region -- so a name is region- and account-portable where an ARN pins both. Pinning the region is what forced the split in the first place: passing the us-west-2 ARN to a us-east-1 Nova job is rejected with "Model package group ARN region 'us-west-2' does not match expected region 'us-east-1'". One name serves both regions and drops a hardcoded account ID from a public repo. Verified: the bare name resolves to the same ARN via DescribeModelPackageGroup in us-west-2, and the us-west-2 recipe path still submits -- SFT, DPO and RLVR minimal-request tests pass against the service (3 passed). Collection unchanged at 100 tests. * test: resolve the shallow suite's CPU training image per region `CPU_IMAGE` hardcoded a us-west-2 URI in the public DLC account. Replace it with `cpu_image(sagemaker_session)`, which resolves the same image in the session's own region through `image_uris.retrieve` -- the resolver the SDK's framework estimators already use, so this is the supported mapping rather than a reconstruction of it. The registry account is not constant, which is what makes the hardcoded form actually wrong rather than merely untidy: it is 763104351884 across the commercial regions but 442386744353 in GovCloud and 727897471807 in China (on .com.cn). A pinned URI is unusable outside one partition, and it fails as an ECR error from the backend's role-assuming validators, which reads like a test bug rather than a hardcoded constant. A function rather than a constant because it needs the session's region; all three call sites already had a session in scope. Verified against AWS: reproduces the previously hardcoded URI byte-for-byte in us-west-2, and returns the correct in-region host (and per-partition registry) in us-east-1, eu-west-1, ap-northeast-1, us-gov-west-1 and cn-north-1. The affected tests pass on a real account -- 10 passed in 88s, covering the ModelTrainer helper, the raw TrainingJob.create path, and the tuner. * fix(train): bring the tuner path inside the shallow concurrency cap `_tuning()` submitted via `tuner.tune()` without acquiring slots, because a tuning job is stopped through `tuner.stop_tuning_job()` rather than `stop_quietly` and so never went through `submitted()`. Meanwhile the `DEFAULT_MAX_CONCURRENT_JOBS` note, the README quota table and `_requested_slots` all described the tuner as being inside the cap. It wasn't. Wrap it in `job_slots()` and drain after stopping. Slots are sized from the tuner's `max_parallel_jobs`, not a compute block: a tuning job occupies instance quota through the child training jobs it launches, which is also why `_requested_slots` cannot size this and `_tuning()` requests its own. The drain matters for the same reason it does elsewhere -- `stop_tuning_job()` returns while the job is still `Stopping` and its children are still tearing down, so releasing there is the release-before-terminal pattern that caused the ~37-concurrent breach. `_STATUS_ATTRS` already carried `hyper_parameter_tuning_job_status`, so the waiter handled this job type already; nothing ever called it with one. Renamed `_wait_until_terminal` -> `wait_until_terminal`. A test module outside the harness now needs it, and no other test imports a private name from there. Real impact today is small and worth saying so: both tuner tests are `max_parallel_jobs=1`, so this is 1 slot each. It is wired up because the cost is one context manager and the failure mode otherwise is silent -- a future test raising `max_parallel_jobs` would consume capacity outside a cap that still claimed to bound it. Verified: 3 unit scenarios (slots held through tune -> stop -> drain and released after; a failing test body still stops and releases; a missing or None `max_parallel_jobs` yields 1 slot, never an unbounded 0), plus a real run -- 2 passed in 28s, both jobs logging "reached Stopped; releasing slot" with no drain timeout. Docs corrected in all three places that overclaimed.
1 parent fd775cd commit 7198091

27 files changed

Lines changed: 3715 additions & 6 deletions

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

Lines changed: 112 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -243,6 +243,118 @@ jobs:
243243
project-name: ${{ github.event.repository.name }}-ci-${{ matrix.submodule }}-integ-tests
244244
source-version-override: 'refs/pull/${{ github.event.pull_request.number }}/head^{${{ github.event.pull_request.head.sha }}}'
245245

246+
# Additive: runs the shallow (submit-then-stop) suite for sagemaker-train
247+
# alongside the existing integ-tests job above, which is unchanged.
248+
#
249+
# Why a separate job rather than folding this into the CodeBuild suite: this
250+
# job's selection is reviewable in the PR that changes it, whereas the
251+
# sagemaker-train CodeBuild buildspec is CDK-managed outside this repo. It also
252+
# reports as its own check, so a shallow failure is distinguishable at a glance
253+
# from a deep-suite failure, and it finishes in minutes -- fast feedback that
254+
# does not wait on the 2XLARGE CodeBuild container.
255+
#
256+
# What runs here: only tests/integ/train/shallow. The client-side tests
257+
# (recipe resolution, data utils, dry-run, log streaming) are deliberately NOT
258+
# repeated -- the CodeBuild suite already runs the whole tests/integ tree, so
259+
# widening this job's scope would duplicate them and double the job creation
260+
# the shallow suite performs.
261+
#
262+
# Why submit-then-stop is worth gating on: CreateTrainingJob returns a
263+
# TrainingJobArn only after the request has cleared public-model validation,
264+
# SigV4, sagemaker:CreateTrainingJob authorization, iam:PassRole, the training
265+
# backend's request validators (including the role-assuming ones that resolve
266+
# S3 and ECR as the customer) and the final duplicate-name write. So a returned
267+
# ARN proves the payload and the caller's permissions are both good -- without
268+
# paying for a training run. The job is stopped immediately.
269+
#
270+
# It asserts nothing about training *behaviour* (artifacts, metrics,
271+
# convergence); that remains the deep suites' job.
272+
fast-integ-tests:
273+
runs-on: ubuntu-latest
274+
needs: [detect-changes]
275+
if: contains(fromJson(needs.detect-changes.outputs.submodules), 'sagemaker-train')
276+
steps:
277+
- uses: actions/checkout@v3
278+
with:
279+
# pull_request_target checks out the base ref by default; these tests
280+
# must run against the PR's code.
281+
ref: 'refs/pull/${{ github.event.pull_request.number }}/head'
282+
283+
- name: Set up Python
284+
uses: actions/setup-python@v5
285+
with:
286+
python-version: '3.12'
287+
288+
- name: Configure AWS Credentials
289+
uses: aws-actions/configure-aws-credentials@v4
290+
with:
291+
role-to-assume: ${{ secrets.CI_AWS_ROLE_ARN }}
292+
aws-region: us-west-2
293+
role-duration-seconds: 10800
294+
295+
- name: Install sagemaker-train and test dependencies
296+
run: |
297+
python -m pip install --upgrade pip
298+
pip install ./sagemaker-core
299+
pip install ./sagemaker-train
300+
pip install -r requirements/extras/test_requirements.txt
301+
302+
- name: Run shallow sagemaker-train integ tests
303+
working-directory: sagemaker-train
304+
env:
305+
AWS_DEFAULT_REGION: us-west-2
306+
# Role resolution goes through iam:SimulatePrincipalPolicy, which is
307+
# low-TPS; adaptive retries keep parallel workers from throttling each
308+
# other.
309+
AWS_RETRY_MODE: adaptive
310+
AWS_MAX_ATTEMPTS: '10'
311+
# Cap the training jobs the *service* counts against the concurrency
312+
# quota, across all xdist workers, so the suite stays inside the
313+
# "concurrent model customization serverless jobs per Region" quota
314+
# (20) with room for the deep integ-tests suite running the same
315+
# account concurrently. The harness holds each slot until the job is
316+
# terminal, not until stop() returns -- the service counts a job for
317+
# ~1-3 min after the stop -- so 10 means "at most 10 jobs counted at
318+
# once", the batches-of-10 behaviour, not "10 stops in flight".
319+
#
320+
# Not redundant with -n 8. -n caps worker processes; this caps what
321+
# the service counts, and with the slot held to terminal those diverge
322+
# sharply (each drain outlives the worker's stop() by minutes). It is
323+
# also what keeps the ceiling stable if -n is raised. A serverful job
324+
# counts one slot per instance.
325+
SHALLOW_MAX_CONCURRENT_JOBS: '10'
326+
run: |
327+
# Scoped to shallow/ only -- see the comment above this job for why the
328+
# rest of tests/integ/train is not repeated here.
329+
#
330+
# 84 of the suite's 100 tests run; the 16 deselected are:
331+
# us_east_1 (5) -- Nova cases; this job holds us-west-2
332+
# credentials only, so they run in the
333+
# integ-tests-us-east-1 job instead.
334+
# gpu_intensive (11) -- the CPT and MTRL classes. These are written in
335+
# the shallow style but cannot be made
336+
# self-contained: CPT submits only via HyperPod
337+
# (a pre-provisioned cluster, not
338+
# CreateTrainingJob) and MTRL needs an agent
339+
# runtime plus an MLflow app. Both become
340+
# gate-eligible by dropping one marker once those
341+
# prerequisites exist in the PR account.
342+
python -m pytest tests/integ/train/shallow \
343+
-m "not gpu_intensive and not us_east_1" \
344+
-n 8 \
345+
--dist loadfile \
346+
-v \
347+
--durations=15 \
348+
--junitxml=shallow-integ-results.xml
349+
350+
- name: Upload test results
351+
if: always()
352+
uses: actions/upload-artifact@v4
353+
with:
354+
name: shallow-integ-test-results
355+
path: sagemaker-train/shallow-integ-results.xml
356+
if-no-files-found: warn
357+
246358
integ-tests-us-east-1:
247359
runs-on: ubuntu-latest
248360
needs: [detect-changes]

‎sagemaker-train/pyproject.toml‎

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

8998
[tool.black]

‎sagemaker-train/tests/integ/conftest.py‎

Lines changed: 98 additions & 4 deletions
Original file line numberDiff line numberDiff line change
@@ -31,10 +31,25 @@
3131
its own). ``adaptive`` mode adds client-side rate limiting so bursts of
3232
``SimulatePrincipalPolicy`` calls ride out transient throttling.
3333
34-
Throttling that still exhausts the adaptive retry budget is deliberately left to
35-
fail the test loudly (rather than being converted to a skip), so a persistent
36-
rate-limit regression stays visible instead of silently disappearing from the
37-
results.
34+
* ``_memoize_role_validation`` (autouse) — retries alone were not enough. A PR-gate
35+
run failed four tests with ``(Throttling) ... SimulatePrincipalPolicy (reached
36+
max retries: 9)``: the adaptive budget was exhausted, not merely stressed. The
37+
cause is volume, not burstiness — ~190 tests each construct a trainer, every
38+
construction calls ``get_role``, and each of those runs a *paginated*
39+
``SimulatePrincipalPolicy`` over ~20 action names. Under ``-n auto`` on a large
40+
CodeBuild container that is thousands of calls against a low, account-wide TPS
41+
limit, so raising the retry budget only trades failures for a slower build.
42+
43+
Since the arguments repeat, the result does too: this memoizes
44+
``resolve_and_validate_role`` per worker, collapsing those calls to one per
45+
distinct ``(provided_role, role_type, region)``. Validation still happens — once,
46+
and its outcome (including a raised ``RoleValidationError``) is what gets reused,
47+
so a genuinely bad role still fails every test that uses it.
48+
49+
Throttling that still exhausts the retry budget after memoization is deliberately
50+
left to fail the test loudly (rather than being converted to a skip), so a
51+
persistent rate-limit regression stays visible instead of silently disappearing
52+
from the results.
3853
"""
3954
from __future__ import absolute_import
4055

@@ -66,3 +81,82 @@ def _configure_boto_adaptive_retries():
6681
os.environ.pop(key, None)
6782
else:
6883
os.environ[key] = value
84+
85+
86+
# Modules that did `from ...iam_role_resolver import resolve_and_validate_role`
87+
# hold their own reference to the original function, so patching only the defining
88+
# module would leave those bindings calling IAM directly. Each importer is patched
89+
# too. Kept as a list of (module path, attribute) so adding a caller is one line.
90+
_ROLE_RESOLVER_CALLERS = (
91+
("sagemaker.core.helper.iam_role_resolver", "resolve_and_validate_role"),
92+
("sagemaker.train.defaults", "resolve_and_validate_role"),
93+
("sagemaker.train.evaluate.base_evaluator", "resolve_and_validate_role"),
94+
)
95+
96+
97+
@pytest.fixture(autouse=True, scope="session")
98+
def _memoize_role_validation():
99+
"""Validate each distinct role once per xdist worker instead of once per test.
100+
101+
See this module's docstring for why retries alone were insufficient. Caches on
102+
``(provided_role, role_type, region)`` -- region is part of the key because the
103+
Nova tests validate the same role against us-east-1, and a role's resolution is
104+
region-scoped. Exceptions are cached alongside successes so a bad role keeps
105+
failing rather than being silently retried per test.
106+
"""
107+
import importlib
108+
109+
patched = []
110+
cache = {}
111+
112+
try:
113+
source = importlib.import_module(_ROLE_RESOLVER_CALLERS[0][0])
114+
except ImportError: # pragma: no cover - SDK layout changed
115+
yield
116+
return
117+
118+
original = source.resolve_and_validate_role
119+
120+
def memoized(provided_role=None, role_type=None, sagemaker_session=None, **kwargs):
121+
region = None
122+
if sagemaker_session is not None:
123+
region = getattr(sagemaker_session, "boto_region_name", None)
124+
key = (provided_role, role_type, region)
125+
126+
if key not in cache:
127+
try:
128+
cache[key] = (
129+
original(
130+
provided_role=provided_role,
131+
role_type=role_type,
132+
sagemaker_session=sagemaker_session,
133+
**kwargs,
134+
),
135+
None,
136+
)
137+
except Exception as exc: # cache the verdict, not just the happy path
138+
cache[key] = (None, exc)
139+
140+
result, error = cache[key]
141+
if error is not None:
142+
raise error
143+
return result
144+
145+
for module_path, attribute in _ROLE_RESOLVER_CALLERS:
146+
try:
147+
module = importlib.import_module(module_path)
148+
except ImportError:
149+
continue # optional/renamed caller; the others still get patched
150+
if getattr(module, attribute, None) is original:
151+
setattr(module, attribute, memoized)
152+
patched.append((module, attribute))
153+
154+
yield
155+
156+
# Restore by checking for `memoized` rather than only undoing what was patched
157+
# above: a caller imported *after* the source module was patched binds the
158+
# memoized function at its own import time, so it needs restoring too even
159+
# though this fixture never set it.
160+
for module, attribute in patched:
161+
if getattr(module, attribute, None) is memoized:
162+
setattr(module, attribute, original)

0 commit comments

Comments
 (0)