Repository navigation
Commit 7198091
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
File tree
- .github/workflows
- sagemaker-train
- tests/integ
- train
- shallow
| Original file line number | Diff line number | Diff line change | |
|---|---|---|---|
| |||
243 | 243 | | |
244 | 244 | | |
245 | 245 | | |
| 246 | + | |
| 247 | + | |
| 248 | + | |
| 249 | + | |
| 250 | + | |
| 251 | + | |
| 252 | + | |
| 253 | + | |
| 254 | + | |
| 255 | + | |
| 256 | + | |
| 257 | + | |
| 258 | + | |
| 259 | + | |
| 260 | + | |
| 261 | + | |
| 262 | + | |
| 263 | + | |
| 264 | + | |
| 265 | + | |
| 266 | + | |
| 267 | + | |
| 268 | + | |
| 269 | + | |
| 270 | + | |
| 271 | + | |
| 272 | + | |
| 273 | + | |
| 274 | + | |
| 275 | + | |
| 276 | + | |
| 277 | + | |
| 278 | + | |
| 279 | + | |
| 280 | + | |
| 281 | + | |
| 282 | + | |
| 283 | + | |
| 284 | + | |
| 285 | + | |
| 286 | + | |
| 287 | + | |
| 288 | + | |
| 289 | + | |
| 290 | + | |
| 291 | + | |
| 292 | + | |
| 293 | + | |
| 294 | + | |
| 295 | + | |
| 296 | + | |
| 297 | + | |
| 298 | + | |
| 299 | + | |
| 300 | + | |
| 301 | + | |
| 302 | + | |
| 303 | + | |
| 304 | + | |
| 305 | + | |
| 306 | + | |
| 307 | + | |
| 308 | + | |
| 309 | + | |
| 310 | + | |
| 311 | + | |
| 312 | + | |
| 313 | + | |
| 314 | + | |
| 315 | + | |
| 316 | + | |
| 317 | + | |
| 318 | + | |
| 319 | + | |
| 320 | + | |
| 321 | + | |
| 322 | + | |
| 323 | + | |
| 324 | + | |
| 325 | + | |
| 326 | + | |
| 327 | + | |
| 328 | + | |
| 329 | + | |
| 330 | + | |
| 331 | + | |
| 332 | + | |
| 333 | + | |
| 334 | + | |
| 335 | + | |
| 336 | + | |
| 337 | + | |
| 338 | + | |
| 339 | + | |
| 340 | + | |
| 341 | + | |
| 342 | + | |
| 343 | + | |
| 344 | + | |
| 345 | + | |
| 346 | + | |
| 347 | + | |
| 348 | + | |
| 349 | + | |
| 350 | + | |
| 351 | + | |
| 352 | + | |
| 353 | + | |
| 354 | + | |
| 355 | + | |
| 356 | + | |
| 357 | + | |
246 | 358 | | |
247 | 359 | | |
248 | 360 | | |
| |||
| Original file line number | Diff line number | Diff line change | |
|---|---|---|---|
| |||
84 | 84 | | |
85 | 85 | | |
86 | 86 | | |
| 87 | + | |
| 88 | + | |
| 89 | + | |
| 90 | + | |
| 91 | + | |
| 92 | + | |
| 93 | + | |
| 94 | + | |
| 95 | + | |
87 | 96 | | |
88 | 97 | | |
89 | 98 | | |
| |||
| Original file line number | Diff line number | Diff line change | |
|---|---|---|---|
| |||
31 | 31 | | |
32 | 32 | | |
33 | 33 | | |
34 | | - | |
35 | | - | |
36 | | - | |
37 | | - | |
| 34 | + | |
| 35 | + | |
| 36 | + | |
| 37 | + | |
| 38 | + | |
| 39 | + | |
| 40 | + | |
| 41 | + | |
| 42 | + | |
| 43 | + | |
| 44 | + | |
| 45 | + | |
| 46 | + | |
| 47 | + | |
| 48 | + | |
| 49 | + | |
| 50 | + | |
| 51 | + | |
| 52 | + | |
38 | 53 | | |
39 | 54 | | |
40 | 55 | | |
| |||
66 | 81 | | |
67 | 82 | | |
68 | 83 | | |
| 84 | + | |
| 85 | + | |
| 86 | + | |
| 87 | + | |
| 88 | + | |
| 89 | + | |
| 90 | + | |
| 91 | + | |
| 92 | + | |
| 93 | + | |
| 94 | + | |
| 95 | + | |
| 96 | + | |
| 97 | + | |
| 98 | + | |
| 99 | + | |
| 100 | + | |
| 101 | + | |
| 102 | + | |
| 103 | + | |
| 104 | + | |
| 105 | + | |
| 106 | + | |
| 107 | + | |
| 108 | + | |
| 109 | + | |
| 110 | + | |
| 111 | + | |
| 112 | + | |
| 113 | + | |
| 114 | + | |
| 115 | + | |
| 116 | + | |
| 117 | + | |
| 118 | + | |
| 119 | + | |
| 120 | + | |
| 121 | + | |
| 122 | + | |
| 123 | + | |
| 124 | + | |
| 125 | + | |
| 126 | + | |
| 127 | + | |
| 128 | + | |
| 129 | + | |
| 130 | + | |
| 131 | + | |
| 132 | + | |
| 133 | + | |
| 134 | + | |
| 135 | + | |
| 136 | + | |
| 137 | + | |
| 138 | + | |
| 139 | + | |
| 140 | + | |
| 141 | + | |
| 142 | + | |
| 143 | + | |
| 144 | + | |
| 145 | + | |
| 146 | + | |
| 147 | + | |
| 148 | + | |
| 149 | + | |
| 150 | + | |
| 151 | + | |
| 152 | + | |
| 153 | + | |
| 154 | + | |
| 155 | + | |
| 156 | + | |
| 157 | + | |
| 158 | + | |
| 159 | + | |
| 160 | + | |
| 161 | + | |
| 162 | + | |
0 commit comments