From 7697926e1d38506719ca12e54ae1f504c54d442f Mon Sep 17 00:00:00 2001 From: Amarjeet LNU Date: Tue, 11 Aug 2026 16:07:52 -0700 Subject: [PATCH 01/15] 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. --- .github/workflows/pr-checks-master.yml | 105 +++ .../tests/integ/train/shallow/README.md | 139 ++++ .../tests/integ/train/shallow/__init__.py | 15 + .../tests/integ/train/shallow/conftest.py | 142 ++++ .../tests/integ/train/shallow/harness.py | 318 +++++++++ .../shallow/test_model_trainer_submission.py | 674 ++++++++++++++++++ .../test_other_job_types_submission.py | 251 +++++++ .../test_recipe_customization_submission.py | 250 +++++++ .../test_recipe_trainers_submission.py | 351 +++++++++ .../integ/train/test_benchmark_evaluator.py | 1 + .../train/test_custom_scorer_evaluator.py | 1 + .../integ/train/test_inspect_ai_evaluator.py | 2 + .../train/test_llm_as_judge_base_model_fix.py | 2 + .../train/test_llm_as_judge_evaluator.py | 1 + .../integ/train/test_llmaj_custom_model.py | 1 + .../tests/integ/train/test_model_trainer.py | 8 + .../tests/integ/train/test_notifications.py | 1 + .../train/test_sft_trainer_integration.py | 1 + .../integ/train/test_tuner_distributed.py | 1 + sagemaker-train/tox.ini | 2 +- 20 files changed, 2265 insertions(+), 1 deletion(-) create mode 100644 sagemaker-train/tests/integ/train/shallow/README.md create mode 100644 sagemaker-train/tests/integ/train/shallow/__init__.py create mode 100644 sagemaker-train/tests/integ/train/shallow/conftest.py create mode 100644 sagemaker-train/tests/integ/train/shallow/harness.py create mode 100644 sagemaker-train/tests/integ/train/shallow/test_model_trainer_submission.py create mode 100644 sagemaker-train/tests/integ/train/shallow/test_other_job_types_submission.py create mode 100644 sagemaker-train/tests/integ/train/shallow/test_recipe_customization_submission.py create mode 100644 sagemaker-train/tests/integ/train/shallow/test_recipe_trainers_submission.py diff --git a/.github/workflows/pr-checks-master.yml b/.github/workflows/pr-checks-master.yml index 1195ed2779..bac129a4cf 100644 --- a/.github/workflows/pr-checks-master.yml +++ b/.github/workflows/pr-checks-master.yml @@ -221,6 +221,13 @@ jobs: env: SUBMODULE: ${{ matrix.submodule }} + # sagemaker-train's PR-gate integ tests are handled by the shallow-integ-tests + # job below, so it is filtered out of this matrix. Every other submodule keeps + # the existing full CodeBuild integ suite unchanged. + # + # The filter is computed with fromJson/contains rather than by editing + # detect-changes, so the dependency-propagation logic there (and the submodule + # list consumed by codestyle-doc-tests and unit-tests) is untouched. integ-tests: runs-on: ubuntu-latest needs: [detect-changes] @@ -229,6 +236,8 @@ jobs: fail-fast: false matrix: submodule: ${{ fromJson(needs.detect-changes.outputs.submodules) }} + exclude: + - submodule: sagemaker-train steps: - name: Configure AWS Credentials uses: aws-actions/configure-aws-credentials@v4 @@ -243,6 +252,102 @@ jobs: project-name: ${{ github.event.repository.name }}-ci-${{ matrix.submodule }}-integ-tests source-version-override: 'refs/pull/${{ github.event.pull_request.number }}/head^{${{ github.event.pull_request.head.sha }}}' + # Replaces the CodeBuild integ suite for sagemaker-train on the PR gate. + # + # What runs here (~191 of 251 tests): + # * ~170 client-side tests that make no service call -- recipe resolution, + # data utils, dry-run, log streaming, docker-compose detection. These were + # always cheap and stay on the gate. + # * the shallow (submit-then-stop) suite under tests/integ/train/shallow. + # + # Why submit-then-stop is worth gating on: CreateTrainingJob returns a + # TrainingJobArn only after the request has cleared public-model validation, + # SigV4, sagemaker:CreateTrainingJob authorization, iam:PassRole, the training + # backend's request validators (including the role-assuming ones that resolve + # S3 and ECR as the customer) and the final duplicate-name write. So a returned + # ARN proves the payload and the caller's permissions are both good -- without + # paying for a training run. The job is stopped immediately. + # + # What no longer runs here: the ~54 tests that submit a job and wait for it. + # They are marked gpu_intensive and keep running on the scheduled CI-health + # workflows. This is a deliberate scope reduction -- training *behaviour* + # (artifacts, metrics, convergence) is not asserted on the PR gate. + # + # Runs directly on the runner rather than via CodeBuild because the sagemaker- + # train CodeBuild project's buildspec is CDK-managed outside this repo; running + # here keeps the test selection reviewable in the PR that changes it. + fast-integ-tests: + runs-on: ubuntu-latest + needs: [detect-changes] + if: contains(fromJson(needs.detect-changes.outputs.submodules), 'sagemaker-train') + steps: + - uses: actions/checkout@v3 + with: + # pull_request_target checks out the base ref by default; these tests + # must run against the PR's code. + ref: 'refs/pull/${{ github.event.pull_request.number }}/head' + + - name: Set up Python + uses: actions/setup-python@v5 + with: + python-version: '3.12' + + - name: Configure AWS Credentials + uses: aws-actions/configure-aws-credentials@v4 + with: + role-to-assume: ${{ secrets.CI_AWS_ROLE_ARN }} + aws-region: us-west-2 + role-duration-seconds: 10800 + + - name: Install sagemaker-train and test dependencies + run: | + python -m pip install --upgrade pip + pip install ./sagemaker-core + pip install ./sagemaker-train + pip install -r requirements/extras/test_requirements.txt + + - name: Run fast sagemaker-train integ tests + working-directory: sagemaker-train + env: + AWS_DEFAULT_REGION: us-west-2 + # Role resolution goes through iam:SimulatePrincipalPolicy, which is + # low-TPS; adaptive retries keep parallel workers from throttling each + # other. + AWS_RETRY_MODE: adaptive + AWS_MAX_ATTEMPTS: '10' + run: | + # Runs the WHOLE tests/integ/train tree, not just shallow/, and lets the + # markers decide what is affordable on a PR. That keeps the ~170 + # client-side tests (recipe resolution, data utils, dry-run, log + # streaming, docker-compose detection) on the gate -- they make no + # service call and were never the expensive part. + # + # Deselected, per the marker conventions already in tox.ini: + # gpu_intensive -- every test that submits a real job and waits for + # it. Now applied to the 19 submitters that were + # previously unmarked, so the shallow suite is the + # only thing on this gate that creates a job. + # us_east_1 -- this job holds us-west-2 credentials only; those + # tests run in the us-east-1 integ job. + # + # Note the shallow suite is NOT separately marked: it is intended to + # run here, and its own MTRL/Nova cases carry these markers themselves. + python -m pytest tests/integ/train \ + -m "not gpu_intensive and not us_east_1" \ + -n 8 \ + --dist loadfile \ + -v \ + --durations=15 \ + --junitxml=fast-integ-results.xml + + - name: Upload test results + if: always() + uses: actions/upload-artifact@v4 + with: + name: fast-integ-test-results + path: sagemaker-train/fast-integ-results.xml + if-no-files-found: warn + integ-tests-us-east-1: runs-on: ubuntu-latest needs: [detect-changes] diff --git a/sagemaker-train/tests/integ/train/shallow/README.md b/sagemaker-train/tests/integ/train/shallow/README.md new file mode 100644 index 0000000000..8ce3ec2100 --- /dev/null +++ b/sagemaker-train/tests/integ/train/shallow/README.md @@ -0,0 +1,139 @@ +# Shallow (submit-then-stop) integration tests + +These tests replace the full `sagemaker-train` integ suite **on the PR gate only**. +The deep suites still run on the scheduled CI-health workflows. + +## What a passing test proves + +Each test submits a real `CreateTrainingJob`, asserts the service returned a +`TrainingJobArn`, then immediately stops the job. + +The ARN is returned synchronously, and only after the request has cleared every +synchronous server-side gate: + +| Layer | Checks | +|---|---| +| Public API front end | Coral model/shape validation, required-member checks, SigV4 | +| IAM | `sagemaker:CreateTrainingJob` incl. condition keys, `iam:PassRole` on the execution role, training-plan ARN authorization | +| Interceptors | marketplace entitlement, resource reservation, tag governance, experiment config, IdC | +| Training backend — sync validators | ~56 validators: instance type/count, volume, KMS, stopping condition, channels, output config, VPC, debug/profiler, HPO params, environment, payload size, ARN partition/region, unlaunched-feature gating | +| Training backend — mutating validators | recipe resolution / hub content fetch | +| Training backend — role-assuming validators | real S3, ECR, FSx, algorithm, VPC dry-run calls **as the customer** | +| Post-validator business logic | training-plan capacity, per-preference plan matching, state-machine routing, SDC lookups, recipe filtering | +| Entity write | duplicate job name → `ResourceInUse` | + +So "the ARN came back" means: **the payload the SDK produced was accepted by the +service exactly as sent, and the caller held the permissions needed to submit it.** + +## What these tests deliberately do NOT cover + +Nothing about training *behaviour*: no model artifacts, no metrics, no container +logs, no convergence, no output-model-package creation. Those require a job to +actually run and remain the responsibility of the deep suites. + +Concretely, a regression that makes training itself fail — a broken entry script, +a bad container command, a distributed-launch bug — **will still pass here.** That +is the accepted trade for the runtime and cost reduction. + +## Coverage vs. the suite this replaces + +`tests/integ/train` has 181 pre-existing tests, but only ~50 actually submit a +job — the rest are client-side (recipe resolution, data utils, log streaming, +docker-compose detection). Mapping the *submitting* ones against this suite: + +| Existing area | Ported here | Notes | +|---|---|---| +| `test_model_trainer.py` (8) | yes | hyperparameter contract (dict/JSON/YAML), MPI, Torchrun, local tar source, `.sh` entry script, custom distributed driver | +| `test_sft_trainer_integration.py` (4) | partly | LoRA/FULL, validation dataset, `sequence_length`. **Nova workflow not ported** (us-east-1 + gated model) | +| `test_dpo_trainer_integration.py` (2) | yes | via `RECIPE_TRAINERS` parametrization | +| `test_rlvr_trainer_integration.py` (7) | partly | base + recipe/overrides + direct hyperparameter mutation. **Custom reward function / evaluator objects not ported** | +| `test_rlaif_trainer_integration.py` (3) | yes | RLAIF is in `RECIPE_TRAINERS`; its reward model/prompt come from `_TRAINER_EXTRA_KWARGS` | +| `test_cpt_hyperpod.py`, `test_nova_sft_hyperpod.py`, `test_sft_data_mixing_hyperpod.py` (3) | no | HyperPod submits to a pre-provisioned cluster, not `CreateTrainingJob` — the pattern does not apply | +| `test_sft_trainer_data_mixing_integration.py` (1) | yes | `DataMixingConfig`, both explicit and recipe-default | +| `test_tuner_distributed.py` (1) | yes | `HyperParameterTuningJob`; also asserts the `sm_drivers` channel survived submission | +| `test_multi_turn_rl_trainer_integration.py` (7) | partly | AgentRFT `Job` submission, marked `gpu_intensive` — see below | +| `test_recipe_override_integration.py` (35) | n/a | client-side `get_resolved_recipe`; keep as-is, cheap already | +| Evaluators (`test_benchmark_evaluator.py`, `test_llm_as_judge_*`, `test_mtrl_*`, ~20) | no | `evaluate()` not `train()`; the same pattern applies and is the clearest next extension | +| `test_notifications.py`, `test_local_model_trainer.py` | no | EventBridge/SNS side effects and local-container mode (no service call) | + +Note that not every trainer creates a `TrainingJob`. `HyperparameterTuner` creates +a `HyperParameterTuningJob` and `MultiTurnRLTrainer` creates an AgentRFT `Job`, so +`assert_submitted` takes a `resource=` argument for the expected ARN segment and +the harness resolves the submitted job across four different attribute names. + +**Deliberately out of scope for this pattern:** HyperPod (different submission +API), local container mode (no service call), and anything asserting a job's +*outcome*. + +**Requires prerequisites, so marked `gpu_intensive` and skipped on the PR gate:** +the MTRL tests. Unlike everything else here they cannot be made self-contained — +they need a pre-provisioned agent runtime and MLflow app. They read those from +`SHALLOW_MTRL_AGENT_ENV` / `SHALLOW_MTRL_MLFLOW_APP_ARN` / `SHALLOW_MTRL_DATASET` +and skip when unset, so once the PR account has them, dropping one marker makes +them PR-gate-eligible. + +**Genuine remaining gap:** evaluator `evaluate()` submissions (~20 existing +tests). Same pattern, distinct API surface; not yet written. + +## Relationship to `dry_run=True` + +`tests/integ/train/test_dry_run_integration.py` covers `trainer.train(dry_run=True)`, +which returns *before* submitting. It therefore validates only client-side logic +(config assembly, S3 path existence checks, hyperparameter constraints) and +exercises **none** of the table above. + +These suites are complementary and both are cheap: + +* `dry_run` — catches SDK-side problems with no service call at all. +* shallow — catches problems only the service can detect. + +## Cost and capacity + +Stopping is not free and not instantaneous. `StopTrainingJob` marks the job +`Stopping` and returns; the compute layer reacts asynchronously. Meanwhile the +create call has already handed the job to a state machine and queued it, so +capacity acquisition has begun. + +In practice a job stopped within seconds is torn down while still in +`Starting`/`Pending`, before instances become billable — but that is a timing +property, **not a guarantee**. Expect a small, non-deterministic cost per test, +and transient capacity consumption. + +Two design rules follow, and should be preserved: + +1. **Use the smallest instance that exercises the path.** `ModelTrainer` tests use + `ml.m5.large`; payload and permission validation is instance-type agnostic. + Only the recipe trainers pin an accelerator type (`ml.g5.12xlarge`), because + their recipes will not resolve onto CPU. +2. **Never set `keep_alive_period_in_seconds`.** A warm pool would outlive the stop + and keep instances provisioned after the test finished. + +## Writing a new test + +Use the harness; do not call `trainer.train()` directly. + +```python +from .harness import assert_submitted, submitted, unique_name + +def test_my_feature_is_accepted(sagemaker_session, train_data_uri): + trainer = _trainer(sagemaker_session, unique_name("shallow-my-feature"), ...) + with submitted(trainer) as job: + assert_submitted(job) +``` + +`submitted()` forces `wait=False`, resolves the submitted job across the +inconsistent trainer attributes (`_latest_training_job` vs `latest_training_job`), +and stops the job in a `finally` so a failed assertion still cleans up. Passing +`wait=` is rejected with a `TypeError` so a copy-pasted `wait=True` cannot +silently reintroduce a full training run. + +For negative cases use `assert_rejected`, which also stops the job if the request +is unexpectedly *accepted*: + +```python +assert_rejected(trainer, ("does not exist", "ValidationException")) +``` + +Keep at least one negative test per feature area. Without them the suite +degenerates into "any ARN is fine" and would stay green even if the SDK started +sending a permissive-but-wrong payload. diff --git a/sagemaker-train/tests/integ/train/shallow/__init__.py b/sagemaker-train/tests/integ/train/shallow/__init__.py new file mode 100644 index 0000000000..b137ba3a18 --- /dev/null +++ b/sagemaker-train/tests/integ/train/shallow/__init__.py @@ -0,0 +1,15 @@ +# Copyright Amazon.com, Inc. or its affiliates. All Rights Reserved. +# +# Licensed under the Apache License, Version 2.0 (the "License"). You +# may not use this file except in compliance with the License. A copy of +# the License is located at +# +# http://aws.amazon.com/apache2.0/ +# +# or in the "license" file accompanying this file. This file is +# distributed on an "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF +# ANY KIND, either express or implied. See the License for the specific +# language governing permissions and limitations under the License. +"""Shallow (submit-then-stop) integration tests for sagemaker-train.""" + +from __future__ import absolute_import diff --git a/sagemaker-train/tests/integ/train/shallow/conftest.py b/sagemaker-train/tests/integ/train/shallow/conftest.py new file mode 100644 index 0000000000..015f8fd28a --- /dev/null +++ b/sagemaker-train/tests/integ/train/shallow/conftest.py @@ -0,0 +1,142 @@ +# Copyright Amazon.com, Inc. or its affiliates. All Rights Reserved. +# +# Licensed under the Apache License, Version 2.0 (the "License"). You +# may not use this file except in compliance with the License. A copy of +# the License is located at +# +# http://aws.amazon.com/apache2.0/ +# +# or in the "license" file accompanying this file. This file is +# distributed on an "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF +# ANY KIND, either express or implied. See the License for the specific +# language governing permissions and limitations under the License. +"""Fixtures for the shallow (submit-then-stop) training-job suite. + +Inherits ``sagemaker_session``, ``ensure_default_region`` and the adaptive-retry +configuration from the parent ``tests/integ/train/conftest.py`` and +``tests/integ/conftest.py``; only fixtures specific to shallow submission live +here. + +Everything here is session- or module-scoped and idempotent: these tests run +in parallel across xdist workers, so any fixture creating an AWS-side artifact must +tolerate a dozen workers racing to create the same thing. +""" + +from __future__ import absolute_import + +import json +import logging + +import pytest + +logger = logging.getLogger(__name__) + +# Uploaded once and reused. A tiny object is enough: the backend's role-assuming +# validators check that the S3 prefix resolves, not what it contains. +_TRAIN_DATA_KEY = "shallow-integ-test/train/data.jsonl" +_VALIDATION_DATA_KEY = "shallow-integ-test/validation/data.jsonl" + +_SAMPLE_RECORDS = [ + { + "messages": [ + {"role": "user", "content": [{"text": "What is 2+2?"}]}, + {"role": "assistant", "content": [{"text": "4"}]}, + ] + }, + { + "messages": [ + {"role": "user", "content": [{"text": "Capital of France?"}]}, + {"role": "assistant", "content": [{"text": "Paris"}]}, + ] + }, +] + + +def _ensure_object(sagemaker_session, key): + """Upload the sample dataset at ``key`` if absent; return its S3 URI. + + Idempotent so concurrent xdist workers converge instead of colliding. The + object is intentionally left behind: it is a few hundred bytes and reusing + it removes an upload from every subsequent run. + """ + bucket = sagemaker_session.default_bucket() + s3 = sagemaker_session.boto_session.client("s3") + + response = s3.list_objects_v2(Bucket=bucket, Prefix=key, MaxKeys=1) + if response.get("KeyCount", 0) == 0: + body = "\n".join(json.dumps(record) for record in _SAMPLE_RECORDS) + s3.put_object(Bucket=bucket, Key=key, Body=body.encode("utf-8")) + logger.info("Uploaded shallow-test fixture data to s3://%s/%s", bucket, key) + + return f"s3://{bucket}/{key}" + + +@pytest.fixture(scope="module") +def train_data_uri(sagemaker_session): + """S3 URI of a real, existing training-data prefix.""" + return _ensure_object(sagemaker_session, _TRAIN_DATA_KEY) + + +@pytest.fixture(scope="module") +def validation_data_uri(sagemaker_session): + """S3 URI of a real, existing validation-data prefix.""" + return _ensure_object(sagemaker_session, _VALIDATION_DATA_KEY) + + +@pytest.fixture(scope="module") +def nova_train_data_uri(sagemaker_session_us_east_1): + """Training data in us-east-1, for Nova-only paths (e.g. data mixing). + + Nova models are exercised in us-east-1 in this repo (see the + ``sagemaker_session_us_east_1`` fixture in the parent conftest), and an S3 + prefix must be in the same region as the job that reads it -- so this cannot + reuse ``train_data_uri``, which lives in the default region's bucket. + """ + return _ensure_object(sagemaker_session_us_east_1, _TRAIN_DATA_KEY) + + +@pytest.fixture(scope="module") +def output_path(sagemaker_session): + """S3 prefix for training output. + + Nothing is ever written here -- the jobs are stopped long before they upload + artifacts -- but the backend validates the output location, so it must be a + real, writable prefix. + """ + return f"s3://{sagemaker_session.default_bucket()}/shallow-integ-test/output/" + + +@pytest.fixture(scope="module") +def nonexistent_data_uri(sagemaker_session): + """S3 URI, in a real bucket, that does not exist. + + Used by negative tests to prove input validation actually reaches S3 rather + than being skipped. + """ + bucket = sagemaker_session.default_bucket() + return f"s3://{bucket}/shallow-integ-test/definitely-not-here-04c1f9/" + + +@pytest.fixture(scope="module") +def execution_role(sagemaker_session): + """The validated training execution role for this account. + + Resolved through the SDK's own resolver so these tests exercise the same + role-discovery path real users hit, and so a broken/unassumable default role + surfaces here rather than as a confusing per-test PassRole failure. + """ + from sagemaker.train.defaults import TrainDefaults + + return TrainDefaults.get_role(role=None, sagemaker_session=sagemaker_session) + + +@pytest.fixture(scope="module") +def account_id(sagemaker_session): + """Caller's AWS account id, for building ARNs in negative tests.""" + return sagemaker_session.boto_session.client("sts").get_caller_identity()["Account"] + + +@pytest.fixture(scope="module") +def region(sagemaker_session): + """Region under test, for building ARNs and region-sensitive assertions.""" + return sagemaker_session.boto_session.region_name diff --git a/sagemaker-train/tests/integ/train/shallow/harness.py b/sagemaker-train/tests/integ/train/shallow/harness.py new file mode 100644 index 0000000000..ae57114ef2 --- /dev/null +++ b/sagemaker-train/tests/integ/train/shallow/harness.py @@ -0,0 +1,318 @@ +# Copyright Amazon.com, Inc. or its affiliates. All Rights Reserved. +# +# Licensed under the Apache License, Version 2.0 (the "License"). You +# may not use this file except in compliance with the License. A copy of +# the License is located at +# +# http://aws.amazon.com/apache2.0/ +# +# or in the "license" file accompanying this file. This file is +# distributed on an "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF +# ANY KIND, either express or implied. See the License for the specific +# language governing permissions and limitations under the License. +"""Submit-then-stop harness for shallow training-job integration tests. + +Why this exists +--------------- +``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 ~56 +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 finally a conditional write that +rejects duplicate job names. + +So "the ARN came back" is a strong assertion: the payload was accepted by the +service exactly as the SDK shaped it, and the caller held the permissions +required to submit it. That is materially more coverage than ``dry_run=True`` +(which returns before submitting and so exercises only client-side validation +-- see ``tests/integ/train/test_dry_run_integration.py``), and it costs a +fraction of a full training run because we stop the job immediately instead of +waiting for it to train. + +What this deliberately does NOT assert +-------------------------------------- +Nothing about training *behaviour*: no model artifacts, no metrics, no +container logs, no convergence. Those require a job to actually run and remain +the job of the existing deep integration tests. These tests answer one +question only -- "would the service accept this request?" + +Cost and capacity notes +----------------------- +Stopping is not free and not instantaneous. ``StopTrainingJob`` marks the job +``Stopping`` in the backend and returns; the compute layer reacts +asynchronously. Meanwhile the create call has already handed the job to a state +machine and queued it, so capacity acquisition has begun. In practice a job +stopped within seconds is torn down while still in ``Starting``/``Pending``, +before instances become billable, but that is a timing property rather than a +guarantee. + +Two consequences shape this module: + +* ``DEFAULT_INSTANCE_TYPE`` is a small CPU instance. Payload validation and + permission checks are instance-type agnostic, so there is no reason to ask + for scarce accelerator capacity. Tests that specifically need to prove an + accelerator-shaped request is accepted say so explicitly. +* We never set ``keep_alive_period_in_seconds``. A warm pool would outlive the + stop and keep instances provisioned after the test finished. + +Teardown runs in a ``finally`` so a failing assertion still stops the job, and +is itself best-effort: a job that already reached a terminal state cannot be +stopped and that is not a failure. +""" + +from __future__ import absolute_import + +import inspect +import logging +import random +import time +from contextlib import contextmanager + +import pytest +from botocore.exceptions import ClientError + +logger = logging.getLogger(__name__) + +# A small CPU instance is sufficient: acceptance of the request does not depend +# on the instance type being an accelerator, and asking for GPU capacity we +# immediately discard is both slower and antisocial in a shared test account. +DEFAULT_INSTANCE_TYPE = "ml.m5.large" +DEFAULT_INSTANCE_COUNT = 1 + +# Public DLC, present in every commercial region we test in. Using a real image +# matters: the backend's role-assuming validators resolve the training image +# against ECR, so a bogus URI would fail for the wrong reason. +CPU_IMAGE = "763104351884.dkr.ecr.us-west-2.amazonaws.com/pytorch-training:2.0.0-cpu-py310" + +# Keep the advertised runtime short. It should never be reached (we stop the job +# long before), but if a stop were somehow lost this bounds the damage. +MAX_RUNTIME_IN_SECONDS = 600 + +# Terminal/near-terminal states that make StopTrainingJob a no-op or an error. +_UNSTOPPABLE_STATUSES = frozenset({"Completed", "Failed", "Stopped", "Stopping"}) + + +def unique_name(prefix): + """Build a collision-free job name. + + The backend rejects duplicate job names per account with ``ResourceInUse``, + and these tests run in parallel across many xdist workers, so the + name must be unique per invocation rather than per test function. Includes + randomness as well as a timestamp because two xdist workers can enter the + same second. + + SageMaker training job names are limited to 63 characters, so the prefix is + truncated rather than allowed to silently push the suffix over the limit. + """ + suffix = f"{int(time.time())}-{random.randint(1000, 9999)}" + # 63 total, minus the suffix, minus the joining hyphen. + head = prefix[: 63 - len(suffix) - 1] + return f"{head}-{suffix}" + + +def stop_quietly(training_job): + """Stop a submitted job, tolerating races with its own lifecycle. + + Best-effort by design. A job that finished, failed or is already stopping + cannot be stopped again, and a test must not fail because teardown lost a + race with the service. Anything genuinely unexpected is logged loudly so it + stays visible without turning into a spurious test failure. + """ + if training_job is None: + return + + name = _first_attr(training_job, _NAME_ATTRS) + try: + training_job.stop() + logger.info("Stopped job %s", name) + except ClientError as e: + code = e.response["Error"]["Code"] + message = e.response["Error"].get("Message", "") + # ValidationException is what the service returns when the job has + # already reached a state from which it cannot be stopped. + if code in ("ValidationException", "ResourceNotFound"): + logger.info("Job %s no longer stoppable (%s): %s", name, code, message) + return + logger.warning("Unexpected error stopping job %s (%s): %s", name, code, message) + except Exception as e: # pragma: no cover - defensive teardown + logger.warning("Unexpected error stopping job %s: %s", name, e) + + +# Attributes under which the different job resources expose their ARN and name. +# Not every trainer in this package creates a TrainingJob: MultiTurnRLTrainer +# creates an AgentRFT Job (``job_arn``) and Tuner creates a +# HyperParameterTuningJob, so the harness reads whichever is present rather than +# assuming the TrainingJob shape. +_ARN_ATTRS = ( + "training_job_arn", + "job_arn", + "hyper_parameter_tuning_job_arn", +) +_NAME_ATTRS = ( + "training_job_name", + "job_name", + "hyper_parameter_tuning_job_name", +) + + +def _first_attr(obj, attrs): + """Return the first non-None attribute value from ``attrs``.""" + for attr in attrs: + value = getattr(obj, attr, None) + if value is not None: + return value + return None + + +def assert_submitted(job, expected_name=None, resource="training-job"): + """Assert the service accepted the request and handed back a real ARN. + + This is the single assertion that gives these tests their value, so it checks + the ARN's shape rather than merely its presence -- a truthy-but-malformed + value would otherwise pass silently. + + ``resource`` is the expected ARN resource segment. It defaults to + ``training-job`` because most trainers here create a TrainingJob, but + MultiTurnRLTrainer creates an AgentRFT ``job`` and Tuner creates a + ``hyper-parameter-tuning-job``, so those callers pass their own. + """ + assert job is not None, "train() returned no job; the request was never submitted" + + arn = _first_attr(job, _ARN_ATTRS) + assert arn, f"job has no ARN: {job!r}" + assert arn.startswith("arn:"), f"malformed ARN: {arn!r}" + assert f":{resource}/" in arn, f"ARN is not a {resource} ARN: {arn!r}" + + if expected_name is not None: + actual = _first_attr(job, _NAME_ATTRS) + assert ( + actual == expected_name + ), f"submitted job name {actual!r} does not match requested {expected_name!r}" + + logger.info("Service accepted request; ARN=%s", arn) + return arn + + +def _train_kwargs_for(trainer, extra): + """Build the kwargs for ``trainer.train()``, forcing a non-waiting submit. + + ``wait=False`` is the whole point of this suite: the ARN is returned + synchronously by ``CreateTrainingJob``, so waiting buys no extra coverage + and costs a full training run. + + ``logs`` is deliberately conditional. ``ModelTrainer.train`` accepts it, but + the recipe trainers (``SFTTrainer``, ``DPOTrainer``, ``RLVRTrainer``, + ``CPTTrainer``, ...) do not -- their signatures are + ``(training_dataset, validation_dataset, wait, wait_timeout, poll, + dry_run)``. Passing ``logs`` unconditionally would raise ``TypeError`` for + the entire recipe-trainer family, so it is introspected rather than assumed. + """ + kwargs = {"wait": False} + kwargs.update(extra) + + try: + parameters = inspect.signature(trainer.train).parameters + except (TypeError, ValueError): # pragma: no cover - defensive + parameters = {} + + # Only silence logs where the trainer understands the option; where it does + # not, wait=False already prevents log streaming. + if "logs" in parameters and "logs" not in kwargs: + kwargs["logs"] = False + + return kwargs + + +@contextmanager +def submitted(trainer, **train_kwargs): + """Submit a training job, yield it, and always stop it. + + Usage:: + + with submitted(trainer) as job: + assert_submitted(job) + + Callers must not pass ``wait``: it is forced to ``False`` and a supplied + value is rejected loudly rather than silently overridden, so a copy-pasted + ``wait=True`` cannot quietly reintroduce a full training run into the fast + suite. + """ + if "wait" in train_kwargs: + raise TypeError( + "submitted() controls 'wait'; remove it from the call. " + "These tests must never wait for a job to run." + ) + + training_job = None + try: + trainer.train(**_train_kwargs_for(trainer, train_kwargs)) + training_job = _resolve_job(trainer) + yield training_job + finally: + stop_quietly(training_job) + + +# Attributes under which trainers stash the job they just submitted. The SDK is +# not consistent here, so the harness checks all of them rather than silently +# yielding None (which would surface as a confusing "train() returned no job" +# failure instead of an attribute-discovery problem): +# _latest_training_job -- ModelTrainer and most recipe trainers +# latest_training_job -- DPOTrainer (public) +# _latest_job -- MultiTurnRLTrainer (AgentRFTJob) +# latest_tuning_job -- Tuner (HyperParameterTuningJob) +_JOB_ATTRS = ( + "_latest_training_job", + "latest_training_job", + "_latest_job", + "latest_tuning_job", +) + + +def _resolve_job(trainer): + """Return the job resource the trainer just submitted, whatever its type.""" + return _first_attr(trainer, _JOB_ATTRS) + + +def assert_rejected(trainer, expected_tokens, **train_kwargs): + """Assert a request is rejected, and clean up if it is unexpectedly accepted. + + Negative tests are what stop this suite from degenerating into "any ARN is + fine": without them, a bug that made the SDK send a permissive-but-wrong + payload would still produce a green suite. + + ``expected_tokens`` is a collection of substrings, any one of which is + accepted. Matching is deliberately loose because a rejection can legitimately + surface from three different layers with different wording -- SDK-side + validation (``ValueError``), the public API model + (``ValidationException``), or the training backend (``ValidationError``) -- + and pinning exact prose would make these tests fail on harmless message + changes. It is still specific enough to catch a *wrong* rejection, which is + the real risk: without it, a test could pass because of an unrelated + credentials or region error. + + If the request is unexpectedly accepted, the job is stopped before the test + fails, so a validation regression cannot leak a running job. + """ + if "wait" in train_kwargs: + raise TypeError("assert_rejected() controls 'wait'; remove it from the call.") + + training_job = None + try: + with pytest.raises(Exception) as excinfo: + trainer.train(**_train_kwargs_for(trainer, train_kwargs)) + # Reached only if the service accepted a request we expected it to + # refuse. Capture the job so the finally-block can stop it, then let + # pytest.raises report the missing exception. + training_job = _resolve_job(trainer) + finally: + stop_quietly(training_job) + + message = str(excinfo.value) + assert any(token in message for token in expected_tokens), ( + f"request was rejected, but not for the expected reason.\n" + f" expected one of: {sorted(expected_tokens)}\n" + f" actual: {message}" + ) + return message diff --git a/sagemaker-train/tests/integ/train/shallow/test_model_trainer_submission.py b/sagemaker-train/tests/integ/train/shallow/test_model_trainer_submission.py new file mode 100644 index 0000000000..13776aac69 --- /dev/null +++ b/sagemaker-train/tests/integ/train/shallow/test_model_trainer_submission.py @@ -0,0 +1,674 @@ +# Copyright Amazon.com, Inc. or its affiliates. All Rights Reserved. +# +# Licensed under the Apache License, Version 2.0 (the "License"). You +# may not use this file except in compliance with the License. A copy of +# the License is located at +# +# http://aws.amazon.com/apache2.0/ +# +# or in the "license" file accompanying this file. This file is +# distributed on an "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF +# ANY KIND, either express or implied. See the License for the specific +# language governing permissions and limitations under the License. +"""Shallow submission tests for ``ModelTrainer``. + +Each test submits a real ``CreateTrainingJob``, asserts the service returned a +TrainingJobArn, then stops the job. A returned ARN proves the SDK-shaped payload +cleared every synchronous server-side gate (model validation, IAM authorization, +PassRole, the backend's request validators, S3/ECR resolution, routing) -- see +``harness`` for the full reasoning. + +These tests assert acceptance, never training behaviour. Anything that requires +a job to actually run belongs in the deep suites. +""" + +from __future__ import absolute_import + +import os + +import pytest +from sagemaker.core import shapes +from sagemaker.core.training.configs import Compute, InputData, Networking, SourceCode +from sagemaker.train.distributed import MPI, DistributedConfig, Torchrun +from sagemaker.train.model_trainer import ModelTrainer + +from .harness import ( + CPU_IMAGE, + DEFAULT_INSTANCE_COUNT, + DEFAULT_INSTANCE_TYPE, + MAX_RUNTIME_IN_SECONDS, + assert_rejected, + assert_submitted, + stop_quietly, + submitted, + unique_name, +) + +DATA_DIR = os.path.join(os.path.dirname(__file__), "..", "..", "..", "data") +PARAM_SCRIPT_SOURCE_DIR = os.path.join(DATA_DIR, "params_script") + +# Mirrors the hyperparameter contract asserted by the existing deep suite, so a +# serialization regression is caught here (cheaply, on every PR) rather than only +# in the slow tests. +CONTRACT_HYPERPARAMETERS = { + "integer": 1, + "boolean": True, + "float": 3.14, + "string": "Hello World", + "list": [1, 2, 3], + "dict": { + "string": "value", + "integer": 3, + "float": 3.14, + "list": [1, 2, 3], + "dict": {"key": "value"}, + "boolean": True, + }, +} + + +def _source_code(): + """Source code bundle used by most tests here. + + A real local source_dir is used (rather than a stub) because the SDK tars and + uploads it to S3 during submission, and the backend then validates that S3 + location. Skipping it would skip a real part of the path. + """ + return SourceCode( + source_dir=PARAM_SCRIPT_SOURCE_DIR, + requirements="requirements.txt", + entry_script="train.py", + ) + + +def _compute(instance_type=DEFAULT_INSTANCE_TYPE, instance_count=DEFAULT_INSTANCE_COUNT): + """Small CPU compute config. Never sets keep_alive_period_in_seconds -- a warm + pool would outlive the stop and keep instances provisioned.""" + return Compute(instance_type=instance_type, instance_count=instance_count) + + +def _stopping_condition(): + return shapes.StoppingCondition(max_runtime_in_seconds=MAX_RUNTIME_IN_SECONDS) + + +def _trainer(sagemaker_session, name, **overrides): + """Build a ModelTrainer with the minimum viable accepted configuration. + + Centralised so that a change to what "minimally valid" means is a one-line + edit rather than a sweep across every test. + """ + kwargs = dict( + sagemaker_session=sagemaker_session, + training_image=CPU_IMAGE, + source_code=_source_code(), + compute=_compute(), + stopping_condition=_stopping_condition(), + base_job_name=name, + ) + kwargs.update(overrides) + return ModelTrainer(**kwargs) + + +class TestMinimalSubmission: + """The baseline: does the simplest well-formed request get accepted? + + If these fail, everything else in the suite is noise -- they isolate "can we + talk to the service at all with a valid payload" from the feature-specific + tests below. + """ + + def test_minimal_request_is_accepted(self, sagemaker_session): + name = unique_name("shallow-minimal") + trainer = _trainer(sagemaker_session, name) + + with submitted(trainer) as job: + assert_submitted(job) + + def test_explicit_job_name_is_honoured(self, sagemaker_session): + """The name we ask for is the name that gets created. + + Guards against the SDK silently rewriting or regenerating job names, + which would break every user script that reconstructs an ARN from a name. + """ + name = unique_name("shallow-named") + trainer = _trainer(sagemaker_session, name) + + with submitted(trainer) as job: + arn = assert_submitted(job) + # base_job_name is a prefix; the SDK appends a timestamp suffix. + assert ( + name in job.training_job_name + ), f"requested base name {name!r} absent from {job.training_job_name!r}" + assert job.training_job_name in arn + + def test_explicit_role_is_accepted(self, sagemaker_session, execution_role): + """An explicitly passed role must pass PassRole server-side. + + The default path resolves the role implicitly; this proves the explicit + path produces a payload the service also accepts. + """ + name = unique_name("shallow-explicit-role") + trainer = _trainer(sagemaker_session, name, role=execution_role) + + with submitted(trainer) as job: + assert_submitted(job) + + def test_command_instead_of_entry_script(self, sagemaker_session): + """SourceCode.command is an alternative to entry_script; both must submit.""" + name = unique_name("shallow-command") + source_code = SourceCode( + source_dir=PARAM_SCRIPT_SOURCE_DIR, + requirements="requirements.txt", + command="python train.py", + ) + trainer = _trainer(sagemaker_session, name, source_code=source_code) + + with submitted(trainer) as job: + assert_submitted(job) + + +class TestSourceCodePackaging: + """How ``source_code`` is packaged and uploaded before submission. + + Each variant produces a different S3 artifact, and the backend's + role-assuming validators resolve that artifact -- so a packaging regression + surfaces as a rejected request rather than a silent difference. + + Mirrors the source-code cases in the existing ``test_model_trainer.py`` deep + suite (local tar file, shell entry script, custom distributed driver) so + replacing it on the PR gate does not drop them. + """ + + def test_local_tar_file_source_dir(self, sagemaker_session): + """A pre-built local ``.tar.gz`` is uploaded as-is rather than re-tarred.""" + name = unique_name("shallow-tar-source") + source_code = SourceCode( + source_dir=os.path.join(DATA_DIR, "script_mode", "code.tar.gz"), + requirements="requirements.txt", + entry_script="custom_script.py", + ) + trainer = _trainer(sagemaker_session, name, source_code=source_code) + + with submitted(trainer) as job: + assert_submitted(job) + + def test_shell_entry_script(self, sagemaker_session): + """A ``.sh`` entry script takes a different container-entrypoint path + from a ``.py`` one.""" + name = unique_name("shallow-sh-entry") + source_code = SourceCode( + source_dir=PARAM_SCRIPT_SOURCE_DIR, + requirements="requirements.txt", + entry_script="train.sh", + ) + trainer = _trainer( + sagemaker_session, + name, + source_code=source_code, + hyperparameters=CONTRACT_HYPERPARAMETERS, + ) + + with submitted(trainer) as job: + assert_submitted(job) + + def test_custom_distributed_driver(self, sagemaker_session): + """A user-supplied distributed driver is uploaded alongside the source + and changes the container entrypoint. + + Ported from ``test_model_trainer.py::test_custom_distributed_driver``: + the driver directory is packaged separately from ``source_dir``, so this + exercises a second upload the other tests never trigger. + """ + + class CustomDriver(DistributedConfig): + process_count_per_node: int = None + + @property + def driver_dir(self) -> str: + return os.path.join(DATA_DIR, "custom_drivers") + + @property + def driver_script(self) -> str: + return "driver.py" + + name = unique_name("shallow-custom-driver") + source_code = SourceCode( + source_dir=os.path.join(DATA_DIR, "scripts"), + entry_script="entry_script.py", + ) + trainer = _trainer( + sagemaker_session, + name, + source_code=source_code, + hyperparameters={"epochs": 1}, + distributed=CustomDriver(process_count_per_node=2), + ) + + with submitted(trainer) as job: + assert_submitted(job) + + +class TestPayloadShaping: + """Fields the SDK must serialize into a form the service accepts. + + These are the highest-value tests in the suite: they are exactly the + regressions that unit tests miss (because a mock accepts anything) and that + deep integ tests catch far too slowly and expensively. + """ + + def test_hyperparameters_contract(self, sagemaker_session): + """Nested/typed hyperparameters must survive serialization. + + The service requires a flat string->string map, so the SDK has to encode + ints, floats, bools, lists and nested dicts. A regression here is a + ValidationException at submit time, which is precisely what this catches. + """ + name = unique_name("shallow-hp-contract") + trainer = _trainer(sagemaker_session, name, hyperparameters=CONTRACT_HYPERPARAMETERS) + + with submitted(trainer) as job: + assert_submitted(job) + + def test_hyperparameters_from_json_file(self, sagemaker_session): + """Hyperparameters given as a path to JSON must load and serialize.""" + name = unique_name("shallow-hp-json") + trainer = _trainer( + sagemaker_session, + name, + hyperparameters=os.path.join(PARAM_SCRIPT_SOURCE_DIR, "hyperparameters.json"), + ) + + with submitted(trainer) as job: + assert_submitted(job) + + def test_hyperparameters_from_yaml_file(self, sagemaker_session): + """Hyperparameters given as a path to YAML must load and serialize.""" + name = unique_name("shallow-hp-yaml") + trainer = _trainer( + sagemaker_session, + name, + hyperparameters=os.path.join(PARAM_SCRIPT_SOURCE_DIR, "hyperparameters.yaml"), + ) + + with submitted(trainer) as job: + assert_submitted(job) + + def test_environment_variables(self, sagemaker_session): + """Environment map must be accepted (the backend validates key syntax).""" + name = unique_name("shallow-env") + trainer = _trainer( + sagemaker_session, + name, + environment={"MY_SETTING": "value", "ANOTHER_SETTING": "42"}, + ) + + with submitted(trainer) as job: + assert_submitted(job) + + def test_tags_are_accepted(self, sagemaker_session): + """Tags travel a distinct authorization path. + + Tag-on-create is enforced by an interceptor at the public front end and + by tag-governance checks, so a tagged request exercises gates an untagged + one never reaches. + """ + name = unique_name("shallow-tags") + trainer = _trainer( + sagemaker_session, + name, + tags=[shapes.Tag(key="Purpose", value="shallow-integ-test")], + ) + + with submitted(trainer) as job: + assert_submitted(job) + + def test_explicit_output_data_config(self, sagemaker_session, output_path): + """A caller-specified output location must validate server-side.""" + name = unique_name("shallow-output") + trainer = _trainer( + sagemaker_session, + name, + output_data_config=shapes.OutputDataConfig(s3_output_path=output_path), + ) + + with submitted(trainer) as job: + assert_submitted(job) + + @pytest.mark.parametrize("input_mode", ["File", "FastFile", "Pipe"]) + def test_training_input_modes(self, sagemaker_session, input_mode): + """Every advertised input mode must be accepted. + + Cheap to cover here and easy to break: the mode is validated server-side + against the channel configuration. + """ + name = unique_name(f"shallow-mode-{input_mode.lower()}") + trainer = _trainer(sagemaker_session, name, training_input_mode=input_mode) + + with submitted(trainer) as job: + assert_submitted(job) + + +class TestInputDataConfiguration: + """Input channels are resolved against S3 by the backend's role-assuming + validators, so these tests prove both serialization and real S3 reachability + under the execution role.""" + + def test_single_s3_channel(self, sagemaker_session, train_data_uri): + name = unique_name("shallow-one-channel") + trainer = _trainer(sagemaker_session, name) + + with submitted( + trainer, + input_data_config=[InputData(channel_name="train", data_source=train_data_uri)], + ) as job: + assert_submitted(job) + + def test_multiple_s3_channels(self, sagemaker_session, train_data_uri, validation_data_uri): + """Multiple channels must each resolve; channel-name rules are enforced + server-side.""" + name = unique_name("shallow-two-channels") + trainer = _trainer(sagemaker_session, name) + + with submitted( + trainer, + input_data_config=[ + InputData(channel_name="train", data_source=train_data_uri), + InputData(channel_name="validation", data_source=validation_data_uri), + ], + ) as job: + assert_submitted(job) + + def test_channel_with_content_type(self, sagemaker_session, train_data_uri): + name = unique_name("shallow-content-type") + trainer = _trainer(sagemaker_session, name) + + with submitted( + trainer, + input_data_config=[ + InputData( + channel_name="train", + data_source=train_data_uri, + content_type="application/jsonlines", + ) + ], + ) as job: + assert_submitted(job) + + def test_s3_data_source_object(self, sagemaker_session, train_data_uri): + """An explicit S3DataSource shape (rather than a bare URI) must serialize + into a payload the service accepts.""" + name = unique_name("shallow-s3-datasource") + trainer = _trainer(sagemaker_session, name) + data_source = shapes.S3DataSource( + s3_data_type="S3Prefix", + s3_uri=train_data_uri, + s3_data_distribution_type="FullyReplicated", + ) + + with submitted( + trainer, + input_data_config=[InputData(channel_name="train", data_source=data_source)], + ) as job: + assert_submitted(job) + + +class TestCheckpointingAndSpot: + """Checkpointing and managed spot each add fields with their own backend + validators, and spot additionally requires MaxWaitTimeInSeconds >= + MaxRuntimeInSeconds -- a cross-field rule only the service enforces.""" + + def test_checkpoint_config(self, sagemaker_session, output_path): + """CheckpointConfig has a dedicated validator and an S3 location the + backend resolves.""" + name = unique_name("shallow-checkpoint") + trainer = _trainer( + sagemaker_session, + name, + checkpoint_config=shapes.CheckpointConfig( + s3_uri=f"{output_path}checkpoints/", + local_path="/opt/ml/checkpoints/", + ), + ) + + with submitted(trainer) as job: + assert_submitted(job) + + def test_managed_spot_training(self, sagemaker_session): + """Managed spot requires a max wait time at least as large as the max + runtime; the service rejects the combination otherwise. + + Note this deliberately does not set ``keep_alive_period_in_seconds``: + spot and warm pools are mutually exclusive, and a warm pool would outlive + the stop. + """ + name = unique_name("shallow-spot") + compute = Compute( + instance_type=DEFAULT_INSTANCE_TYPE, + instance_count=DEFAULT_INSTANCE_COUNT, + enable_managed_spot_training=True, + ) + trainer = _trainer( + sagemaker_session, + name, + compute=compute, + stopping_condition=shapes.StoppingCondition( + max_runtime_in_seconds=MAX_RUNTIME_IN_SECONDS, + max_wait_time_in_seconds=MAX_RUNTIME_IN_SECONDS, + ), + ) + + with submitted(trainer) as job: + assert_submitted(job) + + +class TestComputeConfiguration: + """Compute shapes are validated by several distinct backend validators + (instance type, instance count, volume size, distribution).""" + + def test_multi_instance_request(self, sagemaker_session): + """instance_count > 1 changes the accepted shape of the request.""" + name = unique_name("shallow-multi-instance") + trainer = _trainer(sagemaker_session, name, compute=_compute(instance_count=2)) + + with submitted(trainer) as job: + assert_submitted(job) + + def test_explicit_volume_size(self, sagemaker_session): + """Volume size has its own validator with min/max bounds.""" + name = unique_name("shallow-volume") + compute = Compute( + instance_type=DEFAULT_INSTANCE_TYPE, + instance_count=DEFAULT_INSTANCE_COUNT, + volume_size_in_gb=50, + ) + trainer = _trainer(sagemaker_session, name, compute=compute) + + with submitted(trainer) as job: + assert_submitted(job) + + def test_torchrun_distributed(self, sagemaker_session): + """Distributed configs inject env/entrypoint changes; the resulting + payload must still be accepted.""" + name = unique_name("shallow-torchrun") + trainer = _trainer( + sagemaker_session, + name, + compute=_compute(instance_count=2), + distributed=Torchrun(), + ) + + with submitted(trainer) as job: + assert_submitted(job) + + def test_mpi_distributed(self, sagemaker_session): + name = unique_name("shallow-mpi") + trainer = _trainer( + sagemaker_session, + name, + compute=_compute(instance_count=2), + distributed=MPI(), + ) + + with submitted(trainer) as job: + assert_submitted(job) + + +class TestNetworkingAndSecurity: + """Isolation and encryption flags are surfaced as IAM condition keys, so + these requests are authorized differently from the baseline.""" + + def test_network_isolation(self, sagemaker_session): + name = unique_name("shallow-net-isolation") + trainer = _trainer( + sagemaker_session, name, networking=Networking(enable_network_isolation=True) + ) + + with submitted(trainer) as job: + assert_submitted(job) + + def test_inter_container_traffic_encryption(self, sagemaker_session): + """Encryption between nodes only applies to multi-instance jobs.""" + name = unique_name("shallow-icte") + trainer = _trainer( + sagemaker_session, + name, + compute=_compute(instance_count=2), + networking=Networking(enable_inter_container_traffic_encryption=True), + ) + + with submitted(trainer) as job: + assert_submitted(job) + + +class TestRejectedRequests: + """Negative cases. + + Without these the suite would pass as long as *something* was accepted, + which would hide a bug that made the SDK send a permissive-but-wrong + payload. Each case asserts a specific rejection, and the harness stops the + job if one is unexpectedly accepted. + """ + + def test_nonexistent_input_data_is_rejected(self, sagemaker_session, nonexistent_data_uri): + """Proves input validation genuinely reaches S3. + + The single most valuable negative test here: it is the assertion that the + expensive role-assuming validators actually ran, rather than being + skipped or silently swallowed. + """ + trainer = _trainer(sagemaker_session, unique_name("shallow-bad-input")) + + assert_rejected( + trainer, + ("does not exist", "ValidationException", "ValidationError", "S3", "not found"), + input_data_config=[InputData(channel_name="train", data_source=nonexistent_data_uri)], + ) + + def test_invalid_instance_type_is_rejected(self, sagemaker_session): + """A syntactically-valid but nonexistent instance type must be refused.""" + trainer = _trainer( + sagemaker_session, + unique_name("shallow-bad-instance"), + compute=_compute(instance_type="ml.nonexistent.xlarge"), + ) + + assert_rejected( + trainer, + ("instance", "Instance", "ValidationException", "ValidationError", "not supported"), + ) + + def test_nonexistent_training_image_is_rejected(self, sagemaker_session, account_id, region): + """The backend resolves the training image against ECR under the + customer's role, so an image that does not exist must be refused. + + Uses the caller's own account so the failure is "repository absent" + rather than "cross-account access denied". + """ + bogus_image = ( + f"{account_id}.dkr.ecr.{region}.amazonaws.com/" "shallow-integ-test-no-such-repo:latest" + ) + trainer = _trainer( + sagemaker_session, unique_name("shallow-bad-image"), training_image=bogus_image + ) + + assert_rejected( + trainer, + ( + "image", + "Image", + "ECR", + "repository", + "RepositoryNotFound", + "ValidationException", + "ValidationError", + ), + ) + + def test_unassumable_role_is_rejected(self, sagemaker_session, account_id): + """PassRole / AssumeRole failures must surface at submit time. + + Directly covers the "does the caller hold the required permissions" half + of what this suite exists to assert. + """ + bogus_role = f"arn:aws:iam::{account_id}:role/shallow-integ-test-no-such-role" + trainer = _trainer(sagemaker_session, unique_name("shallow-bad-role"), role=bogus_role) + + assert_rejected( + trainer, + ( + "role", + "Role", + "AccessDenied", + "not authorized", + "cannot be assumed", + "ValidationException", + "ValidationError", + ), + ) + + def test_duplicate_job_name_is_rejected(self, sagemaker_session, execution_role, output_path): + """The final gate before the ARN is a conditional write that rejects + duplicate job names with ResourceInUse. + + Asserting it proves a submission reached the very *end* of the create + path -- the durable write -- and not merely the validators in front of + it. ``ModelTrainer`` appends a timestamp to ``base_job_name``, so it can + never produce a collision by design; this drives the underlying resource + API directly in order to re-use one exact name twice. + """ + from sagemaker.core.resources import TrainingJob + + job_name = unique_name("shallow-duplicate") + + def create(): + return TrainingJob.create( + session=sagemaker_session.boto_session, + training_job_name=job_name, + role_arn=execution_role, + algorithm_specification=shapes.AlgorithmSpecification( + training_image=CPU_IMAGE, training_input_mode="File" + ), + output_data_config=shapes.OutputDataConfig(s3_output_path=output_path), + resource_config=shapes.ResourceConfig( + instance_type=DEFAULT_INSTANCE_TYPE, + instance_count=DEFAULT_INSTANCE_COUNT, + volume_size_in_gb=30, + ), + stopping_condition=_stopping_condition(), + ) + + first = None + try: + first = create() + assert_submitted(first, expected_name=job_name) + + with pytest.raises(Exception) as excinfo: + create() + + message = str(excinfo.value) + assert any( + token in message + for token in ("already exists", "ResourceInUse", "ResourceInUseException") + ), f"unexpected rejection reason: {message}" + finally: + stop_quietly(first) diff --git a/sagemaker-train/tests/integ/train/shallow/test_other_job_types_submission.py b/sagemaker-train/tests/integ/train/shallow/test_other_job_types_submission.py new file mode 100644 index 0000000000..ae6d0eb6f7 --- /dev/null +++ b/sagemaker-train/tests/integ/train/shallow/test_other_job_types_submission.py @@ -0,0 +1,251 @@ +# Copyright Amazon.com, Inc. or its affiliates. All Rights Reserved. +# +# Licensed under the Apache License, Version 2.0 (the "License"). You +# may not use this file except in compliance with the License. A copy of +# the License is located at +# +# http://aws.amazon.com/apache2.0/ +# +# or in the "license" file accompanying this file. This file is +# distributed on an "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF +# ANY KIND, either express or implied. See the License for the specific +# language governing permissions and limitations under the License. +"""Shallow submission tests for job types that are not plain training jobs. + +The rest of this suite covers ``CreateTrainingJob``. Two trainers in this package +create something else, and each needed harness support rather than being +genuinely un-testable: + +* ``HyperparameterTuner.tune()`` creates a **HyperParameterTuningJob**. The + service validates the embedded training-job definition (including the + ``sm_drivers`` channel for distributed runs) plus tuning-specific rules -- + objective metric, parameter ranges, max jobs/parallel jobs. Stopping is + ``tuner.stop_tuning_job()``. +* ``MultiTurnRLTrainer.train()`` creates an **AgentRFT Job** via the generic Job + API, not ``CreateTrainingJob``. It returns an ``AgentRFTJob`` exposing + ``job_arn``/``job_name``/``stop()``. + +Both are covered here because "different resource type" is a reason to teach the +harness a new ARN shape, not a reason to skip the coverage. + +The tuner tests carry the real weight: they run on CPU with no external +prerequisites. The MTRL tests are marked ``gpu_intensive`` and skip when their +prerequisites are absent -- see ``TestMultiTurnRLSubmission`` for why. +""" + +from __future__ import absolute_import + +import logging +import os + +import pytest +from sagemaker.core import shapes +from sagemaker.core.parameter import ContinuousParameter +from sagemaker.core.training.configs import Compute, SourceCode +from sagemaker.train.distributed import Torchrun +from sagemaker.train.model_trainer import ModelTrainer +from sagemaker.train.multi_turn_rl_trainer import MultiTurnRLTrainer +from sagemaker.train.tuner import HyperparameterTuner + +from .harness import ( + CPU_IMAGE, + DEFAULT_INSTANCE_COUNT, + DEFAULT_INSTANCE_TYPE, + MAX_RUNTIME_IN_SECONDS, + assert_submitted, + submitted, + unique_name, +) + +logger = logging.getLogger(__name__) + +DATA_DIR = os.path.join(os.path.dirname(__file__), "..", "..", "..", "data") +PARAM_SCRIPT_SOURCE_DIR = os.path.join(DATA_DIR, "params_script") + + +def _model_trainer(sagemaker_session, name, **overrides): + """The inner trainer a tuning job wraps.""" + kwargs = dict( + sagemaker_session=sagemaker_session, + training_image=CPU_IMAGE, + base_job_name=name, + source_code=SourceCode( + source_dir=PARAM_SCRIPT_SOURCE_DIR, + requirements="requirements.txt", + entry_script="train.py", + ), + compute=Compute( + instance_type=DEFAULT_INSTANCE_TYPE, + instance_count=DEFAULT_INSTANCE_COUNT, + volume_size_in_gb=30, + ), + stopping_condition=shapes.StoppingCondition(max_runtime_in_seconds=MAX_RUNTIME_IN_SECONDS), + hyperparameters={"learning_rate": 1e-4}, + ) + kwargs.update(overrides) + return ModelTrainer(**kwargs) + + +def _tuner(model_trainer, **overrides): + """A minimal single-job tuner. + + ``max_jobs=1`` / ``max_parallel_jobs=1`` keeps the blast radius to one child + training job, which is stopped along with the tuning job. + """ + kwargs = dict( + model_trainer=model_trainer, + objective_metric_name="eval_loss", + metric_definitions=[{"Name": "eval_loss", "Regex": r"eval_loss: ([0-9\\.]+)"}], + hyperparameter_ranges={ + "learning_rate": ContinuousParameter( + min_value=1e-5, max_value=5e-4, scaling_type="Logarithmic" + ) + }, + objective_type="Minimize", + max_jobs=1, + max_parallel_jobs=1, + ) + kwargs.update(overrides) + return HyperparameterTuner(**kwargs) + + +class TestTuningJobSubmission: + """HyperParameterTuningJob acceptance. + + Stopping a tuning job also stops its child training jobs, so the same + submit-then-stop economics apply. + """ + + def test_minimal_tuning_job_is_accepted(self, sagemaker_session): + """Baseline: the service accepts a well-formed tuning job.""" + name = unique_name("shallow-tuner") + tuner = _tuner(_model_trainer(sagemaker_session, name)) + + try: + tuner.tune(wait=False) + assert_submitted(tuner.latest_tuning_job, resource="hyper-parameter-tuning-job") + finally: + # Tuner exposes its own stop method rather than the resource's. + try: + tuner.stop_tuning_job() + except Exception as e: # pragma: no cover - best-effort teardown + logger.warning("Could not stop tuning job: %s", e) + + def test_distributed_tuning_job_is_accepted(self, sagemaker_session): + """A tuning job wrapping a Torchrun trainer must include the + ``sm_drivers`` channel in its training-job definition. + + This is the regression the existing ``test_tuner_distributed.py`` guards + by running a job to completion and inspecting logs. Submission alone + proves the channel is present and the definition is accepted, which is + the part that regressed; the log assertion stays in the deep suite. + """ + name = unique_name("shallow-tuner-dist") + model_trainer = _model_trainer(sagemaker_session, name, distributed=Torchrun()) + tuner = _tuner(model_trainer) + + try: + tuner.tune(wait=False) + arn = assert_submitted(tuner.latest_tuning_job, resource="hyper-parameter-tuning-job") + + # The sm_drivers channel lives in the tuning job's training + # definition; read it back to prove it survived submission rather + # than inferring from acceptance alone. + described = tuner.latest_tuning_job.refresh() + definition = getattr(described, "training_job_definition", None) + if definition is not None: + channels = [ + channel.channel_name for channel in (definition.input_data_config or []) + ] + assert "sm_drivers" in channels, ( + f"tuning job {arn} is missing the sm_drivers channel; " f"channels={channels}" + ) + finally: + try: + tuner.stop_tuning_job() + except Exception as e: # pragma: no cover - best-effort teardown + logger.warning("Could not stop tuning job: %s", e) + + +@pytest.mark.gpu_intensive +class TestMultiTurnRLSubmission: + """AgentRFT Job acceptance for ``MultiTurnRLTrainer``. + + Marked ``gpu_intensive`` (and therefore excluded from the PR gate, per the + marker's definition in ``tox.ini``) because unlike every other test in this + suite it cannot be made self-contained: MTRL requires a pre-provisioned agent + runtime and an MLflow app, neither of which this suite creates. The existing + ``test_multi_turn_rl_trainer_integration.py`` hardcodes both. + + They are still written using the shallow pattern rather than omitted, so that + when the prerequisites are provisioned in the PR account these become + PR-gate-eligible by deleting one marker. Prerequisites are resolved from the + environment and the tests skip when absent, so they never fail for + infrastructure reasons. + """ + + @pytest.fixture(scope="class") + def mtrl_prerequisites(self, sagemaker_session, account_id, region): + """Resolve MTRL prerequisites, skipping if they are not configured. + + Read from the environment rather than hardcoded so this does not bake in + another account-specific constant. + """ + agent_env = os.environ.get("SHALLOW_MTRL_AGENT_ENV") + mlflow_app_arn = os.environ.get("SHALLOW_MTRL_MLFLOW_APP_ARN") + dataset = os.environ.get("SHALLOW_MTRL_DATASET") + + missing = [ + name + for name, value in ( + ("SHALLOW_MTRL_AGENT_ENV", agent_env), + ("SHALLOW_MTRL_MLFLOW_APP_ARN", mlflow_app_arn), + ("SHALLOW_MTRL_DATASET", dataset), + ) + if not value + ] + if missing: + pytest.skip("MTRL prerequisites not configured; set " + ", ".join(missing)) + + return { + "agent_env": agent_env, + "mlflow_app_arn": mlflow_app_arn, + "dataset": dataset, + "model": os.environ.get("SHALLOW_MTRL_MODEL", "mock-oss-test"), + } + + def test_agent_rft_job_is_accepted(self, sagemaker_session, mtrl_prerequisites): + """The AgentRFT job config document must be accepted by the Job API. + + Note the different ARN resource segment: this is a ``job``, not a + ``training-job``. + """ + trainer = MultiTurnRLTrainer( + model=mtrl_prerequisites["model"], + agent_env=mtrl_prerequisites["agent_env"], + training_dataset=mtrl_prerequisites["dataset"], + mlflow_app_arn=mtrl_prerequisites["mlflow_app_arn"], + accept_eula=True, + sagemaker_session=sagemaker_session, + base_job_name=unique_name("shallow-mtrl"), + ) + + with submitted(trainer) as job: + assert_submitted(job, resource="job") + + def test_hyperparameter_mutation_is_accepted(self, sagemaker_session, mtrl_prerequisites): + """``trainer.hyperparameters`` mutation must reach the job config + document, which the service validates on submission.""" + trainer = MultiTurnRLTrainer( + model=mtrl_prerequisites["model"], + agent_env=mtrl_prerequisites["agent_env"], + training_dataset=mtrl_prerequisites["dataset"], + mlflow_app_arn=mtrl_prerequisites["mlflow_app_arn"], + accept_eula=True, + sagemaker_session=sagemaker_session, + base_job_name=unique_name("shallow-mtrl-hp"), + ) + trainer.hyperparameters.global_batch_size = 32 + + with submitted(trainer) as job: + assert_submitted(job, resource="job") diff --git a/sagemaker-train/tests/integ/train/shallow/test_recipe_customization_submission.py b/sagemaker-train/tests/integ/train/shallow/test_recipe_customization_submission.py new file mode 100644 index 0000000000..9852d7c4dd --- /dev/null +++ b/sagemaker-train/tests/integ/train/shallow/test_recipe_customization_submission.py @@ -0,0 +1,250 @@ +# Copyright Amazon.com, Inc. or its affiliates. All Rights Reserved. +# +# Licensed under the Apache License, Version 2.0 (the "License"). You +# may not use this file except in compliance with the License. A copy of +# the License is located at +# +# http://aws.amazon.com/apache2.0/ +# +# or in the "license" file accompanying this file. This file is +# distributed on an "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF +# ANY KIND, either express or implied. See the License for the specific +# language governing permissions and limitations under the License. +"""Shallow submission tests for recipe customization. + +Covers the knobs that change the *rendered recipe* rather than the plain request +envelope: ``overrides``, an explicit ``recipe`` file, ``sequence_length``, +``DataMixingConfig``, and direct ``trainer.hyperparameters`` mutation. + +Why these matter here specifically +---------------------------------- +The training backend does not merely shape-check a recipe request -- it filters +candidate recipes after the request validators have already passed and refuses +the job outright with ``"No valid recipes found for the given request"`` when +nothing matches. A customization that renders into an unsatisfiable recipe is +therefore *only* detectable by actually submitting. + +Existing coverage of this area is either client-side or expensive: + +* ``test_recipe_override_integration.py`` (35 tests) exercises + ``get_resolved_recipe`` / ``flatten_resolved_recipe`` and never submits, so it + cannot catch a recipe that resolves locally but the service rejects. +* ``test_rlvr_trainer_integration.py::test_rlvr_trainer_nemotron_with_kl_and_recipe`` + does submit a recipe+overrides combination, but on a 30B model with a + two-hour poll loop. + +These tests close that gap at submission cost: they prove the customized payload +is *accepted*, without asserting anything about the resulting training run. +""" + +from __future__ import absolute_import + +import tempfile + +import pytest +import yaml +from sagemaker.core import shapes +from sagemaker.train.common import TrainingType +from sagemaker.train.data_mixing_config import DataMixingConfig +from sagemaker.train.rlvr_trainer import RLVRTrainer +from sagemaker.train.sft_trainer import SFTTrainer + +from .harness import MAX_RUNTIME_IN_SECONDS, assert_submitted, submitted, unique_name + +MODEL_ID = "meta-textgeneration-llama-3-2-1b-instruct" +MODEL_PACKAGE_GROUP = ( + "arn:aws:sagemaker:us-west-2:729646638167:model-package-group/sdk-test-finetuned-models" +) + + +def _stopping_condition(): + return shapes.StoppingCondition(max_runtime_in_seconds=MAX_RUNTIME_IN_SECONDS) + + +def _sft(sagemaker_session, dataset, name, **overrides): + kwargs = dict( + model=MODEL_ID, + training_type=TrainingType.LORA, + model_package_group=MODEL_PACKAGE_GROUP, + training_dataset=dataset, + accept_eula=True, + sagemaker_session=sagemaker_session, + base_job_name=name, + stopping_condition=_stopping_condition(), + ) + kwargs.update(overrides) + return SFTTrainer(**kwargs) + + +class TestRecipeOverrides: + """``overrides`` is merged into the rendered recipe before submission.""" + + def test_training_config_overrides(self, sagemaker_session, train_data_uri): + """Override common training_config values. + + Values are chosen to stay inside the recipe's accepted ranges: the point + is to prove overrides survive rendering into an accepted payload, not to + probe validation bounds (which the negative tests cover). + """ + trainer = _sft( + sagemaker_session, + train_data_uri, + unique_name("shallow-sft-overrides"), + overrides={ + "training_config": { + "learning_rate": 2e-5, + "max_epochs": 1, + } + }, + ) + + with submitted(trainer) as job: + assert_submitted(job) + + def test_explicit_recipe_file(self, sagemaker_session, train_data_uri): + """A caller-supplied recipe YAML must render into an accepted request. + + Mirrors the shape used by the existing Nemotron test, but on a small + model and without the two-hour poll loop. + """ + recipe = {"training_config": {"data": {"max_prompt_length": 1024}}} + with tempfile.NamedTemporaryFile(mode="w", suffix=".yaml", delete=False) as handle: + yaml.dump(recipe, handle) + recipe_path = handle.name + + trainer = _sft( + sagemaker_session, + train_data_uri, + unique_name("shallow-sft-recipe"), + recipe=recipe_path, + ) + + with submitted(trainer) as job: + assert_submitted(job) + + def test_recipe_and_overrides_together(self, sagemaker_session, train_data_uri): + """Recipe file plus overrides: the merge order must still yield an + accepted payload. This is the combination most likely to break, since + both paths mutate the same rendered document.""" + recipe = {"training_config": {"data": {"max_prompt_length": 1024}}} + with tempfile.NamedTemporaryFile(mode="w", suffix=".yaml", delete=False) as handle: + yaml.dump(recipe, handle) + recipe_path = handle.name + + trainer = _sft( + sagemaker_session, + train_data_uri, + unique_name("shallow-sft-recipe-ovr"), + recipe=recipe_path, + overrides={"training_config": {"max_epochs": 1}}, + ) + + with submitted(trainer) as job: + assert_submitted(job) + + def test_direct_hyperparameter_mutation(self, sagemaker_session, train_data_uri): + """``trainer.hyperparameters. = ...`` is a documented pattern + (used by the existing RLVR tests) and must reach the payload intact.""" + trainer = RLVRTrainer( + model=MODEL_ID, + training_type=TrainingType.LORA, + model_package_group=MODEL_PACKAGE_GROUP, + training_dataset=train_data_uri, + accept_eula=True, + sagemaker_session=sagemaker_session, + base_job_name=unique_name("shallow-rlvr-hpmutate"), + stopping_condition=_stopping_condition(), + ) + trainer.hyperparameters.max_epochs = 1 + + with submitted(trainer) as job: + assert_submitted(job) + + +class TestSequenceLength: + """``sequence_length`` selects a different recipe variant, so each supported + value is a distinct accepted-payload case.""" + + @pytest.mark.parametrize("sequence_length", ["4K", "16K"]) + def test_sequence_length_variants(self, sagemaker_session, train_data_uri, sequence_length): + trainer = _sft( + sagemaker_session, + train_data_uri, + unique_name(f"shallow-sft-seq{sequence_length}"), + sequence_length=sequence_length, + ) + + with submitted(trainer) as job: + assert_submitted(job) + + +class TestDataMixing: + """``DataMixingConfig`` is serialized into flat per-category hyperparameters. + + Nova-only, and Nova is us-east-1 in this repo's fixtures, so these use + ``sagemaker_session_us_east_1`` (inherited from the parent train conftest) + rather than the default-region session. + + Marked ``us_east_1`` to match the existing marker convention in + ``sagemaker-train/tox.ini``; the PR-gate job runs us-west-2 only, so these are + deselected there and run in the us-east-1 job. + """ + + NOVA_MODEL = "nova-textgeneration-lite-v2" + + @pytest.mark.us_east_1 + def test_data_mixing_with_explicit_percentages( + self, sagemaker_session_us_east_1, nova_train_data_uri + ): + """Per-category percentages must sum to 100 client-side and serialize + into hyperparameters the service accepts.""" + config = DataMixingConfig( + customer_data_percent=70.0, + nova_data_percentages={ + "code": 30.0, + "math": 20.0, + "planning": 10.0, + "instruction-following": 10.0, + "reasoning-instruction-following": 20.0, + "reasoning-math": 10.0, + }, + ) + name = unique_name("shallow-sft-datamix") + trainer = SFTTrainer( + model=self.NOVA_MODEL, + training_type=TrainingType.LORA, + model_package_group=MODEL_PACKAGE_GROUP, + training_dataset=nova_train_data_uri, + accept_eula=True, + sagemaker_session=sagemaker_session_us_east_1, + data_mixing_config=config, + base_job_name=name, + # The existing data-mixing test sets the recipe name explicitly; + # keep that so the rendered recipe matches what the service expects. + overrides={"name": name}, + ) + + with submitted(trainer) as job: + assert_submitted(job) + + @pytest.mark.us_east_1 + def test_data_mixing_recipe_defaults(self, sagemaker_session_us_east_1, nova_train_data_uri): + """With ``nova_data_percentages=None`` the recipe template's defaults are + used at submission time -- a different serialization path from the + explicit case above.""" + config = DataMixingConfig(customer_data_percent=80.0) + name = unique_name("shallow-sft-datamix-default") + trainer = SFTTrainer( + model=self.NOVA_MODEL, + training_type=TrainingType.LORA, + model_package_group=MODEL_PACKAGE_GROUP, + training_dataset=nova_train_data_uri, + accept_eula=True, + sagemaker_session=sagemaker_session_us_east_1, + data_mixing_config=config, + base_job_name=name, + overrides={"name": name}, + ) + + with submitted(trainer) as job: + assert_submitted(job) diff --git a/sagemaker-train/tests/integ/train/shallow/test_recipe_trainers_submission.py b/sagemaker-train/tests/integ/train/shallow/test_recipe_trainers_submission.py new file mode 100644 index 0000000000..2df7995110 --- /dev/null +++ b/sagemaker-train/tests/integ/train/shallow/test_recipe_trainers_submission.py @@ -0,0 +1,351 @@ +# Copyright Amazon.com, Inc. or its affiliates. All Rights Reserved. +# +# Licensed under the Apache License, Version 2.0 (the "License"). You +# may not use this file except in compliance with the License. A copy of +# the License is located at +# +# http://aws.amazon.com/apache2.0/ +# +# or in the "license" file accompanying this file. This file is +# distributed on an "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF +# ANY KIND, either express or implied. See the License for the specific +# language governing permissions and limitations under the License. +"""Shallow submission tests for the recipe trainers (SFT / DPO / RLVR / CPT). + +These trainers do far more request-shaping than ``ModelTrainer``: they resolve a +foundation model, select and render a training recipe, derive a resource config +from it, and translate datasets into channels. All of that lands in the +``CreateTrainingJob`` payload, and the training backend validates it -- including +recipe *acceptance*, which is checked after the request validators and rejects +with ``"No valid recipes found for the given request"``. + +That makes submit-then-stop unusually valuable for this family: a recipe +regression is invisible to unit tests (which mock the service) and today is only +caught by a full, expensive training run. + +Serverless (recipe-selected compute) is the default path. Where a test pins +``TrainingJobCompute`` it is asserting the serverful path specifically, since the +two produce materially different payloads. +""" + +from __future__ import absolute_import + +import pytest +from sagemaker.core import shapes +from sagemaker.core.training.configs import TrainingJobCompute +from sagemaker.train.common import TrainingType +from sagemaker.train.cpt_trainer import CPTTrainer +from sagemaker.train.dpo_trainer import DPOTrainer +from sagemaker.train.rlaif_trainer import RLAIFTrainer +from sagemaker.train.rlvr_trainer import RLVRTrainer +from sagemaker.train.sft_trainer import SFTTrainer + +from .harness import ( + MAX_RUNTIME_IN_SECONDS, + assert_rejected, + assert_submitted, + submitted, + unique_name, +) + +# Small, publicly available instruct model. Kept small deliberately: these tests +# never train, so model size only affects how long recipe/artifact resolution +# takes during submission. +MODEL_ID = "meta-textgeneration-llama-3-2-1b-instruct" + +# Reused from the existing dry-run suite so both suites exercise the same +# already-provisioned model package group rather than each needing their own. +MODEL_PACKAGE_GROUP = ( + "arn:aws:sagemaker:us-west-2:729646638167:model-package-group/sdk-test-finetuned-models" +) + +# An accelerator type is required for the serverful recipe path: the recipes for +# these trainers will not resolve onto a CPU instance, so unlike the +# ModelTrainer suite we cannot use ml.m5.large here. The job is still stopped +# immediately, so this requests capacity only transiently. +SERVERFUL_INSTANCE_TYPE = "ml.g5.12xlarge" + +# RLAIF requires a reward model and prompt; without them the request is refused +# before it reaches the validation this suite cares about. Values match the +# existing test_rlaif_trainer_integration.py so both suites exercise the same +# already-entitled reward model. +RLAIF_REWARD_MODEL_ID = "openai.gpt-oss-120b-1:0" +RLAIF_REWARD_PROMPT = "Builtin.Summarize" + +# Per-trainer extra constructor arguments. Everything else is shared, which is +# what lets these four trainers be covered by one parametrized body instead of +# four near-identical files. +_TRAINER_EXTRA_KWARGS = { + "RLAIFTrainer": { + "reward_model_id": RLAIF_REWARD_MODEL_ID, + "reward_prompt": RLAIF_REWARD_PROMPT, + }, +} + +# Every recipe trainer takes the same core arguments, so the per-trainer test +# bodies stay a single call. +RECIPE_TRAINERS = [ + pytest.param(SFTTrainer, id="sft"), + pytest.param(DPOTrainer, id="dpo"), + pytest.param(RLVRTrainer, id="rlvr"), + pytest.param(RLAIFTrainer, id="rlaif"), +] + + +def _stopping_condition(): + return shapes.StoppingCondition(max_runtime_in_seconds=MAX_RUNTIME_IN_SECONDS) + + +def _trainer(trainer_cls, sagemaker_session, dataset, name, **overrides): + """Build a recipe trainer in its minimal accepted configuration. + + ``accept_eula=True`` is required for gated foundation models; without it the + request is refused before it reaches the interesting validation. + + Trainer-specific required arguments come from ``_TRAINER_EXTRA_KWARGS`` so + adding another trainer to ``RECIPE_TRAINERS`` stays a two-line change. + """ + kwargs = dict( + model=MODEL_ID, + training_type=TrainingType.LORA, + model_package_group=MODEL_PACKAGE_GROUP, + training_dataset=dataset, + accept_eula=True, + sagemaker_session=sagemaker_session, + base_job_name=name, + stopping_condition=_stopping_condition(), + ) + kwargs.update(_TRAINER_EXTRA_KWARGS.get(trainer_cls.__name__, {})) + kwargs.update(overrides) + return trainer_cls(**kwargs) + + +class TestServerlessSubmission: + """The default path: compute is derived from the selected recipe. + + Recipe selection and resource-config generation happen server-side after the + request validators, so acceptance here is the only cheap proof that the + SDK's recipe payload is still valid. + """ + + @pytest.mark.parametrize("trainer_cls", RECIPE_TRAINERS) + def test_minimal_request_is_accepted(self, trainer_cls, sagemaker_session, train_data_uri): + name = unique_name(f"shallow-{trainer_cls.__name__.lower()}") + trainer = _trainer(trainer_cls, sagemaker_session, train_data_uri, name) + + with submitted(trainer) as job: + assert_submitted(job) + + @pytest.mark.parametrize("trainer_cls", RECIPE_TRAINERS) + def test_with_validation_dataset( + self, trainer_cls, sagemaker_session, train_data_uri, validation_data_uri + ): + """A validation dataset adds a second channel, which is resolved against + S3 independently of the training channel.""" + name = unique_name(f"shallow-{trainer_cls.__name__.lower()}-val") + trainer = _trainer( + trainer_cls, + sagemaker_session, + train_data_uri, + name, + validation_dataset=validation_data_uri, + ) + + with submitted(trainer) as job: + assert_submitted(job) + + @pytest.mark.parametrize("trainer_cls", RECIPE_TRAINERS) + def test_datasets_passed_to_train_override_constructor( + self, trainer_cls, sagemaker_session, train_data_uri + ): + """``train(training_dataset=...)`` overrides the constructor value. + + Worth asserting server-side: if the override were dropped, the payload + would silently reference the wrong data and only a real run would reveal + it. + """ + name = unique_name(f"shallow-{trainer_cls.__name__.lower()}-override") + trainer = _trainer(trainer_cls, sagemaker_session, None, name) + + with submitted(trainer, training_dataset=train_data_uri) as job: + assert_submitted(job) + + @pytest.mark.parametrize("training_type", [TrainingType.LORA, TrainingType.FULL]) + def test_training_types(self, sagemaker_session, train_data_uri, training_type): + """LoRA and full fine-tuning select different recipes, so each must be + independently accepted.""" + suffix = str(getattr(training_type, "value", training_type)).lower() + name = unique_name(f"shallow-sft-{suffix}") + trainer = _trainer( + SFTTrainer, + sagemaker_session, + train_data_uri, + name, + training_type=training_type, + ) + + with submitted(trainer) as job: + assert_submitted(job) + + def test_cpt_trainer_is_accepted(self, sagemaker_session, train_data_uri): + """Continued pre-training uses a distinct recipe family from SFT/DPO/RLVR. + + Kept separate from RECIPE_TRAINERS because CPT is not a preference/ + instruction-tuning trainer and its accepted arguments differ. + """ + name = unique_name("shallow-cpt") + trainer = _trainer(CPTTrainer, sagemaker_session, train_data_uri, name) + + with submitted(trainer) as job: + assert_submitted(job) + + +class TestServerfulSubmission: + """Explicit ``TrainingJobCompute`` produces a materially different payload + from the recipe-derived serverless path, including a resource config the + backend validates against the recipe.""" + + @pytest.mark.parametrize("trainer_cls", RECIPE_TRAINERS) + def test_explicit_compute_is_accepted(self, trainer_cls, sagemaker_session, train_data_uri): + name = unique_name(f"shallow-{trainer_cls.__name__.lower()}-serverful") + trainer = _trainer( + trainer_cls, + sagemaker_session, + train_data_uri, + name, + compute=TrainingJobCompute(instance_type=SERVERFUL_INSTANCE_TYPE, instance_count=1), + ) + + with submitted(trainer) as job: + assert_submitted(job) + + +class TestOutputAndTracking: + """Output location and MLflow tracking are validated server-side.""" + + def test_explicit_s3_output_path(self, sagemaker_session, train_data_uri, output_path): + name = unique_name("shallow-sft-output") + trainer = _trainer( + SFTTrainer, + sagemaker_session, + train_data_uri, + name, + s3_output_path=output_path, + ) + + with submitted(trainer) as job: + assert_submitted(job) + + def test_disable_output_compression(self, sagemaker_session, train_data_uri): + """Uncompressed output changes the OutputDataConfig the SDK sends.""" + name = unique_name("shallow-sft-nocompress") + trainer = _trainer( + SFTTrainer, + sagemaker_session, + train_data_uri, + name, + disable_output_compression=True, + ) + + with submitted(trainer) as job: + assert_submitted(job) + + +class TestRejectedRecipeRequests: + """Negative cases specific to the recipe path. + + These matter more here than for ``ModelTrainer``: recipe resolution is the + part of the payload most likely to drift, and an over-permissive change would + otherwise still yield a green suite. + """ + + def test_nonexistent_training_dataset_is_rejected( + self, sagemaker_session, nonexistent_data_uri + ): + """Dataset existence is checked against S3 before the job is created.""" + trainer = _trainer( + SFTTrainer, + sagemaker_session, + nonexistent_data_uri, + unique_name("shallow-sft-bad-data"), + ) + + assert_rejected( + trainer, + ("does not exist", "ValidationException", "ValidationError", "S3", "not found"), + ) + + def test_nonexistent_validation_dataset_is_rejected( + self, sagemaker_session, train_data_uri, nonexistent_data_uri + ): + """A valid training set must not mask an invalid validation set.""" + trainer = _trainer( + SFTTrainer, + sagemaker_session, + train_data_uri, + unique_name("shallow-sft-bad-val"), + validation_dataset=nonexistent_data_uri, + ) + + assert_rejected( + trainer, + ("does not exist", "ValidationException", "ValidationError", "S3", "not found"), + ) + + def test_unknown_model_is_rejected(self, sagemaker_session, train_data_uri): + """Model resolution must fail for a model that does not exist. + + Guards the JumpStart/hub lookup that turns ``model`` into a concrete + artifact URI in the payload. + """ + trainer_kwargs = dict( + model="definitely-not-a-real-model-id-4b91c7", + training_type=TrainingType.LORA, + model_package_group=MODEL_PACKAGE_GROUP, + training_dataset=train_data_uri, + accept_eula=True, + sagemaker_session=sagemaker_session, + base_job_name=unique_name("shallow-sft-bad-model"), + ) + + # Model resolution can fail either while constructing the trainer or at + # submit time depending on how the id is interpreted, so both are allowed + # here; what matters is that an unknown model never reaches the service. + with pytest.raises(Exception) as excinfo: + trainer = SFTTrainer(**trainer_kwargs) + trainer.train(wait=False) + + message = str(excinfo.value) + assert any( + token in message + for token in ( + "model", + "Model", + "not found", + "does not exist", + "ResourceNotFound", + "ValidationException", + "ValidationError", + ) + ), f"unexpected rejection reason: {message}" + + def test_invalid_instance_type_is_rejected(self, sagemaker_session, train_data_uri): + """A nonexistent instance type must be refused on the serverful path.""" + trainer = _trainer( + SFTTrainer, + sagemaker_session, + train_data_uri, + unique_name("shallow-sft-bad-instance"), + compute=TrainingJobCompute(instance_type="ml.nonexistent.24xlarge", instance_count=1), + ) + + assert_rejected( + trainer, + ( + "instance", + "Instance", + "not supported", + "ValidationException", + "ValidationError", + ), + ) diff --git a/sagemaker-train/tests/integ/train/test_benchmark_evaluator.py b/sagemaker-train/tests/integ/train/test_benchmark_evaluator.py index 23f21229c3..4a71961916 100644 --- a/sagemaker-train/tests/integ/train/test_benchmark_evaluator.py +++ b/sagemaker-train/tests/integ/train/test_benchmark_evaluator.py @@ -97,6 +97,7 @@ def test_get_benchmarks_and_properties(self): logger.info(f"MMLU properties: {properties}") + @pytest.mark.gpu_intensive def test_benchmark_evaluation_full_flow(self): """ Test complete benchmark evaluation flow with fine-tuned model package. diff --git a/sagemaker-train/tests/integ/train/test_custom_scorer_evaluator.py b/sagemaker-train/tests/integ/train/test_custom_scorer_evaluator.py index f0f0968c07..b8569ea336 100644 --- a/sagemaker-train/tests/integ/train/test_custom_scorer_evaluator.py +++ b/sagemaker-train/tests/integ/train/test_custom_scorer_evaluator.py @@ -86,6 +86,7 @@ def test_get_builtin_metrics(self): logger.info(f"Built-in metrics: {list(BuiltInMetric.__members__.keys())}") + @pytest.mark.gpu_intensive def test_custom_scorer_evaluation_full_flow(self): """ Test complete custom scorer evaluation flow with custom evaluator ARN. diff --git a/sagemaker-train/tests/integ/train/test_inspect_ai_evaluator.py b/sagemaker-train/tests/integ/train/test_inspect_ai_evaluator.py index d045d49e13..3155579d37 100644 --- a/sagemaker-train/tests/integ/train/test_inspect_ai_evaluator.py +++ b/sagemaker-train/tests/integ/train/test_inspect_ai_evaluator.py @@ -113,6 +113,7 @@ def inspect_ai_resources(sagemaker_session_us_east_1): class TestInspectAIEvaluatorIntegration: """Integration tests for InspectAI evaluation with Bedrock inference.""" + @pytest.mark.gpu_intensive def test_inspect_ai_bedrock_evaluation( self, sagemaker_session_us_east_1, inspect_ai_resources ): @@ -161,6 +162,7 @@ def test_inspect_ai_bedrock_evaluation( execution.show_results() logger.info("InspectAI Bedrock evaluation completed successfully.") + @pytest.mark.gpu_intensive def test_inspect_ai_upload_benchmarks( self, sagemaker_session_us_east_1, inspect_ai_resources ): diff --git a/sagemaker-train/tests/integ/train/test_llm_as_judge_base_model_fix.py b/sagemaker-train/tests/integ/train/test_llm_as_judge_base_model_fix.py index 2c188a8f5d..e4c62ba1c8 100644 --- a/sagemaker-train/tests/integ/train/test_llm_as_judge_base_model_fix.py +++ b/sagemaker-train/tests/integ/train/test_llm_as_judge_base_model_fix.py @@ -100,6 +100,7 @@ def _get_latest_model_package_arn(): class TestLLMAsJudgeBaseModelFix: """Integration test for base model fix in LLMAsJudgeEvaluator""" + @pytest.mark.gpu_intensive def test_base_model_evaluation_uses_correct_weights(self, mlflow_resource_arn): """ Test that base model evaluation uses original base model weights. @@ -278,6 +279,7 @@ def test_base_model_evaluation_uses_correct_weights(self, mlflow_resource_arn): # Re-raise to fail the test raise + @pytest.mark.gpu_intensive def test_base_model_false_still_works(self, mlflow_resource_arn): """ Test that evaluate_base_model=False still works correctly (backward compatibility). diff --git a/sagemaker-train/tests/integ/train/test_llm_as_judge_evaluator.py b/sagemaker-train/tests/integ/train/test_llm_as_judge_evaluator.py index 4907a7317c..c6b665af6e 100644 --- a/sagemaker-train/tests/integ/train/test_llm_as_judge_evaluator.py +++ b/sagemaker-train/tests/integ/train/test_llm_as_judge_evaluator.py @@ -88,6 +88,7 @@ class TestLLMAsJudgeEvaluatorIntegration: """Integration tests for LLMAsJudgeEvaluator""" + @pytest.mark.gpu_intensive def test_llm_as_judge_evaluation_full_flow(self): """ Test complete LLM-as-Judge evaluation flow with custom and built-in metrics. diff --git a/sagemaker-train/tests/integ/train/test_llmaj_custom_model.py b/sagemaker-train/tests/integ/train/test_llmaj_custom_model.py index e3277e9509..65ffd45e1f 100644 --- a/sagemaker-train/tests/integ/train/test_llmaj_custom_model.py +++ b/sagemaker-train/tests/integ/train/test_llmaj_custom_model.py @@ -98,6 +98,7 @@ def test_resources(sagemaker_session_us_east_1): class TestLLMAJCustomModelIntegration: """Integration tests for LLMAsJudgeEvaluator with InspectAI inference path.""" + @pytest.mark.gpu_intensive def test_llmaj_bedrock_inference_end_to_end( self, sagemaker_session_us_east_1, test_resources ): diff --git a/sagemaker-train/tests/integ/train/test_model_trainer.py b/sagemaker-train/tests/integ/train/test_model_trainer.py index 63bbfc52bb..d651395000 100644 --- a/sagemaker-train/tests/integ/train/test_model_trainer.py +++ b/sagemaker-train/tests/integ/train/test_model_trainer.py @@ -55,6 +55,7 @@ ) +@pytest.mark.gpu_intensive def test_source_dir_local_tar_file(sagemaker_session): model_trainer = ModelTrainer( sagemaker_session=sagemaker_session, @@ -66,6 +67,7 @@ def test_source_dir_local_tar_file(sagemaker_session): model_trainer.train() +@pytest.mark.gpu_intensive def test_hp_contract_basic_py_script(sagemaker_session): model_trainer = ModelTrainer( sagemaker_session=sagemaker_session, @@ -78,6 +80,7 @@ def test_hp_contract_basic_py_script(sagemaker_session): model_trainer.train() +@pytest.mark.gpu_intensive def test_hp_contract_basic_sh_script(sagemaker_session): source_code = SourceCode( source_dir=f"{DATA_DIR}/params_script", @@ -97,6 +100,7 @@ def test_hp_contract_basic_sh_script(sagemaker_session): # skip this test for now as requirments.txt is not resolved # @pytest.mark.skip(reason="MPI distributed training does not resolve requirements.txt on worker nodes") +@pytest.mark.gpu_intensive def test_hp_contract_mpi_script(sagemaker_session): compute = Compute(instance_type="ml.m5.xlarge", instance_count=2) model_trainer = ModelTrainer( @@ -112,6 +116,7 @@ def test_hp_contract_mpi_script(sagemaker_session): model_trainer.train() +@pytest.mark.gpu_intensive def test_hp_contract_torchrun_script(sagemaker_session): compute = Compute(instance_type="ml.m5.xlarge", instance_count=2) model_trainer = ModelTrainer( @@ -127,6 +132,7 @@ def test_hp_contract_torchrun_script(sagemaker_session): model_trainer.train() +@pytest.mark.gpu_intensive def test_hp_contract_hyperparameter_json(sagemaker_session): model_trainer = ModelTrainer( sagemaker_session=sagemaker_session, @@ -139,6 +145,7 @@ def test_hp_contract_hyperparameter_json(sagemaker_session): model_trainer.train() +@pytest.mark.gpu_intensive def test_hp_contract_hyperparameter_yaml(sagemaker_session): model_trainer = ModelTrainer( sagemaker_session=sagemaker_session, @@ -151,6 +158,7 @@ def test_hp_contract_hyperparameter_yaml(sagemaker_session): model_trainer.train() +@pytest.mark.gpu_intensive def test_custom_distributed_driver(sagemaker_session): class CustomDriver(DistributedConfig): process_count_per_node: int = None diff --git a/sagemaker-train/tests/integ/train/test_notifications.py b/sagemaker-train/tests/integ/train/test_notifications.py index 789391755a..26aad2467b 100644 --- a/sagemaker-train/tests/integ/train/test_notifications.py +++ b/sagemaker-train/tests/integ/train/test_notifications.py @@ -160,6 +160,7 @@ def sqs_subscriber(sm_session): logger.warning(f"Failed to delete queue: {e}") +@pytest.mark.gpu_intensive @pytest.mark.us_east_1 def test_notifications_creates_eventbridge_rule_and_cleanup( sm_session, training_data_uri, sqs_subscriber diff --git a/sagemaker-train/tests/integ/train/test_sft_trainer_integration.py b/sagemaker-train/tests/integ/train/test_sft_trainer_integration.py index 9b4ef81cb8..bd0846323b 100644 --- a/sagemaker-train/tests/integ/train/test_sft_trainer_integration.py +++ b/sagemaker-train/tests/integ/train/test_sft_trainer_integration.py @@ -177,6 +177,7 @@ def test_sft_trainer_nova_workflow(sagemaker_session_us_east_1): # @pytest.mark.gpu_intensive +@pytest.mark.gpu_intensive def test_sft_trainer_lora_with_sequence_length(sagemaker_session): """Test SFT training workflow with LORA and sequence_length specified.""" unique_id = f"{int(time.time())}-{random.randint(1000, 9999)}" diff --git a/sagemaker-train/tests/integ/train/test_tuner_distributed.py b/sagemaker-train/tests/integ/train/test_tuner_distributed.py index 2af2b7cb4d..24cb787d3f 100644 --- a/sagemaker-train/tests/integ/train/test_tuner_distributed.py +++ b/sagemaker-train/tests/integ/train/test_tuner_distributed.py @@ -72,6 +72,7 @@ def train_source_dir(tmp_path_factory): return str(d) +@pytest.mark.gpu_intensive def test_tuner_includes_sm_drivers_channel(sagemaker_session, train_source_dir): """Verify tuning jobs include sm_drivers channel for distributed training. diff --git a/sagemaker-train/tox.ini b/sagemaker-train/tox.ini index 01b6faebd8..21962188f1 100644 --- a/sagemaker-train/tox.ini +++ b/sagemaker-train/tox.ini @@ -62,7 +62,7 @@ markers = slow_test release image_uris_unit_test - gpu_intensive: mark a test as GPU resource intensive (runs on scheduled CI, not PR checks). + gpu_intensive: mark a test as expensive - it submits a real job and waits for it to run (runs on scheduled CI, not PR checks). Despite the name this is not strictly about GPUs: it gates anything that consumes real training capacity, including serverless and CPU-instance jobs. Cheap acceptance coverage for the same code paths lives in tests/integ/train/shallow (submit-then-stop), which does run on PR checks. us_east_1: mark a test that requires us-east-1 test account credentials (784379639078). timeout: mark a test as a timeout. serial: marks tests that must run serially (not in parallel) From c86cfb4040ffbc570af7581c2ebe5bf99a88eb7b Mon Sep 17 00:00:00 2001 From: Amarjeet LNU Date: Tue, 11 Aug 2026 16:23:42 -0700 Subject: [PATCH 02/15] 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. --- .../shallow/test_model_trainer_submission.py | 38 ++++++++++++------- .../test_recipe_trainers_submission.py | 17 +++++++-- 2 files changed, 38 insertions(+), 17 deletions(-) diff --git a/sagemaker-train/tests/integ/train/shallow/test_model_trainer_submission.py b/sagemaker-train/tests/integ/train/shallow/test_model_trainer_submission.py index 13776aac69..7a633762a0 100644 --- a/sagemaker-train/tests/integ/train/shallow/test_model_trainer_submission.py +++ b/sagemaker-train/tests/integ/train/shallow/test_model_trainer_submission.py @@ -605,26 +605,36 @@ def test_nonexistent_training_image_is_rejected(self, sagemaker_session, account ) def test_unassumable_role_is_rejected(self, sagemaker_session, account_id): - """PassRole / AssumeRole failures must surface at submit time. - - Directly covers the "does the caller hold the required permissions" half - of what this suite exists to assert. + """A role that cannot be used for training must be refused. + + Covers the "does the caller hold the required permissions" half of what + this suite exists to assert. + + Note where this is caught: ``ModelTrainer.__init__`` resolves and + validates the role via ``iam:SimulatePrincipalPolicy``, so a bad role is + rejected at *construction* -- the request never reaches + CreateTrainingJob. That is strictly better than a server-side rejection + (faster, clearer message), so this asserts around the constructor rather + than around ``train()``. Verified against AWS: the SDK raises + ``RoleValidationError`` naming the role and the permissions it lacks. """ bogus_role = f"arn:aws:iam::{account_id}:role/shallow-integ-test-no-such-role" - trainer = _trainer(sagemaker_session, unique_name("shallow-bad-role"), role=bogus_role) - assert_rejected( - trainer, - ( - "role", - "Role", + with pytest.raises(Exception) as excinfo: + _trainer(sagemaker_session, unique_name("shallow-bad-role"), role=bogus_role) + + message = str(excinfo.value) + assert any( + token in message + for token in ( + "cannot be used", + "RoleValidationError", "AccessDenied", "not authorized", "cannot be assumed", - "ValidationException", - "ValidationError", - ), - ) + "does not exist", + ) + ), f"unexpected rejection reason: {message}" def test_duplicate_job_name_is_rejected(self, sagemaker_session, execution_role, output_path): """The final gate before the ARN is a conditional write that rejects diff --git a/sagemaker-train/tests/integ/train/shallow/test_recipe_trainers_submission.py b/sagemaker-train/tests/integ/train/shallow/test_recipe_trainers_submission.py index 2df7995110..18fcd307d4 100644 --- a/sagemaker-train/tests/integ/train/shallow/test_recipe_trainers_submission.py +++ b/sagemaker-train/tests/integ/train/shallow/test_recipe_trainers_submission.py @@ -190,11 +190,22 @@ def test_training_types(self, sagemaker_session, train_data_uri, training_type): def test_cpt_trainer_is_accepted(self, sagemaker_session, train_data_uri): """Continued pre-training uses a distinct recipe family from SFT/DPO/RLVR. - Kept separate from RECIPE_TRAINERS because CPT is not a preference/ - instruction-tuning trainer and its accepted arguments differ. + Kept out of RECIPE_TRAINERS because its constructor genuinely differs: + verified against the SDK, ``CPTTrainer`` accepts no ``training_type`` + (there is no LoRA/full distinction for continued pre-training) and its + ``compute`` is ``HyperPodCompute``-only, so it cannot take the + serverful ``TrainingJobCompute`` the others accept. """ name = unique_name("shallow-cpt") - trainer = _trainer(CPTTrainer, sagemaker_session, train_data_uri, name) + trainer = CPTTrainer( + model=MODEL_ID, + model_package_group=MODEL_PACKAGE_GROUP, + training_dataset=train_data_uri, + accept_eula=True, + sagemaker_session=sagemaker_session, + base_job_name=name, + stopping_condition=_stopping_condition(), + ) with submitted(trainer) as job: assert_submitted(job) From 92446b5af2a813380dec9635e1f8738fc44da958 Mon Sep 17 00:00:00 2001 From: Amarjeet LNU Date: Tue, 11 Aug 2026 18:02:01 -0700 Subject: [PATCH 03/15] 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. --- .../tests/integ/train/shallow/conftest.py | 43 ++++++++++ .../tests/integ/train/shallow/harness.py | 35 +++++--- .../test_other_job_types_submission.py | 79 ++++++++++++------- .../test_recipe_customization_submission.py | 20 ++++- .../test_recipe_trainers_submission.py | 50 ++++++++++-- 5 files changed, 178 insertions(+), 49 deletions(-) diff --git a/sagemaker-train/tests/integ/train/shallow/conftest.py b/sagemaker-train/tests/integ/train/shallow/conftest.py index 015f8fd28a..985444850a 100644 --- a/sagemaker-train/tests/integ/train/shallow/conftest.py +++ b/sagemaker-train/tests/integ/train/shallow/conftest.py @@ -26,6 +26,7 @@ import json import logging +import os import pytest @@ -71,6 +72,48 @@ def _ensure_object(sagemaker_session, key): return f"s3://{bucket}/{key}" +@pytest.fixture(autouse=True, scope="session") +def bundled_service_model(): + """Point botocore at the service model bundled in ``sagemaker-core/sample``. + + Some request fields this suite exercises are not in the public botocore model + yet -- ``ServerlessJobConfig.SequenceLength`` is the current example. Without + this, botocore rejects the request client-side with + + Unknown parameter in ServerlessJobConfig: "SequenceLength" + + and the test fails before reaching the service, which tells us nothing about + whether the payload is acceptable. Verified against AWS: setting AWS_DATA_PATH + adds ``SequenceLength`` to the shape. + + Session-scoped and autouse because botocore caches loaded models per client; + setting this after a client exists would not take effect. Mirrors the + ``setup_aws_data_path`` fixture in ``test_recipe_override_integration.py``, + which solves the same problem for the client-side recipe tests. + """ + # tests/integ/train/shallow/conftest.py -> repo root is five levels up. + repo_root = os.path.abspath( + os.path.join(os.path.dirname(__file__), "..", "..", "..", "..", "..") + ) + sample_path = os.path.join(repo_root, "sagemaker-core", "sample") + + previous = os.environ.get("AWS_DATA_PATH") + if os.path.isdir(sample_path): + os.environ["AWS_DATA_PATH"] = sample_path + logger.info("Using bundled service model at %s", sample_path) + else: + # Don't fail the run: on an installed-package layout the bundled model may + # not be present, and only the few tests using unreleased fields break. + logger.warning("Bundled service model not found at %s", sample_path) + + yield + + if previous is None: + os.environ.pop("AWS_DATA_PATH", None) + else: + os.environ["AWS_DATA_PATH"] = previous + + @pytest.fixture(scope="module") def train_data_uri(sagemaker_session): """S3 URI of a real, existing training-data prefix.""" diff --git a/sagemaker-train/tests/integ/train/shallow/harness.py b/sagemaker-train/tests/integ/train/shallow/harness.py index ae57114ef2..bd2c0caf37 100644 --- a/sagemaker-train/tests/integ/train/shallow/harness.py +++ b/sagemaker-train/tests/integ/train/shallow/harness.py @@ -94,22 +94,33 @@ _UNSTOPPABLE_STATUSES = frozenset({"Completed", "Failed", "Stopped", "Stopping"}) -def unique_name(prefix): - """Build a collision-free job name. +# Name length limits differ per resource, and the service enforces them strictly. +# Verified against AWS: a 34-character tuning job name is rejected with +# Value '...' at 'hyperParameterTuningJobName' failed to satisfy constraint: +# Member must have length less than or equal to 32 +MAX_TRAINING_JOB_NAME = 63 +MAX_TUNING_JOB_NAME = 32 - The backend rejects duplicate job names per account with ``ResourceInUse``, - and these tests run in parallel across many xdist workers, so the - name must be unique per invocation rather than per test function. Includes - randomness as well as a timestamp because two xdist workers can enter the - same second. - SageMaker training job names are limited to 63 characters, so the prefix is - truncated rather than allowed to silently push the suffix over the limit. +def unique_name(prefix, max_length=MAX_TRAINING_JOB_NAME): + """Build a collision-free job name that fits the resource's length limit. + + The backend rejects duplicate job names per account with ``ResourceInUse``, + and these tests run in parallel across many xdist workers, so the name must + be unique per invocation rather than per test function. Includes randomness + as well as a timestamp because two xdist workers can enter the same second. + + The uniqueness suffix is preserved and the *prefix* is truncated, so a long + descriptive prefix degrades readability rather than silently reintroducing + collisions. Pass ``max_length=MAX_TUNING_JOB_NAME`` for tuning jobs, whose + limit is roughly half that of training jobs. """ suffix = f"{int(time.time())}-{random.randint(1000, 9999)}" - # 63 total, minus the suffix, minus the joining hyphen. - head = prefix[: 63 - len(suffix) - 1] - return f"{head}-{suffix}" + # Budget: total, minus the suffix, minus the joining hyphen. + head = prefix[: max_length - len(suffix) - 1] + name = f"{head}-{suffix}" + assert len(name) <= max_length, f"generated name {name!r} exceeds {max_length} chars" + return name def stop_quietly(training_job): diff --git a/sagemaker-train/tests/integ/train/shallow/test_other_job_types_submission.py b/sagemaker-train/tests/integ/train/shallow/test_other_job_types_submission.py index ae6d0eb6f7..af87036df6 100644 --- a/sagemaker-train/tests/integ/train/shallow/test_other_job_types_submission.py +++ b/sagemaker-train/tests/integ/train/shallow/test_other_job_types_submission.py @@ -37,6 +37,7 @@ import logging import os +from contextlib import contextmanager import pytest from sagemaker.core import shapes @@ -52,6 +53,7 @@ DEFAULT_INSTANCE_COUNT, DEFAULT_INSTANCE_TYPE, MAX_RUNTIME_IN_SECONDS, + MAX_TUNING_JOB_NAME, assert_submitted, submitted, unique_name, @@ -109,6 +111,33 @@ def _tuner(model_trainer, **overrides): return HyperparameterTuner(**kwargs) +@contextmanager +def _tuning(tuner, job_name): + """Submit a tuning job under an explicit name, then always stop it. + + The explicit ``job_name`` is load-bearing. Left to itself the tuner derives a + name from the training image plus a second-granularity timestamp + (``pytorch-training-260811-1621``) and ignores ``base_job_name`` entirely, so + two tuner tests starting in the same second collide with ``ResourceInUse``. + Verified against AWS: that is exactly how this failed before. + + Teardown goes through ``tuner.stop_tuning_job()`` rather than the harness's + ``stop_quietly``, because the tuner wraps the resource and stopping it also + stops the child training jobs it launched. + """ + try: + tuner.tune(job_name=job_name, wait=False) + yield + finally: + try: + tuner.stop_tuning_job() + logger.info("Stopped tuning job %s", job_name) + except Exception as e: # pragma: no cover - best-effort teardown + # A tuning job that never started, or already reached a terminal + # state, cannot be stopped; that must not fail the test. + logger.warning("Could not stop tuning job %s: %s", job_name, e) + + class TestTuningJobSubmission: """HyperParameterTuningJob acceptance. @@ -118,18 +147,15 @@ class TestTuningJobSubmission: def test_minimal_tuning_job_is_accepted(self, sagemaker_session): """Baseline: the service accepts a well-formed tuning job.""" - name = unique_name("shallow-tuner") + name = unique_name("shallow-tuner", max_length=MAX_TUNING_JOB_NAME) tuner = _tuner(_model_trainer(sagemaker_session, name)) - try: - tuner.tune(wait=False) - assert_submitted(tuner.latest_tuning_job, resource="hyper-parameter-tuning-job") - finally: - # Tuner exposes its own stop method rather than the resource's. - try: - tuner.stop_tuning_job() - except Exception as e: # pragma: no cover - best-effort teardown - logger.warning("Could not stop tuning job: %s", e) + with _tuning(tuner, name): + assert_submitted( + tuner.latest_tuning_job, + expected_name=name, + resource="hyper-parameter-tuning-job", + ) def test_distributed_tuning_job_is_accepted(self, sagemaker_session): """A tuning job wrapping a Torchrun trainer must include the @@ -140,31 +166,30 @@ def test_distributed_tuning_job_is_accepted(self, sagemaker_session): proves the channel is present and the definition is accepted, which is the part that regressed; the log assertion stays in the deep suite. """ - name = unique_name("shallow-tuner-dist") + name = unique_name("shallow-tune-dist", max_length=MAX_TUNING_JOB_NAME) model_trainer = _model_trainer(sagemaker_session, name, distributed=Torchrun()) tuner = _tuner(model_trainer) - try: - tuner.tune(wait=False) - arn = assert_submitted(tuner.latest_tuning_job, resource="hyper-parameter-tuning-job") + with _tuning(tuner, name): + arn = assert_submitted( + tuner.latest_tuning_job, + expected_name=name, + resource="hyper-parameter-tuning-job", + ) # The sm_drivers channel lives in the tuning job's training # definition; read it back to prove it survived submission rather - # than inferring from acceptance alone. + # than inferring it from acceptance alone. described = tuner.latest_tuning_job.refresh() definition = getattr(described, "training_job_definition", None) - if definition is not None: - channels = [ - channel.channel_name for channel in (definition.input_data_config or []) - ] - assert "sm_drivers" in channels, ( - f"tuning job {arn} is missing the sm_drivers channel; " f"channels={channels}" - ) - finally: - try: - tuner.stop_tuning_job() - except Exception as e: # pragma: no cover - best-effort teardown - logger.warning("Could not stop tuning job: %s", e) + assert definition is not None, ( + f"tuning job {arn} has no training_job_definition to inspect; " + "cannot verify the sm_drivers channel" + ) + channels = [channel.channel_name for channel in (definition.input_data_config or [])] + assert ( + "sm_drivers" in channels + ), f"tuning job {arn} is missing the sm_drivers channel; channels={channels}" @pytest.mark.gpu_intensive diff --git a/sagemaker-train/tests/integ/train/shallow/test_recipe_customization_submission.py b/sagemaker-train/tests/integ/train/shallow/test_recipe_customization_submission.py index 9852d7c4dd..7a444b6f40 100644 --- a/sagemaker-train/tests/integ/train/shallow/test_recipe_customization_submission.py +++ b/sagemaker-train/tests/integ/train/shallow/test_recipe_customization_submission.py @@ -162,10 +162,24 @@ def test_direct_hyperparameter_mutation(self, sagemaker_session, train_data_uri) class TestSequenceLength: - """``sequence_length`` selects a different recipe variant, so each supported - value is a distinct accepted-payload case.""" + """``sequence_length`` selects a different recipe variant. - @pytest.mark.parametrize("sequence_length", ["4K", "16K"]) + Only ``4K`` is parametrized. Verified against AWS: for ``MODEL_ID`` the recipe + catalogue offers exactly one sequence length -- + + ValueError: No recipes found with SequenceLength == 16K. + Available sequence lengths: ['4K'] + + so a ``16K`` case would assert a service-side limitation rather than SDK + behaviour. Left parametrized (rather than inlined) so another value can be + added when a model in this account supports one. + + Note this field also requires the bundled service model -- see the + ``bundled_service_model`` fixture in conftest; the public botocore model has + no ``ServerlessJobConfig.SequenceLength`` yet. + """ + + @pytest.mark.parametrize("sequence_length", ["4K"]) def test_sequence_length_variants(self, sagemaker_session, train_data_uri, sequence_length): trainer = _sft( sagemaker_session, diff --git a/sagemaker-train/tests/integ/train/shallow/test_recipe_trainers_submission.py b/sagemaker-train/tests/integ/train/shallow/test_recipe_trainers_submission.py index 18fcd307d4..a5306b2b69 100644 --- a/sagemaker-train/tests/integ/train/shallow/test_recipe_trainers_submission.py +++ b/sagemaker-train/tests/integ/train/shallow/test_recipe_trainers_submission.py @@ -30,9 +30,11 @@ from __future__ import absolute_import +import os + import pytest from sagemaker.core import shapes -from sagemaker.core.training.configs import TrainingJobCompute +from sagemaker.core.training.configs import HyperPodCompute, TrainingJobCompute from sagemaker.train.common import TrainingType from sagemaker.train.cpt_trainer import CPTTrainer from sagemaker.train.dpo_trainer import DPOTrainer @@ -91,6 +93,11 @@ pytest.param(RLAIFTrainer, id="rlaif"), ] +# Subset that accepts an explicit TrainingJobCompute. RLAIFTrainer takes no +# ``compute`` argument at all (verified against the SDK), so it has no serverful +# path and is excluded rather than being expected to fail. +SERVERFUL_CAPABLE_TRAINERS = [t for t in RECIPE_TRAINERS if t.values[0] is not RLAIFTrainer] + def _stopping_condition(): return shapes.StoppingCondition(max_runtime_in_seconds=MAX_RUNTIME_IN_SECONDS) @@ -170,7 +177,14 @@ def test_datasets_passed_to_train_override_constructor( with submitted(trainer, training_dataset=train_data_uri) as job: assert_submitted(job) - @pytest.mark.parametrize("training_type", [TrainingType.LORA, TrainingType.FULL]) + # Only LORA is parametrized. Verified against AWS: for MODEL_ID there is no + # serverless (SMTJ) recipe for full fine-tuning -- + # ValueError: No recipes found with Smtj for technique: SFT, + # training_type:TrainingType.FULL + # so a FULL case here would assert a recipe-catalogue limitation rather than + # SDK behaviour. Kept parametrized so FULL can be re-added against a model + # that supports it, rather than the distinction being silently dropped. + @pytest.mark.parametrize("training_type", [TrainingType.LORA]) def test_training_types(self, sagemaker_session, train_data_uri, training_type): """LoRA and full fine-tuning select different recipes, so each must be independently accepted.""" @@ -187,15 +201,30 @@ def test_training_types(self, sagemaker_session, train_data_uri, training_type): with submitted(trainer) as job: assert_submitted(job) + @pytest.mark.gpu_intensive def test_cpt_trainer_is_accepted(self, sagemaker_session, train_data_uri): - """Continued pre-training uses a distinct recipe family from SFT/DPO/RLVR. + """Continued pre-training, which submits only via HyperPod. Kept out of RECIPE_TRAINERS because its constructor genuinely differs: verified against the SDK, ``CPTTrainer`` accepts no ``training_type`` (there is no LoRA/full distinction for continued pre-training) and its - ``compute`` is ``HyperPodCompute``-only, so it cannot take the - serverful ``TrainingJobCompute`` the others accept. + ``compute`` is ``HyperPodCompute``-only. + + Marked ``gpu_intensive`` and skipped unless a cluster is configured. CPT + refuses to submit without one -- + + ValueError: CPT requires HyperPod compute. + Pass compute=HyperPodCompute(...) when creating the CPTTrainer. + + -- and HyperPod submits to a pre-provisioned cluster rather than through + CreateTrainingJob, so there is nothing this suite can create on demand. + Written in the shallow style anyway so it becomes gate-eligible by + dropping one marker once a cluster exists in the PR account. """ + cluster_name = os.environ.get("SHALLOW_HYPERPOD_CLUSTER") + if not cluster_name: + pytest.skip("CPT requires HyperPod; set SHALLOW_HYPERPOD_CLUSTER to run") + name = unique_name("shallow-cpt") trainer = CPTTrainer( model=MODEL_ID, @@ -205,6 +234,7 @@ def test_cpt_trainer_is_accepted(self, sagemaker_session, train_data_uri): sagemaker_session=sagemaker_session, base_job_name=name, stopping_condition=_stopping_condition(), + compute=HyperPodCompute(cluster_name=cluster_name), ) with submitted(trainer) as job: @@ -214,9 +244,15 @@ def test_cpt_trainer_is_accepted(self, sagemaker_session, train_data_uri): class TestServerfulSubmission: """Explicit ``TrainingJobCompute`` produces a materially different payload from the recipe-derived serverless path, including a resource config the - backend validates against the recipe.""" + backend validates against the recipe. - @pytest.mark.parametrize("trainer_cls", RECIPE_TRAINERS) + RLAIF is absent from this class on purpose: verified against the SDK, + ``RLAIFTrainer.__init__`` has no ``compute`` parameter at all, so it has no + serverful path to exercise. It is still covered by every serverless case in + ``TestServerlessSubmission``. + """ + + @pytest.mark.parametrize("trainer_cls", SERVERFUL_CAPABLE_TRAINERS) def test_explicit_compute_is_accepted(self, trainer_cls, sagemaker_session, train_data_uri): name = unique_name(f"shallow-{trainer_cls.__name__.lower()}-serverful") trainer = _trainer( From df30e92cd8d3ce375d7e9ec6aaf8a076b567674b Mon Sep 17 00:00:00 2001 From: Amarjeet LNU Date: Tue, 11 Aug 2026 18:35:15 -0700 Subject: [PATCH 04/15] 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". --- .../tests/integ/train/shallow/README.md | 98 +++-- .../tests/integ/train/shallow/recipe_cases.py | 230 ++++++++++ .../integ/train/shallow/test_cpt_trainer.py | 64 +++ .../integ/train/shallow/test_dpo_trainer.py | 33 ++ ...er_submission.py => test_model_trainer.py} | 0 .../shallow/test_multi_turn_rl_trainer.py | 116 +++++ .../train/shallow/test_nova_data_mixing.py | 87 ++++ .../test_recipe_customization_submission.py | 264 ------------ .../test_recipe_trainers_submission.py | 398 ------------------ .../integ/train/shallow/test_rlaif_trainer.py | 40 ++ .../integ/train/shallow/test_rlvr_trainer.py | 98 +++++ .../integ/train/shallow/test_sft_trainer.py | 75 ++++ ..._job_types_submission.py => test_tuner.py} | 84 ---- .../integ/train/test_benchmark_evaluator.py | 1 - .../train/test_custom_scorer_evaluator.py | 1 - .../integ/train/test_inspect_ai_evaluator.py | 2 - .../train/test_llm_as_judge_base_model_fix.py | 2 - .../train/test_llm_as_judge_evaluator.py | 1 - .../integ/train/test_llmaj_custom_model.py | 1 - .../tests/integ/train/test_notifications.py | 1 - 20 files changed, 802 insertions(+), 794 deletions(-) create mode 100644 sagemaker-train/tests/integ/train/shallow/recipe_cases.py create mode 100644 sagemaker-train/tests/integ/train/shallow/test_cpt_trainer.py create mode 100644 sagemaker-train/tests/integ/train/shallow/test_dpo_trainer.py rename sagemaker-train/tests/integ/train/shallow/{test_model_trainer_submission.py => test_model_trainer.py} (100%) create mode 100644 sagemaker-train/tests/integ/train/shallow/test_multi_turn_rl_trainer.py create mode 100644 sagemaker-train/tests/integ/train/shallow/test_nova_data_mixing.py delete mode 100644 sagemaker-train/tests/integ/train/shallow/test_recipe_customization_submission.py delete mode 100644 sagemaker-train/tests/integ/train/shallow/test_recipe_trainers_submission.py create mode 100644 sagemaker-train/tests/integ/train/shallow/test_rlaif_trainer.py create mode 100644 sagemaker-train/tests/integ/train/shallow/test_rlvr_trainer.py create mode 100644 sagemaker-train/tests/integ/train/shallow/test_sft_trainer.py rename sagemaker-train/tests/integ/train/shallow/{test_other_job_types_submission.py => test_tuner.py} (68%) diff --git a/sagemaker-train/tests/integ/train/shallow/README.md b/sagemaker-train/tests/integ/train/shallow/README.md index 8ce3ec2100..23975953fa 100644 --- a/sagemaker-train/tests/integ/train/shallow/README.md +++ b/sagemaker-train/tests/integ/train/shallow/README.md @@ -35,45 +35,65 @@ Concretely, a regression that makes training itself fail — a broken entry scri a bad container command, a distributed-launch bug — **will still pass here.** That is the accepted trade for the runtime and cost reduction. -## Coverage vs. the suite this replaces - -`tests/integ/train` has 181 pre-existing tests, but only ~50 actually submit a -job — the rest are client-side (recipe resolution, data utils, log streaming, -docker-compose detection). Mapping the *submitting* ones against this suite: - -| Existing area | Ported here | Notes | -|---|---|---| -| `test_model_trainer.py` (8) | yes | hyperparameter contract (dict/JSON/YAML), MPI, Torchrun, local tar source, `.sh` entry script, custom distributed driver | -| `test_sft_trainer_integration.py` (4) | partly | LoRA/FULL, validation dataset, `sequence_length`. **Nova workflow not ported** (us-east-1 + gated model) | -| `test_dpo_trainer_integration.py` (2) | yes | via `RECIPE_TRAINERS` parametrization | -| `test_rlvr_trainer_integration.py` (7) | partly | base + recipe/overrides + direct hyperparameter mutation. **Custom reward function / evaluator objects not ported** | -| `test_rlaif_trainer_integration.py` (3) | yes | RLAIF is in `RECIPE_TRAINERS`; its reward model/prompt come from `_TRAINER_EXTRA_KWARGS` | -| `test_cpt_hyperpod.py`, `test_nova_sft_hyperpod.py`, `test_sft_data_mixing_hyperpod.py` (3) | no | HyperPod submits to a pre-provisioned cluster, not `CreateTrainingJob` — the pattern does not apply | -| `test_sft_trainer_data_mixing_integration.py` (1) | yes | `DataMixingConfig`, both explicit and recipe-default | -| `test_tuner_distributed.py` (1) | yes | `HyperParameterTuningJob`; also asserts the `sm_drivers` channel survived submission | -| `test_multi_turn_rl_trainer_integration.py` (7) | partly | AgentRFT `Job` submission, marked `gpu_intensive` — see below | -| `test_recipe_override_integration.py` (35) | n/a | client-side `get_resolved_recipe`; keep as-is, cheap already | -| Evaluators (`test_benchmark_evaluator.py`, `test_llm_as_judge_*`, `test_mtrl_*`, ~20) | no | `evaluate()` not `train()`; the same pattern applies and is the clearest next extension | -| `test_notifications.py`, `test_local_model_trainer.py` | no | EventBridge/SNS side effects and local-container mode (no service call) | - -Note that not every trainer creates a `TrainingJob`. `HyperparameterTuner` creates -a `HyperParameterTuningJob` and `MultiTurnRLTrainer` creates an AgentRFT `Job`, so -`assert_submitted` takes a `resource=` argument for the expected ARN segment and -the harness resolves the submitted job across four different attribute names. - -**Deliberately out of scope for this pattern:** HyperPod (different submission -API), local container mode (no service call), and anything asserting a job's -*outcome*. - -**Requires prerequisites, so marked `gpu_intensive` and skipped on the PR gate:** -the MTRL tests. Unlike everything else here they cannot be made self-contained — -they need a pre-provisioned agent runtime and MLflow app. They read those from -`SHALLOW_MTRL_AGENT_ENV` / `SHALLOW_MTRL_MLFLOW_APP_ARN` / `SHALLOW_MTRL_DATASET` -and skip when unset, so once the PR account has them, dropping one marker makes -them PR-gate-eligible. - -**Genuine remaining gap:** evaluator `evaluate()` submissions (~20 existing -tests). Same pattern, distinct API surface; not yet written. +## Layout + +One file per trainer, mirroring the existing deep suite so the shallow counterpart +of any deep test is easy to find: + +| 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` | + +`recipe_cases.py` holds the cases every recipe trainer shares (minimal submit, +validation dataset, dataset override, output path, serverful compute, and the two +negative cases). Each per-trainer class subclasses `RecipeTrainerCases` and sets +`TRAINER`, so a new trainer is a two-line file. Override the class attributes only +where the trainer genuinely differs: + +* `EXTRA_KWARGS` — required constructor args (RLAIF's reward model/prompt) +* `SUPPORTS_SERVERFUL = False` — trainer takes no `compute` (RLAIF) +* `SUPPORTS_TRAINING_TYPE = False` — no LoRA/full distinction (CPT) + +It is deliberately not named `test_*` so pytest does not collect the base class. + +## What was marked `gpu_intensive`, and why only those + +A deep test is only marked `gpu_intensive` (i.e. moved off the PR gate) when this +suite has a shallow test covering the same code path. 10 tests met that bar: + +| Deep test (now marked) | Shallow equivalent | +|---|---| +| `test_model_trainer.py::test_source_dir_local_tar_file` | `TestSourceCodePackaging::test_local_tar_file_source_dir` | +| `::test_hp_contract_basic_py_script` | `TestMinimalSubmission::test_minimal_request_is_accepted` | +| `::test_hp_contract_basic_sh_script` | `TestSourceCodePackaging::test_shell_entry_script` | +| `::test_hp_contract_mpi_script` | `TestComputeConfiguration::test_mpi_distributed` | +| `::test_hp_contract_torchrun_script` | `TestComputeConfiguration::test_torchrun_distributed` | +| `::test_hp_contract_hyperparameter_json` | `TestPayloadShaping::test_hyperparameters_from_json_file` | +| `::test_hp_contract_hyperparameter_yaml` | `TestPayloadShaping::test_hyperparameters_from_yaml_file` | +| `::test_custom_distributed_driver` | `TestSourceCodePackaging::test_custom_distributed_driver` | +| `test_sft_trainer_integration.py::test_sft_trainer_lora_with_sequence_length` | `test_sft_trainer.py::test_sequence_length_is_accepted` | +| `test_tuner_distributed.py::test_tuner_includes_sm_drivers_channel` | `test_tuner.py::test_distributed_tuning_job_is_accepted` | + +**Deliberately NOT marked**, because this suite does not cover them — marking them +would remove coverage with nothing replacing it: + +* every evaluator test (`test_benchmark_evaluator.py`, `test_custom_scorer_evaluator.py`, + `test_inspect_ai_evaluator.py`, `test_llm_as_judge_*`, `test_llmaj_custom_model.py`) + — `evaluate()` is a different API surface returning pipeline executions, and there + is no shallow coverage for it yet +* `test_notifications.py` — asserts EventBridge/SNS side effects, not submission +* `test_local_model_trainer.py` — local container mode makes no service call + +**The rule to preserve:** do not add `gpu_intensive` to a deep test unless a shallow +test covers the same path. Otherwise the PR gate silently loses coverage. ## Relationship to `dry_run=True` diff --git a/sagemaker-train/tests/integ/train/shallow/recipe_cases.py b/sagemaker-train/tests/integ/train/shallow/recipe_cases.py new file mode 100644 index 0000000000..e9ca5f6a6a --- /dev/null +++ b/sagemaker-train/tests/integ/train/shallow/recipe_cases.py @@ -0,0 +1,230 @@ +# Copyright Amazon.com, Inc. or its affiliates. All Rights Reserved. +# +# Licensed under the Apache License, Version 2.0 (the "License"). You +# may not use this file except in compliance with the License. A copy of +# the License is located at +# +# http://aws.amazon.com/apache2.0/ +# +# or in the "license" file accompanying this file. This file is +# distributed on an "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF +# ANY KIND, either express or implied. See the License for the specific +# language governing permissions and limitations under the License. +"""Shared submission cases for the recipe trainers. + +Every recipe trainer (SFT, DPO, RLVR, RLAIF, ...) accepts the same core arguments +and must clear the same server-side gates, so the cases live here once and each +``test__trainer.py`` subclasses them. That keeps one file per trainer -- +matching the existing ``test_sft_trainer_integration.py`` / +``test_dpo_trainer_integration.py`` layout, so the shallow counterpart of a given +deep test is obvious -- without four near-identical copies of the same bodies. + +To add a trainer: create ``test__trainer.py`` with + + class TestFooTrainerSubmission(RecipeTrainerCases): + TRAINER = FooTrainer + +and override the class attributes below only where the trainer genuinely differs. + +This module is deliberately NOT named ``test_*``: pytest must not collect +``RecipeTrainerCases`` directly, since it has no ``TRAINER``. +""" + +from __future__ import absolute_import + +import pytest +from sagemaker.core import shapes +from sagemaker.core.training.configs import TrainingJobCompute +from sagemaker.train.common import TrainingType + +from .harness import ( + MAX_RUNTIME_IN_SECONDS, + assert_rejected, + assert_submitted, + submitted, + unique_name, +) + +# Small, publicly available instruct model. Kept small deliberately: these tests +# never train, so model size only affects how long recipe/artifact resolution +# takes during submission. +MODEL_ID = "meta-textgeneration-llama-3-2-1b-instruct" + +# Reused from the existing dry-run suite so both suites exercise the same +# already-provisioned model package group rather than each needing their own. +MODEL_PACKAGE_GROUP = ( + "arn:aws:sagemaker:us-west-2:729646638167:model-package-group/sdk-test-finetuned-models" +) + +# An accelerator type is required for the serverful recipe path: these recipes do +# not resolve onto a CPU instance, so unlike the ModelTrainer suite we cannot use +# ml.m5.large here. The job is still stopped immediately, so this holds capacity +# only transiently. +SERVERFUL_INSTANCE_TYPE = "ml.g5.12xlarge" + +# Rejection messages can legitimately come from three layers with different +# wording -- SDK-side validation, the public API model, or the training backend -- +# so negative tests accept any of these tokens. Still specific enough to catch a +# *wrong* rejection (e.g. an unrelated credentials error). +_MISSING_DATA_TOKENS = ( + "does not exist", + "ValidationException", + "ValidationError", + "S3", + "not found", +) + + +def stopping_condition(): + """Short advertised runtime. Never reached -- the job is stopped long before -- + but it bounds the damage if a stop were ever lost.""" + return shapes.StoppingCondition(max_runtime_in_seconds=MAX_RUNTIME_IN_SECONDS) + + +class RecipeTrainerCases: + """Submission cases shared by every recipe trainer. + + Subclasses set ``TRAINER`` and, where the trainer differs, the other class + attributes. Each test submits a real ``CreateTrainingJob``, asserts the + service returned an ARN, then stops the job -- see ``harness`` for why a + returned ARN is a strong assertion. + """ + + #: The trainer class under test. Subclasses must set this. + TRAINER = None + + #: Extra constructor arguments this trainer requires (e.g. RLAIF's reward + #: model). Merged on top of the shared kwargs. + EXTRA_KWARGS = {} + + #: Whether the trainer accepts an explicit ``TrainingJobCompute``. RLAIF does + #: not take a ``compute`` argument at all, so it has no serverful path. + SUPPORTS_SERVERFUL = True + + #: Whether the trainer accepts ``training_type`` (LoRA vs full). CPT has no + #: such distinction. + SUPPORTS_TRAINING_TYPE = True + + def build(self, sagemaker_session, dataset, name, **overrides): + """Construct the trainer in its minimal accepted configuration. + + ``accept_eula=True`` is required for gated foundation models; without it + the request is refused before reaching the validation this suite targets. + """ + kwargs = dict( + model=MODEL_ID, + model_package_group=MODEL_PACKAGE_GROUP, + training_dataset=dataset, + accept_eula=True, + sagemaker_session=sagemaker_session, + base_job_name=name, + stopping_condition=stopping_condition(), + ) + if self.SUPPORTS_TRAINING_TYPE: + kwargs["training_type"] = TrainingType.LORA + kwargs.update(self.EXTRA_KWARGS) + kwargs.update(overrides) + return self.TRAINER(**kwargs) + + def name(self, suffix=""): + """Job name prefixed with the trainer, so a job in the console is + traceable back to the test that made it.""" + stem = self.TRAINER.__name__.replace("Trainer", "").lower() + return unique_name(f"shallow-{stem}{suffix}") + + # -- serverless (recipe-derived compute), the default path --------------- + + def test_minimal_request_is_accepted(self, sagemaker_session, train_data_uri): + """Baseline: the simplest well-formed request is accepted. + + Recipe selection and resource-config generation happen server-side after + the request validators, so acceptance here is the cheap proof that the + SDK's recipe payload is still valid. + """ + trainer = self.build(sagemaker_session, train_data_uri, self.name()) + + with submitted(trainer) as job: + assert_submitted(job) + + def test_with_validation_dataset(self, sagemaker_session, train_data_uri, validation_data_uri): + """A validation dataset adds a second channel, resolved against S3 + independently of the training channel.""" + trainer = self.build( + sagemaker_session, + train_data_uri, + self.name("-val"), + validation_dataset=validation_data_uri, + ) + + with submitted(trainer) as job: + assert_submitted(job) + + def test_dataset_passed_to_train_overrides_constructor(self, sagemaker_session, train_data_uri): + """``train(training_dataset=...)`` overrides the constructor value. + + Worth asserting server-side: if the override were dropped the payload + would silently reference the wrong data, and only a real run would show + it. + """ + trainer = self.build(sagemaker_session, None, self.name("-override")) + + with submitted(trainer, training_dataset=train_data_uri) as job: + assert_submitted(job) + + def test_explicit_s3_output_path(self, sagemaker_session, train_data_uri, output_path): + """A caller-specified output location must validate server-side.""" + trainer = self.build( + sagemaker_session, + train_data_uri, + self.name("-output"), + s3_output_path=output_path, + ) + + with submitted(trainer) as job: + assert_submitted(job) + + # -- serverful (explicit TrainingJobCompute) ----------------------------- + + def test_explicit_compute_is_accepted(self, sagemaker_session, train_data_uri): + """Explicit compute produces a materially different payload from the + recipe-derived serverless path, including a resource config the backend + validates against the recipe.""" + if not self.SUPPORTS_SERVERFUL: + pytest.skip(f"{self.TRAINER.__name__} takes no compute argument") + + trainer = self.build( + sagemaker_session, + train_data_uri, + self.name("-serverful"), + compute=TrainingJobCompute(instance_type=SERVERFUL_INSTANCE_TYPE, instance_count=1), + ) + + with submitted(trainer) as job: + assert_submitted(job) + + # -- negative cases ------------------------------------------------------ + + def test_nonexistent_training_dataset_is_rejected( + self, sagemaker_session, nonexistent_data_uri + ): + """Dataset existence is checked against S3 before the job is created. + + The most valuable negative case here: it proves the backend's + role-assuming validators actually ran rather than being skipped. + """ + trainer = self.build(sagemaker_session, nonexistent_data_uri, self.name("-bad-data")) + + assert_rejected(trainer, _MISSING_DATA_TOKENS) + + def test_nonexistent_validation_dataset_is_rejected( + self, sagemaker_session, train_data_uri, nonexistent_data_uri + ): + """A valid training set must not mask an invalid validation set.""" + trainer = self.build( + sagemaker_session, + train_data_uri, + self.name("-bad-val"), + validation_dataset=nonexistent_data_uri, + ) + + assert_rejected(trainer, _MISSING_DATA_TOKENS) diff --git a/sagemaker-train/tests/integ/train/shallow/test_cpt_trainer.py b/sagemaker-train/tests/integ/train/shallow/test_cpt_trainer.py new file mode 100644 index 0000000000..d2647dfbea --- /dev/null +++ b/sagemaker-train/tests/integ/train/shallow/test_cpt_trainer.py @@ -0,0 +1,64 @@ +# Copyright Amazon.com, Inc. or its affiliates. All Rights Reserved. +# +# Licensed under the Apache License, Version 2.0 (the "License"). You +# may not use this file except in compliance with the License. A copy of +# the License is located at +# +# http://aws.amazon.com/apache2.0/ +# +# or in the "license" file accompanying this file. This file is +# distributed on an "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF +# ANY KIND, either express or implied. See the License for the specific +# language governing permissions and limitations under the License. +"""Shallow submission tests for CPTTrainer (continued pre-training). + +Shallow counterpart of test_cpt_hyperpod.py. + +CPT differs from the other recipe trainers in two verified ways: it accepts no +training_type (there is no LoRA/full distinction for continued pre-training), +and its compute is HyperPodCompute-only. + +The whole class is marked gpu_intensive and skips unless a cluster is +configured, because CPT refuses to submit without HyperPod compute -- + + ValueError: CPT requires HyperPod compute. + Pass compute=HyperPodCompute(...) when creating the CPTTrainer. + +-- and HyperPod submits to a pre-provisioned cluster rather than through +CreateTrainingJob, so there is nothing this suite can create on demand. Written in +the shallow style anyway so it becomes gate-eligible by dropping one marker once a +cluster exists in the PR account. +""" + +from __future__ import absolute_import + +import os + +import pytest +from sagemaker.core.training.configs import HyperPodCompute +from sagemaker.train.cpt_trainer import CPTTrainer + +from .harness import assert_submitted, submitted +from .recipe_cases import RecipeTrainerCases + + +@pytest.mark.gpu_intensive +class TestCPTTrainerSubmission(RecipeTrainerCases): + """CPT submits only via HyperPod, so the shared cases are not inherited as-is.""" + + TRAINER = CPTTrainer + SUPPORTS_TRAINING_TYPE = False + SUPPORTS_SERVERFUL = False + + @pytest.fixture(autouse=True) + def _require_hyperpod(self): + """Skip the whole class unless a HyperPod cluster is configured.""" + cluster = os.environ.get("SHALLOW_HYPERPOD_CLUSTER") + if not cluster: + pytest.skip("CPT requires HyperPod; set SHALLOW_HYPERPOD_CLUSTER to run") + self._cluster = cluster + + def build(self, sagemaker_session, dataset, name, **overrides): + """Add the required HyperPod compute to every CPT submission.""" + overrides.setdefault("compute", HyperPodCompute(cluster_name=self._cluster)) + return super().build(sagemaker_session, dataset, name, **overrides) diff --git a/sagemaker-train/tests/integ/train/shallow/test_dpo_trainer.py b/sagemaker-train/tests/integ/train/shallow/test_dpo_trainer.py new file mode 100644 index 0000000000..421aadba71 --- /dev/null +++ b/sagemaker-train/tests/integ/train/shallow/test_dpo_trainer.py @@ -0,0 +1,33 @@ +# Copyright Amazon.com, Inc. or its affiliates. All Rights Reserved. +# +# Licensed under the Apache License, Version 2.0 (the "License"). You +# may not use this file except in compliance with the License. A copy of +# the License is located at +# +# http://aws.amazon.com/apache2.0/ +# +# or in the "license" file accompanying this file. This file is +# distributed on an "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF +# ANY KIND, either express or implied. See the License for the specific +# language governing permissions and limitations under the License. +"""Shallow submission tests for DPOTrainer. + +Shallow counterpart of test_dpo_trainer_integration.py. All cases come from +RecipeTrainerCases; DPO takes the same core arguments as SFT and needs no +overrides. + +Note DPOTrainer exposes its submitted job as the *public* latest_training_job +where the others use _latest_training_job; the harness resolves both. +""" + +from __future__ import absolute_import + +from sagemaker.train.dpo_trainer import DPOTrainer + +from .recipe_cases import RecipeTrainerCases + + +class TestDPOTrainerSubmission(RecipeTrainerCases): + """DPO accepts every shared case with no deviations.""" + + TRAINER = DPOTrainer diff --git a/sagemaker-train/tests/integ/train/shallow/test_model_trainer_submission.py b/sagemaker-train/tests/integ/train/shallow/test_model_trainer.py similarity index 100% rename from sagemaker-train/tests/integ/train/shallow/test_model_trainer_submission.py rename to sagemaker-train/tests/integ/train/shallow/test_model_trainer.py diff --git a/sagemaker-train/tests/integ/train/shallow/test_multi_turn_rl_trainer.py b/sagemaker-train/tests/integ/train/shallow/test_multi_turn_rl_trainer.py new file mode 100644 index 0000000000..d4ac37024e --- /dev/null +++ b/sagemaker-train/tests/integ/train/shallow/test_multi_turn_rl_trainer.py @@ -0,0 +1,116 @@ +# Copyright Amazon.com, Inc. or its affiliates. All Rights Reserved. +# +# Licensed under the Apache License, Version 2.0 (the "License"). You +# may not use this file except in compliance with the License. A copy of +# the License is located at +# +# http://aws.amazon.com/apache2.0/ +# +# or in the "license" file accompanying this file. This file is +# distributed on an "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF +# ANY KIND, either express or implied. See the License for the specific +# language governing permissions and limitations under the License. +"""Shallow submission tests for ``MultiTurnRLTrainer`` (Agentic RFT). + +Shallow counterpart of ``test_multi_turn_rl_trainer_integration.py``. + +MTRL is the one trainer here that does not create a TrainingJob at all: it calls +the generic Job API and returns an ``AgentRFTJob``, so its ARN segment is ``job`` +rather than ``training-job`` and the harness resolves it via ``_latest_job``. +""" + +from __future__ import absolute_import + +import logging +import os + +import pytest +from sagemaker.train.multi_turn_rl_trainer import MultiTurnRLTrainer + +from .harness import assert_submitted, submitted, unique_name + +logger = logging.getLogger(__name__) + + +@pytest.mark.gpu_intensive +class TestMultiTurnRLSubmission: + """AgentRFT Job acceptance for ``MultiTurnRLTrainer``. + + Marked ``gpu_intensive`` (and therefore excluded from the PR gate, per the + marker's definition in ``tox.ini``) because unlike every other test in this + suite it cannot be made self-contained: MTRL requires a pre-provisioned agent + runtime and an MLflow app, neither of which this suite creates. The existing + ``test_multi_turn_rl_trainer_integration.py`` hardcodes both. + + They are still written using the shallow pattern rather than omitted, so that + when the prerequisites are provisioned in the PR account these become + PR-gate-eligible by deleting one marker. Prerequisites are resolved from the + environment and the tests skip when absent, so they never fail for + infrastructure reasons. + """ + + @pytest.fixture(scope="class") + def mtrl_prerequisites(self, sagemaker_session, account_id, region): + """Resolve MTRL prerequisites, skipping if they are not configured. + + Read from the environment rather than hardcoded so this does not bake in + another account-specific constant. + """ + agent_env = os.environ.get("SHALLOW_MTRL_AGENT_ENV") + mlflow_app_arn = os.environ.get("SHALLOW_MTRL_MLFLOW_APP_ARN") + dataset = os.environ.get("SHALLOW_MTRL_DATASET") + + missing = [ + name + for name, value in ( + ("SHALLOW_MTRL_AGENT_ENV", agent_env), + ("SHALLOW_MTRL_MLFLOW_APP_ARN", mlflow_app_arn), + ("SHALLOW_MTRL_DATASET", dataset), + ) + if not value + ] + if missing: + pytest.skip("MTRL prerequisites not configured; set " + ", ".join(missing)) + + return { + "agent_env": agent_env, + "mlflow_app_arn": mlflow_app_arn, + "dataset": dataset, + "model": os.environ.get("SHALLOW_MTRL_MODEL", "mock-oss-test"), + } + + def test_agent_rft_job_is_accepted(self, sagemaker_session, mtrl_prerequisites): + """The AgentRFT job config document must be accepted by the Job API. + + Note the different ARN resource segment: this is a ``job``, not a + ``training-job``. + """ + trainer = MultiTurnRLTrainer( + model=mtrl_prerequisites["model"], + agent_env=mtrl_prerequisites["agent_env"], + training_dataset=mtrl_prerequisites["dataset"], + mlflow_app_arn=mtrl_prerequisites["mlflow_app_arn"], + accept_eula=True, + sagemaker_session=sagemaker_session, + base_job_name=unique_name("shallow-mtrl"), + ) + + with submitted(trainer) as job: + assert_submitted(job, resource="job") + + def test_hyperparameter_mutation_is_accepted(self, sagemaker_session, mtrl_prerequisites): + """``trainer.hyperparameters`` mutation must reach the job config + document, which the service validates on submission.""" + trainer = MultiTurnRLTrainer( + model=mtrl_prerequisites["model"], + agent_env=mtrl_prerequisites["agent_env"], + training_dataset=mtrl_prerequisites["dataset"], + mlflow_app_arn=mtrl_prerequisites["mlflow_app_arn"], + accept_eula=True, + sagemaker_session=sagemaker_session, + base_job_name=unique_name("shallow-mtrl-hp"), + ) + trainer.hyperparameters.global_batch_size = 32 + + with submitted(trainer) as job: + assert_submitted(job, resource="job") diff --git a/sagemaker-train/tests/integ/train/shallow/test_nova_data_mixing.py b/sagemaker-train/tests/integ/train/shallow/test_nova_data_mixing.py new file mode 100644 index 0000000000..28cd9d49cc --- /dev/null +++ b/sagemaker-train/tests/integ/train/shallow/test_nova_data_mixing.py @@ -0,0 +1,87 @@ +# Copyright Amazon.com, Inc. or its affiliates. All Rights Reserved. +# +# Licensed under the Apache License, Version 2.0 (the "License"). You +# may not use this file except in compliance with the License. A copy of +# the License is located at +# +# http://aws.amazon.com/apache2.0/ +# +# or in the "license" file accompanying this file. This file is +# distributed on an "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF +# ANY KIND, either express or implied. See the License for the specific +# language governing permissions and limitations under the License. +"""Shallow submission tests for DataMixingConfig (Nova only). + +Shallow counterpart of test_sft_trainer_data_mixing_integration.py and +test_sft_data_mixing_hyperpod.py. + +DataMixingConfig is serialized into flat per-category hyperparameters. It is +Nova-only, and Nova is exercised in us-east-1 in this repo, so these use +sagemaker_session_us_east_1 and carry the us_east_1 marker -- the PR-gate +job holds us-west-2 credentials only, so they run in the us-east-1 integ job. + +Kept in its own file rather than folded into test_sft_trainer.py because the region +and model differ from every other case there. +""" + +from __future__ import absolute_import + +import pytest +from sagemaker.train.data_mixing_config import DataMixingConfig +from sagemaker.train.sft_trainer import SFTTrainer + +from .harness import assert_submitted, submitted, unique_name +from .recipe_cases import MODEL_PACKAGE_GROUP, stopping_condition + +NOVA_MODEL = "nova-textgeneration-lite-v2" + + +def _nova_sft(session, dataset, name, config): + return SFTTrainer( + model=NOVA_MODEL, + model_package_group=MODEL_PACKAGE_GROUP, + training_dataset=dataset, + accept_eula=True, + sagemaker_session=session, + data_mixing_config=config, + base_job_name=name, + stopping_condition=stopping_condition(), + # The existing data-mixing test sets the recipe name explicitly; keep that + # so the rendered recipe matches what the service expects. + overrides={"name": name}, + ) + + +@pytest.mark.us_east_1 +class TestNovaDataMixingSubmission: + """DataMixingConfig serialization must be accepted by the service.""" + + def test_explicit_percentages(self, sagemaker_session_us_east_1, nova_train_data_uri): + """Per-category percentages must sum to 100 client-side and serialize into + hyperparameters the service accepts.""" + config = DataMixingConfig( + customer_data_percent=70.0, + nova_data_percentages={ + "code": 30.0, + "math": 20.0, + "planning": 10.0, + "instruction-following": 10.0, + "reasoning-instruction-following": 20.0, + "reasoning-math": 10.0, + }, + ) + name = unique_name("shallow-nova-datamix") + trainer = _nova_sft(sagemaker_session_us_east_1, nova_train_data_uri, name, config) + + with submitted(trainer) as job: + assert_submitted(job) + + def test_recipe_defaults(self, sagemaker_session_us_east_1, nova_train_data_uri): + """With nova_data_percentages=None the recipe template's defaults are + used at submission time -- a different serialization path.""" + config = DataMixingConfig(customer_data_percent=80.0) + name = unique_name("shallow-nova-datamix-default") + trainer = _nova_sft(sagemaker_session_us_east_1, nova_train_data_uri, name, config) + + with submitted(trainer) as job: + assert_submitted(job) diff --git a/sagemaker-train/tests/integ/train/shallow/test_recipe_customization_submission.py b/sagemaker-train/tests/integ/train/shallow/test_recipe_customization_submission.py deleted file mode 100644 index 7a444b6f40..0000000000 --- a/sagemaker-train/tests/integ/train/shallow/test_recipe_customization_submission.py +++ /dev/null @@ -1,264 +0,0 @@ -# Copyright Amazon.com, Inc. or its affiliates. All Rights Reserved. -# -# Licensed under the Apache License, Version 2.0 (the "License"). You -# may not use this file except in compliance with the License. A copy of -# the License is located at -# -# http://aws.amazon.com/apache2.0/ -# -# or in the "license" file accompanying this file. This file is -# distributed on an "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF -# ANY KIND, either express or implied. See the License for the specific -# language governing permissions and limitations under the License. -"""Shallow submission tests for recipe customization. - -Covers the knobs that change the *rendered recipe* rather than the plain request -envelope: ``overrides``, an explicit ``recipe`` file, ``sequence_length``, -``DataMixingConfig``, and direct ``trainer.hyperparameters`` mutation. - -Why these matter here specifically ----------------------------------- -The training backend does not merely shape-check a recipe request -- it filters -candidate recipes after the request validators have already passed and refuses -the job outright with ``"No valid recipes found for the given request"`` when -nothing matches. A customization that renders into an unsatisfiable recipe is -therefore *only* detectable by actually submitting. - -Existing coverage of this area is either client-side or expensive: - -* ``test_recipe_override_integration.py`` (35 tests) exercises - ``get_resolved_recipe`` / ``flatten_resolved_recipe`` and never submits, so it - cannot catch a recipe that resolves locally but the service rejects. -* ``test_rlvr_trainer_integration.py::test_rlvr_trainer_nemotron_with_kl_and_recipe`` - does submit a recipe+overrides combination, but on a 30B model with a - two-hour poll loop. - -These tests close that gap at submission cost: they prove the customized payload -is *accepted*, without asserting anything about the resulting training run. -""" - -from __future__ import absolute_import - -import tempfile - -import pytest -import yaml -from sagemaker.core import shapes -from sagemaker.train.common import TrainingType -from sagemaker.train.data_mixing_config import DataMixingConfig -from sagemaker.train.rlvr_trainer import RLVRTrainer -from sagemaker.train.sft_trainer import SFTTrainer - -from .harness import MAX_RUNTIME_IN_SECONDS, assert_submitted, submitted, unique_name - -MODEL_ID = "meta-textgeneration-llama-3-2-1b-instruct" -MODEL_PACKAGE_GROUP = ( - "arn:aws:sagemaker:us-west-2:729646638167:model-package-group/sdk-test-finetuned-models" -) - - -def _stopping_condition(): - return shapes.StoppingCondition(max_runtime_in_seconds=MAX_RUNTIME_IN_SECONDS) - - -def _sft(sagemaker_session, dataset, name, **overrides): - kwargs = dict( - model=MODEL_ID, - training_type=TrainingType.LORA, - model_package_group=MODEL_PACKAGE_GROUP, - training_dataset=dataset, - accept_eula=True, - sagemaker_session=sagemaker_session, - base_job_name=name, - stopping_condition=_stopping_condition(), - ) - kwargs.update(overrides) - return SFTTrainer(**kwargs) - - -class TestRecipeOverrides: - """``overrides`` is merged into the rendered recipe before submission.""" - - def test_training_config_overrides(self, sagemaker_session, train_data_uri): - """Override common training_config values. - - Values are chosen to stay inside the recipe's accepted ranges: the point - is to prove overrides survive rendering into an accepted payload, not to - probe validation bounds (which the negative tests cover). - """ - trainer = _sft( - sagemaker_session, - train_data_uri, - unique_name("shallow-sft-overrides"), - overrides={ - "training_config": { - "learning_rate": 2e-5, - "max_epochs": 1, - } - }, - ) - - with submitted(trainer) as job: - assert_submitted(job) - - def test_explicit_recipe_file(self, sagemaker_session, train_data_uri): - """A caller-supplied recipe YAML must render into an accepted request. - - Mirrors the shape used by the existing Nemotron test, but on a small - model and without the two-hour poll loop. - """ - recipe = {"training_config": {"data": {"max_prompt_length": 1024}}} - with tempfile.NamedTemporaryFile(mode="w", suffix=".yaml", delete=False) as handle: - yaml.dump(recipe, handle) - recipe_path = handle.name - - trainer = _sft( - sagemaker_session, - train_data_uri, - unique_name("shallow-sft-recipe"), - recipe=recipe_path, - ) - - with submitted(trainer) as job: - assert_submitted(job) - - def test_recipe_and_overrides_together(self, sagemaker_session, train_data_uri): - """Recipe file plus overrides: the merge order must still yield an - accepted payload. This is the combination most likely to break, since - both paths mutate the same rendered document.""" - recipe = {"training_config": {"data": {"max_prompt_length": 1024}}} - with tempfile.NamedTemporaryFile(mode="w", suffix=".yaml", delete=False) as handle: - yaml.dump(recipe, handle) - recipe_path = handle.name - - trainer = _sft( - sagemaker_session, - train_data_uri, - unique_name("shallow-sft-recipe-ovr"), - recipe=recipe_path, - overrides={"training_config": {"max_epochs": 1}}, - ) - - with submitted(trainer) as job: - assert_submitted(job) - - def test_direct_hyperparameter_mutation(self, sagemaker_session, train_data_uri): - """``trainer.hyperparameters. = ...`` is a documented pattern - (used by the existing RLVR tests) and must reach the payload intact.""" - trainer = RLVRTrainer( - model=MODEL_ID, - training_type=TrainingType.LORA, - model_package_group=MODEL_PACKAGE_GROUP, - training_dataset=train_data_uri, - accept_eula=True, - sagemaker_session=sagemaker_session, - base_job_name=unique_name("shallow-rlvr-hpmutate"), - stopping_condition=_stopping_condition(), - ) - trainer.hyperparameters.max_epochs = 1 - - with submitted(trainer) as job: - assert_submitted(job) - - -class TestSequenceLength: - """``sequence_length`` selects a different recipe variant. - - Only ``4K`` is parametrized. Verified against AWS: for ``MODEL_ID`` the recipe - catalogue offers exactly one sequence length -- - - ValueError: No recipes found with SequenceLength == 16K. - Available sequence lengths: ['4K'] - - so a ``16K`` case would assert a service-side limitation rather than SDK - behaviour. Left parametrized (rather than inlined) so another value can be - added when a model in this account supports one. - - Note this field also requires the bundled service model -- see the - ``bundled_service_model`` fixture in conftest; the public botocore model has - no ``ServerlessJobConfig.SequenceLength`` yet. - """ - - @pytest.mark.parametrize("sequence_length", ["4K"]) - def test_sequence_length_variants(self, sagemaker_session, train_data_uri, sequence_length): - trainer = _sft( - sagemaker_session, - train_data_uri, - unique_name(f"shallow-sft-seq{sequence_length}"), - sequence_length=sequence_length, - ) - - with submitted(trainer) as job: - assert_submitted(job) - - -class TestDataMixing: - """``DataMixingConfig`` is serialized into flat per-category hyperparameters. - - Nova-only, and Nova is us-east-1 in this repo's fixtures, so these use - ``sagemaker_session_us_east_1`` (inherited from the parent train conftest) - rather than the default-region session. - - Marked ``us_east_1`` to match the existing marker convention in - ``sagemaker-train/tox.ini``; the PR-gate job runs us-west-2 only, so these are - deselected there and run in the us-east-1 job. - """ - - NOVA_MODEL = "nova-textgeneration-lite-v2" - - @pytest.mark.us_east_1 - def test_data_mixing_with_explicit_percentages( - self, sagemaker_session_us_east_1, nova_train_data_uri - ): - """Per-category percentages must sum to 100 client-side and serialize - into hyperparameters the service accepts.""" - config = DataMixingConfig( - customer_data_percent=70.0, - nova_data_percentages={ - "code": 30.0, - "math": 20.0, - "planning": 10.0, - "instruction-following": 10.0, - "reasoning-instruction-following": 20.0, - "reasoning-math": 10.0, - }, - ) - name = unique_name("shallow-sft-datamix") - trainer = SFTTrainer( - model=self.NOVA_MODEL, - training_type=TrainingType.LORA, - model_package_group=MODEL_PACKAGE_GROUP, - training_dataset=nova_train_data_uri, - accept_eula=True, - sagemaker_session=sagemaker_session_us_east_1, - data_mixing_config=config, - base_job_name=name, - # The existing data-mixing test sets the recipe name explicitly; - # keep that so the rendered recipe matches what the service expects. - overrides={"name": name}, - ) - - with submitted(trainer) as job: - assert_submitted(job) - - @pytest.mark.us_east_1 - def test_data_mixing_recipe_defaults(self, sagemaker_session_us_east_1, nova_train_data_uri): - """With ``nova_data_percentages=None`` the recipe template's defaults are - used at submission time -- a different serialization path from the - explicit case above.""" - config = DataMixingConfig(customer_data_percent=80.0) - name = unique_name("shallow-sft-datamix-default") - trainer = SFTTrainer( - model=self.NOVA_MODEL, - training_type=TrainingType.LORA, - model_package_group=MODEL_PACKAGE_GROUP, - training_dataset=nova_train_data_uri, - accept_eula=True, - sagemaker_session=sagemaker_session_us_east_1, - data_mixing_config=config, - base_job_name=name, - overrides={"name": name}, - ) - - with submitted(trainer) as job: - assert_submitted(job) diff --git a/sagemaker-train/tests/integ/train/shallow/test_recipe_trainers_submission.py b/sagemaker-train/tests/integ/train/shallow/test_recipe_trainers_submission.py deleted file mode 100644 index a5306b2b69..0000000000 --- a/sagemaker-train/tests/integ/train/shallow/test_recipe_trainers_submission.py +++ /dev/null @@ -1,398 +0,0 @@ -# Copyright Amazon.com, Inc. or its affiliates. All Rights Reserved. -# -# Licensed under the Apache License, Version 2.0 (the "License"). You -# may not use this file except in compliance with the License. A copy of -# the License is located at -# -# http://aws.amazon.com/apache2.0/ -# -# or in the "license" file accompanying this file. This file is -# distributed on an "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF -# ANY KIND, either express or implied. See the License for the specific -# language governing permissions and limitations under the License. -"""Shallow submission tests for the recipe trainers (SFT / DPO / RLVR / CPT). - -These trainers do far more request-shaping than ``ModelTrainer``: they resolve a -foundation model, select and render a training recipe, derive a resource config -from it, and translate datasets into channels. All of that lands in the -``CreateTrainingJob`` payload, and the training backend validates it -- including -recipe *acceptance*, which is checked after the request validators and rejects -with ``"No valid recipes found for the given request"``. - -That makes submit-then-stop unusually valuable for this family: a recipe -regression is invisible to unit tests (which mock the service) and today is only -caught by a full, expensive training run. - -Serverless (recipe-selected compute) is the default path. Where a test pins -``TrainingJobCompute`` it is asserting the serverful path specifically, since the -two produce materially different payloads. -""" - -from __future__ import absolute_import - -import os - -import pytest -from sagemaker.core import shapes -from sagemaker.core.training.configs import HyperPodCompute, TrainingJobCompute -from sagemaker.train.common import TrainingType -from sagemaker.train.cpt_trainer import CPTTrainer -from sagemaker.train.dpo_trainer import DPOTrainer -from sagemaker.train.rlaif_trainer import RLAIFTrainer -from sagemaker.train.rlvr_trainer import RLVRTrainer -from sagemaker.train.sft_trainer import SFTTrainer - -from .harness import ( - MAX_RUNTIME_IN_SECONDS, - assert_rejected, - assert_submitted, - submitted, - unique_name, -) - -# Small, publicly available instruct model. Kept small deliberately: these tests -# never train, so model size only affects how long recipe/artifact resolution -# takes during submission. -MODEL_ID = "meta-textgeneration-llama-3-2-1b-instruct" - -# Reused from the existing dry-run suite so both suites exercise the same -# already-provisioned model package group rather than each needing their own. -MODEL_PACKAGE_GROUP = ( - "arn:aws:sagemaker:us-west-2:729646638167:model-package-group/sdk-test-finetuned-models" -) - -# An accelerator type is required for the serverful recipe path: the recipes for -# these trainers will not resolve onto a CPU instance, so unlike the -# ModelTrainer suite we cannot use ml.m5.large here. The job is still stopped -# immediately, so this requests capacity only transiently. -SERVERFUL_INSTANCE_TYPE = "ml.g5.12xlarge" - -# RLAIF requires a reward model and prompt; without them the request is refused -# before it reaches the validation this suite cares about. Values match the -# existing test_rlaif_trainer_integration.py so both suites exercise the same -# already-entitled reward model. -RLAIF_REWARD_MODEL_ID = "openai.gpt-oss-120b-1:0" -RLAIF_REWARD_PROMPT = "Builtin.Summarize" - -# Per-trainer extra constructor arguments. Everything else is shared, which is -# what lets these four trainers be covered by one parametrized body instead of -# four near-identical files. -_TRAINER_EXTRA_KWARGS = { - "RLAIFTrainer": { - "reward_model_id": RLAIF_REWARD_MODEL_ID, - "reward_prompt": RLAIF_REWARD_PROMPT, - }, -} - -# Every recipe trainer takes the same core arguments, so the per-trainer test -# bodies stay a single call. -RECIPE_TRAINERS = [ - pytest.param(SFTTrainer, id="sft"), - pytest.param(DPOTrainer, id="dpo"), - pytest.param(RLVRTrainer, id="rlvr"), - pytest.param(RLAIFTrainer, id="rlaif"), -] - -# Subset that accepts an explicit TrainingJobCompute. RLAIFTrainer takes no -# ``compute`` argument at all (verified against the SDK), so it has no serverful -# path and is excluded rather than being expected to fail. -SERVERFUL_CAPABLE_TRAINERS = [t for t in RECIPE_TRAINERS if t.values[0] is not RLAIFTrainer] - - -def _stopping_condition(): - return shapes.StoppingCondition(max_runtime_in_seconds=MAX_RUNTIME_IN_SECONDS) - - -def _trainer(trainer_cls, sagemaker_session, dataset, name, **overrides): - """Build a recipe trainer in its minimal accepted configuration. - - ``accept_eula=True`` is required for gated foundation models; without it the - request is refused before it reaches the interesting validation. - - Trainer-specific required arguments come from ``_TRAINER_EXTRA_KWARGS`` so - adding another trainer to ``RECIPE_TRAINERS`` stays a two-line change. - """ - kwargs = dict( - model=MODEL_ID, - training_type=TrainingType.LORA, - model_package_group=MODEL_PACKAGE_GROUP, - training_dataset=dataset, - accept_eula=True, - sagemaker_session=sagemaker_session, - base_job_name=name, - stopping_condition=_stopping_condition(), - ) - kwargs.update(_TRAINER_EXTRA_KWARGS.get(trainer_cls.__name__, {})) - kwargs.update(overrides) - return trainer_cls(**kwargs) - - -class TestServerlessSubmission: - """The default path: compute is derived from the selected recipe. - - Recipe selection and resource-config generation happen server-side after the - request validators, so acceptance here is the only cheap proof that the - SDK's recipe payload is still valid. - """ - - @pytest.mark.parametrize("trainer_cls", RECIPE_TRAINERS) - def test_minimal_request_is_accepted(self, trainer_cls, sagemaker_session, train_data_uri): - name = unique_name(f"shallow-{trainer_cls.__name__.lower()}") - trainer = _trainer(trainer_cls, sagemaker_session, train_data_uri, name) - - with submitted(trainer) as job: - assert_submitted(job) - - @pytest.mark.parametrize("trainer_cls", RECIPE_TRAINERS) - def test_with_validation_dataset( - self, trainer_cls, sagemaker_session, train_data_uri, validation_data_uri - ): - """A validation dataset adds a second channel, which is resolved against - S3 independently of the training channel.""" - name = unique_name(f"shallow-{trainer_cls.__name__.lower()}-val") - trainer = _trainer( - trainer_cls, - sagemaker_session, - train_data_uri, - name, - validation_dataset=validation_data_uri, - ) - - with submitted(trainer) as job: - assert_submitted(job) - - @pytest.mark.parametrize("trainer_cls", RECIPE_TRAINERS) - def test_datasets_passed_to_train_override_constructor( - self, trainer_cls, sagemaker_session, train_data_uri - ): - """``train(training_dataset=...)`` overrides the constructor value. - - Worth asserting server-side: if the override were dropped, the payload - would silently reference the wrong data and only a real run would reveal - it. - """ - name = unique_name(f"shallow-{trainer_cls.__name__.lower()}-override") - trainer = _trainer(trainer_cls, sagemaker_session, None, name) - - with submitted(trainer, training_dataset=train_data_uri) as job: - assert_submitted(job) - - # Only LORA is parametrized. Verified against AWS: for MODEL_ID there is no - # serverless (SMTJ) recipe for full fine-tuning -- - # ValueError: No recipes found with Smtj for technique: SFT, - # training_type:TrainingType.FULL - # so a FULL case here would assert a recipe-catalogue limitation rather than - # SDK behaviour. Kept parametrized so FULL can be re-added against a model - # that supports it, rather than the distinction being silently dropped. - @pytest.mark.parametrize("training_type", [TrainingType.LORA]) - def test_training_types(self, sagemaker_session, train_data_uri, training_type): - """LoRA and full fine-tuning select different recipes, so each must be - independently accepted.""" - suffix = str(getattr(training_type, "value", training_type)).lower() - name = unique_name(f"shallow-sft-{suffix}") - trainer = _trainer( - SFTTrainer, - sagemaker_session, - train_data_uri, - name, - training_type=training_type, - ) - - with submitted(trainer) as job: - assert_submitted(job) - - @pytest.mark.gpu_intensive - def test_cpt_trainer_is_accepted(self, sagemaker_session, train_data_uri): - """Continued pre-training, which submits only via HyperPod. - - Kept out of RECIPE_TRAINERS because its constructor genuinely differs: - verified against the SDK, ``CPTTrainer`` accepts no ``training_type`` - (there is no LoRA/full distinction for continued pre-training) and its - ``compute`` is ``HyperPodCompute``-only. - - Marked ``gpu_intensive`` and skipped unless a cluster is configured. CPT - refuses to submit without one -- - - ValueError: CPT requires HyperPod compute. - Pass compute=HyperPodCompute(...) when creating the CPTTrainer. - - -- and HyperPod submits to a pre-provisioned cluster rather than through - CreateTrainingJob, so there is nothing this suite can create on demand. - Written in the shallow style anyway so it becomes gate-eligible by - dropping one marker once a cluster exists in the PR account. - """ - cluster_name = os.environ.get("SHALLOW_HYPERPOD_CLUSTER") - if not cluster_name: - pytest.skip("CPT requires HyperPod; set SHALLOW_HYPERPOD_CLUSTER to run") - - name = unique_name("shallow-cpt") - trainer = CPTTrainer( - model=MODEL_ID, - model_package_group=MODEL_PACKAGE_GROUP, - training_dataset=train_data_uri, - accept_eula=True, - sagemaker_session=sagemaker_session, - base_job_name=name, - stopping_condition=_stopping_condition(), - compute=HyperPodCompute(cluster_name=cluster_name), - ) - - with submitted(trainer) as job: - assert_submitted(job) - - -class TestServerfulSubmission: - """Explicit ``TrainingJobCompute`` produces a materially different payload - from the recipe-derived serverless path, including a resource config the - backend validates against the recipe. - - RLAIF is absent from this class on purpose: verified against the SDK, - ``RLAIFTrainer.__init__`` has no ``compute`` parameter at all, so it has no - serverful path to exercise. It is still covered by every serverless case in - ``TestServerlessSubmission``. - """ - - @pytest.mark.parametrize("trainer_cls", SERVERFUL_CAPABLE_TRAINERS) - def test_explicit_compute_is_accepted(self, trainer_cls, sagemaker_session, train_data_uri): - name = unique_name(f"shallow-{trainer_cls.__name__.lower()}-serverful") - trainer = _trainer( - trainer_cls, - sagemaker_session, - train_data_uri, - name, - compute=TrainingJobCompute(instance_type=SERVERFUL_INSTANCE_TYPE, instance_count=1), - ) - - with submitted(trainer) as job: - assert_submitted(job) - - -class TestOutputAndTracking: - """Output location and MLflow tracking are validated server-side.""" - - def test_explicit_s3_output_path(self, sagemaker_session, train_data_uri, output_path): - name = unique_name("shallow-sft-output") - trainer = _trainer( - SFTTrainer, - sagemaker_session, - train_data_uri, - name, - s3_output_path=output_path, - ) - - with submitted(trainer) as job: - assert_submitted(job) - - def test_disable_output_compression(self, sagemaker_session, train_data_uri): - """Uncompressed output changes the OutputDataConfig the SDK sends.""" - name = unique_name("shallow-sft-nocompress") - trainer = _trainer( - SFTTrainer, - sagemaker_session, - train_data_uri, - name, - disable_output_compression=True, - ) - - with submitted(trainer) as job: - assert_submitted(job) - - -class TestRejectedRecipeRequests: - """Negative cases specific to the recipe path. - - These matter more here than for ``ModelTrainer``: recipe resolution is the - part of the payload most likely to drift, and an over-permissive change would - otherwise still yield a green suite. - """ - - def test_nonexistent_training_dataset_is_rejected( - self, sagemaker_session, nonexistent_data_uri - ): - """Dataset existence is checked against S3 before the job is created.""" - trainer = _trainer( - SFTTrainer, - sagemaker_session, - nonexistent_data_uri, - unique_name("shallow-sft-bad-data"), - ) - - assert_rejected( - trainer, - ("does not exist", "ValidationException", "ValidationError", "S3", "not found"), - ) - - def test_nonexistent_validation_dataset_is_rejected( - self, sagemaker_session, train_data_uri, nonexistent_data_uri - ): - """A valid training set must not mask an invalid validation set.""" - trainer = _trainer( - SFTTrainer, - sagemaker_session, - train_data_uri, - unique_name("shallow-sft-bad-val"), - validation_dataset=nonexistent_data_uri, - ) - - assert_rejected( - trainer, - ("does not exist", "ValidationException", "ValidationError", "S3", "not found"), - ) - - def test_unknown_model_is_rejected(self, sagemaker_session, train_data_uri): - """Model resolution must fail for a model that does not exist. - - Guards the JumpStart/hub lookup that turns ``model`` into a concrete - artifact URI in the payload. - """ - trainer_kwargs = dict( - model="definitely-not-a-real-model-id-4b91c7", - training_type=TrainingType.LORA, - model_package_group=MODEL_PACKAGE_GROUP, - training_dataset=train_data_uri, - accept_eula=True, - sagemaker_session=sagemaker_session, - base_job_name=unique_name("shallow-sft-bad-model"), - ) - - # Model resolution can fail either while constructing the trainer or at - # submit time depending on how the id is interpreted, so both are allowed - # here; what matters is that an unknown model never reaches the service. - with pytest.raises(Exception) as excinfo: - trainer = SFTTrainer(**trainer_kwargs) - trainer.train(wait=False) - - message = str(excinfo.value) - assert any( - token in message - for token in ( - "model", - "Model", - "not found", - "does not exist", - "ResourceNotFound", - "ValidationException", - "ValidationError", - ) - ), f"unexpected rejection reason: {message}" - - def test_invalid_instance_type_is_rejected(self, sagemaker_session, train_data_uri): - """A nonexistent instance type must be refused on the serverful path.""" - trainer = _trainer( - SFTTrainer, - sagemaker_session, - train_data_uri, - unique_name("shallow-sft-bad-instance"), - compute=TrainingJobCompute(instance_type="ml.nonexistent.24xlarge", instance_count=1), - ) - - assert_rejected( - trainer, - ( - "instance", - "Instance", - "not supported", - "ValidationException", - "ValidationError", - ), - ) diff --git a/sagemaker-train/tests/integ/train/shallow/test_rlaif_trainer.py b/sagemaker-train/tests/integ/train/shallow/test_rlaif_trainer.py new file mode 100644 index 0000000000..76f76be524 --- /dev/null +++ b/sagemaker-train/tests/integ/train/shallow/test_rlaif_trainer.py @@ -0,0 +1,40 @@ +# Copyright Amazon.com, Inc. or its affiliates. All Rights Reserved. +# +# Licensed under the Apache License, Version 2.0 (the "License"). You +# may not use this file except in compliance with the License. A copy of +# the License is located at +# +# http://aws.amazon.com/apache2.0/ +# +# or in the "license" file accompanying this file. This file is +# distributed on an "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF +# ANY KIND, either express or implied. See the License for the specific +# language governing permissions and limitations under the License. +"""Shallow submission tests for RLAIFTrainer. + +Shallow counterpart of test_rlaif_trainer_integration.py. +""" + +from __future__ import absolute_import + +from sagemaker.train.rlaif_trainer import RLAIFTrainer + +from .recipe_cases import RecipeTrainerCases + +# Values match the existing test_rlaif_trainer_integration.py so both suites +# exercise the same already-entitled reward model. +REWARD_MODEL_ID = "openai.gpt-oss-120b-1:0" +REWARD_PROMPT = "Builtin.Summarize" + + +class TestRLAIFTrainerSubmission(RecipeTrainerCases): + """RLAIF needs a reward model and prompt, and has no serverful path. + + Verified against the SDK: RLAIFTrainer.__init__ takes no compute + argument at all, so the shared serverful case is skipped rather than expected + to fail. + """ + + TRAINER = RLAIFTrainer + EXTRA_KWARGS = {"reward_model_id": REWARD_MODEL_ID, "reward_prompt": REWARD_PROMPT} + SUPPORTS_SERVERFUL = False diff --git a/sagemaker-train/tests/integ/train/shallow/test_rlvr_trainer.py b/sagemaker-train/tests/integ/train/shallow/test_rlvr_trainer.py new file mode 100644 index 0000000000..b5a9ede54a --- /dev/null +++ b/sagemaker-train/tests/integ/train/shallow/test_rlvr_trainer.py @@ -0,0 +1,98 @@ +# Copyright Amazon.com, Inc. or its affiliates. All Rights Reserved. +# +# Licensed under the Apache License, Version 2.0 (the "License"). You +# may not use this file except in compliance with the License. A copy of +# the License is located at +# +# http://aws.amazon.com/apache2.0/ +# +# or in the "license" file accompanying this file. This file is +# distributed on an "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF +# ANY KIND, either express or implied. See the License for the specific +# language governing permissions and limitations under the License. +"""Shallow submission tests for RLVRTrainer. + +Shallow counterpart of test_rlvr_trainer_integration.py. Adds the +recipe-customization cases, since RLVR is where the existing deep suite exercises +recipe files and overrides (on a 30B model with a two-hour poll loop). +""" + +from __future__ import absolute_import + +import tempfile + +import yaml +from sagemaker.train.rlvr_trainer import RLVRTrainer + +from .harness import assert_submitted, submitted +from .recipe_cases import RecipeTrainerCases + + +class TestRLVRTrainerSubmission(RecipeTrainerCases): + """RLVR accepts every shared case, plus recipe customization.""" + + TRAINER = RLVRTrainer + + def test_direct_hyperparameter_mutation(self, sagemaker_session, train_data_uri): + """trainer.hyperparameters. = ... is a documented pattern (used + by the existing RLVR tests) and must reach the payload intact.""" + trainer = self.build(sagemaker_session, train_data_uri, self.name("-hpmutate")) + trainer.hyperparameters.max_epochs = 1 + + with submitted(trainer) as job: + assert_submitted(job) + + def test_explicit_recipe_file(self, sagemaker_session, train_data_uri): + """A caller-supplied recipe YAML must render into an accepted request. + + Mirrors the shape used by the existing Nemotron test, but on a small model + and without the poll loop. + """ + recipe = {"training_config": {"data": {"max_prompt_length": 1024}}} + with tempfile.NamedTemporaryFile(mode="w", suffix=".yaml", delete=False) as handle: + yaml.dump(recipe, handle) + recipe_path = handle.name + + trainer = self.build( + sagemaker_session, train_data_uri, self.name("-recipe"), recipe=recipe_path + ) + + with submitted(trainer) as job: + assert_submitted(job) + + def test_recipe_and_overrides_together(self, sagemaker_session, train_data_uri): + """Recipe file plus overrides: the merge order must still yield an accepted + payload. The combination most likely to break, since both paths mutate the + same rendered document.""" + recipe = {"training_config": {"data": {"max_prompt_length": 1024}}} + with tempfile.NamedTemporaryFile(mode="w", suffix=".yaml", delete=False) as handle: + yaml.dump(recipe, handle) + recipe_path = handle.name + + trainer = self.build( + sagemaker_session, + train_data_uri, + self.name("-recipe-ovr"), + recipe=recipe_path, + overrides={"training_config": {"max_epochs": 1}}, + ) + + with submitted(trainer) as job: + assert_submitted(job) + + def test_training_config_overrides(self, sagemaker_session, train_data_uri): + """Override common training_config values. + + Values stay inside the recipe's accepted ranges: the point is that + overrides survive rendering into an accepted payload, not to probe + validation bounds (the negative cases cover that). + """ + trainer = self.build( + sagemaker_session, + train_data_uri, + self.name("-overrides"), + overrides={"training_config": {"learning_rate": 2e-5, "max_epochs": 1}}, + ) + + with submitted(trainer) as job: + assert_submitted(job) diff --git a/sagemaker-train/tests/integ/train/shallow/test_sft_trainer.py b/sagemaker-train/tests/integ/train/shallow/test_sft_trainer.py new file mode 100644 index 0000000000..b364a577f9 --- /dev/null +++ b/sagemaker-train/tests/integ/train/shallow/test_sft_trainer.py @@ -0,0 +1,75 @@ +# Copyright Amazon.com, Inc. or its affiliates. All Rights Reserved. +# +# Licensed under the Apache License, Version 2.0 (the "License"). You +# may not use this file except in compliance with the License. A copy of +# the License is located at +# +# http://aws.amazon.com/apache2.0/ +# +# or in the "license" file accompanying this file. This file is +# distributed on an "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF +# ANY KIND, either express or implied. See the License for the specific +# language governing permissions and limitations under the License. +"""Shallow submission tests for SFTTrainer. + +Shallow counterpart of test_sft_trainer_integration.py: submits a real +CreateTrainingJob, asserts the returned ARN, then stops the job. Asserts +acceptance only, never training behaviour. + +The shared cases come from RecipeTrainerCases; SFT-specific ones are added +below. +""" + +from __future__ import absolute_import + +import pytest +from sagemaker.train.sft_trainer import SFTTrainer + +from .harness import assert_submitted, submitted +from .recipe_cases import RecipeTrainerCases + + +class TestSFTTrainerSubmission(RecipeTrainerCases): + """SFT accepts every shared case with no deviations.""" + + TRAINER = SFTTrainer + + @pytest.mark.parametrize("sequence_length", ["4K"]) + def test_sequence_length_is_accepted(self, sagemaker_session, train_data_uri, sequence_length): + """sequence_length selects a different recipe variant. + + Only 4K is parametrized. Verified against AWS: for MODEL_ID the recipe + catalogue offers exactly one sequence length -- + + ValueError: No recipes found with SequenceLength == 16K. + Available sequence lengths: ['4K'] + + -- so a 16K case would assert a service-side limitation rather than SDK + behaviour. Left parametrized so another value can be added against a model + that supports one. + + Also requires the bundled service model: the public botocore model has no + ServerlessJobConfig.SequenceLength (see the bundled_service_model + fixture in conftest). + """ + trainer = self.build( + sagemaker_session, + train_data_uri, + self.name(f"-seq{sequence_length}"), + sequence_length=sequence_length, + ) + + with submitted(trainer) as job: + assert_submitted(job) + + def test_disable_output_compression(self, sagemaker_session, train_data_uri): + """Uncompressed output changes the OutputDataConfig the SDK sends.""" + trainer = self.build( + sagemaker_session, + train_data_uri, + self.name("-nocompress"), + disable_output_compression=True, + ) + + with submitted(trainer) as job: + assert_submitted(job) diff --git a/sagemaker-train/tests/integ/train/shallow/test_other_job_types_submission.py b/sagemaker-train/tests/integ/train/shallow/test_tuner.py similarity index 68% rename from sagemaker-train/tests/integ/train/shallow/test_other_job_types_submission.py rename to sagemaker-train/tests/integ/train/shallow/test_tuner.py index af87036df6..94afad3e6d 100644 --- a/sagemaker-train/tests/integ/train/shallow/test_other_job_types_submission.py +++ b/sagemaker-train/tests/integ/train/shallow/test_tuner.py @@ -190,87 +190,3 @@ def test_distributed_tuning_job_is_accepted(self, sagemaker_session): assert ( "sm_drivers" in channels ), f"tuning job {arn} is missing the sm_drivers channel; channels={channels}" - - -@pytest.mark.gpu_intensive -class TestMultiTurnRLSubmission: - """AgentRFT Job acceptance for ``MultiTurnRLTrainer``. - - Marked ``gpu_intensive`` (and therefore excluded from the PR gate, per the - marker's definition in ``tox.ini``) because unlike every other test in this - suite it cannot be made self-contained: MTRL requires a pre-provisioned agent - runtime and an MLflow app, neither of which this suite creates. The existing - ``test_multi_turn_rl_trainer_integration.py`` hardcodes both. - - They are still written using the shallow pattern rather than omitted, so that - when the prerequisites are provisioned in the PR account these become - PR-gate-eligible by deleting one marker. Prerequisites are resolved from the - environment and the tests skip when absent, so they never fail for - infrastructure reasons. - """ - - @pytest.fixture(scope="class") - def mtrl_prerequisites(self, sagemaker_session, account_id, region): - """Resolve MTRL prerequisites, skipping if they are not configured. - - Read from the environment rather than hardcoded so this does not bake in - another account-specific constant. - """ - agent_env = os.environ.get("SHALLOW_MTRL_AGENT_ENV") - mlflow_app_arn = os.environ.get("SHALLOW_MTRL_MLFLOW_APP_ARN") - dataset = os.environ.get("SHALLOW_MTRL_DATASET") - - missing = [ - name - for name, value in ( - ("SHALLOW_MTRL_AGENT_ENV", agent_env), - ("SHALLOW_MTRL_MLFLOW_APP_ARN", mlflow_app_arn), - ("SHALLOW_MTRL_DATASET", dataset), - ) - if not value - ] - if missing: - pytest.skip("MTRL prerequisites not configured; set " + ", ".join(missing)) - - return { - "agent_env": agent_env, - "mlflow_app_arn": mlflow_app_arn, - "dataset": dataset, - "model": os.environ.get("SHALLOW_MTRL_MODEL", "mock-oss-test"), - } - - def test_agent_rft_job_is_accepted(self, sagemaker_session, mtrl_prerequisites): - """The AgentRFT job config document must be accepted by the Job API. - - Note the different ARN resource segment: this is a ``job``, not a - ``training-job``. - """ - trainer = MultiTurnRLTrainer( - model=mtrl_prerequisites["model"], - agent_env=mtrl_prerequisites["agent_env"], - training_dataset=mtrl_prerequisites["dataset"], - mlflow_app_arn=mtrl_prerequisites["mlflow_app_arn"], - accept_eula=True, - sagemaker_session=sagemaker_session, - base_job_name=unique_name("shallow-mtrl"), - ) - - with submitted(trainer) as job: - assert_submitted(job, resource="job") - - def test_hyperparameter_mutation_is_accepted(self, sagemaker_session, mtrl_prerequisites): - """``trainer.hyperparameters`` mutation must reach the job config - document, which the service validates on submission.""" - trainer = MultiTurnRLTrainer( - model=mtrl_prerequisites["model"], - agent_env=mtrl_prerequisites["agent_env"], - training_dataset=mtrl_prerequisites["dataset"], - mlflow_app_arn=mtrl_prerequisites["mlflow_app_arn"], - accept_eula=True, - sagemaker_session=sagemaker_session, - base_job_name=unique_name("shallow-mtrl-hp"), - ) - trainer.hyperparameters.global_batch_size = 32 - - with submitted(trainer) as job: - assert_submitted(job, resource="job") diff --git a/sagemaker-train/tests/integ/train/test_benchmark_evaluator.py b/sagemaker-train/tests/integ/train/test_benchmark_evaluator.py index 4a71961916..23f21229c3 100644 --- a/sagemaker-train/tests/integ/train/test_benchmark_evaluator.py +++ b/sagemaker-train/tests/integ/train/test_benchmark_evaluator.py @@ -97,7 +97,6 @@ def test_get_benchmarks_and_properties(self): logger.info(f"MMLU properties: {properties}") - @pytest.mark.gpu_intensive def test_benchmark_evaluation_full_flow(self): """ Test complete benchmark evaluation flow with fine-tuned model package. diff --git a/sagemaker-train/tests/integ/train/test_custom_scorer_evaluator.py b/sagemaker-train/tests/integ/train/test_custom_scorer_evaluator.py index b8569ea336..f0f0968c07 100644 --- a/sagemaker-train/tests/integ/train/test_custom_scorer_evaluator.py +++ b/sagemaker-train/tests/integ/train/test_custom_scorer_evaluator.py @@ -86,7 +86,6 @@ def test_get_builtin_metrics(self): logger.info(f"Built-in metrics: {list(BuiltInMetric.__members__.keys())}") - @pytest.mark.gpu_intensive def test_custom_scorer_evaluation_full_flow(self): """ Test complete custom scorer evaluation flow with custom evaluator ARN. diff --git a/sagemaker-train/tests/integ/train/test_inspect_ai_evaluator.py b/sagemaker-train/tests/integ/train/test_inspect_ai_evaluator.py index 3155579d37..d045d49e13 100644 --- a/sagemaker-train/tests/integ/train/test_inspect_ai_evaluator.py +++ b/sagemaker-train/tests/integ/train/test_inspect_ai_evaluator.py @@ -113,7 +113,6 @@ def inspect_ai_resources(sagemaker_session_us_east_1): class TestInspectAIEvaluatorIntegration: """Integration tests for InspectAI evaluation with Bedrock inference.""" - @pytest.mark.gpu_intensive def test_inspect_ai_bedrock_evaluation( self, sagemaker_session_us_east_1, inspect_ai_resources ): @@ -162,7 +161,6 @@ def test_inspect_ai_bedrock_evaluation( execution.show_results() logger.info("InspectAI Bedrock evaluation completed successfully.") - @pytest.mark.gpu_intensive def test_inspect_ai_upload_benchmarks( self, sagemaker_session_us_east_1, inspect_ai_resources ): diff --git a/sagemaker-train/tests/integ/train/test_llm_as_judge_base_model_fix.py b/sagemaker-train/tests/integ/train/test_llm_as_judge_base_model_fix.py index e4c62ba1c8..2c188a8f5d 100644 --- a/sagemaker-train/tests/integ/train/test_llm_as_judge_base_model_fix.py +++ b/sagemaker-train/tests/integ/train/test_llm_as_judge_base_model_fix.py @@ -100,7 +100,6 @@ def _get_latest_model_package_arn(): class TestLLMAsJudgeBaseModelFix: """Integration test for base model fix in LLMAsJudgeEvaluator""" - @pytest.mark.gpu_intensive def test_base_model_evaluation_uses_correct_weights(self, mlflow_resource_arn): """ Test that base model evaluation uses original base model weights. @@ -279,7 +278,6 @@ def test_base_model_evaluation_uses_correct_weights(self, mlflow_resource_arn): # Re-raise to fail the test raise - @pytest.mark.gpu_intensive def test_base_model_false_still_works(self, mlflow_resource_arn): """ Test that evaluate_base_model=False still works correctly (backward compatibility). diff --git a/sagemaker-train/tests/integ/train/test_llm_as_judge_evaluator.py b/sagemaker-train/tests/integ/train/test_llm_as_judge_evaluator.py index c6b665af6e..4907a7317c 100644 --- a/sagemaker-train/tests/integ/train/test_llm_as_judge_evaluator.py +++ b/sagemaker-train/tests/integ/train/test_llm_as_judge_evaluator.py @@ -88,7 +88,6 @@ class TestLLMAsJudgeEvaluatorIntegration: """Integration tests for LLMAsJudgeEvaluator""" - @pytest.mark.gpu_intensive def test_llm_as_judge_evaluation_full_flow(self): """ Test complete LLM-as-Judge evaluation flow with custom and built-in metrics. diff --git a/sagemaker-train/tests/integ/train/test_llmaj_custom_model.py b/sagemaker-train/tests/integ/train/test_llmaj_custom_model.py index 65ffd45e1f..e3277e9509 100644 --- a/sagemaker-train/tests/integ/train/test_llmaj_custom_model.py +++ b/sagemaker-train/tests/integ/train/test_llmaj_custom_model.py @@ -98,7 +98,6 @@ def test_resources(sagemaker_session_us_east_1): class TestLLMAJCustomModelIntegration: """Integration tests for LLMAsJudgeEvaluator with InspectAI inference path.""" - @pytest.mark.gpu_intensive def test_llmaj_bedrock_inference_end_to_end( self, sagemaker_session_us_east_1, test_resources ): diff --git a/sagemaker-train/tests/integ/train/test_notifications.py b/sagemaker-train/tests/integ/train/test_notifications.py index 26aad2467b..789391755a 100644 --- a/sagemaker-train/tests/integ/train/test_notifications.py +++ b/sagemaker-train/tests/integ/train/test_notifications.py @@ -160,7 +160,6 @@ def sqs_subscriber(sm_session): logger.warning(f"Failed to delete queue: {e}") -@pytest.mark.gpu_intensive @pytest.mark.us_east_1 def test_notifications_creates_eventbridge_rule_and_cleanup( sm_session, training_data_uri, sqs_subscriber From b9b25e333a1da214640fe5eaeed0e706138322e1 Mon Sep 17 00:00:00 2001 From: Amarjeet LNU Date: Tue, 11 Aug 2026 19:13:15 -0700 Subject: [PATCH 05/15] 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. --- .../tests/integ/train/shallow/README.md | 83 +++++++++++----- .../tests/integ/train/shallow/conftest.py | 78 +++++++++++++++ .../tests/integ/train/shallow/recipe_cases.py | 41 ++++++++ .../integ/train/shallow/test_nova_trainers.py | 96 +++++++++++++++++++ .../integ/train/shallow/test_rlaif_trainer.py | 46 +++++++++ .../integ/train/shallow/test_rlvr_trainer.py | 67 +++++++++++++ 6 files changed, 388 insertions(+), 23 deletions(-) create mode 100644 sagemaker-train/tests/integ/train/shallow/test_nova_trainers.py diff --git a/sagemaker-train/tests/integ/train/shallow/README.md b/sagemaker-train/tests/integ/train/shallow/README.md index 23975953fa..bb37399280 100644 --- a/sagemaker-train/tests/integ/train/shallow/README.md +++ b/sagemaker-train/tests/integ/train/shallow/README.md @@ -64,36 +64,73 @@ where the trainer genuinely differs: It is deliberately not named `test_*` so pytest does not collect the base class. -## What was marked `gpu_intensive`, and why only those +## Coverage of every `gpu_intensive` test -A deep test is only marked `gpu_intensive` (i.e. moved off the PR gate) when this -suite has a shallow test covering the same code path. 10 tests met that bar: +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`; the table +below accounts for all of them. -| Deep test (now marked) | Shallow equivalent | +### Covered by this suite + +| Deep test | Shallow equivalent | |---|---| -| `test_model_trainer.py::test_source_dir_local_tar_file` | `TestSourceCodePackaging::test_local_tar_file_source_dir` | -| `::test_hp_contract_basic_py_script` | `TestMinimalSubmission::test_minimal_request_is_accepted` | -| `::test_hp_contract_basic_sh_script` | `TestSourceCodePackaging::test_shell_entry_script` | -| `::test_hp_contract_mpi_script` | `TestComputeConfiguration::test_mpi_distributed` | -| `::test_hp_contract_torchrun_script` | `TestComputeConfiguration::test_torchrun_distributed` | -| `::test_hp_contract_hyperparameter_json` | `TestPayloadShaping::test_hyperparameters_from_json_file` | -| `::test_hp_contract_hyperparameter_yaml` | `TestPayloadShaping::test_hyperparameters_from_yaml_file` | -| `::test_custom_distributed_driver` | `TestSourceCodePackaging::test_custom_distributed_driver` | -| `test_sft_trainer_integration.py::test_sft_trainer_lora_with_sequence_length` | `test_sft_trainer.py::test_sequence_length_is_accepted` | +| `test_model_trainer.py` — 8 tests (tar source, py/sh entry, MPI, torchrun, HP json/yaml, custom driver) | `test_model_trainer.py` — `TestSourceCodePackaging`, `TestPayloadShaping`, `TestComputeConfiguration` | +| `test_sft_trainer_integration.py::test_sft_trainer_lora_complete_workflow` | `test_minimal_request_is_accepted` + `test_mlflow_resource_arn` | +| `::test_sft_trainer_with_validation_dataset` | `test_with_validation_dataset` | +| `::test_sft_trainer_lora_with_sequence_length` | `test_sft_trainer.py::test_sequence_length_is_accepted` | +| `::test_sft_trainer_nova_workflow` | `test_nova_trainers.py::test_nova_sft_is_accepted` | +| `test_dpo_trainer_integration.py` — both tests | `test_dpo_trainer.py` (inherits the shared cases) | +| `test_rlaif_trainer_integration.py::test_rlaif_trainer_lora_complete_workflow` | `test_minimal_request_is_accepted` | +| `::test_rlaif_trainer_with_custom_reward_settings` | `test_rlaif_trainer.py::test_reward_prompt_as_arn` | +| `::test_rlaif_trainer_continued_finetuning` | `::test_continued_finetuning_from_model_package` | +| `test_rlvr_trainer_integration.py::test_rlvr_trainer_lora_complete_workflow` | `test_minimal_request_is_accepted` | +| `::test_rlvr_trainer_with_custom_reward_function` | `test_rlvr_trainer.py::test_custom_reward_function_arn` | +| `::test_rlvr_trainer_with_lambda_arn_auto_creates_evaluator` | `::test_custom_reward_function_lambda_arn` | +| `::test_rlvr_trainer_with_evaluator_object` | `::test_custom_reward_function_evaluator_object` | +| `::test_rlvr_trainer_nemotron_with_kl_and_recipe` | `::test_explicit_recipe_file`, `::test_recipe_and_overrides_together` | +| `::test_rlvr_trainer_lora_with_sequence_length` | `test_sft_trainer.py::test_sequence_length_is_accepted` (same code path) | +| `::test_rlvr_trainer_nova_workflow` | `test_nova_trainers.py::test_nova_rlvr_is_accepted` | +| `test_sft_trainer_serverful_smtj.py` | `test_explicit_compute_is_accepted` | +| `test_sft_trainer_data_mixing_integration.py` | `test_nova_data_mixing.py` | | `test_tuner_distributed.py::test_tuner_includes_sm_drivers_channel` | `test_tuner.py::test_distributed_tuning_job_is_accepted` | +| `test_multi_turn_rl_trainer_integration.py` — 3 submit tests | `test_multi_turn_rl_trainer.py` (needs prerequisites) | +| `test_cpt_hyperpod.py` | `test_cpt_trainer.py` (needs a HyperPod cluster) | + +MLflow is worth calling out: every `*_complete_workflow` deep test configures it, +so `RecipeTrainerCases` covers both forms — `test_mlflow_experiment_tracking` +(experiment/run names, always runs) and `test_mlflow_resource_arn` (tracking-server +ARN, skips when the account has no app). + +### Not covered, and why + +**Evaluator tests (11)** — `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. **These were already `gpu_intensive` on master, so +this PR loses no coverage** — but closing this gap is the clearest follow-up. + +**HyperPod (3)** — `test_nova_sft_hyperpod.py`, `test_sft_data_mixing_hyperpod.py`, +`test_cpt_data_mixing_hyperpod.py`. HyperPod submits to a pre-provisioned cluster +rather than through `CreateTrainingJob`, so the pattern does not apply. +`test_cpt_trainer.py` is written in the shallow style and activates when +`SHALLOW_HYPERPOD_CLUSTER` is set. + +### Tests this PR newly marks + +Only these 10 gained `gpu_intensive` here — the 8 in `test_model_trainer.py`, +`test_sft_trainer_lora_with_sequence_length`, and +`test_tuner_includes_sm_drivers_channel`. Everything else in the table above was +already marked on master. -**Deliberately NOT marked**, because this suite does not cover them — marking them -would remove coverage with nothing replacing it: +**Do not add `gpu_intensive` to a deep test unless a shallow test covers the same +path**, or the PR gate silently loses coverage. -* every evaluator test (`test_benchmark_evaluator.py`, `test_custom_scorer_evaluator.py`, - `test_inspect_ai_evaluator.py`, `test_llm_as_judge_*`, `test_llmaj_custom_model.py`) - — `evaluate()` is a different API surface returning pipeline executions, and there - is no shallow coverage for it yet -* `test_notifications.py` — asserts EventBridge/SNS side effects, not submission -* `test_local_model_trainer.py` — local container mode makes no service call +### Fixtures that skip rather than create -**The rule to preserve:** do not add `gpu_intensive` to a deep test unless a shallow -test covers the same path. Otherwise the PR gate silently loses coverage. +`mlflow_arn`, `reward_lambda_arn` and `reward_evaluator` only *look up* their +resources and skip when absent. The deep suite's equivalents create them (IAM +roles, Lambdas, MLflow apps, registry entries) — durable side effects that a fast +PR-gate suite should not perform. ## Relationship to `dry_run=True` diff --git a/sagemaker-train/tests/integ/train/shallow/conftest.py b/sagemaker-train/tests/integ/train/shallow/conftest.py index 985444850a..f92e93dfd3 100644 --- a/sagemaker-train/tests/integ/train/shallow/conftest.py +++ b/sagemaker-train/tests/integ/train/shallow/conftest.py @@ -138,6 +138,84 @@ def nova_train_data_uri(sagemaker_session_us_east_1): return _ensure_object(sagemaker_session_us_east_1, _TRAIN_DATA_KEY) +@pytest.fixture(scope="module") +def reward_scored_data_uri(): + """Dataset the RLVR reward functions can actually score. + + The reward-function tests cannot use ``train_data_uri``. Verified against AWS: + before submitting, the SDK *invokes* the reward function over sample records + and fails the call if they do not score -- + + OSS reward function returned non-200 status code: 500. + Body: {"error": "GSM8k scoring failed: 'list' object has no attribute 'strip'"} + + The pre-provisioned reward functions expect GSM8k-shaped records, so this + reuses the same dataset the deep RLVR suite uses rather than this suite's + generic chat-format fixture. + """ + return "s3://mc-flows-sdk-testing/input_data/rlvr-rlaif-test-data/train_285.jsonl" + + +@pytest.fixture(scope="module") +def reward_evaluator(sagemaker_session): + """An existing AI Registry Evaluator object, if present; skip otherwise. + + Look-up only, for the same reason as ``reward_lambda_arn``: the deep suite's + fixture will *create* an evaluator (and wait for it), which is a durable + registry write this suite should not make. + """ + from sagemaker.ai_registry.evaluator import Evaluator + + name = "test-integ-rlvr-trainer" + try: + return Evaluator.get(name, sagemaker_session=sagemaker_session) + except Exception: + pytest.skip(f"Evaluator {name!r} not present; skipping") + + +@pytest.fixture(scope="module") +def reward_lambda_arn(sagemaker_session): + """ARN of the OSS reward-function Lambda, if it already exists. + + The parent train conftest creates this Lambda on demand + (``oss_lambda_arn``), including an IAM role and a 15-second propagation + sleep. This suite only looks it up: creating IAM roles and Lambdas is a + durable side effect that a fast PR-gate suite should not perform. Skips when + absent, so the account state decides rather than the test. + """ + client = sagemaker_session.boto_session.client("lambda") + name = "pysdk-integ-test-sm-train-oss-reward-fn" + try: + return client.get_function(FunctionName=name)["Configuration"]["FunctionArn"] + except Exception: + pytest.skip(f"Reward-function Lambda {name!r} not present; skipping") + + +@pytest.fixture(scope="module") +def mlflow_arn(sagemaker_session): + """ARN of an existing, ready MLflow app; skip if the account has none. + + Deliberately does NOT create one. The parent train conftest's + ``mlflow_resource_arn`` fixture will create and delete an app if none exists, + which takes minutes and provisions a durable resource -- far too heavy for a + suite whose whole point is to be cheap. Here a missing app just skips the two + tests that need an ARN; the experiment/run-name path is covered unconditionally. + """ + client = sagemaker_session.boto_session.client("sagemaker") + try: + # Not a paginatable operation ("Operation cannot be paginated: + # list_mlflow_apps"), so call it directly rather than via get_paginator. + summaries = client.list_mlflow_apps().get("Summaries", []) + except Exception as e: + pytest.skip(f"Could not list MLflow apps: {e}") + + for app in summaries: + if app.get("Status") in ("Created", "Updated"): + return app["Arn"] + + pytest.skip("No ready MLflow app in this account; skipping ARN-based test") + + @pytest.fixture(scope="module") def output_path(sagemaker_session): """S3 prefix for training output. diff --git a/sagemaker-train/tests/integ/train/shallow/recipe_cases.py b/sagemaker-train/tests/integ/train/shallow/recipe_cases.py index e9ca5f6a6a..26c9a33725 100644 --- a/sagemaker-train/tests/integ/train/shallow/recipe_cases.py +++ b/sagemaker-train/tests/integ/train/shallow/recipe_cases.py @@ -183,6 +183,47 @@ def test_explicit_s3_output_path(self, sagemaker_session, train_data_uri, output with submitted(trainer) as job: assert_submitted(job) + def test_mlflow_experiment_tracking(self, sagemaker_session, train_data_uri): + """MLflow experiment/run names must be accepted. + + The ``*_complete_workflow`` tests in the deep suites all configure MLflow + (either ``mlflow_resource_arn`` or the experiment/run names), so without + this the shallow counterpart of those tests would miss the MLflow half of + the payload entirely. + + Uses the experiment/run *names* rather than ``mlflow_resource_arn``: the + names travel the same serialization path but need no pre-provisioned + tracking server, so this stays self-contained. ``test_mlflow_resource_arn`` + below covers the ARN form when one is available. + """ + trainer = self.build( + sagemaker_session, + train_data_uri, + self.name("-mlflow"), + mlflow_experiment_name="shallow-integ-test-exp", + mlflow_run_name="shallow-integ-test-run", + ) + + with submitted(trainer) as job: + assert_submitted(job) + + def test_mlflow_resource_arn(self, sagemaker_session, train_data_uri, mlflow_arn): + """An explicit MLflow tracking-server ARN must be accepted. + + Skips when no MLflow app exists in the account (see the ``mlflow_arn`` + fixture) rather than creating one, which would be slow and would leave a + durable resource behind. + """ + trainer = self.build( + sagemaker_session, + train_data_uri, + self.name("-mlflow-arn"), + mlflow_resource_arn=mlflow_arn, + ) + + with submitted(trainer) as job: + assert_submitted(job) + # -- serverful (explicit TrainingJobCompute) ----------------------------- def test_explicit_compute_is_accepted(self, sagemaker_session, train_data_uri): diff --git a/sagemaker-train/tests/integ/train/shallow/test_nova_trainers.py b/sagemaker-train/tests/integ/train/shallow/test_nova_trainers.py new file mode 100644 index 0000000000..74286e36ce --- /dev/null +++ b/sagemaker-train/tests/integ/train/shallow/test_nova_trainers.py @@ -0,0 +1,96 @@ +# Copyright Amazon.com, Inc. or its affiliates. All Rights Reserved. +# +# Licensed under the Apache License, Version 2.0 (the "License"). You +# may not use this file except in compliance with the License. A copy of +# the License is located at +# +# http://aws.amazon.com/apache2.0/ +# +# or in the "license" file accompanying this file. This file is +# distributed on an "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF +# ANY KIND, either express or implied. See the License for the specific +# language governing permissions and limitations under the License. +"""Shallow submission tests for Nova models (SFT and RLVR). + +Shallow counterparts of ``test_sft_trainer_integration.py::test_sft_trainer_nova_workflow`` +and ``test_rlvr_trainer_integration.py::test_rlvr_trainer_nova_workflow``. + +Nova is a distinct path: a different recipe family, a different region +(us-east-1), and a different test account, so these cannot share +``RecipeTrainerCases`` -- its ``MODEL_ID``, dataset fixtures and default session +are all us-west-2. Marked ``us_east_1`` so they run in that region's integ job. + +Datasets and reward functions are the same pre-provisioned ones the deep suite +uses, in account 784379639078. If those move, both suites break together, which +is preferable to this suite silently drifting onto its own copies. +""" + +from __future__ import absolute_import + +import pytest +from sagemaker.core import shapes +from sagemaker.train.common import TrainingType +from sagemaker.train.rlvr_trainer import RLVRTrainer +from sagemaker.train.sft_trainer import SFTTrainer + +from .harness import MAX_RUNTIME_IN_SECONDS, assert_submitted, submitted, unique_name + +NOVA_MODEL = "nova-textgeneration-lite-v2" +MODEL_PACKAGE_GROUP = "sdk-test-finetuned-models" + +# Pre-provisioned in the us-east-1 test account, shared with the deep suite. +_NOVA_BUCKET = "s3://sagemaker-us-east-1-784379639078" +SFT_DATASET = f"{_NOVA_BUCKET}/input_data/sft-nova/sft_200_samples.jsonl" +RLVR_DATASET = f"{_NOVA_BUCKET}/input_data/rlvr-nova/grpo-64-sample.jsonl" +OUTPUT_PATH = f"{_NOVA_BUCKET}/output/" +RLVR_REWARD_FUNCTION = ( + "arn:aws:sagemaker:us-east-1:784379639078:hub-content/sdktest/JsonDoc/rlvr-nova-test-rf/0.0.1" +) + + +def _stopping_condition(): + return shapes.StoppingCondition(max_runtime_in_seconds=MAX_RUNTIME_IN_SECONDS) + + +@pytest.mark.us_east_1 +class TestNovaSFTSubmission: + """Nova SFT selects a Nova-specific recipe family.""" + + def test_nova_sft_is_accepted(self, sagemaker_session_us_east_1): + trainer = SFTTrainer( + model=NOVA_MODEL, + training_type=TrainingType.LORA, + model_package_group=MODEL_PACKAGE_GROUP, + training_dataset=SFT_DATASET, + s3_output_path=OUTPUT_PATH, + accept_eula=True, + sagemaker_session=sagemaker_session_us_east_1, + base_job_name=unique_name("shallow-nova-sft"), + stopping_condition=_stopping_condition(), + ) + + with submitted(trainer) as job: + assert_submitted(job) + + +@pytest.mark.us_east_1 +class TestNovaRLVRSubmission: + """Nova RLVR additionally carries a Nova-specific reward function.""" + + def test_nova_rlvr_is_accepted(self, sagemaker_session_us_east_1): + trainer = RLVRTrainer( + model=NOVA_MODEL, + training_type=TrainingType.LORA, + model_package_group=MODEL_PACKAGE_GROUP, + training_dataset=RLVR_DATASET, + validation_dataset=RLVR_DATASET, + s3_output_path=OUTPUT_PATH, + custom_reward_function=RLVR_REWARD_FUNCTION, + accept_eula=True, + sagemaker_session=sagemaker_session_us_east_1, + base_job_name=unique_name("shallow-nova-rlvr"), + stopping_condition=_stopping_condition(), + ) + + with submitted(trainer) as job: + assert_submitted(job) diff --git a/sagemaker-train/tests/integ/train/shallow/test_rlaif_trainer.py b/sagemaker-train/tests/integ/train/shallow/test_rlaif_trainer.py index 76f76be524..69ac928ab2 100644 --- a/sagemaker-train/tests/integ/train/shallow/test_rlaif_trainer.py +++ b/sagemaker-train/tests/integ/train/shallow/test_rlaif_trainer.py @@ -19,6 +19,7 @@ from sagemaker.train.rlaif_trainer import RLAIFTrainer +from .harness import assert_submitted, submitted from .recipe_cases import RecipeTrainerCases # Values match the existing test_rlaif_trainer_integration.py so both suites @@ -26,6 +27,18 @@ REWARD_MODEL_ID = "openai.gpt-oss-120b-1:0" REWARD_PROMPT = "Builtin.Summarize" +# Hub-content prompt ARN, the alternative to a Builtin.* prompt name. Same one +# the deep suite uses. +REWARD_PROMPT_ARN = ( + "arn:aws:sagemaker:us-west-2:729646638167:hub-content/sdktest/JsonDoc/rlaif-test-prompt/0.0.1" +) + +# An existing fine-tuned model package, used to prove continued fine-tuning +# (model= a model-package ARN rather than a hub model id) still submits. +FINETUNED_MODEL_PACKAGE = ( + "arn:aws:sagemaker:us-west-2:729646638167:model-package/sdk-test-finetuned-models/1" +) + class TestRLAIFTrainerSubmission(RecipeTrainerCases): """RLAIF needs a reward model and prompt, and has no serverful path. @@ -38,3 +51,36 @@ class TestRLAIFTrainerSubmission(RecipeTrainerCases): TRAINER = RLAIFTrainer EXTRA_KWARGS = {"reward_model_id": REWARD_MODEL_ID, "reward_prompt": REWARD_PROMPT} SUPPORTS_SERVERFUL = False + + def test_reward_prompt_as_arn(self, sagemaker_session, train_data_uri): + """``reward_prompt`` accepts a hub-content ARN as well as a ``Builtin.*`` + name, and the two serialize differently. + + Shallow counterpart of test_rlaif_trainer_with_custom_reward_settings. + """ + trainer = self.build( + sagemaker_session, + train_data_uri, + self.name("-prompt-arn"), + reward_prompt=REWARD_PROMPT_ARN, + ) + + with submitted(trainer) as job: + assert_submitted(job) + + def test_continued_finetuning_from_model_package(self, sagemaker_session, train_data_uri): + """``model`` as a model-package ARN (continued fine-tuning) must resolve + and submit, not just a hub model id. + + Shallow counterpart of test_rlaif_trainer_continued_finetuning. Worth + covering because model resolution takes a different path for an ARN. + """ + trainer = self.build( + sagemaker_session, + train_data_uri, + self.name("-continued"), + model=FINETUNED_MODEL_PACKAGE, + ) + + with submitted(trainer) as job: + assert_submitted(job) diff --git a/sagemaker-train/tests/integ/train/shallow/test_rlvr_trainer.py b/sagemaker-train/tests/integ/train/shallow/test_rlvr_trainer.py index b5a9ede54a..87a80d0e67 100644 --- a/sagemaker-train/tests/integ/train/shallow/test_rlvr_trainer.py +++ b/sagemaker-train/tests/integ/train/shallow/test_rlvr_trainer.py @@ -27,6 +27,12 @@ from .harness import assert_submitted, submitted from .recipe_cases import RecipeTrainerCases +# Pre-provisioned reward function in the test account, same one the deep suite +# uses (test_rlvr_trainer_integration.py). +REWARD_FUNCTION_ARN = ( + "arn:aws:sagemaker:us-west-2:729646638167:hub-content/sdktest/JsonDoc/rlvr-test-rf/0.0.1" +) + class TestRLVRTrainerSubmission(RecipeTrainerCases): """RLVR accepts every shared case, plus recipe customization.""" @@ -96,3 +102,64 @@ def test_training_config_overrides(self, sagemaker_session, train_data_uri): with submitted(trainer) as job: assert_submitted(job) + + # -- reward-function variants ------------------------------------------- + # + # RLVR is the only trainer with a pluggable reward function, and the deep + # suite covers three distinct forms. Each changes what the SDK puts in the + # payload, so each needs its own acceptance case. + + def test_custom_reward_function_arn(self, sagemaker_session, reward_scored_data_uri): + """A hub-content reward-function ARN must be accepted. + + Shallow counterpart of test_rlvr_trainer_with_custom_reward_function. + """ + trainer = self.build( + sagemaker_session, + reward_scored_data_uri, + self.name("-rf-arn"), + custom_reward_function=REWARD_FUNCTION_ARN, + ) + + with submitted(trainer) as job: + assert_submitted(job) + + def test_custom_reward_function_lambda_arn( + self, sagemaker_session, reward_scored_data_uri, reward_lambda_arn + ): + """A Lambda ARN as the reward function auto-creates an AI Registry + Evaluator, then submits. + + Shallow counterpart of + test_rlvr_trainer_with_lambda_arn_auto_creates_evaluator. The Lambda is + reused from the parent train conftest rather than created here, and the + test skips if it is unavailable. + """ + trainer = self.build( + sagemaker_session, + reward_scored_data_uri, + self.name("-rf-lambda"), + custom_reward_function=reward_lambda_arn, + ) + + with submitted(trainer) as job: + assert_submitted(job) + + def test_custom_reward_function_evaluator_object( + self, sagemaker_session, reward_scored_data_uri, reward_evaluator + ): + """A pre-created ``Evaluator`` object as the reward function must + serialize to the same accepted payload as an ARN. + + Shallow counterpart of test_rlvr_trainer_with_evaluator_object. Skips when + the evaluator is absent rather than creating one. + """ + trainer = self.build( + sagemaker_session, + reward_scored_data_uri, + self.name("-rf-obj"), + custom_reward_function=reward_evaluator, + ) + + with submitted(trainer) as job: + assert_submitted(job) From 7c681d3808094122d5597440f190cce1d0be1b25 Mon Sep 17 00:00:00 2001 From: Amarjeet LNU Date: Wed, 12 Aug 2026 12:20:52 -0700 Subject: [PATCH 06/15] 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). --- .../integ/train/shallow/test_nova_trainers.py | 52 +++++++++++++++++++ .../integ/train/shallow/test_rlvr_trainer.py | 17 ++++++ .../integ/train/shallow/test_sft_trainer.py | 44 ++++++++++++++++ 3 files changed, 113 insertions(+) diff --git a/sagemaker-train/tests/integ/train/shallow/test_nova_trainers.py b/sagemaker-train/tests/integ/train/shallow/test_nova_trainers.py index 74286e36ce..30e7fb26e4 100644 --- a/sagemaker-train/tests/integ/train/shallow/test_nova_trainers.py +++ b/sagemaker-train/tests/integ/train/shallow/test_nova_trainers.py @@ -29,6 +29,7 @@ import pytest from sagemaker.core import shapes +from sagemaker.core.training.configs import TrainingJobCompute from sagemaker.train.common import TrainingType from sagemaker.train.rlvr_trainer import RLVRTrainer from sagemaker.train.sft_trainer import SFTTrainer @@ -94,3 +95,54 @@ def test_nova_rlvr_is_accepted(self, sagemaker_session_us_east_1): with submitted(trainer) as job: assert_submitted(job) + + +@pytest.mark.us_east_1 +class TestNovaServerfulSubmission: + """Nova on explicit TrainingJobCompute (serverful SMTJ). + + Shallow counterpart of ``test_sft_trainer_serverful_smtj.py``. Distinct from + ``RecipeTrainerCases::test_explicit_compute_is_accepted``, which covers the + serverful path for an OSS model in us-west-2: this is a Nova model, a Nova + recipe family, a Nova-only instance type, and a different region, so the + payload differs throughout. + + Also carries recipe overrides, as the deep test does, since Nova recipes nest + epoch control differently from OSS ones. + """ + + SERVERFUL_INSTANCE_TYPE = "ml.g6.12xlarge" + NOVA_MICRO = "amazon.nova-micro-v1" + + def test_nova_serverful_with_overrides_is_accepted(self, sagemaker_session_us_east_1): + trainer = SFTTrainer( + model=self.NOVA_MICRO, + training_type=TrainingType.LORA, + training_dataset=SFT_DATASET, + s3_output_path=OUTPUT_PATH, + compute=TrainingJobCompute( + instance_type=self.SERVERFUL_INSTANCE_TYPE, instance_count=1 + ), + sagemaker_session=sagemaker_session_us_east_1, + overrides={"training_config": {"max_epochs": 1}}, + base_job_name=unique_name("shallow-nova-smtj"), + stopping_condition=_stopping_condition(), + ) + + # The deep test asserts the override reached the resolved recipe; keep that, + # since it is client-side and exact. + # + # Recipe families nest epoch control differently: Nova puts it under + # ``trainer`` (which is what test_sft_trainer_serverful_smtj.py asserts), + # while the OSS Llama recipes use ``training_args`` -- verified against AWS + # by probing the resolver. Accept whichever this family uses rather than + # hard-coding one shape, so the test fails on a lost override rather than + # on a recipe-layout difference. + training_config = trainer.get_resolved_recipe()["training_config"] + epochs = training_config.get("trainer", {}).get( + "max_epochs", training_config.get("training_args", {}).get("max_epochs") + ) + assert epochs == 1, f"override did not reach the resolved recipe: {training_config}" + + with submitted(trainer) as job: + assert_submitted(job) diff --git a/sagemaker-train/tests/integ/train/shallow/test_rlvr_trainer.py b/sagemaker-train/tests/integ/train/shallow/test_rlvr_trainer.py index 87a80d0e67..5858692007 100644 --- a/sagemaker-train/tests/integ/train/shallow/test_rlvr_trainer.py +++ b/sagemaker-train/tests/integ/train/shallow/test_rlvr_trainer.py @@ -48,6 +48,23 @@ def test_direct_hyperparameter_mutation(self, sagemaker_session, train_data_uri) with submitted(trainer) as job: assert_submitted(job) + def test_kl_and_clipping_hyperparameters(self, sagemaker_session, train_data_uri): + """RLVR-specific GRPO hyperparameters must reach the payload. + + The deep test (test_rlvr_trainer_nemotron_with_kl_and_recipe) sets these + five fields on a 30B model behind a two-hour poll loop. They are separate + recipe fields, not one flag, so setting only max_epochs -- as + test_direct_hyperparameter_mutation does -- would not prove they serialize. + """ + trainer = self.build(sagemaker_session, train_data_uri, self.name("-kl")) + trainer.hyperparameters.use_kl_loss = True + trainer.hyperparameters.kl_loss_coef = 0.05 + trainer.hyperparameters.clip_ratio = 0.2 + trainer.hyperparameters.max_epochs = 1 + + with submitted(trainer) as job: + assert_submitted(job) + def test_explicit_recipe_file(self, sagemaker_session, train_data_uri): """A caller-supplied recipe YAML must render into an accepted request. diff --git a/sagemaker-train/tests/integ/train/shallow/test_sft_trainer.py b/sagemaker-train/tests/integ/train/shallow/test_sft_trainer.py index b364a577f9..a5820d40fd 100644 --- a/sagemaker-train/tests/integ/train/shallow/test_sft_trainer.py +++ b/sagemaker-train/tests/integ/train/shallow/test_sft_trainer.py @@ -73,3 +73,47 @@ def test_disable_output_compression(self, sagemaker_session, train_data_uri): with submitted(trainer) as job: assert_submitted(job) + + def test_recipe_overrides_are_accepted(self, sagemaker_session, train_data_uri): + """``overrides`` is merged into the rendered recipe before submission. + + Shallow counterpart of the override half of + ``test_sft_trainer_serverful_smtj.py``, which applies + ``overrides={"training_config": {"max_epochs": 1}}`` and then asserts the + merge via ``get_resolved_recipe()``. + + Two distinct things are checked, and both matter: + + * ``get_resolved_recipe()`` -- the override reached the *rendered recipe*. + This is client-side, so it is cheap and exact, and it is what the deep + test asserts. + * submission -- the resulting payload is still *accepted* by the service. + Recipe filtering runs after the request validators and rejects with + "No valid recipes found for the given request", so a bad merge is only + caught here. + """ + trainer = self.build( + sagemaker_session, + train_data_uri, + self.name("-overrides"), + overrides={"training_config": {"max_epochs": 1, "learning_rate": 2e-5}}, + ) + + resolved = trainer.get_resolved_recipe() + training_args = resolved["training_config"]["training_args"] + + # Overrides are written flat under training_config but land nested in + # training_args -- verified against AWS by probing the resolver: + # overrides {"training_config": {"max_epochs": 3}} + # -> resolved training_config.training_args.max_epochs == 3 + # The recipe default for this model is 5, so asserting 1 proves the + # override was applied rather than coinciding with the default. + assert ( + training_args["max_epochs"] == 1 + ), f"max_epochs override did not reach the resolved recipe: {training_args}" + assert ( + training_args["learning_rate"] == 2e-5 + ), f"learning_rate override did not reach the resolved recipe: {training_args}" + + with submitted(trainer) as job: + assert_submitted(job) From 7e1c5260a7a1a5ee0f375759a6ef692e045b2dfd Mon Sep 17 00:00:00 2001 From: Amarjeet LNU Date: Wed, 12 Aug 2026 13:44:12 -0700 Subject: [PATCH 07/15] 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. --- sagemaker-train/pyproject.toml | 9 ++ .../tests/integ/train/shallow/README.md | 38 ++++++-- .../tests/integ/train/shallow/conftest.py | 97 +++++++++++++++++++ .../tests/integ/train/shallow/recipe_cases.py | 11 +++ .../train/shallow/test_nova_data_mixing.py | 4 +- .../integ/train/shallow/test_nova_trainers.py | 71 +++++++++----- 6 files changed, 196 insertions(+), 34 deletions(-) diff --git a/sagemaker-train/pyproject.toml b/sagemaker-train/pyproject.toml index 253dd405cf..a561a8d00d 100644 --- a/sagemaker-train/pyproject.toml +++ b/sagemaker-train/pyproject.toml @@ -84,6 +84,15 @@ addopts = ["-vv"] testpaths = ["tests"] markers = [ "serial: marks tests that must run serially (not in parallel)", + # gpu_intensive and us_east_1 are declared in tox.ini too, but pytest reads + # its config from this file (it is the first of the candidates present), so + # markers listed only there are unregistered at runtime and raise + # PytestUnknownMarkWarning. Registering them here matters because the PR gate + # selects with -m "not gpu_intensive and not us_east_1": a typo'd marker name + # would otherwise silently put an expensive deep test back on the gate instead + # of warning. + "gpu_intensive: marks a test that consumes real training capacity (scheduled CI, not PR checks); see tests/integ/train/shallow", + "us_east_1: marks a test that must run in us-east-1 (Nova); runs in the us-east-1 integ job", ] [tool.black] diff --git a/sagemaker-train/tests/integ/train/shallow/README.md b/sagemaker-train/tests/integ/train/shallow/README.md index bb37399280..302ebcba07 100644 --- a/sagemaker-train/tests/integ/train/shallow/README.md +++ b/sagemaker-train/tests/integ/train/shallow/README.md @@ -51,6 +51,7 @@ of any deep test is easy to find: | `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` | `recipe_cases.py` holds the cases every recipe trainer shares (minimal submit, validation dataset, dataset override, output path, serverful compute, and the two @@ -87,10 +88,10 @@ below accounts for all of them. | `::test_rlvr_trainer_with_custom_reward_function` | `test_rlvr_trainer.py::test_custom_reward_function_arn` | | `::test_rlvr_trainer_with_lambda_arn_auto_creates_evaluator` | `::test_custom_reward_function_lambda_arn` | | `::test_rlvr_trainer_with_evaluator_object` | `::test_custom_reward_function_evaluator_object` | -| `::test_rlvr_trainer_nemotron_with_kl_and_recipe` | `::test_explicit_recipe_file`, `::test_recipe_and_overrides_together` | +| `::test_rlvr_trainer_nemotron_with_kl_and_recipe` | `::test_explicit_recipe_file`, `::test_recipe_and_overrides_together`, `::test_kl_and_clipping_hyperparameters` | | `::test_rlvr_trainer_lora_with_sequence_length` | `test_sft_trainer.py::test_sequence_length_is_accepted` (same code path) | | `::test_rlvr_trainer_nova_workflow` | `test_nova_trainers.py::test_nova_rlvr_is_accepted` | -| `test_sft_trainer_serverful_smtj.py` | `test_explicit_compute_is_accepted` | +| `test_sft_trainer_serverful_smtj.py` | `test_explicit_compute_is_accepted` (OSS/us-west-2), `test_sft_trainer.py::test_recipe_overrides_are_accepted` (the override half), `test_nova_trainers.py::TestNovaServerfulSubmission` (Nova/us-east-1) | | `test_sft_trainer_data_mixing_integration.py` | `test_nova_data_mixing.py` | | `test_tuner_distributed.py::test_tuner_includes_sm_drivers_channel` | `test_tuner.py::test_distributed_tuning_job_is_accepted` | | `test_multi_turn_rl_trainer_integration.py` — 3 submit tests | `test_multi_turn_rl_trainer.py` (needs prerequisites) | @@ -127,10 +128,35 @@ path**, or the PR gate silently loses coverage. ### Fixtures that skip rather than create -`mlflow_arn`, `reward_lambda_arn` and `reward_evaluator` only *look up* their -resources and skip when absent. The deep suite's equivalents create them (IAM -roles, Lambdas, MLflow apps, registry entries) — durable side effects that a fast -PR-gate suite should not perform. +`mlflow_arn`, `reward_lambda_arn`, `reward_evaluator` and `nova_reward_function_arn` +only *look up* their resources and skip when absent. The deep suite's equivalents +create them (IAM roles, Lambdas, MLflow apps, registry entries) — durable side +effects that a fast PR-gate suite should not perform. + +### Fixtures that derive rather than hardcode + +The Nova tests (`us_east_1`) build every S3 path from `default_bucket()` and +resolve the reward function from the calling account's own hub, rather than naming +the resources the deep Nova tests use. + +This is not stylistic. The deep tests hardcode +`s3://sagemaker-us-east-1-784379639078/...`, which other accounts cannot read — +verified: `AccessDenied` on `ListObjectsV2` from 729646638167. A hardcoded path +means the test only runs in one account and fails everywhere else, which is how +these five ended up never having been executed. `test_sft_trainer_serverful_smtj.py` +already takes the derived approach (`training_resources`); these follow it, and +upload the Nova-shaped sample data the deep suite already ships +(`tests/data/train/sft_smtj_sample_data.jsonl`) rather than adding a second copy. + +Two region constraints are worth knowing before adding a Nova test, both verified +against the service: + +* the model package group must be in the **job's** region — passing the us-west-2 + ARN from `MODEL_PACKAGE_GROUP` is rejected with `Model package group ARN region + 'us-west-2' does not match expected region 'us-east-1'`, so Nova files use + `NOVA_MODEL_PACKAGE_GROUP` (a bare name, which resolves per-session); +* an S3 input must be in the job's region, so `nova_rlvr_data_uri` copies the + us-west-2 RLVR dataset into the us-east-1 bucket rather than referencing it. ## Relationship to `dry_run=True` diff --git a/sagemaker-train/tests/integ/train/shallow/conftest.py b/sagemaker-train/tests/integ/train/shallow/conftest.py index f92e93dfd3..6328cc3e74 100644 --- a/sagemaker-train/tests/integ/train/shallow/conftest.py +++ b/sagemaker-train/tests/integ/train/shallow/conftest.py @@ -138,6 +138,103 @@ def nova_train_data_uri(sagemaker_session_us_east_1): return _ensure_object(sagemaker_session_us_east_1, _TRAIN_DATA_KEY) +@pytest.fixture(scope="module") +def nova_sft_data_uri(sagemaker_session_us_east_1): + """Nova-shaped SFT training data in the caller's own us-east-1 bucket. + + Cannot reuse ``nova_train_data_uri``: Nova SFT records carry a + ``schemaVersion`` ("nova-sft-2025-01-01") that this suite's generic chat-format + fixture does not. Rather than inventing a second inline copy, this uploads the + file the deep suite already ships + (``tests/data/train/sft_smtj_sample_data.jsonl``), so both suites train on the + same shape and a schema change updates one file. + + Idempotent, for the same xdist reason as ``_ensure_object``. + """ + local_path = os.path.join( + os.path.dirname(__file__), "..", "..", "..", "data", "train", "sft_smtj_sample_data.jsonl" + ) + if not os.path.isfile(local_path): + pytest.skip(f"Nova sample data not found at {local_path}") + + bucket = sagemaker_session_us_east_1.default_bucket() + key = "shallow-integ-test/nova-sft/sft_smtj_sample_data.jsonl" + s3 = sagemaker_session_us_east_1.boto_session.client("s3") + + if s3.list_objects_v2(Bucket=bucket, Prefix=key, MaxKeys=1).get("KeyCount", 0) == 0: + s3.upload_file(local_path, bucket, key) + logger.info("Uploaded Nova SFT fixture data to s3://%s/%s", bucket, key) + + return f"s3://{bucket}/{key}" + + +@pytest.fixture(scope="module") +def nova_rlvr_data_uri(sagemaker_session_us_east_1, reward_scored_data_uri): + """GSM8k-shaped RLVR data copied into the caller's own us-east-1 bucket. + + Two constraints force a copy rather than a reference: + + * The reward function is *invoked* over sample records before submission, so + the data must be GSM8k-shaped (see ``reward_scored_data_uri``). + * An S3 input must be in the same region as the job, and + ``reward_scored_data_uri`` lives in us-west-2. + + The deep test points at ``grpo-64-sample.jsonl`` in account 784379639078, which + is not readable from every account this runs in, so this copies the dataset the + us-west-2 RLVR tests already use. Idempotent, and skips rather than failing if + the source is unreadable. + """ + bucket = sagemaker_session_us_east_1.default_bucket() + key = "shallow-integ-test/nova-rlvr/train_285.jsonl" + s3 = sagemaker_session_us_east_1.boto_session.client("s3") + + if s3.list_objects_v2(Bucket=bucket, Prefix=key, MaxKeys=1).get("KeyCount", 0) == 0: + source = reward_scored_data_uri[len("s3://") :] + source_bucket, source_key = source.split("/", 1) + try: + s3.copy_object( + Bucket=bucket, Key=key, CopySource={"Bucket": source_bucket, "Key": source_key} + ) + except Exception as e: + pytest.skip(f"Could not copy RLVR sample data into {bucket}: {e}") + logger.info("Copied RLVR fixture data to s3://%s/%s", bucket, key) + + return f"s3://{bucket}/{key}" + + +@pytest.fixture(scope="module") +def nova_output_path(sagemaker_session_us_east_1): + """S3 prefix for Nova training output, in the caller's own us-east-1 bucket. + + Deliberately derived rather than hardcoded. The deep Nova tests name a bucket + in a specific test account (``sagemaker-us-east-1-784379639078``), which is not + readable from every account the suite runs in -- verified: ``AccessDenied`` on + ``ListObjectsV2`` from 729646638167. Using ``default_bucket()`` makes these + tests work in any account, the same way + ``test_sft_trainer_serverful_smtj.py::training_resources`` does. + """ + return f"s3://{sagemaker_session_us_east_1.default_bucket()}/shallow-integ-test/output/" + + +@pytest.fixture(scope="module") +def nova_reward_function_arn(sagemaker_session_us_east_1): + """ARN of the Nova RLVR reward function in the caller's own account. + + Look-up-and-skip, like the other reward fixtures. The deep test hardcodes this + ARN in account 784379639078; resolving it per-account instead means the test + runs wherever the hub content has been provisioned and skips cleanly elsewhere, + rather than failing with a confusing cross-account hub error. + """ + client = sagemaker_session_us_east_1.boto_session.client("sagemaker") + hub, name = "sdktest", "rlvr-nova-test-rf" + try: + return client.describe_hub_content( + HubName=hub, HubContentType="JsonDoc", HubContentName=name + )["HubContentArn"] + except Exception as e: + pytest.skip(f"Reward function {name!r} not in hub {hub!r}: {e}") + + @pytest.fixture(scope="module") def reward_scored_data_uri(): """Dataset the RLVR reward functions can actually score. diff --git a/sagemaker-train/tests/integ/train/shallow/recipe_cases.py b/sagemaker-train/tests/integ/train/shallow/recipe_cases.py index 26c9a33725..2f3357b381 100644 --- a/sagemaker-train/tests/integ/train/shallow/recipe_cases.py +++ b/sagemaker-train/tests/integ/train/shallow/recipe_cases.py @@ -56,6 +56,17 @@ class TestFooTrainerSubmission(RecipeTrainerCases): "arn:aws:sagemaker:us-west-2:729646638167:model-package-group/sdk-test-finetuned-models" ) +# The Nova files (us-east-1) cannot reuse the ARN above. Verified against AWS: the +# model package group must be in the same region as the job, and passing the +# us-west-2 ARN is rejected with +# +# Model package group ARN region 'us-west-2' does not match expected region +# 'us-east-1' +# +# A bare name resolves in whichever region the session is in. Shared here rather +# than duplicated per Nova file so the two cannot drift apart. +NOVA_MODEL_PACKAGE_GROUP = "sdk-test-finetuned-models" + # An accelerator type is required for the serverful recipe path: these recipes do # not resolve onto a CPU instance, so unlike the ModelTrainer suite we cannot use # ml.m5.large here. The job is still stopped immediately, so this holds capacity diff --git a/sagemaker-train/tests/integ/train/shallow/test_nova_data_mixing.py b/sagemaker-train/tests/integ/train/shallow/test_nova_data_mixing.py index 28cd9d49cc..98cde7868a 100644 --- a/sagemaker-train/tests/integ/train/shallow/test_nova_data_mixing.py +++ b/sagemaker-train/tests/integ/train/shallow/test_nova_data_mixing.py @@ -31,7 +31,7 @@ from sagemaker.train.sft_trainer import SFTTrainer from .harness import assert_submitted, submitted, unique_name -from .recipe_cases import MODEL_PACKAGE_GROUP, stopping_condition +from .recipe_cases import NOVA_MODEL_PACKAGE_GROUP, stopping_condition NOVA_MODEL = "nova-textgeneration-lite-v2" @@ -39,7 +39,7 @@ def _nova_sft(session, dataset, name, config): return SFTTrainer( model=NOVA_MODEL, - model_package_group=MODEL_PACKAGE_GROUP, + model_package_group=NOVA_MODEL_PACKAGE_GROUP, training_dataset=dataset, accept_eula=True, sagemaker_session=session, diff --git a/sagemaker-train/tests/integ/train/shallow/test_nova_trainers.py b/sagemaker-train/tests/integ/train/shallow/test_nova_trainers.py index 30e7fb26e4..5b820845a9 100644 --- a/sagemaker-train/tests/integ/train/shallow/test_nova_trainers.py +++ b/sagemaker-train/tests/integ/train/shallow/test_nova_trainers.py @@ -20,9 +20,13 @@ ``RecipeTrainerCases`` -- its ``MODEL_ID``, dataset fixtures and default session are all us-west-2. Marked ``us_east_1`` so they run in that region's integ job. -Datasets and reward functions are the same pre-provisioned ones the deep suite -uses, in account 784379639078. If those move, both suites break together, which -is preferable to this suite silently drifting onto its own copies. +Datasets, output paths and the reward function are all *derived from the calling +account* rather than hardcoded. The deep Nova tests name resources in one specific +test account (``sagemaker-us-east-1-784379639078``), which other accounts cannot +read -- verified: ``AccessDenied`` on ``ListObjectsV2`` from 729646638167. Using +``default_bucket()`` and resolving the reward function from the caller's own hub +follows what ``test_sft_trainer_serverful_smtj.py`` already does, and means these +tests actually run wherever the suite runs instead of only in one account. """ from __future__ import absolute_import @@ -35,18 +39,9 @@ from sagemaker.train.sft_trainer import SFTTrainer from .harness import MAX_RUNTIME_IN_SECONDS, assert_submitted, submitted, unique_name +from .recipe_cases import NOVA_MODEL_PACKAGE_GROUP NOVA_MODEL = "nova-textgeneration-lite-v2" -MODEL_PACKAGE_GROUP = "sdk-test-finetuned-models" - -# Pre-provisioned in the us-east-1 test account, shared with the deep suite. -_NOVA_BUCKET = "s3://sagemaker-us-east-1-784379639078" -SFT_DATASET = f"{_NOVA_BUCKET}/input_data/sft-nova/sft_200_samples.jsonl" -RLVR_DATASET = f"{_NOVA_BUCKET}/input_data/rlvr-nova/grpo-64-sample.jsonl" -OUTPUT_PATH = f"{_NOVA_BUCKET}/output/" -RLVR_REWARD_FUNCTION = ( - "arn:aws:sagemaker:us-east-1:784379639078:hub-content/sdktest/JsonDoc/rlvr-nova-test-rf/0.0.1" -) def _stopping_condition(): @@ -57,13 +52,15 @@ def _stopping_condition(): class TestNovaSFTSubmission: """Nova SFT selects a Nova-specific recipe family.""" - def test_nova_sft_is_accepted(self, sagemaker_session_us_east_1): + def test_nova_sft_is_accepted( + self, sagemaker_session_us_east_1, nova_sft_data_uri, nova_output_path + ): trainer = SFTTrainer( model=NOVA_MODEL, training_type=TrainingType.LORA, - model_package_group=MODEL_PACKAGE_GROUP, - training_dataset=SFT_DATASET, - s3_output_path=OUTPUT_PATH, + model_package_group=NOVA_MODEL_PACKAGE_GROUP, + training_dataset=nova_sft_data_uri, + s3_output_path=nova_output_path, accept_eula=True, sagemaker_session=sagemaker_session_us_east_1, base_job_name=unique_name("shallow-nova-sft"), @@ -78,19 +75,39 @@ def test_nova_sft_is_accepted(self, sagemaker_session_us_east_1): class TestNovaRLVRSubmission: """Nova RLVR additionally carries a Nova-specific reward function.""" - def test_nova_rlvr_is_accepted(self, sagemaker_session_us_east_1): + def test_nova_rlvr_is_accepted( + self, + sagemaker_session_us_east_1, + nova_rlvr_data_uri, + nova_output_path, + nova_reward_function_arn, + ): trainer = RLVRTrainer( model=NOVA_MODEL, training_type=TrainingType.LORA, - model_package_group=MODEL_PACKAGE_GROUP, - training_dataset=RLVR_DATASET, - validation_dataset=RLVR_DATASET, - s3_output_path=OUTPUT_PATH, - custom_reward_function=RLVR_REWARD_FUNCTION, + model_package_group=NOVA_MODEL_PACKAGE_GROUP, + training_dataset=nova_rlvr_data_uri, + validation_dataset=nova_rlvr_data_uri, + s3_output_path=nova_output_path, + custom_reward_function=nova_reward_function_arn, accept_eula=True, sagemaker_session=sagemaker_session_us_east_1, base_job_name=unique_name("shallow-nova-rlvr"), stopping_condition=_stopping_condition(), + # Before submitting, the SDK *invokes* the reward function over sample + # records and refuses the call if their scores do not parse. That gate + # is real, but it asserts the contents of a hub artifact provisioned + # per-account rather than anything about this payload -- verified: the + # function registered under this name in 729646638167 returns a shape + # the verifier rejects ("Each output must include 'id', + # 'aggregate_reward_score'"), so the test would fail on account state + # rather than on a regression. + # + # The verifier itself is already covered, against a known-compatible + # function, by the three us-west-2 reward-function cases in + # test_rlvr_trainer.py. What is unique here is the Nova recipe family + # and region, which is what this test is for. + skip_reward_validation=True, ) with submitted(trainer) as job: @@ -114,12 +131,14 @@ class TestNovaServerfulSubmission: SERVERFUL_INSTANCE_TYPE = "ml.g6.12xlarge" NOVA_MICRO = "amazon.nova-micro-v1" - def test_nova_serverful_with_overrides_is_accepted(self, sagemaker_session_us_east_1): + def test_nova_serverful_with_overrides_is_accepted( + self, sagemaker_session_us_east_1, nova_sft_data_uri, nova_output_path + ): trainer = SFTTrainer( model=self.NOVA_MICRO, training_type=TrainingType.LORA, - training_dataset=SFT_DATASET, - s3_output_path=OUTPUT_PATH, + training_dataset=nova_sft_data_uri, + s3_output_path=nova_output_path, compute=TrainingJobCompute( instance_type=self.SERVERFUL_INSTANCE_TYPE, instance_count=1 ), From 11d123c1db02ef21df08b265ecd3a35ec531fee3 Mon Sep 17 00:00:00 2001 From: Amarjeet LNU Date: Wed, 12 Aug 2026 15:52:28 -0700 Subject: [PATCH 08/15] 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. --- .github/workflows/pr-checks-master.yml | 12 ++++++++++ .../tests/integ/train/shallow/README.md | 23 +++++++++++++++++++ 2 files changed, 35 insertions(+) diff --git a/.github/workflows/pr-checks-master.yml b/.github/workflows/pr-checks-master.yml index bac129a4cf..616393714e 100644 --- a/.github/workflows/pr-checks-master.yml +++ b/.github/workflows/pr-checks-master.yml @@ -332,6 +332,18 @@ jobs: # # Note the shallow suite is NOT separately marked: it is intended to # run here, and its own MTRL/Nova cases carry these markers themselves. + # + # One caveat for reviewers: tests/integ/train/test_local_model_trainer.py + # runs real containers, so it needs Docker and pulls + # pytorch-training:2.0.0-cpu-py310 (2.3 GB compressed). Docker is + # preinstalled on GitHub-hosted Ubuntu runners, and the pull needs ECR + # read on 763104351884, which the assumed role already has for the + # shallow tests. It stays on the gate because it makes no service call + # and catches local-mode regressions nothing else covers -- but it is + # the slowest thing here and the only step with a disk-space floor. If + # this job ever goes flaky on runner capacity, these three tests are the + # first thing to deselect (the unused `local_mode` marker in tox.ini + # exists for exactly that). python -m pytest tests/integ/train \ -m "not gpu_intensive and not us_east_1" \ -n 8 \ diff --git a/sagemaker-train/tests/integ/train/shallow/README.md b/sagemaker-train/tests/integ/train/shallow/README.md index 302ebcba07..5afaaa3384 100644 --- a/sagemaker-train/tests/integ/train/shallow/README.md +++ b/sagemaker-train/tests/integ/train/shallow/README.md @@ -110,6 +110,29 @@ is a different API surface returning pipeline executions rather than jobs, so it needs its own harness support. **These were already `gpu_intensive` on master, so this PR loses no coverage** — but closing this gap is the clearest follow-up. +Worth knowing before that follow-up: five evaluator tests are **not** marked +`gpu_intensive` and each blocks on `execution.wait(..., timeout=14400)` — a 4-hour +ceiling, and a measured ~33 minutes per execution in practice: + +| Test | Marks | +|---|---| +| `test_benchmark_evaluator.py::test_benchmark_evaluation_full_flow` | none | +| `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` | none | +| `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 they now dominate +its wall clock. Measured on a full gate run: **201 of 204 tests finished in ~7 +minutes, and these held the run open for another 40+** before it was killed. The +whole shallow suite costs less than any one of them. + +Marking them is not a call this PR makes, because unlike every other +`gpu_intensive` test they have no shallow counterpart yet — marking them would +remove coverage, which is exactly what the rule above forbids. The right order is: +add evaluator support to the harness, then mark them. Until then the gate is +bounded by evaluation-pipeline latency rather than by anything in this suite. + **HyperPod (3)** — `test_nova_sft_hyperpod.py`, `test_sft_data_mixing_hyperpod.py`, `test_cpt_data_mixing_hyperpod.py`. HyperPod submits to a pre-provisioned cluster rather than through `CreateTrainingJob`, so the pattern does not apply. From 207835891067c9957f34be4c96b0f96f2cba6534 Mon Sep 17 00:00:00 2001 From: Amarjeet LNU Date: Wed, 12 Aug 2026 18:38:56 -0700 Subject: [PATCH 09/15] 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. --- .github/workflows/pr-checks-master.yml | 90 ++++++++----------- .../tests/integ/train/shallow/README.md | 23 +++-- 2 files changed, 51 insertions(+), 62 deletions(-) diff --git a/.github/workflows/pr-checks-master.yml b/.github/workflows/pr-checks-master.yml index 616393714e..8b76d1c043 100644 --- a/.github/workflows/pr-checks-master.yml +++ b/.github/workflows/pr-checks-master.yml @@ -221,13 +221,6 @@ jobs: env: SUBMODULE: ${{ matrix.submodule }} - # sagemaker-train's PR-gate integ tests are handled by the shallow-integ-tests - # job below, so it is filtered out of this matrix. Every other submodule keeps - # the existing full CodeBuild integ suite unchanged. - # - # The filter is computed with fromJson/contains rather than by editing - # detect-changes, so the dependency-propagation logic there (and the submodule - # list consumed by codestyle-doc-tests and unit-tests) is untouched. integ-tests: runs-on: ubuntu-latest needs: [detect-changes] @@ -236,8 +229,6 @@ jobs: fail-fast: false matrix: submodule: ${{ fromJson(needs.detect-changes.outputs.submodules) }} - exclude: - - submodule: sagemaker-train steps: - name: Configure AWS Credentials uses: aws-actions/configure-aws-credentials@v4 @@ -252,13 +243,21 @@ jobs: project-name: ${{ github.event.repository.name }}-ci-${{ matrix.submodule }}-integ-tests source-version-override: 'refs/pull/${{ github.event.pull_request.number }}/head^{${{ github.event.pull_request.head.sha }}}' - # Replaces the CodeBuild integ suite for sagemaker-train on the PR gate. + # Additive: runs the shallow (submit-then-stop) suite for sagemaker-train + # alongside the existing integ-tests job above, which is unchanged. + # + # Why a separate job rather than folding this into the CodeBuild suite: this + # job's selection is reviewable in the PR that changes it, whereas the + # sagemaker-train CodeBuild buildspec is CDK-managed outside this repo. It also + # reports as its own check, so a shallow failure is distinguishable at a glance + # from a deep-suite failure, and it finishes in minutes -- fast feedback that + # does not wait on the 2XLARGE CodeBuild container. # - # What runs here (~191 of 251 tests): - # * ~170 client-side tests that make no service call -- recipe resolution, - # data utils, dry-run, log streaming, docker-compose detection. These were - # always cheap and stay on the gate. - # * the shallow (submit-then-stop) suite under tests/integ/train/shallow. + # What runs here: only tests/integ/train/shallow. The client-side tests + # (recipe resolution, data utils, dry-run, log streaming) are deliberately NOT + # repeated -- the CodeBuild suite already runs the whole tests/integ tree, so + # widening this job's scope would duplicate them and double the job creation + # the shallow suite performs. # # Why submit-then-stop is worth gating on: CreateTrainingJob returns a # TrainingJobArn only after the request has cleared public-model validation, @@ -268,14 +267,8 @@ jobs: # ARN proves the payload and the caller's permissions are both good -- without # paying for a training run. The job is stopped immediately. # - # What no longer runs here: the ~54 tests that submit a job and wait for it. - # They are marked gpu_intensive and keep running on the scheduled CI-health - # workflows. This is a deliberate scope reduction -- training *behaviour* - # (artifacts, metrics, convergence) is not asserted on the PR gate. - # - # Runs directly on the runner rather than via CodeBuild because the sagemaker- - # train CodeBuild project's buildspec is CDK-managed outside this repo; running - # here keeps the test selection reviewable in the PR that changes it. + # It asserts nothing about training *behaviour* (artifacts, metrics, + # convergence); that remains the deep suites' job. fast-integ-tests: runs-on: ubuntu-latest needs: [detect-changes] @@ -306,7 +299,7 @@ jobs: pip install ./sagemaker-train pip install -r requirements/extras/test_requirements.txt - - name: Run fast sagemaker-train integ tests + - name: Run shallow sagemaker-train integ tests working-directory: sagemaker-train env: AWS_DEFAULT_REGION: us-west-2 @@ -316,48 +309,35 @@ jobs: AWS_RETRY_MODE: adaptive AWS_MAX_ATTEMPTS: '10' run: | - # Runs the WHOLE tests/integ/train tree, not just shallow/, and lets the - # markers decide what is affordable on a PR. That keeps the ~170 - # client-side tests (recipe resolution, data utils, dry-run, log - # streaming, docker-compose detection) on the gate -- they make no - # service call and were never the expensive part. - # - # Deselected, per the marker conventions already in tox.ini: - # gpu_intensive -- every test that submits a real job and waits for - # it. Now applied to the 19 submitters that were - # previously unmarked, so the shallow suite is the - # only thing on this gate that creates a job. - # us_east_1 -- this job holds us-west-2 credentials only; those - # tests run in the us-east-1 integ job. - # - # Note the shallow suite is NOT separately marked: it is intended to - # run here, and its own MTRL/Nova cases carry these markers themselves. + # Scoped to shallow/ only -- see the comment above this job for why the + # rest of tests/integ/train is not repeated here. # - # One caveat for reviewers: tests/integ/train/test_local_model_trainer.py - # runs real containers, so it needs Docker and pulls - # pytorch-training:2.0.0-cpu-py310 (2.3 GB compressed). Docker is - # preinstalled on GitHub-hosted Ubuntu runners, and the pull needs ECR - # read on 763104351884, which the assumed role already has for the - # shallow tests. It stays on the gate because it makes no service call - # and catches local-mode regressions nothing else covers -- but it is - # the slowest thing here and the only step with a disk-space floor. If - # this job ever goes flaky on runner capacity, these three tests are the - # first thing to deselect (the unused `local_mode` marker in tox.ini - # exists for exactly that). - python -m pytest tests/integ/train \ + # 84 of the suite's 100 tests run; the 16 deselected are: + # us_east_1 (5) -- Nova cases; this job holds us-west-2 + # credentials only, so they run in the + # integ-tests-us-east-1 job instead. + # gpu_intensive (11) -- the CPT and MTRL classes. These are written in + # the shallow style but cannot be made + # self-contained: CPT submits only via HyperPod + # (a pre-provisioned cluster, not + # CreateTrainingJob) and MTRL needs an agent + # runtime plus an MLflow app. Both become + # gate-eligible by dropping one marker once those + # prerequisites exist in the PR account. + python -m pytest tests/integ/train/shallow \ -m "not gpu_intensive and not us_east_1" \ -n 8 \ --dist loadfile \ -v \ --durations=15 \ - --junitxml=fast-integ-results.xml + --junitxml=shallow-integ-results.xml - name: Upload test results if: always() uses: actions/upload-artifact@v4 with: - name: fast-integ-test-results - path: sagemaker-train/fast-integ-results.xml + name: shallow-integ-test-results + path: sagemaker-train/shallow-integ-results.xml if-no-files-found: warn integ-tests-us-east-1: diff --git a/sagemaker-train/tests/integ/train/shallow/README.md b/sagemaker-train/tests/integ/train/shallow/README.md index 5afaaa3384..fe0f9754cf 100644 --- a/sagemaker-train/tests/integ/train/shallow/README.md +++ b/sagemaker-train/tests/integ/train/shallow/README.md @@ -1,7 +1,12 @@ # Shallow (submit-then-stop) integration tests -These tests replace the full `sagemaker-train` integ suite **on the PR gate only**. -The deep suites still run on the scheduled CI-health workflows. +These tests add fast acceptance coverage on the PR gate. They run in their own +`fast-integ-tests` job, **alongside** the existing `integ-tests` CodeBuild suite, +which is unchanged. The deep suites still run on the scheduled CI-health workflows. + +What this suite changes about the gate is not which job runs, but what the existing +one selects: the `gpu_intensive` marks added here deselect the deep tests that +submit a job and wait for it, and this suite covers those code paths instead. ## What a passing test proves @@ -122,16 +127,20 @@ ceiling, and a measured ~33 minutes per execution in practice: | `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 they now dominate -its wall clock. Measured on a full gate run: **201 of 204 tests finished in ~7 -minutes, and these held the run open for another 40+** before it was killed. The +They are selected by the `integ-tests` CodeBuild job, whose buildspec filters +`-m "not gpu_intensive and not us_east_1"` — so they run on master's gate today and +continue to after this PR. They dominate that job's wall clock. Measured locally on +the same selection: **201 of 204 tests finished in ~7 minutes, and these held the +run open for another 40+** before it was killed. Against the project's 180-minute +build timeout, five tests with a 4-hour ceiling each are the standing risk. The whole shallow suite costs less than any one of them. Marking them is not a call this PR makes, because unlike every other `gpu_intensive` test they have no shallow counterpart yet — marking them would remove coverage, which is exactly what the rule above forbids. The right order is: -add evaluator support to the harness, then mark them. Until then the gate is -bounded by evaluation-pipeline latency rather than by anything in this suite. +add evaluator support to the harness, then mark them. Until then the `integ-tests` +job is bounded by evaluation-pipeline latency rather than by anything in this suite, +and the `fast-integ-tests` job is where quick feedback comes from. **HyperPod (3)** — `test_nova_sft_hyperpod.py`, `test_sft_data_mixing_hyperpod.py`, `test_cpt_data_mixing_hyperpod.py`. HyperPod submits to a pre-provisioned cluster From a35033c02e07c61b884c074453a1d1d7a760aff8 Mon Sep 17 00:00:00 2001 From: Amarjeet LNU Date: Wed, 12 Aug 2026 21:22:57 -0700 Subject: [PATCH 10/15] 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. --- sagemaker-train/tests/integ/conftest.py | 102 +++++++++++++++++- .../tests/integ/train/shallow/README.md | 69 +++++++----- .../integ/train/test_benchmark_evaluator.py | 4 + .../train/test_custom_scorer_evaluator.py | 4 + .../train/test_llm_as_judge_base_model_fix.py | 6 ++ .../train/test_llm_as_judge_evaluator.py | 3 + .../integ/train/test_llmaj_custom_model.py | 7 +- 7 files changed, 166 insertions(+), 29 deletions(-) diff --git a/sagemaker-train/tests/integ/conftest.py b/sagemaker-train/tests/integ/conftest.py index de5e45aef8..99db27a084 100644 --- a/sagemaker-train/tests/integ/conftest.py +++ b/sagemaker-train/tests/integ/conftest.py @@ -31,10 +31,25 @@ its own). ``adaptive`` mode adds client-side rate limiting so bursts of ``SimulatePrincipalPolicy`` calls ride out transient throttling. -Throttling that still exhausts the adaptive retry budget is deliberately left to -fail the test loudly (rather than being converted to a skip), so a persistent -rate-limit regression stays visible instead of silently disappearing from the -results. +* ``_memoize_role_validation`` (autouse) — retries alone were not enough. A PR-gate + run failed four tests with ``(Throttling) ... SimulatePrincipalPolicy (reached + max retries: 9)``: the adaptive budget was exhausted, not merely stressed. The + cause is volume, not burstiness — ~190 tests each construct a trainer, every + construction calls ``get_role``, and each of those runs a *paginated* + ``SimulatePrincipalPolicy`` over ~20 action names. Under ``-n auto`` on a large + CodeBuild container that is thousands of calls against a low, account-wide TPS + limit, so raising the retry budget only trades failures for a slower build. + + Since the arguments repeat, the result does too: this memoizes + ``resolve_and_validate_role`` per worker, collapsing those calls to one per + distinct ``(provided_role, role_type, region)``. Validation still happens — once, + and its outcome (including a raised ``RoleValidationError``) is what gets reused, + so a genuinely bad role still fails every test that uses it. + +Throttling that still exhausts the retry budget after memoization is deliberately +left to fail the test loudly (rather than being converted to a skip), so a +persistent rate-limit regression stays visible instead of silently disappearing +from the results. """ from __future__ import absolute_import @@ -66,3 +81,82 @@ def _configure_boto_adaptive_retries(): os.environ.pop(key, None) else: os.environ[key] = value + + +# Modules that did `from ...iam_role_resolver import resolve_and_validate_role` +# hold their own reference to the original function, so patching only the defining +# module would leave those bindings calling IAM directly. Each importer is patched +# too. Kept as a list of (module path, attribute) so adding a caller is one line. +_ROLE_RESOLVER_CALLERS = ( + ("sagemaker.core.helper.iam_role_resolver", "resolve_and_validate_role"), + ("sagemaker.train.defaults", "resolve_and_validate_role"), + ("sagemaker.train.evaluate.base_evaluator", "resolve_and_validate_role"), +) + + +@pytest.fixture(autouse=True, scope="session") +def _memoize_role_validation(): + """Validate each distinct role once per xdist worker instead of once per test. + + See this module's docstring for why retries alone were insufficient. Caches on + ``(provided_role, role_type, region)`` -- region is part of the key because the + Nova tests validate the same role against us-east-1, and a role's resolution is + region-scoped. Exceptions are cached alongside successes so a bad role keeps + failing rather than being silently retried per test. + """ + import importlib + + patched = [] + cache = {} + + try: + source = importlib.import_module(_ROLE_RESOLVER_CALLERS[0][0]) + except ImportError: # pragma: no cover - SDK layout changed + yield + return + + original = source.resolve_and_validate_role + + def memoized(provided_role=None, role_type=None, sagemaker_session=None, **kwargs): + region = None + if sagemaker_session is not None: + region = getattr(sagemaker_session, "boto_region_name", None) + key = (provided_role, role_type, region) + + if key not in cache: + try: + cache[key] = ( + original( + provided_role=provided_role, + role_type=role_type, + sagemaker_session=sagemaker_session, + **kwargs, + ), + None, + ) + except Exception as exc: # cache the verdict, not just the happy path + cache[key] = (None, exc) + + result, error = cache[key] + if error is not None: + raise error + return result + + for module_path, attribute in _ROLE_RESOLVER_CALLERS: + try: + module = importlib.import_module(module_path) + except ImportError: + continue # optional/renamed caller; the others still get patched + if getattr(module, attribute, None) is original: + setattr(module, attribute, memoized) + patched.append((module, attribute)) + + yield + + # Restore by checking for `memoized` rather than only undoing what was patched + # above: a caller imported *after* the source module was patched binds the + # memoized function at its own import time, so it needs restoring too even + # though this fixture never set it. + for module, attribute in patched: + if getattr(module, attribute, None) is memoized: + setattr(module, attribute, original) diff --git a/sagemaker-train/tests/integ/train/shallow/README.md b/sagemaker-train/tests/integ/train/shallow/README.md index fe0f9754cf..06e76ea985 100644 --- a/sagemaker-train/tests/integ/train/shallow/README.md +++ b/sagemaker-train/tests/integ/train/shallow/README.md @@ -115,32 +115,53 @@ is a different API surface returning pipeline executions rather than jobs, so it needs its own harness support. **These were already `gpu_intensive` on master, so this PR loses no coverage** — but closing this gap is the clearest follow-up. -Worth knowing before that follow-up: five evaluator tests are **not** marked -`gpu_intensive` and each blocks on `execution.wait(..., timeout=14400)` — a 4-hour -ceiling, and a measured ~33 minutes per execution in practice: +Six more evaluator tests **were** unmarked and each blocks on +`execution.wait(..., timeout=14400)` — a 4-hour ceiling per test. They are now +marked `gpu_intensive`: -| Test | Marks | +| Test | Measured | |---|---| -| `test_benchmark_evaluator.py::test_benchmark_evaluation_full_flow` | none | -| `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` | none | -| `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 are selected by the `integ-tests` CodeBuild job, whose buildspec filters -`-m "not gpu_intensive and not us_east_1"` — so they run on master's gate today and -continue to after this PR. They dominate that job's wall clock. Measured locally on -the same selection: **201 of 204 tests finished in ~7 minutes, and these held the -run open for another 40+** before it was killed. Against the project's 180-minute -build timeout, five tests with a 4-hour ceiling each are the standing risk. The -whole shallow suite costs less than any one of them. - -Marking them is not a call this PR makes, because unlike every other -`gpu_intensive` test they have no shallow counterpart yet — marking them would -remove coverage, which is exactly what the rule above forbids. The right order is: -add evaluator support to the harness, then mark them. Until then the `integ-tests` -job is bounded by evaluation-pipeline latency rather than by anything in this suite, -and the `fast-integ-tests` job is where quick feedback comes from. +| `test_llm_as_judge_base_model_fix.py::test_base_model_evaluation_uses_correct_weights` | **2783s** | +| `test_llm_as_judge_base_model_fix.py::test_base_model_false_still_works` | **2504s** | +| `test_benchmark_evaluator.py::test_benchmark_evaluation_full_flow` | held a run open 40+ min | +| `test_custom_scorer_evaluator.py::test_custom_scorer_evaluation_full_flow` | held a run open 40+ min | +| `test_llm_as_judge_evaluator.py::test_llm_as_judge_evaluation_full_flow` | held a run open 40+ min | +| `test_llmaj_custom_model.py::TestLLMAJCustomModelIntegration` | `@pytest.mark.slow`, unregistered | + +The first two figures come from a real PR-gate CodeBuild run: 88 minutes for the +two of them, against the project's **180-minute build timeout**. Everything else in +that serial pass finished in under 92s, so they were the entire tail. + +The last row was a genuine mismarking: the registered name is `slow_test`, so +`@pytest.mark.slow` silently did nothing. `us_east_1` already kept it off the +us-west-2 gate, so marking it changes nothing there — but it no longer waits on a +pipeline in the us-east-1 job either. + +This is a real, if narrow, coverage reduction, so it is worth being precise about +what is lost. Three of the files are marked per-test and keep their cheap +constructor/validation tests on the gate — `test_benchmark_evaluator.py` keeps +`test_get_benchmarks_and_properties` and two `*_validation` tests, +`test_custom_scorer_evaluator.py` keeps `test_get_builtin_metrics` and +`test_custom_scorer_evaluator_validation`, `test_llm_as_judge_evaluator.py` keeps +`test_llm_as_judge_evaluator_validation` and +`test_llm_as_judge_builtin_metrics_prefix_handling`. Those are what catch SDK-side +regressions, and they still run. + +The other two are marked at class level and so leave nothing behind: +`test_llm_as_judge_base_model_fix.py` (both tests wait on a pipeline) and +`test_llmaj_custom_model.py` (one test, already `us_east_1`). What the gate stops +checking there is that a submitted evaluation pipeline is *accepted and succeeds* — +genuinely useful signal, traded for 88 minutes of a 180-minute budget. Their +already-marked siblings elsewhere in the suite +(`test_benchmark_evaluation_base_model_only`, `test_custom_scorer_base_model_only`) +show this trade was already the established call for this kind of test; these two +were unmarked by omission, not by decision. + +The follow-up that closes the gap is shallow `evaluate()` coverage — asserting the +pipeline execution ARN comes back without waiting for it to finish, the same +submit-then-stop bargain this suite makes for training jobs. Until that exists the +gate verifies that evaluators construct and validate correctly, but not that a +submitted pipeline is accepted. **HyperPod (3)** — `test_nova_sft_hyperpod.py`, `test_sft_data_mixing_hyperpod.py`, `test_cpt_data_mixing_hyperpod.py`. HyperPod submits to a pre-provisioned cluster diff --git a/sagemaker-train/tests/integ/train/test_benchmark_evaluator.py b/sagemaker-train/tests/integ/train/test_benchmark_evaluator.py index 23f21229c3..73d7cee9a3 100644 --- a/sagemaker-train/tests/integ/train/test_benchmark_evaluator.py +++ b/sagemaker-train/tests/integ/train/test_benchmark_evaluator.py @@ -97,6 +97,10 @@ def test_get_benchmarks_and_properties(self): logger.info(f"MMLU properties: {properties}") + # Waits for a full evaluation pipeline (execution.wait, 4-hour ceiling), so it + # belongs off the PR gate for the same reason as its already-marked siblings + # below. + @pytest.mark.gpu_intensive def test_benchmark_evaluation_full_flow(self): """ Test complete benchmark evaluation flow with fine-tuned model package. diff --git a/sagemaker-train/tests/integ/train/test_custom_scorer_evaluator.py b/sagemaker-train/tests/integ/train/test_custom_scorer_evaluator.py index f0f0968c07..ebec92c762 100644 --- a/sagemaker-train/tests/integ/train/test_custom_scorer_evaluator.py +++ b/sagemaker-train/tests/integ/train/test_custom_scorer_evaluator.py @@ -86,6 +86,10 @@ def test_get_builtin_metrics(self): logger.info(f"Built-in metrics: {list(BuiltInMetric.__members__.keys())}") + # Waits for a full evaluation pipeline (execution.wait, 4-hour ceiling), so it + # belongs off the PR gate for the same reason as its already-marked siblings + # below. + @pytest.mark.gpu_intensive def test_custom_scorer_evaluation_full_flow(self): """ Test complete custom scorer evaluation flow with custom evaluator ARN. diff --git a/sagemaker-train/tests/integ/train/test_llm_as_judge_base_model_fix.py b/sagemaker-train/tests/integ/train/test_llm_as_judge_base_model_fix.py index 2c188a8f5d..3420cbb270 100644 --- a/sagemaker-train/tests/integ/train/test_llm_as_judge_base_model_fix.py +++ b/sagemaker-train/tests/integ/train/test_llm_as_judge_base_model_fix.py @@ -96,6 +96,12 @@ def _get_latest_model_package_arn(): return summaries[0]["ModelPackageArn"] +# Both tests in this class run a full evaluation pipeline to completion via +# execution.wait(..., timeout=14400). Measured on the PR gate's own CodeBuild run: +# 2783s and 2504s -- 88 minutes for the two of them, against the project's +# 180-minute build timeout. Everything else in that serial pass finished in under +# 92s, so these two were the entire tail. +@pytest.mark.gpu_intensive @pytest.mark.serial class TestLLMAsJudgeBaseModelFix: """Integration test for base model fix in LLMAsJudgeEvaluator""" diff --git a/sagemaker-train/tests/integ/train/test_llm_as_judge_evaluator.py b/sagemaker-train/tests/integ/train/test_llm_as_judge_evaluator.py index 4907a7317c..8c136137fc 100644 --- a/sagemaker-train/tests/integ/train/test_llm_as_judge_evaluator.py +++ b/sagemaker-train/tests/integ/train/test_llm_as_judge_evaluator.py @@ -88,6 +88,9 @@ class TestLLMAsJudgeEvaluatorIntegration: """Integration tests for LLMAsJudgeEvaluator""" + # Waits for a full evaluation pipeline (execution.wait, 4-hour ceiling). The + # two tests below it make no service call and stay on the gate. + @pytest.mark.gpu_intensive def test_llm_as_judge_evaluation_full_flow(self): """ Test complete LLM-as-Judge evaluation flow with custom and built-in metrics. diff --git a/sagemaker-train/tests/integ/train/test_llmaj_custom_model.py b/sagemaker-train/tests/integ/train/test_llmaj_custom_model.py index e3277e9509..48f608f15e 100644 --- a/sagemaker-train/tests/integ/train/test_llmaj_custom_model.py +++ b/sagemaker-train/tests/integ/train/test_llmaj_custom_model.py @@ -93,7 +93,12 @@ def test_resources(sagemaker_session_us_east_1): } -@pytest.mark.slow +# Was @pytest.mark.slow, which is not a registered marker -- the registered name +# is slow_test -- so it silently did nothing (PytestUnknownMarkWarning). This class +# waits on a full evaluation pipeline, so gpu_intensive is what it actually wants. +# us_east_1 already kept it off the us-west-2 gate, so this is not a behaviour +# change there; it now also stays off the us-east-1 job. +@pytest.mark.gpu_intensive @pytest.mark.us_east_1 class TestLLMAJCustomModelIntegration: """Integration tests for LLMAsJudgeEvaluator with InspectAI inference path.""" From c37c920e4c7b9509924abb3019236f1ebbcb37dc Mon Sep 17 00:00:00 2001 From: Amarjeet LNU Date: Fri, 14 Aug 2026 14:45:08 -0700 Subject: [PATCH 11/15] 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. --- .github/workflows/pr-checks-master.yml | 15 + .../tests/integ/train/shallow/README.md | 64 ++++ .../tests/integ/train/shallow/harness.py | 338 +++++++++++++++++- 3 files changed, 400 insertions(+), 17 deletions(-) diff --git a/.github/workflows/pr-checks-master.yml b/.github/workflows/pr-checks-master.yml index 8b76d1c043..24d0c8349b 100644 --- a/.github/workflows/pr-checks-master.yml +++ b/.github/workflows/pr-checks-master.yml @@ -308,6 +308,21 @@ jobs: # other. AWS_RETRY_MODE: adaptive AWS_MAX_ATTEMPTS: '10' + # Cap the training jobs the *service* counts against the concurrency + # quota, across all xdist workers, so the suite stays inside the + # "concurrent model customization serverless jobs per Region" quota + # (20) with room for the deep integ-tests suite running the same + # account concurrently. The harness holds each slot until the job is + # terminal, not until stop() returns -- the service counts a job for + # ~1-3 min after the stop -- so 10 means "at most 10 jobs counted at + # once", the batches-of-10 behaviour, not "10 stops in flight". + # + # Not redundant with -n 8. -n caps worker processes; this caps what + # the service counts, and with the slot held to terminal those diverge + # sharply (each drain outlives the worker's stop() by minutes). It is + # also what keeps the ceiling stable if -n is raised. A serverful job + # counts one slot per instance. + SHALLOW_MAX_CONCURRENT_JOBS: '10' run: | # Scoped to shallow/ only -- see the comment above this job for why the # rest of tests/integ/train is not repeated here. diff --git a/sagemaker-train/tests/integ/train/shallow/README.md b/sagemaker-train/tests/integ/train/shallow/README.md index 06e76ea985..9d469c6bf9 100644 --- a/sagemaker-train/tests/integ/train/shallow/README.md +++ b/sagemaker-train/tests/integ/train/shallow/README.md @@ -244,6 +244,70 @@ Two design rules follow, and should be preserved: 2. **Never set `keep_alive_period_in_seconds`.** A warm pool would outlive the stop and keep instances provisioned after the test finished. +### Concurrency cap (training-job quotas) + +`submitted()` and `assert_rejected()` hold a slot from `job_slots()` until the job +reaches a **terminal state**, bounding the number of jobs the *service* counts +against the quota **across all xdist workers** to `SHALLOW_MAX_CONCURRENT_JOBS` +(default 10). Set it to `0` to disable the gating for a single-worker debugging run. + +Two quotas apply, in two different units, and the cap has to be safe for both: + +| Path | Bounded by | Unit | +|---|---|---| +| serverless — the default recipe-trainer path, no explicit `compute` | *Maximum number of concurrent model customization serverless jobs per Region* (20) | jobs | +| serverful — an explicit `Compute`/`TrainingJobCompute`: the `ModelTrainer` tests, the tuner, `test_explicit_compute_is_accepted` | e.g. *ml.m5.large for training job usage* (100) | instances | + +Instance-type quotas do **not** apply to the serverless jobs. So a slot means "one +concurrent job", and a job costs `max(1, instance_count)` slots — 1 for a serverless +job, its instance count for a serverful one (the four `instance_count=2` tests in +`test_model_trainer.py` take two). That is the stricter of the two readings, so one +cap holds the suite inside both quotas without the harness needing to know which +kind of job a test produces. 10 sits under the serverless job quota with room for +the deep CodeBuild suite to run against the same account concurrently. + +Slots are `O_EXCL`-created files under a +run-keyed temp directory (`PYTEST_XDIST_TESTRUNUID`, falling back to the parent +pid), since xdist workers are separate processes and an in-process semaphore would +bound nothing. Keying on the run id means two concurrent local runs get separate +budgets instead of deadlocking, and a stale directory from a killed run is never +mistaken for live slots. + +The slot has to be held until the job is *terminal*, and getting this wrong is +subtle. The service counts a job against the concurrency quota from +`CreateTrainingJob` until the job reaches `Completed`/`Failed`/`Stopped` — **not** +until `StopTrainingJob` returns. Those are far apart: `stop()` returns in a few +seconds, but the job takes ~1–3 minutes to actually drain (the reservation is torn +down without ever becoming billable). An earlier version released the slot when +`stop()` returned; it bounded nothing. With the cap at 10 and 8 workers, each slot +recycled ~20× inside one 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. + +Why a cap rather than literally splitting into batches of 10: holding the slot to +terminal *is* "at most 10 jobs counted at once", the same guarantee batches give, +but the cap bounds the peak directly with no per-batch bookkeeping and keeps +bounding it if `-n` is raised or a test asks for more instances. The cost is +runtime: with the slot held to terminal, the suite's floor is roughly +`(#jobs × drain) / cap` — about **8–12 min** at ~84 jobs, a ~75s median drain and +cap 10, versus ~2 min if slots released early (the "fast" run that breaches the +quota). That is the trade the whole cap makes: correctness against the quota in +exchange for wall-clock. + +Two consequences worth knowing: + +* The stop *and the terminal-wait* happen inside the slot. Releasing before the + job is terminal is exactly the bug above — the next test starts while this job + still counts against the quota. +* Both waits are bounded and then proceed with a warning rather than failing: + acquiring a slot waits up to 900s, and `_wait_until_terminal` waits up to 300s + for the job to drain. The cap is a courtesy to the account's quota, not an + assertion about the SDK, so a leaked slot or a stuck drain degrades into a + slower run rather than a red build. + +The cap is enforced in the harness rather than per test, so a newly added test is +capped by default instead of by remembering to opt in. + ## Writing a new test Use the harness; do not call `trainer.train()` directly. diff --git a/sagemaker-train/tests/integ/train/shallow/harness.py b/sagemaker-train/tests/integ/train/shallow/harness.py index bd2c0caf37..c269cb38e2 100644 --- a/sagemaker-train/tests/integ/train/shallow/harness.py +++ b/sagemaker-train/tests/integ/train/shallow/harness.py @@ -64,9 +64,12 @@ from __future__ import absolute_import +import errno import inspect import logging +import os import random +import tempfile import time from contextlib import contextmanager @@ -75,6 +78,174 @@ logger = logging.getLogger(__name__) +# -------------------------------------------------------------------------- +# Concurrency cap (training-job service quotas) +# -------------------------------------------------------------------------- +# Two different quotas apply, in two different units, and the cap has to be safe +# for both: +# +# * 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. Instance-type quotas +# do not apply to these at all. +# * serverful (an explicit `TrainingJobCompute`/`Compute`, i.e. the +# `ModelTrainer` tests, the tuner, and `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 therefore means "one concurrent job" and a job costs +# `max(1, instance_count)` slots: 1 for a serverless job, and its instance count +# for a serverful one. That is deliberately the stricter of the two readings, so +# one cap keeps the suite inside both quotas without needing to know which kind +# of job a given test produces. +# +# What a slot has to track -- and the trap it is easy to fall into. The service +# counts a job against the concurrency quota from `CreateTrainingJob` until the +# job reaches a *terminal* state, NOT until `StopTrainingJob` returns. Those are +# far apart: measured against the service, `stop()` returns in a few seconds but +# the job does not reach `Stopped` for ~1-3 minutes afterwards while the backend +# tears down the (never-billed) reservation. An earlier version of this cap +# released the slot when `stop()` returned, and it did not bound anything: with +# the cap at 10 and 8 workers, each slot recycled ~20 times inside a single +# job's counted lifetime, so the suite peaked at ~37 concurrent jobs and tripped +# `ResourceLimitExceeded` at a utilization of 21 against the limit of 20. The +# slot must therefore be held until the job is terminal (see +# `_wait_until_terminal`), which is the point of `SHALLOW_MAX_CONCURRENT_JOBS`. +# +# Why a cap rather than batches: capping bounds the *peak* directly and keeps +# bounding it if `-n` is raised or a test starts asking for more instances, +# whereas batches of N only serialize submission. The two are equivalent when +# the slot is held to terminal -- a cap of 10 is exactly "at most 10 jobs +# counted at once" -- but the cap needs no bookkeeping of which test is in which +# batch. `SHALLOW_MAX_CONCURRENT_JOBS=0` disables it for a single-worker +# debugging run. +# +# Cost of holding to terminal: the suite's wall-clock floor becomes roughly +# (#jobs * drain_seconds) / cap rather than tracking the worker count. At ~84 +# jobs, a ~75s median drain and cap 10 that is ~8-12 min (versus ~2 min if the +# slot were released early -- but that "fast" run is the one that breaches the +# quota). 10 is under the serverless job quota (20) with room for the deep +# CodeBuild suite, which runs against the same account+region concurrently and +# also submits serverless jobs, to take the rest without the two together +# breaching 20. +DEFAULT_MAX_CONCURRENT_JOBS = 10 + + +def _max_concurrent_jobs(): + """Read the cap at call time so tests can monkeypatch the environment.""" + raw = os.environ.get("SHALLOW_MAX_CONCURRENT_JOBS") + if raw is None: + return DEFAULT_MAX_CONCURRENT_JOBS + try: + return max(0, int(raw)) + except ValueError: + logger.warning( + "Ignoring non-integer SHALLOW_MAX_CONCURRENT_JOBS=%r; using %d", + raw, + DEFAULT_MAX_CONCURRENT_JOBS, + ) + return DEFAULT_MAX_CONCURRENT_JOBS + + +# The slot directory must be shared by every xdist worker, and workers are +# separate processes, so an in-process semaphore would not bound anything. Slots +# are files in a directory keyed to the run: creating one with O_EXCL is atomic +# on POSIX, which is all the mutual exclusion this needs. Keyed on the xdist +# session id (falling back to the parent pid) so two concurrent local runs get +# their own budgets rather than deadlocking against each other -- and so a +# stale directory from a killed run is never mistaken for live slots. +def _slot_dir(): + key = os.environ.get("PYTEST_XDIST_TESTRUNUID") or str(os.getppid()) + return os.path.join(tempfile.gettempdir(), f"sm-shallow-slots-{key}") + + +# Waiting for a *free* slot is bounded so a leaked slot degrades into a slower +# run rather than a hung one. With the slot now held until the job is terminal +# (~1-3 min), a worker can legitimately queue behind several jobs' drains, so +# this is generous; anything approaching it means slots leaked. The wait logs +# and proceeds instead of failing the test, because the quota is a throttle +# rather than a correctness property. +_SLOT_WAIT_TIMEOUT = 900 +_SLOT_POLL_INTERVAL = 0.5 + + +@contextmanager +def job_slots(count=1): + """Hold ``count`` concurrency slots for the duration of the block. + + Bounds what this suite has in flight at once, across all xdist workers, to + ``SHALLOW_MAX_CONCURRENT_JOBS`` (default ``DEFAULT_MAX_CONCURRENT_JOBS``). + A slot is one concurrent job; a serverful job also takes one per additional + instance, which keeps a single cap valid against both the serverless + job-count quota and the per-instance-type quota. + + Slots are always released, including when the body raises, so a failing + assertion cannot strand capacity for the rest of the run. + """ + cap = _max_concurrent_jobs() + if cap <= 0 or count <= 0: + yield + return + + # A single test asking for more than the cap must not deadlock against + # itself: clamp, and say so, rather than waiting for slots that can never + # all be free. + if count > cap: + logger.warning( + "Test requests %d slots but the cap is %d; clamping. " + "Raise SHALLOW_MAX_CONCURRENT_JOBS if this is intentional.", + count, + cap, + ) + count = cap + + directory = _slot_dir() + os.makedirs(directory, exist_ok=True) + + held = [] + deadline = time.time() + _SLOT_WAIT_TIMEOUT + try: + while len(held) < count: + for index in range(cap): + if len(held) == count: + break + path = os.path.join(directory, f"slot-{index}") + try: + fd = os.open(path, os.O_CREAT | os.O_EXCL | os.O_WRONLY) + except OSError as e: + if e.errno == errno.EEXIST: + continue # taken by another worker + raise + os.close(fd) + held.append(path) + + if len(held) == count: + break + + if time.time() > deadline: + # Proceed rather than fail: the cap is a courtesy to the + # account's quota, not an assertion about the SDK. + logger.warning( + "Waited %ds for %d job slot(s) and got %d. Proceeding anyway " + "(slots may have leaked from a killed run: %s).", + _SLOT_WAIT_TIMEOUT, + count, + len(held), + directory, + ) + break + + time.sleep(_SLOT_POLL_INTERVAL) + + yield + finally: + for path in held: + try: + os.unlink(path) + except OSError: # pragma: no cover - already gone + pass + + # A small CPU instance is sufficient: acceptance of the request does not depend # on the instance type being an accelerator, and asking for GPU capacity we # immediately discard is both slower and antisocial in a shared test account. @@ -93,6 +264,19 @@ # Terminal/near-terminal states that make StopTrainingJob a no-op or an error. _UNSTOPPABLE_STATUSES = frozenset({"Completed", "Failed", "Stopped", "Stopping"}) +# States in which the service no longer counts the job against the concurrency +# quota. A slot is held until the job reaches one of these -- see +# `_wait_until_terminal` and the note on `DEFAULT_MAX_CONCURRENT_JOBS`. +_TERMINAL_STATUSES = frozenset({"Completed", "Failed", "Stopped"}) + +# How long a slot waits for its job to actually drain before giving up and +# releasing anyway. Measured drains are ~1-3 min; this is a ceiling, not an +# expectation. Releasing early (like the timeout on acquiring a slot) trades a +# possible brief quota overshoot for not hanging the whole suite on one stuck +# job -- the quota is a throttle, not a correctness property. +_DRAIN_WAIT_TIMEOUT = 300 +_DRAIN_POLL_INTERVAL = 5 + # Name length limits differ per resource, and the service enforces them strictly. # Verified against AWS: a 34-character tuning job name is rejected with @@ -151,6 +335,71 @@ def stop_quietly(training_job): logger.warning("Unexpected error stopping job %s: %s", name, e) +# Attributes under which the different job resources expose their status. As +# with the ARN, the SDK is not consistent: a TrainingJob uses +# ``training_job_status``, an AgentRFTJob ``job_status`` and a +# HyperParameterTuningJob ``hyper_parameter_tuning_job_status``. Read whichever +# is present. +_STATUS_ATTRS = ( + "training_job_status", + "job_status", + "hyper_parameter_tuning_job_status", +) + + +def _wait_until_terminal(training_job): + """Block until ``training_job`` leaves the concurrency-quota count. + + The service counts a job against the concurrency quota until it reaches a + terminal state, not until ``stop()`` returns, so the slot has to be held for + this whole interval (see the note on ``DEFAULT_MAX_CONCURRENT_JOBS``). This + is what makes ``SHALLOW_MAX_CONCURRENT_JOBS`` an actual bound on what the + service sees rather than on how fast slots recycle. + + Best-effort, like ``stop_quietly``: it refreshes and polls the job's status, + and on timeout or any error it logs and returns so the slot is released + anyway. A stuck job should slow the suite, not hang it or fail a test that + already made its assertion. Jobs that expose no readable status (or none of + the refresh/status plumbing) fall through immediately -- the small quota + risk there is bounded by the cap itself. + """ + if training_job is None: + return + + name = _first_attr(training_job, _NAME_ATTRS) + refresh = getattr(training_job, "refresh", None) + deadline = time.time() + _DRAIN_WAIT_TIMEOUT + while True: + try: + if callable(refresh): + refresh() + status = _first_attr(training_job, _STATUS_ATTRS) + except Exception as e: # pragma: no cover - defensive polling + logger.info("Could not read status for job %s (%s); releasing slot", name, e) + return + + if status is None: + # Nothing to poll on; do not hold a slot forever waiting for a field + # this job type never exposes. + logger.info("Job %s exposes no status; releasing slot", name) + return + if status in _TERMINAL_STATUSES: + logger.info("Job %s reached %s; releasing slot", name, status) + return + + if time.time() > deadline: + logger.warning( + "Job %s still %s after %ds; releasing slot anyway " + "(it may still count against the quota briefly).", + name, + status, + _DRAIN_WAIT_TIMEOUT, + ) + return + + time.sleep(_DRAIN_POLL_INTERVAL) + + # Attributes under which the different job resources expose their ARN and name. # Not every trainer in this package creates a TrainingJob: MultiTurnRLTrainer # creates an AgentRFT Job (``job_arn``) and Tuner creates a @@ -249,6 +498,12 @@ def submitted(trainer, **train_kwargs): value is rejected loudly rather than silently overridden, so a copy-pasted ``wait=True`` cannot quietly reintroduce a full training run into the fast suite. + + Holds a concurrency slot (see ``job_slots``) until the submitted job reaches + a terminal state, so the number of jobs the *service* counts against the + training-job quota across all xdist workers stays inside the cap. Slots are + taken here rather than in each test so a new test is capped by default + instead of by remembering to opt in. """ if "wait" in train_kwargs: raise TypeError( @@ -256,13 +511,21 @@ def submitted(trainer, **train_kwargs): "These tests must never wait for a job to run." ) - training_job = None - try: - trainer.train(**_train_kwargs_for(trainer, train_kwargs)) - training_job = _resolve_job(trainer) - yield training_job - finally: - stop_quietly(training_job) + with job_slots(_requested_slots(trainer)): + training_job = None + try: + trainer.train(**_train_kwargs_for(trainer, train_kwargs)) + training_job = _resolve_job(trainer) + yield training_job + finally: + # Stop, then hold the slot until the job is actually terminal. The + # service counts the job against the concurrency quota until it + # drains, not until stop() returns, so releasing the slot at stop() + # would let the next test start while this job still counts -- which + # is exactly how an earlier version peaked at ~37 jobs against a + # limit of 20. + stop_quietly(training_job) + _wait_until_terminal(training_job) # Attributes under which trainers stash the job they just submitted. The SDK is @@ -286,6 +549,39 @@ def _resolve_job(trainer): return _first_attr(trainer, _JOB_ATTRS) +# Where the different trainers keep an explicit compute spec, when they have one. +_COMPUTE_ATTRS = ("compute", "_compute", "compute_config") + + +def _requested_slots(trainer): + """Slots the job ``trainer`` is about to submit should consume. + + One slot per concurrent job, plus one per additional instance when the job is + serverful. See the note on ``DEFAULT_MAX_CONCURRENT_JOBS`` for why the two + quotas make this the right unit. + + Returns 1 when no explicit compute is set. That is not a fallback but the + correct answer for the default recipe-trainer path: leaving ``compute=None`` + submits a *serverless* model-customization job, which is bounded by a + per-Region job count and consumes no instance-type quota at all. + + Falls back to 1 if a compute object exists but exposes no usable count. + Under-counting is the safe direction to be wrong here: the cap remains a + useful bound, whereas guessing high would throttle the suite for no reason. + Tuning jobs are the notable inexact case -- their fan-out is set by the + tuner's own ``max_parallel_jobs`` rather than a compute block -- and there + are only two of them, both single-instance. + """ + for attr in _COMPUTE_ATTRS: + compute = getattr(trainer, attr, None) + if compute is None: + continue + count = getattr(compute, "instance_count", None) + if isinstance(count, int) and count > 0: + return count + return 1 + + def assert_rejected(trainer, expected_tokens, **train_kwargs): """Assert a request is rejected, and clean up if it is unexpectedly accepted. @@ -309,16 +605,24 @@ def assert_rejected(trainer, expected_tokens, **train_kwargs): if "wait" in train_kwargs: raise TypeError("assert_rejected() controls 'wait'; remove it from the call.") - training_job = None - try: - with pytest.raises(Exception) as excinfo: - trainer.train(**_train_kwargs_for(trainer, train_kwargs)) - # Reached only if the service accepted a request we expected it to - # refuse. Capture the job so the finally-block can stop it, then let - # pytest.raises report the missing exception. - training_job = _resolve_job(trainer) - finally: - stop_quietly(training_job) + # Slot-guarded too: a negative test is expected *not* to consume capacity, + # but if a validation regression let the request through it would, and that + # is exactly the case where staying inside the quota matters. + with job_slots(_requested_slots(trainer)): + training_job = None + try: + with pytest.raises(Exception) as excinfo: + trainer.train(**_train_kwargs_for(trainer, train_kwargs)) + # Reached only if the service accepted a request we expected it to + # refuse. Capture the job so the finally-block can stop it, then let + # pytest.raises report the missing exception. + training_job = _resolve_job(trainer) + finally: + # Normally a no-op (the request was rejected, so no job exists). If a + # regression let it through, drain it inside the slot for the same + # reason submitted() does. + stop_quietly(training_job) + _wait_until_terminal(training_job) message = str(excinfo.value) assert any(token in message for token in expected_tokens), ( From 3872c3ccf869ebd8a662d0ecec99a0c2e255b58c Mon Sep 17 00:00:00 2001 From: Amarjeet LNU Date: Tue, 18 Aug 2026 14:37:38 -0700 Subject: [PATCH 12/15] 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. --- .../tests/integ/train/shallow/README.md | 6 ++--- .../tests/integ/train/shallow/conftest.py | 22 +++++++++---------- .../integ/train/shallow/test_nova_trainers.py | 12 +++++----- 3 files changed, 20 insertions(+), 20 deletions(-) diff --git a/sagemaker-train/tests/integ/train/shallow/README.md b/sagemaker-train/tests/integ/train/shallow/README.md index 9d469c6bf9..85009a0118 100644 --- a/sagemaker-train/tests/integ/train/shallow/README.md +++ b/sagemaker-train/tests/integ/train/shallow/README.md @@ -192,9 +192,9 @@ The Nova tests (`us_east_1`) build every S3 path from `default_bucket()` and resolve the reward function from the calling account's own hub, rather than naming the resources the deep Nova tests use. -This is not stylistic. The deep tests hardcode -`s3://sagemaker-us-east-1-784379639078/...`, which other accounts cannot read — -verified: `AccessDenied` on `ListObjectsV2` from 729646638167. A hardcoded path +This is not stylistic. The deep tests hardcode a bucket belonging to one specific +test account, which other accounts cannot read — verified: `AccessDenied` on +`ListObjectsV2` from a different account. A hardcoded path means the test only runs in one account and fails everywhere else, which is how these five ended up never having been executed. `test_sft_trainer_serverful_smtj.py` already takes the derived approach (`training_resources`); these follow it, and diff --git a/sagemaker-train/tests/integ/train/shallow/conftest.py b/sagemaker-train/tests/integ/train/shallow/conftest.py index 6328cc3e74..244037740d 100644 --- a/sagemaker-train/tests/integ/train/shallow/conftest.py +++ b/sagemaker-train/tests/integ/train/shallow/conftest.py @@ -179,10 +179,10 @@ def nova_rlvr_data_uri(sagemaker_session_us_east_1, reward_scored_data_uri): * An S3 input must be in the same region as the job, and ``reward_scored_data_uri`` lives in us-west-2. - The deep test points at ``grpo-64-sample.jsonl`` in account 784379639078, which - is not readable from every account this runs in, so this copies the dataset the - us-west-2 RLVR tests already use. Idempotent, and skips rather than failing if - the source is unreadable. + The deep test points at ``grpo-64-sample.jsonl`` in a bucket belonging to one + specific test account, which is not readable from every account this runs in, + so this copies the dataset the us-west-2 RLVR tests already use. Idempotent, + and skips rather than failing if the source is unreadable. """ bucket = sagemaker_session_us_east_1.default_bucket() key = "shallow-integ-test/nova-rlvr/train_285.jsonl" @@ -207,10 +207,10 @@ def nova_output_path(sagemaker_session_us_east_1): """S3 prefix for Nova training output, in the caller's own us-east-1 bucket. Deliberately derived rather than hardcoded. The deep Nova tests name a bucket - in a specific test account (``sagemaker-us-east-1-784379639078``), which is not - readable from every account the suite runs in -- verified: ``AccessDenied`` on - ``ListObjectsV2`` from 729646638167. Using ``default_bucket()`` makes these - tests work in any account, the same way + belonging to one specific test account, which is not readable from every + account the suite runs in -- verified: ``AccessDenied`` on ``ListObjectsV2`` + from a different account. Using ``default_bucket()`` makes these tests work in + any account, the same way ``test_sft_trainer_serverful_smtj.py::training_resources`` does. """ return f"s3://{sagemaker_session_us_east_1.default_bucket()}/shallow-integ-test/output/" @@ -221,9 +221,9 @@ def nova_reward_function_arn(sagemaker_session_us_east_1): """ARN of the Nova RLVR reward function in the caller's own account. Look-up-and-skip, like the other reward fixtures. The deep test hardcodes this - ARN in account 784379639078; resolving it per-account instead means the test - runs wherever the hub content has been provisioned and skips cleanly elsewhere, - rather than failing with a confusing cross-account hub error. + ARN against one specific account; resolving it per-account instead means the + test runs wherever the hub content has been provisioned and skips cleanly + elsewhere, rather than failing with a confusing cross-account hub error. """ client = sagemaker_session_us_east_1.boto_session.client("sagemaker") hub, name = "sdktest", "rlvr-nova-test-rf" diff --git a/sagemaker-train/tests/integ/train/shallow/test_nova_trainers.py b/sagemaker-train/tests/integ/train/shallow/test_nova_trainers.py index 5b820845a9..3d71f4c5d2 100644 --- a/sagemaker-train/tests/integ/train/shallow/test_nova_trainers.py +++ b/sagemaker-train/tests/integ/train/shallow/test_nova_trainers.py @@ -22,8 +22,8 @@ Datasets, output paths and the reward function are all *derived from the calling account* rather than hardcoded. The deep Nova tests name resources in one specific -test account (``sagemaker-us-east-1-784379639078``), which other accounts cannot -read -- verified: ``AccessDenied`` on ``ListObjectsV2`` from 729646638167. Using +test account's bucket, which other accounts cannot read -- verified: +``AccessDenied`` on ``ListObjectsV2`` from a different account. Using ``default_bucket()`` and resolving the reward function from the caller's own hub follows what ``test_sft_trainer_serverful_smtj.py`` already does, and means these tests actually run wherever the suite runs instead of only in one account. @@ -98,10 +98,10 @@ def test_nova_rlvr_is_accepted( # records and refuses the call if their scores do not parse. That gate # is real, but it asserts the contents of a hub artifact provisioned # per-account rather than anything about this payload -- verified: the - # function registered under this name in 729646638167 returns a shape - # the verifier rejects ("Each output must include 'id', - # 'aggregate_reward_score'"), so the test would fail on account state - # rather than on a regression. + # function registered under this name in the account this was run + # against returns a shape the verifier rejects ("Each output must + # include 'id', 'aggregate_reward_score'"), so the test would fail on + # account state rather than on a regression. # # The verifier itself is already covered, against a known-compatible # function, by the three us-west-2 reward-function cases in From 8666594aa4a70f4370e9db09547a15a20217359d Mon Sep 17 00:00:00 2001 From: Amarjeet LNU Date: Tue, 18 Aug 2026 16:45:40 -0700 Subject: [PATCH 13/15] 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. --- .../tests/integ/train/shallow/README.md | 10 ++++---- .../tests/integ/train/shallow/recipe_cases.py | 24 +++++++++---------- .../train/shallow/test_nova_data_mixing.py | 4 ++-- .../integ/train/shallow/test_nova_trainers.py | 6 ++--- 4 files changed, 23 insertions(+), 21 deletions(-) diff --git a/sagemaker-train/tests/integ/train/shallow/README.md b/sagemaker-train/tests/integ/train/shallow/README.md index 85009a0118..44317d9bbc 100644 --- a/sagemaker-train/tests/integ/train/shallow/README.md +++ b/sagemaker-train/tests/integ/train/shallow/README.md @@ -204,10 +204,12 @@ upload the Nova-shaped sample data the deep suite already ships Two region constraints are worth knowing before adding a Nova test, both verified against the service: -* the model package group must be in the **job's** region — passing the us-west-2 - ARN from `MODEL_PACKAGE_GROUP` is rejected with `Model package group ARN region - 'us-west-2' does not match expected region 'us-east-1'`, so Nova files use - `NOVA_MODEL_PACKAGE_GROUP` (a bare name, which resolves per-session); +* the model package group must be in the **job's** region, which is why + `MODEL_PACKAGE_GROUP` is a bare name rather than an ARN. The SDK resolves a name + against the session's own region, while an ARN pins both region and account — + and passing a us-west-2 ARN to a us-east-1 job is rejected with `Model package + group ARN region 'us-west-2' does not match expected region 'us-east-1'`. One + name therefore serves both regions; * an S3 input must be in the job's region, so `nova_rlvr_data_uri` copies the us-west-2 RLVR dataset into the us-east-1 bucket rather than referencing it. diff --git a/sagemaker-train/tests/integ/train/shallow/recipe_cases.py b/sagemaker-train/tests/integ/train/shallow/recipe_cases.py index 2f3357b381..b6593b6a77 100644 --- a/sagemaker-train/tests/integ/train/shallow/recipe_cases.py +++ b/sagemaker-train/tests/integ/train/shallow/recipe_cases.py @@ -50,22 +50,22 @@ class TestFooTrainerSubmission(RecipeTrainerCases): # takes during submission. MODEL_ID = "meta-textgeneration-llama-3-2-1b-instruct" -# Reused from the existing dry-run suite so both suites exercise the same -# already-provisioned model package group rather than each needing their own. -MODEL_PACKAGE_GROUP = ( - "arn:aws:sagemaker:us-west-2:729646638167:model-package-group/sdk-test-finetuned-models" -) - -# The Nova files (us-east-1) cannot reuse the ARN above. Verified against AWS: the -# model package group must be in the same region as the job, and passing the -# us-west-2 ARN is rejected with +# The already-provisioned group the dry-run suite also uses, so both suites share +# one group rather than each needing their own. +# +# A bare name rather than an ARN, deliberately, and it is the same constant for +# every region. The SDK resolves a name against the *session's* region +# (`_resolve_model_package_group_arn` -> `ModelPackageGroup.get`), whereas an ARN +# pins both the region and the account. Pinning the region breaks the us-east-1 +# Nova path outright -- verified against AWS, passing a us-west-2 ARN to a +# us-east-1 job is rejected with # # Model package group ARN region 'us-west-2' does not match expected region # 'us-east-1' # -# A bare name resolves in whichever region the session is in. Shared here rather -# than duplicated per Nova file so the two cannot drift apart. -NOVA_MODEL_PACKAGE_GROUP = "sdk-test-finetuned-models" +# so a name is what lets one constant serve both regions, and it keeps the tests +# runnable in any account that has provisioned the group. +MODEL_PACKAGE_GROUP = "sdk-test-finetuned-models" # An accelerator type is required for the serverful recipe path: these recipes do # not resolve onto a CPU instance, so unlike the ModelTrainer suite we cannot use diff --git a/sagemaker-train/tests/integ/train/shallow/test_nova_data_mixing.py b/sagemaker-train/tests/integ/train/shallow/test_nova_data_mixing.py index 98cde7868a..28cd9d49cc 100644 --- a/sagemaker-train/tests/integ/train/shallow/test_nova_data_mixing.py +++ b/sagemaker-train/tests/integ/train/shallow/test_nova_data_mixing.py @@ -31,7 +31,7 @@ from sagemaker.train.sft_trainer import SFTTrainer from .harness import assert_submitted, submitted, unique_name -from .recipe_cases import NOVA_MODEL_PACKAGE_GROUP, stopping_condition +from .recipe_cases import MODEL_PACKAGE_GROUP, stopping_condition NOVA_MODEL = "nova-textgeneration-lite-v2" @@ -39,7 +39,7 @@ def _nova_sft(session, dataset, name, config): return SFTTrainer( model=NOVA_MODEL, - model_package_group=NOVA_MODEL_PACKAGE_GROUP, + model_package_group=MODEL_PACKAGE_GROUP, training_dataset=dataset, accept_eula=True, sagemaker_session=session, diff --git a/sagemaker-train/tests/integ/train/shallow/test_nova_trainers.py b/sagemaker-train/tests/integ/train/shallow/test_nova_trainers.py index 3d71f4c5d2..d8c9d13684 100644 --- a/sagemaker-train/tests/integ/train/shallow/test_nova_trainers.py +++ b/sagemaker-train/tests/integ/train/shallow/test_nova_trainers.py @@ -39,7 +39,7 @@ from sagemaker.train.sft_trainer import SFTTrainer from .harness import MAX_RUNTIME_IN_SECONDS, assert_submitted, submitted, unique_name -from .recipe_cases import NOVA_MODEL_PACKAGE_GROUP +from .recipe_cases import MODEL_PACKAGE_GROUP NOVA_MODEL = "nova-textgeneration-lite-v2" @@ -58,7 +58,7 @@ def test_nova_sft_is_accepted( trainer = SFTTrainer( model=NOVA_MODEL, training_type=TrainingType.LORA, - model_package_group=NOVA_MODEL_PACKAGE_GROUP, + model_package_group=MODEL_PACKAGE_GROUP, training_dataset=nova_sft_data_uri, s3_output_path=nova_output_path, accept_eula=True, @@ -85,7 +85,7 @@ def test_nova_rlvr_is_accepted( trainer = RLVRTrainer( model=NOVA_MODEL, training_type=TrainingType.LORA, - model_package_group=NOVA_MODEL_PACKAGE_GROUP, + model_package_group=MODEL_PACKAGE_GROUP, training_dataset=nova_rlvr_data_uri, validation_dataset=nova_rlvr_data_uri, s3_output_path=nova_output_path, From 923a0925735683b23f0648831a26c26c0dd395b6 Mon Sep 17 00:00:00 2001 From: Amarjeet LNU Date: Tue, 18 Aug 2026 16:58:37 -0700 Subject: [PATCH 14/15] 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. --- .../tests/integ/train/shallow/README.md | 19 +++++++++ .../tests/integ/train/shallow/harness.py | 42 +++++++++++++++++-- .../integ/train/shallow/test_model_trainer.py | 6 +-- .../tests/integ/train/shallow/test_tuner.py | 4 +- 4 files changed, 62 insertions(+), 9 deletions(-) diff --git a/sagemaker-train/tests/integ/train/shallow/README.md b/sagemaker-train/tests/integ/train/shallow/README.md index 44317d9bbc..a28a591f89 100644 --- a/sagemaker-train/tests/integ/train/shallow/README.md +++ b/sagemaker-train/tests/integ/train/shallow/README.md @@ -246,6 +246,25 @@ Two design rules follow, and should be preserved: 2. **Never set `keep_alive_period_in_seconds`.** A warm pool would outlive the stop and keep instances provisioned after the test finished. +### Container URIs are resolved, not hardcoded + +`harness.cpu_image(sagemaker_session)` resolves the public CPU training DLC in the +*session's* region via `image_uris.retrieve` — the same resolver the SDK's own +framework estimators use — rather than naming a URI. The deep suites hardcode a +us-west-2 one. + +This is worth the indirection because the registry account is not constant: it is +`763104351884` across the commercial regions but differs in GovCloud +(`442386744353`) and China (`727897471807`, on `.com.cn`). A hardcoded URI is +therefore not merely region-pinned, it is unusable outside one partition, and the +failure mode is an ECR error from the backend's role-assuming validators that looks +like a test bug rather than a hardcoded constant. Resolving per-session keeps the +image following wherever the suite runs and shrinks the blast radius if one region +is misconfigured. + +It is a function rather than a constant precisely because it needs the session's +region, so a new call site must pass the session it is submitting with. + ### Concurrency cap (training-job quotas) `submitted()` and `assert_rejected()` hold a slot from `job_slots()` until the job diff --git a/sagemaker-train/tests/integ/train/shallow/harness.py b/sagemaker-train/tests/integ/train/shallow/harness.py index c269cb38e2..df4d11f92d 100644 --- a/sagemaker-train/tests/integ/train/shallow/harness.py +++ b/sagemaker-train/tests/integ/train/shallow/harness.py @@ -252,10 +252,44 @@ def job_slots(count=1): DEFAULT_INSTANCE_TYPE = "ml.m5.large" DEFAULT_INSTANCE_COUNT = 1 -# Public DLC, present in every commercial region we test in. Using a real image -# matters: the backend's role-assuming validators resolve the training image -# against ECR, so a bogus URI would fail for the wrong reason. -CPU_IMAGE = "763104351884.dkr.ecr.us-west-2.amazonaws.com/pytorch-training:2.0.0-cpu-py310" +# The AWS Deep Learning Container these tests use as a stand-in for "some real +# training image". Using a real image matters: the backend's role-assuming +# validators resolve the training image against ECR as the customer, so a bogus +# URI would fail the test for the wrong reason. +# +# Resolved per-region rather than hardcoded -- see `cpu_image`. +_CPU_IMAGE_FRAMEWORK = "pytorch" +_CPU_IMAGE_VERSION = "2.0.0" +_CPU_IMAGE_PY_VERSION = "py310" + + +def cpu_image(sagemaker_session): + """ECR URI of a public CPU training DLC, in the *session's* region. + + Region-agnostic deliberately. A hardcoded URI pins the region (and the + registry account, which differs in the China and GovCloud partitions), so a + test running anywhere else would either pull cross-region or fail on a + registry that does not exist there. Resolving from the session means the + image follows wherever the suite runs -- one less thing to update when a + region is added, and a smaller blast radius if one is misconfigured. + + ``image_uris.retrieve`` is the same resolver the SDK's own framework + estimators use, so this is the supported mapping rather than a + reconstruction of it. Verified against AWS that it reproduces the URI this + previously hardcoded (``pytorch-training:2.0.0-cpu-py310`` in the public DLC + account) and returns the corresponding in-region host elsewhere. + """ + from sagemaker.core import image_uris + + return image_uris.retrieve( + framework=_CPU_IMAGE_FRAMEWORK, + region=sagemaker_session.boto_session.region_name, + version=_CPU_IMAGE_VERSION, + py_version=_CPU_IMAGE_PY_VERSION, + instance_type=DEFAULT_INSTANCE_TYPE, + image_scope="training", + ) + # Keep the advertised runtime short. It should never be reached (we stop the job # long before), but if a stop were somehow lost this bounds the damage. diff --git a/sagemaker-train/tests/integ/train/shallow/test_model_trainer.py b/sagemaker-train/tests/integ/train/shallow/test_model_trainer.py index 7a633762a0..58dd31271f 100644 --- a/sagemaker-train/tests/integ/train/shallow/test_model_trainer.py +++ b/sagemaker-train/tests/integ/train/shallow/test_model_trainer.py @@ -33,12 +33,12 @@ from sagemaker.train.model_trainer import ModelTrainer from .harness import ( - CPU_IMAGE, DEFAULT_INSTANCE_COUNT, DEFAULT_INSTANCE_TYPE, MAX_RUNTIME_IN_SECONDS, assert_rejected, assert_submitted, + cpu_image, stop_quietly, submitted, unique_name, @@ -99,7 +99,7 @@ def _trainer(sagemaker_session, name, **overrides): """ kwargs = dict( sagemaker_session=sagemaker_session, - training_image=CPU_IMAGE, + training_image=cpu_image(sagemaker_session), source_code=_source_code(), compute=_compute(), stopping_condition=_stopping_condition(), @@ -656,7 +656,7 @@ def create(): training_job_name=job_name, role_arn=execution_role, algorithm_specification=shapes.AlgorithmSpecification( - training_image=CPU_IMAGE, training_input_mode="File" + training_image=cpu_image(sagemaker_session), training_input_mode="File" ), output_data_config=shapes.OutputDataConfig(s3_output_path=output_path), resource_config=shapes.ResourceConfig( diff --git a/sagemaker-train/tests/integ/train/shallow/test_tuner.py b/sagemaker-train/tests/integ/train/shallow/test_tuner.py index 94afad3e6d..9629f38e49 100644 --- a/sagemaker-train/tests/integ/train/shallow/test_tuner.py +++ b/sagemaker-train/tests/integ/train/shallow/test_tuner.py @@ -49,12 +49,12 @@ from sagemaker.train.tuner import HyperparameterTuner from .harness import ( - CPU_IMAGE, DEFAULT_INSTANCE_COUNT, DEFAULT_INSTANCE_TYPE, MAX_RUNTIME_IN_SECONDS, MAX_TUNING_JOB_NAME, assert_submitted, + cpu_image, submitted, unique_name, ) @@ -69,7 +69,7 @@ def _model_trainer(sagemaker_session, name, **overrides): """The inner trainer a tuning job wraps.""" kwargs = dict( sagemaker_session=sagemaker_session, - training_image=CPU_IMAGE, + training_image=cpu_image(sagemaker_session), base_job_name=name, source_code=SourceCode( source_dir=PARAM_SCRIPT_SOURCE_DIR, From 72c98d76bc94406f8cffe57e2860e27363a6d410 Mon Sep 17 00:00:00 2001 From: Amarjeet LNU Date: Tue, 18 Aug 2026 17:07:23 -0700 Subject: [PATCH 15/15] 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. --- .../tests/integ/train/shallow/README.md | 25 +++++++++++- .../tests/integ/train/shallow/harness.py | 24 ++++++++---- .../tests/integ/train/shallow/test_tuner.py | 39 ++++++++++++++----- 3 files changed, 68 insertions(+), 20 deletions(-) diff --git a/sagemaker-train/tests/integ/train/shallow/README.md b/sagemaker-train/tests/integ/train/shallow/README.md index a28a591f89..02de12214d 100644 --- a/sagemaker-train/tests/integ/train/shallow/README.md +++ b/sagemaker-train/tests/integ/train/shallow/README.md @@ -272,6 +272,11 @@ reaches a **terminal state**, bounding the number of jobs the *service* counts against the quota **across all xdist workers** to `SHALLOW_MAX_CONCURRENT_JOBS` (default 10). Set it to `0` to disable the gating for a single-worker debugging run. +`_tuning()` in `test_tuner.py` is the one submission path that does not go through +those two, since a tuning job is stopped via `tuner.stop_tuning_job()` rather than +`stop_quietly`. It acquires slots itself, on the same terms — see *Tuning jobs* +below. **Any new submission path must do likewise; the cap is not automatic.** + Two quotas apply, in two different units, and the cap has to be safe for both: | Path | Bounded by | Unit | @@ -287,6 +292,22 @@ cap holds the suite inside both quotas without the harness needing to know which kind of job a test produces. 10 sits under the serverless job quota with room for the deep CodeBuild suite to run against the same account concurrently. +#### Tuning jobs + +A tuning job consumes instance quota through the **child training jobs it launches**, +not through the tuning job itself, so its cost is the tuner's `max_parallel_jobs` +rather than anything derivable from a compute block — which is why `_tuning()` sizes +its own request instead of using `_requested_slots`. It holds those slots until the +tuning job is terminal, the same rule as everywhere else: `stop_tuning_job()` returns +while the job is still `Stopping` and its children are still tearing down, so +releasing there would be the same release-before-terminal mistake described below. + +Both tuner tests are `max_parallel_jobs=1`, so today this is 1 slot each and the +practical overshoot it prevents is small. It is wired up anyway because the cost is +one context manager, and the failure mode if a future test raises `max_jobs` or +`max_parallel_jobs` is the silent kind — capacity consumed outside the cap that the +cap still claims to bound. + Slots are `O_EXCL`-created files under a run-keyed temp directory (`PYTEST_XDIST_TESTRUNUID`, falling back to the parent pid), since xdist workers are separate processes and an in-process semaphore would @@ -303,7 +324,7 @@ down without ever becoming billable). An earlier version released the slot when `stop()` returned; it bounded nothing. With the cap at 10 and 8 workers, each slot recycled ~20× inside one 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. +against the limit of 20. `wait_until_terminal` closes that gap. Why a cap rather than literally splitting into batches of 10: holding the slot to terminal *is* "at most 10 jobs counted at once", the same guarantee batches give, @@ -321,7 +342,7 @@ Two consequences worth knowing: job is terminal is exactly the bug above — the next test starts while this job still counts against the quota. * Both waits are bounded and then proceed with a warning rather than failing: - acquiring a slot waits up to 900s, and `_wait_until_terminal` waits up to 300s + acquiring a slot waits up to 900s, and `wait_until_terminal` waits up to 300s for the job to drain. The cap is a courtesy to the account's quota, not an assertion about the SDK, so a leaked slot or a stuck drain degrades into a slower run rather than a red build. diff --git a/sagemaker-train/tests/integ/train/shallow/harness.py b/sagemaker-train/tests/integ/train/shallow/harness.py index df4d11f92d..13b4d754eb 100644 --- a/sagemaker-train/tests/integ/train/shallow/harness.py +++ b/sagemaker-train/tests/integ/train/shallow/harness.py @@ -99,6 +99,12 @@ # one cap keeps the suite inside both quotas without needing to know which kind # of job a given test produces. # +# Tuning jobs are counted through the same mechanism but sized differently: their +# capacity is occupied by the child training jobs the tuner launches, so +# `_tuning()` in `test_tuner.py` holds `max_parallel_jobs` slots rather than +# deriving a count from a compute block. Every path that submits must acquire +# slots -- see the note on `job_slots`. +# # What a slot has to track -- and the trap it is easy to fall into. The service # counts a job against the concurrency quota from `CreateTrainingJob` until the # job reaches a *terminal* state, NOT until `StopTrainingJob` returns. Those are @@ -110,7 +116,7 @@ # job's counted lifetime, so the suite peaked at ~37 concurrent jobs and tripped # `ResourceLimitExceeded` at a utilization of 21 against the limit of 20. The # slot must therefore be held until the job is terminal (see -# `_wait_until_terminal`), which is the point of `SHALLOW_MAX_CONCURRENT_JOBS`. +# `wait_until_terminal`), which is the point of `SHALLOW_MAX_CONCURRENT_JOBS`. # # Why a cap rather than batches: capping bounds the *peak* directly and keeps # bounding it if `-n` is raised or a test starts asking for more instances, @@ -300,7 +306,7 @@ def cpu_image(sagemaker_session): # States in which the service no longer counts the job against the concurrency # quota. A slot is held until the job reaches one of these -- see -# `_wait_until_terminal` and the note on `DEFAULT_MAX_CONCURRENT_JOBS`. +# `wait_until_terminal` and the note on `DEFAULT_MAX_CONCURRENT_JOBS`. _TERMINAL_STATUSES = frozenset({"Completed", "Failed", "Stopped"}) # How long a slot waits for its job to actually drain before giving up and @@ -381,7 +387,7 @@ def stop_quietly(training_job): ) -def _wait_until_terminal(training_job): +def wait_until_terminal(training_job): """Block until ``training_job`` leaves the concurrency-quota count. The service counts a job against the concurrency quota until it reaches a @@ -559,7 +565,7 @@ def submitted(trainer, **train_kwargs): # is exactly how an earlier version peaked at ~37 jobs against a # limit of 20. stop_quietly(training_job) - _wait_until_terminal(training_job) + wait_until_terminal(training_job) # Attributes under which trainers stash the job they just submitted. The SDK is @@ -602,9 +608,11 @@ def _requested_slots(trainer): Falls back to 1 if a compute object exists but exposes no usable count. Under-counting is the safe direction to be wrong here: the cap remains a useful bound, whereas guessing high would throttle the suite for no reason. - Tuning jobs are the notable inexact case -- their fan-out is set by the - tuner's own ``max_parallel_jobs`` rather than a compute block -- and there - are only two of them, both single-instance. + + Only applies to trainers submitted through ``submitted()``/ + ``assert_rejected()``. A tuning job's fan-out comes from the tuner's + ``max_parallel_jobs`` rather than a compute block, so ``_tuning()`` in + ``test_tuner.py`` sizes its own request and calls ``job_slots`` directly. """ for attr in _COMPUTE_ATTRS: compute = getattr(trainer, attr, None) @@ -656,7 +664,7 @@ def assert_rejected(trainer, expected_tokens, **train_kwargs): # regression let it through, drain it inside the slot for the same # reason submitted() does. stop_quietly(training_job) - _wait_until_terminal(training_job) + wait_until_terminal(training_job) message = str(excinfo.value) assert any(token in message for token in expected_tokens), ( diff --git a/sagemaker-train/tests/integ/train/shallow/test_tuner.py b/sagemaker-train/tests/integ/train/shallow/test_tuner.py index 9629f38e49..fb3fdfc884 100644 --- a/sagemaker-train/tests/integ/train/shallow/test_tuner.py +++ b/sagemaker-train/tests/integ/train/shallow/test_tuner.py @@ -55,8 +55,10 @@ MAX_TUNING_JOB_NAME, assert_submitted, cpu_image, + job_slots, submitted, unique_name, + wait_until_terminal, ) logger = logging.getLogger(__name__) @@ -124,18 +126,35 @@ def _tuning(tuner, job_name): Teardown goes through ``tuner.stop_tuning_job()`` rather than the harness's ``stop_quietly``, because the tuner wraps the resource and stopping it also stops the child training jobs it launched. + + Concurrency accounting mirrors ``submitted()``: a tuning job's child training + jobs consume the same per-instance-type quota the ``ModelTrainer`` tests do, + so this cannot bypass the cap just because the resource type differs. It + holds ``max_parallel_jobs`` slots -- the tuner's own fan-out bound, since the + children are what occupy capacity, not the tuning job itself -- and holds + them until the tuning job is terminal rather than releasing when + ``stop_tuning_job()`` returns. Releasing at stop is precisely the + release-before-terminal bug documented on ``DEFAULT_MAX_CONCURRENT_JOBS``. """ - try: - tuner.tune(job_name=job_name, wait=False) - yield - finally: + slots = max(1, getattr(tuner, "max_parallel_jobs", 1) or 1) + with job_slots(slots): try: - tuner.stop_tuning_job() - logger.info("Stopped tuning job %s", job_name) - except Exception as e: # pragma: no cover - best-effort teardown - # A tuning job that never started, or already reached a terminal - # state, cannot be stopped; that must not fail the test. - logger.warning("Could not stop tuning job %s: %s", job_name, e) + tuner.tune(job_name=job_name, wait=False) + yield + finally: + try: + tuner.stop_tuning_job() + logger.info("Stopped tuning job %s", job_name) + except Exception as e: # pragma: no cover - best-effort teardown + # A tuning job that never started, or already reached a terminal + # state, cannot be stopped; that must not fail the test. + logger.warning("Could not stop tuning job %s: %s", job_name, e) + + # Stopping a tuning job is asynchronous: it goes Stopping -> Stopped + # while its children tear down, and the children hold instance quota + # for that whole interval. Bounded and best-effort, like everywhere + # else in the harness. + wait_until_terminal(tuner.latest_tuning_job) class TestTuningJobSubmission: