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..bb37399280 --- /dev/null +++ b/sagemaker-train/tests/integ/train/shallow/README.md @@ -0,0 +1,196 @@ +# 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. + +## 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. + +## Coverage of every `gpu_intensive` test + +The rule: **a deep test belongs off the PR gate only if this suite covers the same +code path.** There are 46 `gpu_intensive` tests in `tests/integ/train`; the table +below accounts for all of them. + +### Covered by this suite + +| Deep test | Shallow equivalent | +|---|---| +| `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. + +**Do not add `gpu_intensive` to a deep test unless a shallow test covers the same +path**, or the PR gate silently loses coverage. + +### Fixtures that skip rather than create + +`mlflow_arn`, `reward_lambda_arn` and `reward_evaluator` only *look up* their +resources and skip when absent. The deep suite's equivalents create them (IAM +roles, Lambdas, MLflow apps, registry entries) — durable side effects that a fast +PR-gate suite should not perform. + +## 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..f92e93dfd3 --- /dev/null +++ b/sagemaker-train/tests/integ/train/shallow/conftest.py @@ -0,0 +1,263 @@ +# 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 os + +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(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.""" + 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 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. + + 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..bd2c0caf37 --- /dev/null +++ b/sagemaker-train/tests/integ/train/shallow/harness.py @@ -0,0 +1,329 @@ +# 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"}) + + +# 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 + + +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)}" + # 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): + """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/recipe_cases.py b/sagemaker-train/tests/integ/train/shallow/recipe_cases.py new file mode 100644 index 0000000000..26c9a33725 --- /dev/null +++ b/sagemaker-train/tests/integ/train/shallow/recipe_cases.py @@ -0,0 +1,271 @@ +# 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) + + 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): + """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.py b/sagemaker-train/tests/integ/train/shallow/test_model_trainer.py new file mode 100644 index 0000000000..7a633762a0 --- /dev/null +++ b/sagemaker-train/tests/integ/train/shallow/test_model_trainer.py @@ -0,0 +1,684 @@ +# 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): + """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" + + 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", + "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 + 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_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_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 new file mode 100644 index 0000000000..69ac928ab2 --- /dev/null +++ b/sagemaker-train/tests/integ/train/shallow/test_rlaif_trainer.py @@ -0,0 +1,86 @@ +# 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 .harness import assert_submitted, submitted +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" + +# 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. + + 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 + + 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 new file mode 100644 index 0000000000..87a80d0e67 --- /dev/null +++ b/sagemaker-train/tests/integ/train/shallow/test_rlvr_trainer.py @@ -0,0 +1,165 @@ +# 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 + +# 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.""" + + 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) + + # -- 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) 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_tuner.py b/sagemaker-train/tests/integ/train/shallow/test_tuner.py new file mode 100644 index 0000000000..94afad3e6d --- /dev/null +++ b/sagemaker-train/tests/integ/train/shallow/test_tuner.py @@ -0,0 +1,192 @@ +# 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 +from contextlib import contextmanager + +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, + MAX_TUNING_JOB_NAME, + 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) + + +@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. + + 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", max_length=MAX_TUNING_JOB_NAME) + tuner = _tuner(_model_trainer(sagemaker_session, name)) + + 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 + ``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-tune-dist", max_length=MAX_TUNING_JOB_NAME) + model_trainer = _model_trainer(sagemaker_session, name, distributed=Torchrun()) + tuner = _tuner(model_trainer) + + 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 it from acceptance alone. + described = tuner.latest_tuning_job.refresh() + definition = getattr(described, "training_job_definition", None) + 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}" 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_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)