From 7697926e1d38506719ca12e54ae1f504c54d442f Mon Sep 17 00:00:00 2001 From: Amarjeet LNU Date: Tue, 11 Aug 2026 16:07:52 -0700 Subject: [PATCH 1/5] 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 2/5] 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 3/5] 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 4/5] 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 5/5] 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)