From 7cabe46dd42ff55cdf701304819090ae6214d89b Mon Sep 17 00:00:00 2001 From: Xichao Wang Date: Tue, 13 Oct 2026 00:40:21 +0000 Subject: [PATCH] feat(train): add the interactive Training Sessions client surface Adds the interactive Training Sessions client surface to sagemaker-train. * ServiceClient / TrainingClient / SamplingClient, exposed lazily from sagemaker.train * Typed request and result models, futures, and exception hierarchy * Session lifecycle management with refcounted handles and idempotent stop * Reference docs for the new modules under docs/api/sagemaker_train.rst * A "tokenizer" extra for the optional transformers dependency that get_tokenizer() needs The clients call the Training Sessions control plane and runtime through boto3, both of which are available in the public botocore release, so no vendored service models ship in the wheel. --- docs/api/sagemaker_train.rst | 38 + sagemaker-train/pyproject.toml | 4 +- .../src/sagemaker/train/__init__.py | 12 + .../train/training_session/README.md | 995 ++++++++ .../train/training_session/__init__.py | 149 ++ .../train/training_session/_clients.py | 167 ++ .../train/training_session/_endpoints.py | 119 + .../train/training_session/_session_handle.py | 247 ++ .../train/training_session/_submitter.py | 256 ++ .../train/training_session/_utils.py | 1492 +++++++++++ .../train/training_session/exceptions.py | 212 ++ .../train/training_session/futures.py | 865 +++++++ .../train/training_session/sampling_client.py | 476 ++++ .../train/training_session/service_client.py | 1313 ++++++++++ .../train/training_session/training_client.py | 1090 ++++++++ .../sagemaker/train/training_session/types.py | 1470 +++++++++++ .../unit/train/training_session/__init__.py | 0 .../session_lifecycle_grid.py | 185 ++ .../train/training_session/test_clients.py | 440 ++++ .../train/training_session/test_exceptions.py | 172 ++ .../train/training_session/test_futures.py | 1556 ++++++++++++ .../training_session/test_sampling_client.py | 821 ++++++ .../test_script_compatibility.py | 618 +++++ .../training_session/test_service_client.py | 2190 +++++++++++++++++ .../test_session_lifecycle.py | 586 +++++ .../test_session_lifecycle_grid.py | 98 + .../train/training_session/test_submitter.py | 453 ++++ .../training_session/test_training_client.py | 1170 +++++++++ .../unit/train/training_session/test_types.py | 1511 ++++++++++++ .../unit/train/training_session/test_utils.py | 1745 +++++++++++++ .../training_session_example.ipynb | 536 ++++ 31 files changed, 20985 insertions(+), 1 deletion(-) create mode 100644 sagemaker-train/src/sagemaker/train/training_session/README.md create mode 100644 sagemaker-train/src/sagemaker/train/training_session/__init__.py create mode 100644 sagemaker-train/src/sagemaker/train/training_session/_clients.py create mode 100644 sagemaker-train/src/sagemaker/train/training_session/_endpoints.py create mode 100644 sagemaker-train/src/sagemaker/train/training_session/_session_handle.py create mode 100644 sagemaker-train/src/sagemaker/train/training_session/_submitter.py create mode 100644 sagemaker-train/src/sagemaker/train/training_session/_utils.py create mode 100644 sagemaker-train/src/sagemaker/train/training_session/exceptions.py create mode 100644 sagemaker-train/src/sagemaker/train/training_session/futures.py create mode 100644 sagemaker-train/src/sagemaker/train/training_session/sampling_client.py create mode 100644 sagemaker-train/src/sagemaker/train/training_session/service_client.py create mode 100644 sagemaker-train/src/sagemaker/train/training_session/training_client.py create mode 100644 sagemaker-train/src/sagemaker/train/training_session/types.py create mode 100644 sagemaker-train/tests/unit/train/training_session/__init__.py create mode 100644 sagemaker-train/tests/unit/train/training_session/session_lifecycle_grid.py create mode 100644 sagemaker-train/tests/unit/train/training_session/test_clients.py create mode 100644 sagemaker-train/tests/unit/train/training_session/test_exceptions.py create mode 100644 sagemaker-train/tests/unit/train/training_session/test_futures.py create mode 100644 sagemaker-train/tests/unit/train/training_session/test_sampling_client.py create mode 100644 sagemaker-train/tests/unit/train/training_session/test_script_compatibility.py create mode 100644 sagemaker-train/tests/unit/train/training_session/test_service_client.py create mode 100644 sagemaker-train/tests/unit/train/training_session/test_session_lifecycle.py create mode 100644 sagemaker-train/tests/unit/train/training_session/test_session_lifecycle_grid.py create mode 100644 sagemaker-train/tests/unit/train/training_session/test_submitter.py create mode 100644 sagemaker-train/tests/unit/train/training_session/test_training_client.py create mode 100644 sagemaker-train/tests/unit/train/training_session/test_types.py create mode 100644 sagemaker-train/tests/unit/train/training_session/test_utils.py create mode 100644 v3-examples/training-examples/training_session_example.ipynb diff --git a/docs/api/sagemaker_train.rst b/docs/api/sagemaker_train.rst index 44cdea2053..5c6e55b71d 100644 --- a/docs/api/sagemaker_train.rst +++ b/docs/api/sagemaker_train.rst @@ -39,6 +39,44 @@ Multi-Turn RL Training :undoc-members: :show-inheritance: +Interactive Training Sessions +----------------------------- + +.. automodule:: sagemaker.train.training_session + :members: + :undoc-members: + :show-inheritance: + +.. automodule:: sagemaker.train.training_session.service_client + :members: + :undoc-members: + :show-inheritance: + +.. automodule:: sagemaker.train.training_session.training_client + :members: + :undoc-members: + :show-inheritance: + +.. automodule:: sagemaker.train.training_session.sampling_client + :members: + :undoc-members: + :show-inheritance: + +.. automodule:: sagemaker.train.training_session.types + :members: + :undoc-members: + :show-inheritance: + +.. automodule:: sagemaker.train.training_session.futures + :members: + :undoc-members: + :show-inheritance: + +.. automodule:: sagemaker.train.training_session.exceptions + :members: + :undoc-members: + :show-inheritance: + Model Evaluation ---------------- diff --git a/sagemaker-train/pyproject.toml b/sagemaker-train/pyproject.toml index 3382363588..1fa3c095c7 100644 --- a/sagemaker-train/pyproject.toml +++ b/sagemaker-train/pyproject.toml @@ -53,6 +53,7 @@ Homepage = "https://github.com/aws/sagemaker-python-sdk" [project.optional-dependencies] rft-strands = ["strands-agents"] rft-langchain = ["httpx>=0.24.0"] +tokenizer = ["transformers"] test = [ "pytest", "pytest-cov", @@ -63,7 +64,8 @@ test = [ "scipy", "omegaconf", "graphene", - "IPython" + "IPython", + "transformers", ] notebook = [ "ipywidgets>=8.0.0", diff --git a/sagemaker-train/src/sagemaker/train/__init__.py b/sagemaker-train/src/sagemaker/train/__init__.py index 57e947b73d..e32812f15e 100644 --- a/sagemaker-train/src/sagemaker/train/__init__.py +++ b/sagemaker-train/src/sagemaker/train/__init__.py @@ -158,4 +158,16 @@ def __getattr__(name): from sagemaker.train.common_utils.finetune_utils import list_hyperparameters return list_hyperparameters + elif name == "ServiceClient": + from sagemaker.train.training_session import ServiceClient + + return ServiceClient + elif name == "TrainingClient": + from sagemaker.train.training_session import TrainingClient + + return TrainingClient + elif name == "SamplingClient": + from sagemaker.train.training_session import SamplingClient + + return SamplingClient raise AttributeError(f"module '{__name__}' has no attribute '{name}'") diff --git a/sagemaker-train/src/sagemaker/train/training_session/README.md b/sagemaker-train/src/sagemaker/train/training_session/README.md new file mode 100644 index 0000000000..0068ca02c7 --- /dev/null +++ b/sagemaker-train/src/sagemaker/train/training_session/README.md @@ -0,0 +1,995 @@ +# sagemaker.train.training_session + +Client-driven interactive training sessions on Amazon SageMaker. A +`ServiceClient` creates a session and hands back a `TrainingClient` that exposes +low-level training primitives (`forward_backward`, `optim_step`, `forward`, +`save_state`, `load_state`, `save_weights_for_sampler`) you orchestrate in a +custom training loop, with automatic operation pipelining via sequence numbers. + +The client surface: + +| Client | Role | +|-------------------|-----------------------------------------------------| +| `ServiceClient` | Entry point / factory | +| `TrainingClient` | Drives the training loop | +| `SamplingClient` | Sample from a saved checkpoint | + +## Requirements + +- **Python >= 3.10** +- **AWS credentials** configured (e.g., via `aws configure` or an IAM role) + +## Installation + +```bash +pip install sagemaker-train +``` + +What a training session needs is a botocore whose service description carries +the training-session operations, and that is checked directly: a botocore +without them makes creating a client raise, naming what is missing. + +The built-in tokenizer (`get_tokenizer()` and string prompts/batches) needs +`transformers`, which ships as an optional extra rather than a core dependency. +Install it only if you pass text: + +```bash +pip install sagemaker-train[tokenizer] +``` + +Callers who tokenize client-side can pass token IDs directly and skip the extra. + +## Example Notebook + +A runnable, end-to-end walkthrough (create, train, save/load state, sample, +delete) is in +[`v3-examples/training-examples/training_session_example.ipynb`](../../../../../v3-examples/training-examples/training_session_example.ipynb). + +## Quick Start + +`role_arn` is the only required argument. Everything else has a default, listed +under [ServiceClient](#serviceclient). + +```python +from sagemaker.train.training_session import ( + AdamParams, + Datum, + ModelInput, + ServiceClient, + TensorData, + TensorDtype, +) + +svc = ServiceClient( + region="us-west-2", + role_arn="arn:aws:iam::123456789012:role/SageMakerRole", +) +print(svc.list_base_models()) # base models you can fine-tune in this region +tc = svc.create_lora_training_client( + base_model="", # resolved to a Hub ARN in your region +) + +for tokens in your_token_id_batches: + # Every datum carries the loss inputs its loss function needs, one value per + # token; see Loss function inputs below. + batch = [ + Datum( + model_input=ModelInput.from_ints(tokens), + loss_fn_inputs={ + "target_tokens": TensorData( + data=tokens[1:] + [0], dtype=TensorDtype.INT64 + ), + "weights": TensorData(data=[1.0] * len(tokens)), + }, + ) + ] + + # Each call returns an APIFuture right away, so optim_step is queued behind + # the forward_backward without waiting for it to land. Resolve both once + # they are submitted. + fwd_bwd_op = tc.forward_backward(batch, loss_fn="cross_entropy") + optim_op = tc.optim_step(AdamParams(learning_rate=1e-4)) + + print(fwd_bwd_op.result().metrics, optim_op.result().metrics) + +print(tc.save_state().result().model_package_arn) +``` + +See [Loss function inputs](#loss-function-inputs) for which keys each loss +function needs and how to build them for reinforcement learning. + +Set the account-level plumbing (model package group, S3 output, session name +prefix, and the optional `kms_key_arn` / `vpc_config`) once on the +`ServiceClient` and every client it creates inherits it. Each is also accepted +per session on `create_lora_training_client` when one run needs something +different. + +```python +svc = ServiceClient( + region="us-west-2", + role_arn="arn:aws:iam::123456789012:role/SageMakerRole", + model_package_group_arn="arn:aws:sagemaker:us-west-2:123456789012:model-package-group/my-models", + s3_output_path="s3://my-bucket/output/", + session_name_prefix="my-experiment", # sessions get named "my-experiment-<8 hex>" +) +``` + +## Prerequisites + +### 1. Create a Model Package Group + +```bash +aws sagemaker create-model-package-group \ + --model-package-group-name "" \ + --region us-west-2 +``` + +### 2. Create an S3 bucket for outputs + +```bash +aws s3 mb s3:// --region us-west-2 +``` + +### 3. Create an execution role + +SageMaker assumes this role to write checkpoints to your bucket and register them +in your model package group. You pass its ARN to `ServiceClient`. + +Attach the `AmazonSageMakerTrainingSessionFullAccess` managed policy: + +```bash +aws iam attach-role-policy \ + --role-name "" \ + --policy-arn arn:aws:iam::aws:policy/AmazonSageMakerTrainingSessionFullAccess +``` + +Trust policy: + +```json +{ + "Version": "2012-10-17", + "Statement": [ + { + "Effect": "Allow", + "Principal": { "Service": "training-session.sagemaker.amazonaws.com" }, + "Action": ["sts:AssumeRole", "sts:TagSession"] + } + ] +} +``` + +### 4. Grant your own identity permission to call the SDK + +Attach the `AmazonSageMakerTrainingSessionAccess` managed policy to whichever +identity runs your training script. It covers the session and operation calls the +SDK makes, and the `iam:PassRole` that SageMaker checks when you create a +session with the role from step 3. + +```bash +aws iam attach-role-policy \ + --role-name "" \ + --policy-arn arn:aws:iam::aws:policy/AmazonSageMakerTrainingSessionAccess +``` + +Its S3 permissions cover buckets whose names begin with `sagemaker-`, so name the +bucket from step 2 accordingly if you call +`delete_model_package(delete_artifacts=True)`, which removes a checkpoint's +objects from your bucket. + +## API Reference + +### ServiceClient + +The entry point. Holds the control-plane and runtime clients. + +`role_arn` is required: it is the execution role the session assumes. Session +defaults set here (`model_package_group_arn`, `s3_output_path`, +`session_name_prefix`, `kms_key_arn`, `vpc_config`) apply to every client this +`ServiceClient` creates, and each can be overridden per session on +`create_lora_training_client`. + +| argument | required | default when omitted | +|---|---|---| +| `role_arn` | yes | n/a | +| `region` | no | `AWS_REGION`, then the standard boto3 chain; an error if neither names one | +| `model_package_group_arn` | no | the account's `sagemaker-default` group, created on first use | +| `s3_output_path` | no | `s3://{SageMaker default bucket}/{session name}/output` | +| `session_name_prefix` | no | `sagemaker-training-session` | +| `kms_key_arn` | no | Amazon S3 managed encryption | +| `vpc_config` | no | checkpoint transfers do not traverse your VPC | + +```python +svc = ServiceClient( + region="us-west-2", + role_arn="arn:aws:iam::123456789012:role/SageMakerRole", + model_package_group_arn="arn:aws:sagemaker:us-west-2:123456789012:model-package-group/my-models", + s3_output_path="s3://my-bucket/output/", + session_name_prefix="my-experiment", +) +``` + +#### create_lora_training_client + +Create a `TrainingSession` resource and return a `TrainingClient`. Blocks until +the session reaches `Running`. + +`base_model` is the only required argument and takes either a model id from +`list_base_models()` or a full Hub model ARN. An id resolves to the latest Hub +ARN in your region; pass the ARN to pin a specific model version. + +```python +tc = svc.create_lora_training_client( + base_model="", +) +``` + +| argument | required | default when omitted | +|---|---|---| +| `base_model` | yes | n/a | +| `training_session_name` | no | `{session_name_prefix}-{8 hex}` | +| `rank` | no | `32` | +| `alpha` | no | `64` | +| `model_package_group_arn`, `s3_output_path`, `kms_key_arn`, `vpc_config` | no | the `ServiceClient`'s | + +```python +tc = svc.create_lora_training_client( + base_model="", + training_session_name="my-session", # optional; names it yourself instead + rank=8, + alpha=16, + s3_output_path="s3://another-bucket/run-7/", # optional; overrides the ServiceClient's +) +``` + +The `*_async` twin is a coroutine, for creating the session from async code: + +```python +tc = await svc.create_lora_training_client_async( + base_model="", rank=8, alpha=16 +) +``` + +#### Encrypting outputs with your own KMS key + +Pass `kms_key_arn` to encrypt the session's outputs (saved weights and +checkpoints) in Amazon S3 with your key instead of Amazon S3 managed +encryption. Set it on `ServiceClient` as a default for every session, or on +`create_lora_training_client` per session. The execution role needs +`kms:DescribeKey`, `kms:GenerateDataKey`, and `kms:Decrypt` on the key. + +```python +svc = ServiceClient( + region="us-west-2", + role_arn="...", + kms_key_arn="arn:aws:kms:us-west-2:123456789012:key/1234abcd-...", +) +``` + +#### Routing checkpoint transfers through your VPC + +Pass `vpc_config` (a `VpcConfig` or a dict with `subnets` and +`security_group_ids`) to have checkpoint transfers between the service and your +account go through your VPC, so that traffic reaches your S3 bucket over your +own network path rather than the public internet. It applies to that transfer +only; the training and sampling operations themselves are unaffected. + +Make sure the subnets can reach Amazon S3 (for example via an S3 gateway +endpoint on the subnets' route tables) and the security groups allow outbound +HTTPS, or checkpoint saves fail. The execution role needs the standard Amazon +EC2 networking permissions (`ec2:DescribeVpcs`, `ec2:DescribeSubnets`, +`ec2:DescribeSecurityGroups`, `ec2:DescribeDhcpOptions`, and the +`ec2:*NetworkInterface*` actions). + +```python +from sagemaker.train.training_session import VpcConfig + +svc = ServiceClient( + region="us-west-2", + role_arn="...", + vpc_config=VpcConfig( + subnets=["subnet-0123456789abcdef0", "subnet-0fedcba9876543210"], + security_group_ids=["sg-0123456789abcdef0"], + ), +) +``` + +These settings also apply when `create_sampling_client` creates a training +session on your behalf. + +#### create_training_client_from_state / create_training_client_from_state_with_optimizer + +Resume from a `save_state` checkpoint in a **new** session: the one-call form of +`create_lora_training_client(...)` followed by `load_state(...)`. As with the +`load_state` pair, the method name decides what comes back from the checkpoint: + +| factory | restores | optimizer | +|---|---|---| +| `create_training_client_from_state(...)` | weights only | resets | +| `create_training_client_from_state_with_optimizer(...)` | weights + optimizer | continues where it left off | + +Both read the same checkpoint. Any other keyword argument is forwarded to +`create_lora_training_client`, so `rank`, `s3_output_path`, `vpc_config` and the +rest work here too. Omit `base_model` and it is read from the checkpoint's model +package metadata, pinned to the version the weights were trained on. + +Both return only once the restore has landed, so the client you get back is ready +to train. If the restore fails, its session is stopped and the error propagates. + +```python +from sagemaker.train.training_session import ModelPackageKind + +ckpt = svc.get_latest_model_package( + ModelPackageKind.CHECKPOINT, + model_package_group_arn="arn:aws:sagemaker:us-west-2:123456789012:model-package-group/my-models", +) + +# Exact resume, optimizer momentum included +tc = svc.create_training_client_from_state_with_optimizer(ckpt.model_package_arn) + +# Weights only, e.g. to start a new schedule from trained weights +tc = svc.create_training_client_from_state(ckpt.model_package_arn, rank=8) +``` + +Both have `*_async` coroutine twins with the same contract, awaitable without +holding up the event loop: + +```python +tc = await svc.create_training_client_from_state_with_optimizer_async( + ckpt.model_package_arn +) +``` + +#### list_base_models + +Query the SageMaker Public Hub and return the base models available for +fine-tuning in the client's region (models are region-specific). Each is named by +its own model id, spelled as its provider publishes it wherever the Hub carries +one, so an id normally reads the same here as it does anywhere else the model is +named. Pass one as `base_model=`. + +```python +svc.list_base_models() +``` + +The set changes as models are added, so call it rather than hard-coding an id. + +#### create_sampling_client + +Get a `SamplingClient` without running a training loop, which is mainly +evaluation work: scoring a saved checkpoint after training finishes, comparing +two checkpoints or a checkpoint against the base model, or reranking with +`compute_logprobs`. The SDK creates a training session on your behalf to serve +the requests. Specify exactly one of `base_model` (sample the base model) or +`model_package_arn` (sample saved weights; the base model is resolved from the +package metadata). + +Because the SDK created that session for this client, **the client owns it**: +use it as a context manager, or call `sc.stop()`, so the session is released at +a point you choose. + +```python +with svc.create_sampling_client(model_package_arn=ckpt_arn) as sc: + out = sc.sample([1, 2, 3]).result() +# session stopped here + +with svc.create_sampling_client(base_model="") as base_sc: + ... +``` + +If you neither scope nor stop it, the SDK stops the session for you once the +last client referencing it becomes unreachable. See +[Session lifecycle](#session-lifecycle). + +> **Note:** To sample from the client that is actively training, use +> `TrainingClient.create_sampling_client(...)` / +> `TrainingClient.save_weights_and_get_sampling_client(...)` instead. Those +> share the `TrainingClient`'s session, so they do **not** own it: leaving a +> `with` block on one of those clients drains its queued submits but leaves the +> training session running. + +#### stop_session + +Stop a running session, transitioning it toward `Stopped`. + +```python +svc.stop_session("my-session") +``` + +#### delete_session + +Delete a session. The service only deletes sessions in a terminal state +(`Stopped`, `Failed`, or `DeletionFailed`) and removal is asynchronous (the +session flips to `Deleting` and a background worker finishes the cleanup). Pass +`stop_if_active=True` to stop a still-running session and wait for a terminal +state before deleting. + +```python +svc.delete_session("my-session") # session already terminal +svc.delete_session("my-session", stop_if_active=True) # stop, wait, then delete +``` + +#### delete_model_package + +Delete a saved checkpoint or set of weights. The model package carries the only +reference to the S3 location its weights were written to, so by default this +deletes the package and leaves the objects, which keep costing storage with +nothing pointing at them. Keep the returned `s3_uri` if you plan to clean up +later. + +```python +result = svc.delete_model_package(ckpt.model_package_arn) +print("weights still at", result.s3_uri) +``` + +Pass `delete_artifacts=True` to delete the objects too. They go first and the +package second, so a failure part-way leaves the package still pointing at +whatever survived and the call can be retried. That case raises +`ArtifactDeletionError` with the surviving keys on `.failed_keys`. + +```python +result = svc.delete_model_package(ckpt.model_package_arn, delete_artifacts=True) +print("removed", result.objects_deleted, "objects") +``` + +> **Note:** `create_sampling_client(model_package_arn=...)` and +> `load_state(model_package_arn=...)` read a package's artifacts, so once they are +> deleted do not rely on this package for either. + +Deleting artifacts signs with your own credentials, not the session's `role_arn`, +so you need `s3:ListBucket` on the bucket and `s3:DeleteObject` on the objects +under it. Where SageMaker stores a package's weights on your behalf rather than +in your own bucket there is nothing for you to delete: the call reports them +retained, and deleting the package releases them. + +The call is idempotent. A package that is already gone comes back with +`model_package_deleted=False` rather than raising. To sweep a group: + +```python +for mp in svc.list_model_packages(max_results=None): + svc.delete_model_package(mp.model_package_arn, delete_artifacts=True) +``` + +Deleting every version leaves the (empty) model package group in place; removing +the group is a separate SageMaker call. + +#### list_sessions + +List training sessions, auto-paginating. Optionally filter by name substring or +status, and cap the total with `max_results`. + +```python +svc.list_sessions() # all sessions +svc.list_sessions(status_equals="Running") # only running +svc.list_sessions(name_contains="exp", max_results=10) +``` + +Each entry is a `TrainingSessionSummary` with `training_session_arn`, `status`, +`creation_time` and `last_modified_time`, plus a `training_session_name` read off +the ARN. `status` is a `SessionStatus`, so it compares against the enum rather +than a raw string. `stop_session` and `delete_session` accept either a name or an +ARN, so either attribute can be passed straight through. + +Sweeping up every session still running, e.g. after an interrupted experiment: + +```python +from sagemaker.train.training_session import SessionStatus + +for summary in svc.list_sessions(status_equals=SessionStatus.RUNNING): + print("stopping", summary.training_session_name, summary.creation_time) + svc.stop_session(summary.training_session_arn) +``` + +Pass `name_contains=` to narrow the sweep to one experiment's sessions. + +### TrainingClient + +Drives the training loop over a session. + +Every operation has two forms, and both return an `APIFuture`: + +| form | returns | resolve it with | +|---|---|---| +| `tc.forward_backward(...)` | `APIFuture` | `.result()`, or `await` the future | +| `await tc.forward_backward_async(...)` | `APIFuture` | `await` the future | + +The base method queues its submit on a background thread and hands back the +future immediately; a submit error surfaces when the future resolves, not at the +call site. The `*_async` twin is a coroutine, so the enqueue yields to the event +loop instead of running on the calling thread, and resolves to the same future. + +Training operations carry sequence numbers and must stay ordered, so the enqueue +and the result are separate steps: the future is your handle on a submit that +already holds its place in the sequence. + +```python +result = tc.forward_backward(batch).result() # sync code +result = await tc.forward_backward(batch) # async code, base method + +future = await tc.forward_backward_async(batch) # async code, coroutine twin +result = await future +``` + +The examples below mostly use the base form. + +#### Bounding how long you wait + +`result()` waits indefinitely by default, because a training operation that +succeeds should not fail just because it ran past an arbitrary clock. Pass +`timeout` (seconds) to bound a specific wait. When set, it bounds the *whole* +wait, including the submit round trip and every chunk of a batch large enough to +be split across several operations, so `result(timeout=60)` is a ~60-second +call, not 60 seconds per stage. (It bounds when a new wait may *start*, so a +status check already in flight can overrun it by that request's latency.) + +`await future` takes no argument, so it uses the default; `result_async` is the +await form that accepts one. + +```python +result = tc.forward_backward(batch).result(timeout=300) + +future = tc.forward_backward(batch) +result = await future # waits indefinitely +result = await future.result_async(timeout=300) # bounded +``` + +A `TimeoutError` means the wait ran out, not that the operation was cancelled. +It keeps running server-side, and resolving the same future again picks it back +up. + +#### Loss function inputs + +Every datum needs loss inputs, and which keys are required depends on the loss: + +| `loss_fn` | Required keys | +|---|---| +| `cross_entropy` | `target_tokens`, `weights` | +| `ppo` / `importance_sampling` / `cispo` | `target_tokens`, `weights`, `logprobs`, `advantages` | + +Each is **full-sequence**: one value per token of the datum, prompt included. +`weights` is what selects the tokens the loss applies to: `0.0` over the prompt, +`1.0` over the response. Pass them as `TensorData` (which also accepts a torch +tensor or numpy array) so they carry a dtype. + +For reinforcement learning, build them per datum: `advantages` differ from one +datum to the next, so each `Datum` carries its own. This is the form a training +loop usually yields. + +```python +from sagemaker.train.training_session import Datum, ModelInput, TensorData + +def make_datum(tokens, prompt_len, logprobs, advantages): + """One datum whose loss inputs are full-sequence.""" + return Datum( + model_input=ModelInput.from_ints(tokens), + loss_fn_inputs={ + "target_tokens": TensorData(data=tokens[1:] + [0], dtype="int64"), + "weights": TensorData(data=[0.0] * prompt_len + [1.0] * (len(tokens) - prompt_len)), + "logprobs": TensorData(data=logprobs), # from sampling, one per token + "advantages": TensorData(data=advantages), # one per token + }, + ) + +result = tc.forward_backward(batch=[make_datum(...)], loss_fn="ppo").result() +``` + +The `loss_fn_inputs` argument on `forward_backward` / `forward` is the shortcut +for the other case: one set of values shared by the whole batch. It is filled in +only for datums that do not already carry their own, so the two forms mix freely. + +```python +tokens, prompt_len = [1, 2, 3, 4, 5, 6], 3 + +result = tc.forward_backward( + batch=[tokens], + loss_fn="cross_entropy", + loss_fn_inputs={ + "target_tokens": TensorData(data=tokens[1:] + [0], dtype="int64"), + "weights": TensorData(data=[0.0] * prompt_len + [1.0] * (len(tokens) - prompt_len)), + }, +).result() +``` + +Every datum must end up with the inputs its loss function requires, and two +mistakes are caught before the request leaves. A datum missing any of them +raises, naming the datum, the loss function and the keys it needs. Giving +response-length values instead of full-sequence ones raises too, naming the +datum and the key. + +#### forward_backward + +Run a forward-backward pass. Accepts token ID lists, strings (with a tokenizer), +or raw datum dicts. + +```python +from sagemaker.train.training_session import TensorData, TensorDtype + +tokens, prompt_len = [1, 2, 3, 4, 5, 6], 3 +response_len = len(tokens) - prompt_len + +result = tc.forward_backward( + batch=[tokens], + loss_fn="importance_sampling", # or "ppo", "cispo", "cross_entropy" + loss_fn_inputs={ + "target_tokens": TensorData(data=tokens[1:] + [0], dtype=TensorDtype.INT64), + "weights": TensorData(data=[0.0] * prompt_len + [1.0] * response_len), + "logprobs": TensorData(data=sampled_logprobs), # from the sampling step + "advantages": TensorData(data=[reward] * len(tokens)), + }, +).result() +print(result.metrics) # {"actor/pg_loss": ..., "global_token_num": ...} + +# Loss-function tensor outputs, one entry per datum: +for datum_outputs in result.loss_fn_outputs: + print(datum_outputs) # {"actor/pg_loss": ..., "actor/grad_norm": ...} +``` + +#### forward + +Run a forward-only pass (no gradient accumulation), useful for computing +loss-function outputs or logprobs without updating the model. Takes the same +arguments as `forward_backward`. + +```python +result = tc.forward( + batch=[tokens], + loss_fn="cross_entropy", + loss_fn_inputs={ + "target_tokens": TensorData(data=tokens[1:] + [0], dtype="int64"), + "weights": TensorData(data=[0.0] * prompt_len + [1.0] * (len(tokens) - prompt_len)), + }, +).result() +for datum_outputs in result.loss_fn_outputs: + print(datum_outputs) +``` + +#### optim_step + +Run an optimizer step. + +```python +result = tc.optim_step(AdamParams(learning_rate=1e-4)).result() +``` + +#### save_weights_for_sampler + +Save model weights for inference, registering a **model** package. Pass +`merge_weights=True` to merge the adapter into the base weights, so the result +can be deployed without the adapter. + +```python +print(tc.save_weights_for_sampler().result().model_package_arn) + +# or, in async code +future = await tc.save_weights_for_sampler_async(merge_weights=True) +print((await future).model_package_arn) +``` + +To save and immediately get a client bound to the result, use +`save_weights_and_get_sampling_client()` instead. It returns a `SamplingClient`, +not a future. + +#### save_state + +Save a resumable checkpoint (weights + optimizer state). Always creates a model +package that can be passed to `load_state` to resume training. + +```python +model_package_arn = tc.save_state().result().model_package_arn +print(model_package_arn) +``` + +#### load_state / load_state_with_optimizer + +Load a previously saved `save_state` checkpoint to resume training. Two peer +methods, no flag: the method name decides what is restored from the checkpoint. + +| method | restores | optimizer | +|---|---|---| +| `load_state(arn)` | weights only | resets | +| `load_state_with_optimizer(arn)` | weights + optimizer | continues where it left off | + +Both read the same checkpoint (which always contains both). + +```python +# Weights only, optimizer resets +tc.load_state(model_package_arn=model_package_arn).result() + +# Weights + optimizer, exact resume +tc.load_state_with_optimizer(model_package_arn=model_package_arn).result() +``` + +#### Model package discovery + +Every `save_state` / `save_weights_for_sampler` registers a **model package** +version in the session's group. A `save_state` one is a **checkpoint** (weights + +optimizer state, for `load_state`); a `save_weights_for_sampler` one is a +deployable **model** (weights only, for inference or a new run). +`ModelPackageKind` names that split, and one accessor takes it as an argument. + +`tc.get_latest_model_package(kind)` answers from an in-memory cache of what +**this client** saved in this process, so there is no API call and no exposure to +the brief lag before a freshly-saved package is describable. It returns `None` +when this client hasn't saved that kind (e.g. a fresh or resumed session): a +client deliberately does not imply a package from a session it did not drive. To +discover packages a *different* client saved, use the `ServiceClient` twin, which +queries the group: `svc.get_latest_model_package(kind, model_package_group_arn=...)`. +`tc.list_model_packages(...)` forwards to the `ServiceClient` with this session's +group and returns full metadata. + +```python +from sagemaker.train.training_session import ModelPackageKind + +tc.save_state().result() +ckpt = tc.get_latest_model_package(ModelPackageKind.CHECKPOINT) # cache; resume with load_state* +tc.load_state_with_optimizer(model_package_arn=ckpt.model_package_arn).result() + +tc.save_weights_for_sampler().result() +model = tc.get_latest_model_package(ModelPackageKind.MODEL) # cache; for inference / sampling + +tc.list_model_packages(max_results=20) # full metadata, newest first + +# Across processes, query the group instead of a client's cache +ckpt = svc.get_latest_model_package( + ModelPackageKind.CHECKPOINT, + model_package_group_arn="arn:aws:sagemaker:us-west-2:123456789012:model-package-group/my-models", +) +``` + +To remove one, see +[`ServiceClient.delete_model_package`](#delete_model_package). A model package +outlives the session that saved it, so nothing cleans it up when the session +goes. + +#### stop / delete + +Stop or delete this client's session without repeating its name. + +```python +tc.stop() # stop this session +tc.delete(stop_if_active=True) # stop, wait for terminal, then delete +``` + +### Session lifecycle + +A training session is live compute, so the SDK does not leave one `Running` +just because you forgot to stop it. Once the last client referencing a session +becomes unreachable, the SDK stops that session and emits a `ResourceWarning`: + +```python +tc = svc.create_lora_training_client(base_model="") +... +tc = None # or the script ends, or the notebook cell is re-run + # -> session stopped +``` + +The automatic teardown **stops** the session. A stopped session is still +describable, so its history stays available to audit and a run that went wrong +can still be inspected; deleting one is an explicit `tc.delete(...)` / +`svc.delete_session(...)` call. + +A session can be shared by more than one client, a `TrainingClient` and any +`SamplingClient` derived from it, and it is the **last** of them to go away that +triggers the stop, not any particular one. So dropping one client while another +is still in use leaves the session running: + +```python +tc = svc.create_lora_training_client(base_model="") +sp = tc.create_sampling_client() # shares tc's session; 2 clients now hold it + +del tc # 1 left (sp) -> session stays Running +sp.sample([1, 2, 3]).result() # still works + +del sp # 0 left -> session stops here +``` + +Call `stop()` to release the session at a moment you choose. It flushes queued +operations, then stops without waiting for them to finish, so call it once you +have the results you need. + +Save a checkpoint for anything you want to keep: model packages are durable on +their own, so trained weights outlive the process and a fresh session can be +created from one whenever you need it. + +```python +tc = svc.create_lora_training_client(base_model="") +... # train +mp = tc.save_weights_for_sampler().result().model_package_arn +# later, in any other process: +sp = svc.create_sampling_client(model_package_arn=mp) +``` + +### SamplingClient + +Sample from a saved checkpoint (or the base model). Get one from a +`TrainingClient` via `save_weights_and_get_sampling_client(...)` (saves weights, +then binds to the resulting checkpoint) or `create_sampling_client(...)`, or +without a `TrainingClient` from `ServiceClient.create_sampling_client(...)`. + +Only the `ServiceClient` factory creates the backing session, so only its client +stops that session on `with`-block exit or `stop()`. Clients handed out by a +`TrainingClient` share the training session and leave it running, though they do +keep it alive; see [Session lifecycle](#session-lifecycle). + +Like every method on the training client, `sample` and `compute_logprobs` queue +their submit on a background thread and return an `APIFuture` immediately. Call +`.result()` on it, or `await` it, to get the parsed response. Sample submits are +unordered, so many requests can be queued concurrently and resolved in any +order. + +Because there is no sequence to hold a place in, the sampling coroutines return +the response itself rather than a future: + +| form | returns | +|---|---| +| `sc.sample(...)` | `APIFuture` resolving to `SampleResponse` | +| `await sc.sample_async(...)` | `SampleResponse` | +| `sc.compute_logprobs(...)` | `APIFuture` resolving to `List[Optional[float]]` | +| `await sc.compute_logprobs_async(...)` | `List[Optional[float]]` | + +```python +from sagemaker.train.training_session import SamplingParams + +sc = tc.save_weights_and_get_sampling_client() +future = sc.sample( + [1, 2, 3], # token IDs, a string (with a tokenizer), or a ModelInput + num_samples=4, + sampling_params=SamplingParams(max_tokens=1024), +) +result = future.result() +for seq in result.sequences: + print(seq.tokens, seq.stop_reason) +``` + +Queue a batch of prompts and gather them as they land; hundreds can be in flight +at once: + +```python +futures = [sc.sample(p, num_samples=4) for p in prompts] # all queued up front +for f in futures: + print(f.result().sequences[0].tokens) + +# async: await the futures, or call the coroutine twin +results = await asyncio.gather(*(sc.sample(p, num_samples=4) for p in prompts)) +one = await sc.sample_async(prompts[0], num_samples=4) # -> SampleResponse +``` + +#### compute_logprobs + +Score a prompt without generating: returns one log probability per prompt token +(useful for evaluation, reranking, or importance weighting). + +```python +logprobs = sc.compute_logprobs([1, 2, 3]).result() # or a string, with a tokenizer +print(logprobs) # one logprob per prompt token (first entry is None) + +logprobs = await sc.compute_logprobs_async([1, 2, 3]) # -> List[Optional[float]] +``` + +#### get_tokenizer + +The runtime accepts only token IDs. To pass text prompts (or string training +batches), call `get_tokenizer()` to load the base model's HuggingFace tokenizer; +the SDK then tokenizes strings for you. It is loaded once and cached. Available +on both `TrainingClient` and `SamplingClient`. + +```python +tokenizer = sc.get_tokenizer() # raises if the model has no published tokenizer +result = sc.sample("Hello, world", num_samples=1).result() +``` + +Sampled sequences come back as token IDs, so use the same tokenizer to read them. +It is an ordinary HuggingFace tokenizer, so `encode` and `decode` behave as they +do anywhere else: + +```python +tokenizer = sc.get_tokenizer() + +prompt_ids = tokenizer.encode("What is 12 x 7?") +result = sc.sample(prompt_ids, num_samples=1).result() + +print(tokenizer.decode(result.sequences[0].tokens)) +``` + +Passing the string directly, as above, tokenizes it the same way. Encode yourself +when you need the IDs for something else too, such as measuring prompt length or +building a `ModelInput` with several chunks. + +## Errors + +Service errors reach you as raised by botocore, so `botocore.exceptions.ClientError` +and its subclasses carry the error code, message, and response metadata you would +see from any other AWS client. Catch them the way you already do, and read +`error.response["Error"]["Code"]` to tell them apart: + +| error code | meaning | +|---|---| +| `ValidationException` | the request was rejected; resending it gets the same answer | +| `ResourceNotFoundException` | the session, operation, or checkpoint does not exist | +| `ConflictException` | the session is not in a state that accepts the call | +| `ThrottlingException` | too many requests; retryable | +| `InternalServerException` | the service could not serve the request; retryable | + +Session lifecycle calls go to the SageMaker control plane, which reports +`ResourceNotFound`, `ResourceInUse`, and `ResourceLimitExceeded` under those +shorter names rather than the `*Exception` spellings above. + +On top of those, three errors describe conditions no service error covers: + +| error | raised when | +|---|---| +| `OperationError` | the service reports a submitted operation as failed; carries `operation_id` | +| `ResultMismatchError` | a result does not correspond to the request it was awaited for | +| `ArtifactDeletionError` | `delete_model_package(delete_artifacts=True)` partially completed; carries `failed_keys` | + +All three subclass `TrainingSessionError`, so a single +`except TrainingSessionError` catches every error this package defines. A wait +that runs out of time is the one exception: `result(timeout=...)` raises the +built-in `TimeoutError`, the same as any other future, so catch that alongside +`TrainingSessionError` to cover both. + +`is_retryable(error)` classifies these alongside botocore's own errors, so a +retry policy does not have to match on error types: + +```python +from sagemaker.train.training_session import is_retryable + +try: + result = tc.forward_backward(batch, loss_fn="ppo").result() +except Exception as error: + if not is_retryable(error): + raise + result = tc.forward_backward(batch, loss_fn="ppo").result() +``` + +## Pipelining + +Every base method already pipelines: it returns an `APIFuture` immediately +and the SDK stamps sequence numbers at call time, so you can submit +`optim_step` before the preceding `forward_backward` resolves and both execute +in the same clock cycle. Submits ride a background sender, so queueing a +large batch of operations takes seconds, not one wire round trip per op: + +```python +for batch in data_loader: + fwd_bwd_op = tc.forward_backward(batch, loss_fn="ppo") + optim_op = tc.optim_step() + + fwd_bwd_result = fwd_bwd_op.result() + optim_result = optim_op.result() +``` + +The same loop in async code: `await` the futures instead of calling `.result()`, +so the event loop stays free while the ops run. + +```python +for batch in data_loader: + fwd_bwd_op = tc.forward_backward(batch, loss_fn="ppo") + optim_op = tc.optim_step() + + fwd_bwd_result = await fwd_bwd_op + optim_result = await optim_op +``` + +Use the `*_async` coroutine twins when the enqueue should yield to the event loop +as well; each resolves to the same `APIFuture`: + +```python +fwd_bwd_op = await tc.forward_backward_async(batch, loss_fn="ppo") +optim_op = await tc.optim_step_async() +fwd_bwd_result, optim_result = await asyncio.gather(fwd_bwd_op, optim_op) +``` + +## Concurrent Access + +A training session must only be accessed by one `TrainingClient` at a time. If +two clients reference the same session simultaneously, they issue conflicting +sequence numbers and the service rejects the later one with a +`ConflictException`. + +## Running the Example + +See +[`v3-examples/training-examples/training_session_example.ipynb`](../../../../../v3-examples/training-examples/training_session_example.ipynb) +for a runnable end-to-end walkthrough. Update the configuration cell with your +own model package group, execution role, and S3 output path before running; the +final cell deletes the session. diff --git a/sagemaker-train/src/sagemaker/train/training_session/__init__.py b/sagemaker-train/src/sagemaker/train/training_session/__init__.py new file mode 100644 index 0000000000..dc35ca0bd2 --- /dev/null +++ b/sagemaker-train/src/sagemaker/train/training_session/__init__.py @@ -0,0 +1,149 @@ +# 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. +"""SageMaker training session SDK: client surface for interactive training. + +Create a ``ServiceClient``, derive a ``TrainingClient``, run the training loop, +then derive a ``SamplingClient`` for inference:: + + from sagemaker.train.training_session import ( + AdamParams, + Datum, + ModelInput, + SamplingParams, + ServiceClient, + TensorData, + TensorDtype, + ) + + svc = ServiceClient( + region="us-west-2", + role_arn="arn:aws:iam::123456789012:role/SageMakerTrainingSessionsRole", + model_package_group_arn="arn:aws:sagemaker:us-west-2:123456789012:model-package-group/my-group", + s3_output_path="s3://my-bucket/output", + ) + tc = svc.create_lora_training_client(base_model=svc.list_base_models()[0]) + for tokens in your_token_id_batches: + # Every datum carries the loss inputs its loss function needs, one value + # per token of the datum. + batch = [ + Datum( + model_input=ModelInput.from_ints(tokens), + loss_fn_inputs={ + "target_tokens": TensorData( + data=tokens[1:] + [0], dtype=TensorDtype.INT64 + ), + "weights": TensorData(data=[1.0] * len(tokens)), + }, + ) + ] + # Each call returns an APIFuture right away, so the optim_step is queued + # behind the forward_backward without waiting for it to land. + fwd_bwd = tc.forward_backward(batch, loss_fn="cross_entropy") + optim = tc.optim_step(AdamParams(learning_rate=1e-4)) + print(fwd_bwd.result().metrics, optim.result().metrics) + sc = tc.save_weights_and_get_sampling_client() + sc.sample( + "What is the capital of France?", + num_samples=8, + sampling_params=SamplingParams(max_tokens=512), + ).result() + +From async code, ``await`` the same futures instead of calling ``.result()``:: + + fwd_bwd = tc.forward_backward(batch, loss_fn="cross_entropy") + optim = tc.optim_step(AdamParams(learning_rate=1e-4)) + fwd_bwd_result, optim_result = await asyncio.gather(fwd_bwd, optim) +""" + +from __future__ import annotations + +from sagemaker.train.training_session.futures import APIFuture, RequestStatus +from sagemaker.train.training_session.service_client import ( + DEFAULT_MODEL_PACKAGE_GROUP_NAME, + DEFAULT_SESSION_NAME_PREFIX, + ServiceClient, +) +from sagemaker.train.training_session.training_client import TrainingClient +from sagemaker.train.training_session.sampling_client import SamplingClient +from sagemaker.train.training_session.types import ( + AdamParams, + Datum, + EncodedTextChunk, + ForwardBackwardOutput, + LoadStateResponse, + LossFunction, + ModelInput, + ModelInputChunk, + ModelPackageDeletion, + ModelPackageInfo, + ModelPackageKind, + OptimStepResponse, + SampledSequence, + SampleResponse, + SamplingParams, + SaveStateResponse, + SaveWeightsResponse, + SessionStatus, + TrainingSessionSummary, + StopReason, + TensorData, + TensorDtype, + TokenLogprob, + VpcConfig, +) +from sagemaker.train.training_session.exceptions import ( + ArtifactDeletionError, + OperationError, + ResultMismatchError, + TrainingSessionError, + is_retryable, +) + +__all__ = [ + "ServiceClient", + "TrainingClient", + "SamplingClient", + "APIFuture", + "RequestStatus", + "DEFAULT_MODEL_PACKAGE_GROUP_NAME", + "DEFAULT_SESSION_NAME_PREFIX", + "AdamParams", + "Datum", + "EncodedTextChunk", + "ForwardBackwardOutput", + "LoadStateResponse", + "LossFunction", + "ModelInput", + "ModelInputChunk", + "ModelPackageDeletion", + "ModelPackageInfo", + "ModelPackageKind", + "OptimStepResponse", + "SampledSequence", + "SampleResponse", + "SamplingParams", + "SaveStateResponse", + "SaveWeightsResponse", + "SessionStatus", + "TrainingSessionSummary", + "StopReason", + "TensorData", + "TensorDtype", + "TokenLogprob", + "VpcConfig", + "ArtifactDeletionError", + "OperationError", + "ResultMismatchError", + "TrainingSessionError", + "is_retryable", +] diff --git a/sagemaker-train/src/sagemaker/train/training_session/_clients.py b/sagemaker-train/src/sagemaker/train/training_session/_clients.py new file mode 100644 index 0000000000..7c8375dc10 --- /dev/null +++ b/sagemaker-train/src/sagemaker/train/training_session/_clients.py @@ -0,0 +1,167 @@ +# 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. +"""boto3 clients for the training-session control plane and runtime. + +Both clients are ordinary boto3 clients: botocore serializes the request, +signs it, applies its own retry policy, and parses the response. Neither +endpoint is reached by hand-signed HTTP. + +The control-plane operations (``CreateTrainingSession`` and friends) are part of +the ``sagemaker`` service; the runtime is a separate service. Both models come +from the installed botocore, so a botocore that predates either surface fails +with an actionable message rather than an opaque lookup error. + +Neither client is given an endpoint to call. Both resolve it from their Region +through botocore's endpoint rules, so a configuration that changes which host a +Region answers on is honoured for free. A URL in the environment (see +:mod:`sagemaker.train.training_session._endpoints`) overrides that resolution +for the one endpoint it names. +""" + +from __future__ import annotations + +import logging +from typing import Any, Optional + +import boto3 +import botocore +from botocore.config import Config +from botocore.exceptions import UnknownServiceError + +from sagemaker.train.training_session._endpoints import ( + _endpoint_overrides, + _resolve_region, +) + +logger = logging.getLogger(__name__) + +#: botocore's name for the runtime service. +_RUNTIME_SERVICE_NAME = "sagemakertrainingsessionruntime" + +#: An operation that only a botocore carrying the training-session API has, used +#: to check the installed version without trusting its version string. +_CONTROL_PLANE_PROBE_OPERATION = "CreateTrainingSession" + +#: Control-plane retries per operation, on top of the first try, in botocore's +#: ``standard`` retry mode: throttling, 5xx and connection errors are retried +#: with exponential backoff, and everything else, ``ValidationException`` and +#: ``ResourceNotFound`` above all, is raised on the first try. Three retries +#: means four attempts in total. +_CONTROL_PLANE_RETRIES = 3 + +#: Runtime calls are sent once. Submitting an operation and polling for its +#: result are both already wrapped in retry loops that understand which failures +#: an in-flight operation can survive, so a second retry policy underneath them +#: would multiply the wait for no gain. +_RUNTIME_RETRIES = 0 + +#: Connections kept in the runtime pool. Sampling fans out far wider than the +#: botocore default of 10 (128 concurrent episodes * 8 samples = 1024 in-flight +#: polls), and a pool smaller than the fan-out serializes the polls behind it. +_DEFAULT_MAX_POOL_CONNECTIONS = 512 + +#: Seconds to wait for a connection to be established. Reaching the endpoint is +#: either quick or blocked, so a connection that has not opened by now will not, +#: and waiting out botocore's default of 60 only delays the retry that recovers. +_CONNECT_TIMEOUT = 10 + +#: Seconds to wait for a response once the request is sent. Every operation here +#: returns promptly (work is submitted and polled for, never awaited on the +#: connection), so this bounds a stalled socket rather than a slow operation. +#: Left at botocore's default, stated so it is a decision rather than inherited. +_READ_TIMEOUT = 60 + + +def _control_plane_client( + region: Optional[str] = None, + *, + boto_session: Optional[boto3.Session] = None, +) -> Any: + """Return a ``sagemaker`` client for the training-session control plane. + + Raises: + RuntimeError: The installed botocore has no training-session operations. + """ + resolved_region = _resolve_region(region, boto_session=boto_session) + endpoint, _ = _endpoint_overrides() + session = boto_session or boto3.Session() + client = session.client( + "sagemaker", + region_name=resolved_region, + endpoint_url=endpoint, # None leaves the endpoint to botocore's rules. + config=Config( + retries={"mode": "standard", "max_attempts": _CONTROL_PLANE_RETRIES}, + connect_timeout=_CONNECT_TIMEOUT, + read_timeout=_READ_TIMEOUT, + ), + ) + _require_control_plane_operations(client) + return client + + +def _require_control_plane_operations(client: Any) -> None: + """Fail with an actionable message when botocore predates the API. + + Asks the loaded model what it supports, not what its version string claims, + which answers the question directly. + """ + if _CONTROL_PLANE_PROBE_OPERATION in client.meta.service_model.operation_names: + return + raise RuntimeError( + "The installed botocore does not describe the SageMaker training-session " + f"operations (botocore {botocore.__version__}), so there is nothing for " + "this client to call. Install a botocore that carries them, then create " + "the client again." + ) + + +def _runtime_client( + region: Optional[str] = None, + *, + boto_session: Optional[boto3.Session] = None, + max_pool_connections: int = _DEFAULT_MAX_POOL_CONNECTIONS, +) -> Any: + """Return a client for the training-session runtime. + + Args: + region: AWS region; resolved from the environment when not supplied. + boto_session: Session whose credentials sign the calls. Defaults to a + fresh session using the standard credential chain. + max_pool_connections: Connections kept in the HTTP pool. Raise it for + sampling workloads that fan out past the default. + + Raises: + RuntimeError: The installed botocore does not model the runtime. + """ + resolved_region = _resolve_region(region, boto_session=boto_session) + _, endpoint = _endpoint_overrides() + session = boto_session or boto3.Session() + try: + return session.client( + _RUNTIME_SERVICE_NAME, + region_name=resolved_region, + endpoint_url=endpoint, # None leaves the endpoint to botocore's rules. + config=Config( + retries={"mode": "standard", "max_attempts": _RUNTIME_RETRIES}, + max_pool_connections=max_pool_connections, + connect_timeout=_CONNECT_TIMEOUT, + read_timeout=_READ_TIMEOUT, + ), + ) + except UnknownServiceError as error: + raise RuntimeError( + "The installed botocore does not describe the SageMaker training-session " + f"runtime (botocore {botocore.__version__}), so there is nothing for " + "this client to call. Install a botocore that carries it, then create " + "the client again." + ) from error diff --git a/sagemaker-train/src/sagemaker/train/training_session/_endpoints.py b/sagemaker-train/src/sagemaker/train/training_session/_endpoints.py new file mode 100644 index 0000000000..ee382ae036 --- /dev/null +++ b/sagemaker-train/src/sagemaker/train/training_session/_endpoints.py @@ -0,0 +1,119 @@ +# 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. +"""Region and endpoint-override resolution for the training-session clients. + +Neither endpoint is computed here. Both clients built in +:mod:`sagemaker.train.training_session._clients` resolve their endpoint the way +every other boto3 client does, from the region through botocore's endpoint +rules, so a configuration that changes which host a Region answers on is +honoured without this package having to know about it. + +To send calls somewhere else, an interface VPC endpoint for example, set the full +URL through an environment variable, the env-driven form of boto3's own +``endpoint_url``: + + SAGEMAKER_TRAINING_SESSION_ENDPOINT full endpoint URL + SAGEMAKER_TRAINING_SESSION_RUNTIME_ENDPOINT full runtime endpoint URL + +The two resolve independently, so a call is sent to a given URL only when the +variable naming that endpoint is set. +""" + +from __future__ import annotations + +import logging +import os + +import boto3 + +logger = logging.getLogger(__name__) + +# Endpoint overrides, read from the environment and not taken as client +# parameters. +_SAGEMAKER_ENDPOINT_ENV = "SAGEMAKER_TRAINING_SESSION_ENDPOINT" +_RUNTIME_ENDPOINT_ENV = "SAGEMAKER_TRAINING_SESSION_RUNTIME_ENDPOINT" + + +def _resolve_region( + region: str | None = None, + *, + boto_session: boto3.Session | None = None, +) -> str: + """Resolve the AWS Region to run in. + + In order: the Region named outright, then the Region of a session the + caller supplied, then AWS_REGION, then the standard chain. + + ``Session.region_name`` covers AWS_DEFAULT_REGION and the shared config + file, the same chain sagemaker-core resolves through, but does not read + AWS_REGION, so that one is checked explicitly. + + A supplied session answers on its own, so its clients and these ones agree + on where they are pointed even when the environment names somewhere else. + + Args: + region: Region to use as given, resolved no further. + boto_session: Session whose Region to prefer. ``None`` means the + caller named no session, and the standard chain answers instead. + + Raises: + ValueError: No region could be determined. Raised rather than defaulted, + so an unconfigured environment cannot quietly create sessions, + and bill for their GPUs, in a Region the caller never named. + """ + if boto_session is not None: + # The session answers on its own. Falling through to the standard chain + # would take the Region from a different profile's configuration while + # the credentials came from this one. + resolved = region or boto_session.region_name + else: + resolved = region or os.environ.get("AWS_REGION") or boto3.Session().region_name + if not resolved: + raise ValueError( + "No AWS Region configured. Pass region= to the client, or set up a " + "local AWS configuration with a Region that offers training " + "sessions (AWS_REGION, AWS_DEFAULT_REGION, or the region of your " + "AWS profile)." + ) + logger.info("Resolved AWS region: %s", resolved) + return resolved + + +def _endpoint_overrides() -> tuple[str | None, str | None]: + """Return ``(sagemaker_endpoint, runtime_endpoint)`` from the environment. + + ``None`` for an endpoint no variable names, which leaves that client to + resolve the endpoint from its Region. + """ + return ( + _checked_override(_SAGEMAKER_ENDPOINT_ENV), + _checked_override(_RUNTIME_ENDPOINT_ENV), + ) + + +def _checked_override(variable: str) -> str | None: + """Read an endpoint override, warning if it would send requests in the clear. + + A plain-HTTP endpoint is honoured rather than rejected, since a local test + server is a reasonable thing to point at, but it is worth saying out loud: + the request carries signed credentials and the training data itself. + """ + url = os.environ.get(variable) or None + if url is not None and url.startswith("http://"): + logger.warning( + "%s is an http:// URL, so requests to it, including their " + "credentials and training data, are sent unencrypted. Use https:// " + "unless this is a local endpoint you control.", + variable, + ) + return url diff --git a/sagemaker-train/src/sagemaker/train/training_session/_session_handle.py b/sagemaker-train/src/sagemaker/train/training_session/_session_handle.py new file mode 100644 index 0000000000..b2ae3a3c0c --- /dev/null +++ b/sagemaker-train/src/sagemaker/train/training_session/_session_handle.py @@ -0,0 +1,247 @@ +# 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. +"""Refcounted, stop-once ownership of one TrainingSession's lifetime. + +A TrainingSession is live compute, but a client object is +just a handle to it. Those two lifetimes are not the same: several clients can +share one session (``TrainingClient.create_sampling_client`` hands out a +``SamplingClient`` bound to the training session), so no single client can be +the one that stops it. + +:class:`_SessionHandle` is the session's lifetime, factored out of the clients. +Every client bound to a session holds a strong reference to the *same* handle, +so the handle is what CPython collects last. Its finalizer therefore runs +exactly when the final client referencing the session goes away: no ownership +flags, no parent/child bookkeeping, and no need to transfer responsibility from +a dropped parent to a surviving child. + +The automatic teardown stops the session; deleting one stays an explicit call. +See :func:`_auto_stop`. +""" + +from __future__ import annotations + +import logging +import sys +import threading +import warnings +import weakref +from typing import Any, List, Optional + +logger = logging.getLogger(__name__) + + +def _drain_submitters(submitters: List[Any]) -> None: + """Close every registered background submitter, best effort. + + Draining before the stop request means operations already queued reach the + service first. ``BackgroundSubmitter.close`` is idempotent, so it does not + matter whether a client's own finalizer already drained it. + """ + for submitter in submitters: + try: + submitter.close() + except Exception: # pragma: no cover - shutdown-time best effort + logger.debug("submitter drain failed during cleanup", exc_info=True) + + +#: Error codes meaning the session no longer exists. Absorbed only on the +#: automatic path; see :func:`_session_gone`. +_SESSION_GONE_CODES = frozenset({"ResourceNotFound", "ResourceNotFoundException"}) + + +def _session_gone(error: Exception) -> bool: + """True when a stop failed because there is no such session any more. + + For a best-effort teardown that is the goal state, not a failure: the caller + deleted the session out of band and the automatic stop simply arrived after + them. Duck-typed on ``response`` so this module needs no ``botocore`` import. + + :meth:`ServiceClient.stop_session` deliberately does *not* absorb this: + an explicit stop against a name that does not exist is a caller mistake and + must surface. + """ + response = getattr(error, "response", None) + if not isinstance(response, dict): + return False + return (response.get("Error") or {}).get("Code") in _SESSION_GONE_CODES + + +def _warn_not_stopped_explicitly(session_name: str) -> None: + """Tell the caller a session they still held had to be released for them.""" + try: + warnings.warn( + f"training session {session_name!r} was not stopped explicitly; " + f"stopping it now. Use the client as a context manager, or call " + f".stop(), to control when the session is released.", + ResourceWarning, + ) + except Exception: # pragma: no cover - warnings machinery may be torn down + pass + + +def _auto_stop( + session_name: str, + service_client: Any, + submitters: List[Any], + explicit_stop_failed: List[bool], +) -> None: + """Stop ``session_name``: the last client referencing it became unreachable. + + Module-level on purpose. :func:`weakref.finalize` stores its arguments, so a + callback that closed over the :class:`_SessionHandle` would keep that handle + permanently alive and the cleanup would never fire at all. + + Stop is the teardown action here, which is what makes an automatic teardown + safe to do on the caller's behalf: + + * A stopped session is still describable, so its history stays available to + audit and a run that went wrong can still be inspected. + * :meth:`ServiceClient.stop_session` is idempotent, so this stop is safe + whether or not the caller already stopped the session themselves. + * Stopping already releases the compute the session was using; what remains + is metadata. + + ``TrainingClient.delete`` / ``ServiceClient.delete_session`` stay + explicit-call-only and are unreachable from here. + + Stays quiet when there was nothing to stop. Both the ``ResourceWarning`` and + the log line report an action, so they are emitted only once one has + happened: a caller who already tore the session down, or whose session the + service has since removed, should not be told they leaked something. + + ``explicit_stop_failed`` non-empty means the caller did call ``stop()`` and + the request raised, so this stop retries theirs instead of cleaning up after + them: the stop still happens, but the ``ResourceWarning`` telling them to + call ``stop()`` would be wrong and is suppressed. + """ + # Logged, not acted on: whether a signed HTTPS call still works this late in + # interpreter teardown is CPython-version and host dependent, so the attempt + # below is made either way. + finalizing = sys.is_finalizing() + + # Before the stop either way, so queued operations reach the service first. + _drain_submitters(submitters) + + try: + stopped = service_client._stop_session_if_active(session_name) + if not stopped: + logger.debug( + "Unreferenced training session %r was already stopping or " + "terminal; nothing to stop (interpreter finalizing=%s).", + session_name, + finalizing, + ) + return + if not explicit_stop_failed: + _warn_not_stopped_explicitly(session_name) + logger.info( + "Auto-stopped unreferenced training session %r (interpreter " "finalizing=%s)", + session_name, + finalizing, + ) + except Exception as e: + if _session_gone(e): + logger.debug( + "Unreferenced training session %r no longer exists; nothing to " + "stop (interpreter finalizing=%s).", + session_name, + finalizing, + ) + return + logger.warning( + "Auto-stop of unreferenced training session %r failed (interpreter " + "finalizing=%s); it will be stopped automatically after its " + "inactivity timeout.", + session_name, + finalizing, + exc_info=True, + ) + + +class _SessionHandle: + """The lifetime of one TrainingSession, shared by every client bound to it. + + Args: + session_name: Name of the backing TrainingSession resource. + service_client: :class:`ServiceClient` used to issue the stop. When + ``None`` the handle cannot stop anything, so the automatic stop is + disabled. + + The stop is unconditional whenever a ``service_client`` is available, so a + caller gets the session released without having to track its lifetime. + Durability comes from the *checkpoint*, a model package durable on its + own, so trained weights outlive the process and a new session can be + created from one at any time. + """ + + def __init__( + self, + session_name: str, + service_client: Any = None, + ) -> None: + self.session_name = session_name + self._service_client = service_client + self._lock = threading.Lock() + # Held strongly: the drain must still be possible after the owning + # clients are gone, and a BackgroundSubmitter's daemon thread would + # otherwise outlive its client with queued submits never dispatched. + self._submitters: List[Any] = [] + # A shared list, not a bool attribute, for the reason _auto_stop is + # module-level: an argument reaching back to this handle would keep it + # alive and the cleanup would never fire. + self._explicit_stop_failed: List[bool] = [] + self._finalizer: Optional[weakref.finalize] = None + if service_client is not None: + self._finalizer = weakref.finalize( + self, + _auto_stop, + session_name, + service_client, + self._submitters, + self._explicit_stop_failed, + ) + + @property + def auto_stop_armed(self) -> bool: + """True while an unreferenced-session stop is still pending.""" + return self._finalizer is not None and self._finalizer.alive + + def register_submitter(self, submitter: Any) -> None: + """Attach a client's background submitter to this session's teardown.""" + with self._lock: + self._submitters.append(submitter) + + def note_explicit_stop_failed(self) -> None: + """Record that the caller's own ``stop()`` was issued and raised. + + The automatic stop stays armed so it can still release the session, but + it no longer warns that the caller never stopped it: they did, and the + request failed. Idempotent. + """ + if not self._explicit_stop_failed: + self._explicit_stop_failed.append(True) + + def disarm(self) -> None: + """Cancel the automatic stop: the session is already torn down. + + Called once an explicit ``stop()`` has succeeded, and before ``delete()``, + which removes the session name the automatic stop would otherwise use. + Saves a second network call and keeps the ``ResourceWarning`` from + reaching a caller who *did* clean up. + + Not called when a stop *fails*: detaching is one-way, so disarming there + would drop the only client-side retry the session has left. Idempotent. + """ + if self._finalizer is not None: + self._finalizer.detach() diff --git a/sagemaker-train/src/sagemaker/train/training_session/_submitter.py b/sagemaker-train/src/sagemaker/train/training_session/_submitter.py new file mode 100644 index 0000000000..e0245e046c --- /dev/null +++ b/sagemaker-train/src/sagemaker/train/training_session/_submitter.py @@ -0,0 +1,256 @@ +# 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. +"""Background submission machinery for non-blocking client methods. + +Base methods stamp submission order at call time, hand the network submit to a +dedicated background event loop, and return a future immediately, so callers +never wait on a submit round trip. Ordered submits (training operations) are +released strictly in call order; unordered submits (sampling operations) run +concurrently. A failed submit raises when the caller resolves the future, not at +call time. +""" + +from __future__ import annotations + +import asyncio +import concurrent.futures +import threading +from typing import Any, Awaitable, Callable, Optional + +from sagemaker.train.training_session.exceptions import TrainingSessionError + +#: Raised when a client is asked for another operation after it has been closed. +#: Worded for the caller: what closes a client is stopping or deleting its +#: session, or leaving its ``with`` block. +_CLOSED_MESSAGE = ( + "This client is closed and accepts no new operations. Stopping or deleting " + "the session closes its client, as does leaving the client's `with` block. " + "Create a new client to continue." +) + +#: Raised on an operation that had been submitted but had not finished when its +#: client closed. Worded for the caller: the operation's outcome is unknown, and +#: what to do about it depends on the operation. +_ABANDONED_MESSAGE = ( + "This operation was still in flight when its client closed, so its result is " + "unknown. Keep the client open until the operations you submitted have " + "resolved." +) + + +class BackgroundSubmitter: + """Owns a daemon thread running an asyncio loop for submits.""" + + def __init__(self, name: str = "sagemaker-training-submitter") -> None: + self._name = name + self._lock = threading.Lock() + self._loop: Optional[asyncio.AbstractEventLoop] = None + self._thread: Optional[threading.Thread] = None + self._next_seq = 1 + self._release_seq = 1 + #: Sequence number -> the future that releases it. Touched only on the + #: submitter's loop thread, so it needs no lock. + self._waiters: "dict[int, asyncio.Future[None]]" = {} + #: Sequences that will never run, recorded by whichever of them was + #: cancelled before reaching its turn. The holder skips past these when it + #: advances, so a cancellation cannot let a later submit overtake a + #: running one. Drained as the turn passes them, so it stays small. + self._skipped: "set[int]" = set() + self._closed = False + #: Futures for submits that have been accepted and not yet finished. + #: :meth:`close` waits on exactly these, rather than on whatever tasks the + #: loop happens to hold at the moment it looks. + self._pending: "set[concurrent.futures.Future[Any]]" = set() + + def submit_ordered( + self, coro_fn: Callable[[], Awaitable[Any]] + ) -> "concurrent.futures.Future[Any]": + """Schedule coro_fn so submits run strictly in call order. + + Returns a concurrent.futures.Future resolving to the coroutine result. + + The sequence slot is taken and scheduled under one lock, which + :meth:`close` also takes before stopping the loop. A slot can therefore + never be handed out to a coroutine that the loop stops before running: + such a slot would never be released, and every later ordered submit + waits on its turn forever. + """ + with self._lock: + loop = self._ensure_loop_locked() + seq = self._next_seq + coro = self._run_in_turn(seq, coro_fn) + try: + future = asyncio.run_coroutine_threadsafe(coro, loop) + except RuntimeError: + # Nothing was scheduled, so the slot must stay unclaimed. + coro.close() + raise + self._next_seq += 1 + self._pending.add(future) + future.add_done_callback(self._forget) + return future + + def submit_unordered( + self, coro_fn: Callable[[], Awaitable[Any]] + ) -> "concurrent.futures.Future[Any]": + """Schedule coro_fn with no ordering constraint.""" + with self._lock: + loop = self._ensure_loop_locked() + future = asyncio.run_coroutine_threadsafe(self._run_unordered(coro_fn), loop) + self._pending.add(future) + future.add_done_callback(self._forget) + return future + + def close(self, timeout: float = 5.0) -> None: + """Let accepted submits finish, stop the loop, close it, join the thread. + + Idempotent. A submit after close raises :class:`TrainingSessionError`. + + Marking the submitter closed comes first, so the set of accepted submits + is final by the time it is read, and each is then waited on directly. The + wait cannot start before a submit that was already accepted becomes + visible: if the loop stopped ahead of one, that submit would never run + and the caller would wait on its result forever. + + A submit still unfinished when the wait runs out is failed rather than + left alone. Stopping the loop freezes it where it is, so its result would + never arrive, and a caller waiting on that result with the default + unbounded timeout would wait for the life of the process. + """ + with self._lock: + self._closed = True + loop = self._loop + thread = self._thread + pending = list(self._pending) + if loop is None: + # No submit ever ran, so there is no loop and nothing to drain. + return + if loop.is_closed(): + # An earlier call already finished the shutdown, and this one must + # stay quiet: ``call_soon_threadsafe`` on a closed loop raises. + return + + # Deliberately not gated on ``is_running()``. A loop whose thread has + # started but has not yet entered ``run_forever`` reports False, and + # returning on that skips both the drain and the stop, leaving the loop + # running with submits on it. ``call_soon_threadsafe`` is safe either + # way: a loop that has not started yet runs the callback when it does. + concurrent.futures.wait(pending, timeout=timeout) + for future in pending: + try: + future.set_exception(TrainingSessionError(_ABANDONED_MESSAGE)) + except concurrent.futures.InvalidStateError: + # The loop thread completed this submit between the wait ending + # and this line. Its real result stands, and letting the error + # escape here would skip the stop and join below, leaking the + # thread this is meant to shut down. + pass + loop.call_soon_threadsafe(loop.stop) + if thread is not None: + thread.join(timeout=timeout) + if thread.is_alive(): + # Join ran out with the loop still spinning; closing it now + # would raise. Leave it open for a later close() to finish. + return + # Stopping a loop only parks it: it keeps its readiness poller and its + # self-pipe pair until it is collected, which is whenever the last + # reference to this submitter goes. Close so they come back here. + loop.close() + + def _forget(self, future: "concurrent.futures.Future[Any]") -> None: + """Drop a finished submit, so the pending set does not grow all session. + + Registered as a done callback outside the lock: an already-finished future + runs the callback on the registering thread, which would otherwise + deadlock on the lock that thread still holds. + """ + with self._lock: + self._pending.discard(future) + + def _ensure_loop_locked(self) -> asyncio.AbstractEventLoop: + if self._closed: + raise TrainingSessionError(_CLOSED_MESSAGE) + if self._loop is None: + self._loop = asyncio.new_event_loop() + self._thread = threading.Thread( + target=self._loop.run_forever, name=self._name, daemon=True + ) + self._thread.start() + return self._loop + + async def _run_in_turn(self, seq: int, coro_fn: Callable[[], Awaitable[Any]]) -> Any: + """Wait for ``seq``'s turn, run the submit, then release the next one. + + Every path out of here retires ``seq``, cancellation included: the retire + is in a ``finally`` that only touches state and never awaits. Awaiting + there would be a second cancellation point, and a sequence that is never + retired strands every later submit on its turn. Cancelling a task only + requests cancellation, so this body always reaches its ``try`` and + therefore always runs its ``finally``, even when the cancel is requested + before the task is first scheduled. + + Waiters are plain futures rather than a condition variable for the same + reason: waking the next sequence takes no lock, so it cannot be + interrupted. The counter, the waiter map and the skipped set are touched + only on the submitter's single loop thread, so they need no locking. + """ + loop = asyncio.get_running_loop() + try: + if self._release_seq != seq: + waiter = loop.create_future() + self._waiters[seq] = waiter + try: + await waiter + finally: + self._waiters.pop(seq, None) + return await coro_fn() + finally: + self._retire(seq) + + def _retire(self, seq: int) -> None: + """Retire ``seq``, advancing the turn only if ``seq`` was holding it. + + Advancing unconditionally would let a sequence cancelled while queued + hand the turn on for a slot it never held, waking the next submit while + the current one is still on the wire. So a sequence that never reached + its turn is only recorded, and the holder skips past it on its way out. + """ + if seq != self._release_seq: + self._skipped.add(seq) + return + self._release_seq += 1 + while self._release_seq in self._skipped: + self._skipped.discard(self._release_seq) + self._release_seq += 1 + successor = self._waiters.get(self._release_seq) + if successor is not None and not successor.done(): + successor.set_result(None) + + async def _run_unordered(self, coro_fn: Callable[[], Awaitable[Any]]) -> Any: + return await coro_fn() + + +def run_blocking(fn: Callable[[], Any]) -> Callable[[], Awaitable[Any]]: + """Wrap a blocking callable as a coroutine function for ``submit_*``. + + The callable runs on the loop's default thread pool, so a slow wire call + never stalls the submitter loop: unordered submits keep overlapping, and + an ordered submit still holds its turn until the call completes (wire + order stays equal to call order). + """ + + async def _coro() -> Any: + loop = asyncio.get_running_loop() + return await loop.run_in_executor(None, fn) + + return _coro diff --git a/sagemaker-train/src/sagemaker/train/training_session/_utils.py b/sagemaker-train/src/sagemaker/train/training_session/_utils.py new file mode 100644 index 0000000000..a69ad271f2 --- /dev/null +++ b/sagemaker-train/src/sagemaker/train/training_session/_utils.py @@ -0,0 +1,1492 @@ +# 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. +"""Internal helpers for translating user-facing types to service API shapes. + +Not part of the public API surface. +""" + +from __future__ import annotations + +import logging +import warnings +import weakref +from typing import Any, Dict, List, Optional + +from sagemaker.train.training_session.types import ( + Datum, + EncodedTextChunk, + ModelInput, + ModelInputChunk, + _serialize_loss_fn_inputs, +) + +logger = logging.getLogger(__name__) + +_PUBLIC_HUB_NAME = "SageMakerPublicHub" + +#: Hub content search-keyword prefix identifying a base model this client can +#: fine-tune. Each model's recipes are tagged +#: ``@recipe:{type}_{technique}_{strategy}`` (lowercase) on the Public Hub; +#: matching on the prefix selects models carrying a compatible fine-tuning +#: recipe regardless of the parameter-efficient strategy it uses. +_TRAINABLE_RECIPE_KEYWORD_PREFIX = "@recipe:finetuning_mtrl_" + +#: Hub content search-keyword prefix carrying a model's published id, spelled +#: ``@huggingface-id:{provider}/{model}``. Free on the list summary, but keyword +#: values are lowercased; the model's Hub document has the exact case. +_HF_ID_KEYWORD_PREFIX = "@huggingface-id:" + +#: Model-family names that :func:`list_base_models_in_region` does not report and +#: ``base_model`` does not accept, for models this client does not support end to +#: end. Matched against the dash-separated parts of a Hub content name. +_EXCLUDED_MODEL_NAMES = ("nova",) + +#: Hub client -> {model id: Hub content name} for the fine-tunable set, filled by +#: a Hub scan. Only the mapping is cached, never the versioned ARN: that is +#: described on every resolution, so a session still starts on the model's +#: current version. +#: +#: Kept under the client itself, because the client is what decides the answer: +#: its Region, its endpoint, and its credentials all change which models come +#: back, and none of them is reliably readable to build a key from. The reference +#: is weak, so a catalogue lives exactly as long as the client it describes. +_HUB_MODEL_INDEX: "weakref.WeakKeyDictionary[Any, Dict[str, str]]" = weakref.WeakKeyDictionary() + +#: Hub client -> {Hub content name: model id}. Survives a catalogue refresh, +#: which exists to find new models, not to re-read ids that do not change under +#: a published model. Successful reads only, so a failed one is retried rather +#: than remembered. +#: +#: That holds for a model publishing an id. For one that does not, the fallback +#: is memoised too, so a client keeps reporting the fallback for its lifetime +#: even if the model starts publishing an id. Bounded to one client, and the +#: fallback resolves, so the cost is a stale spelling rather than a stale model. +_MODEL_ID_BY_HUB_CONTENT: "weakref.WeakKeyDictionary[Any, Dict[str, str]]" = ( + weakref.WeakKeyDictionary() +) + +#: Ceiling on concurrent ``DescribeHubContent`` calls for one Hub scan. +_MAX_HUB_DESCRIBE_WORKERS = 8 + + +#: Consequence clause for ``retry_config``, shared by every factory returning a +#: SamplingClient. The retry policy is the SDK's: throttling and 5xx failures are +#: retried on a fixed budget on both submit and poll, and any other error +#: surfaces immediately. +_RETRY_CONFIG_CONSEQUENCE = ( + "the SDK's built-in transient-retry policy applies instead (fixed budget on " + "throttling and 5xx; other errors surface immediately)" +) + + +def _warn_unsupported_args(consequence: str, **passed: Any) -> None: + """Log a warning for each accepted-but-unhonoured keyword argument. + + These parameters are accepted but carry no behaviour. The warning names each + one the caller actually set, so the log says plainly which knobs the run did + not use. + + Only a value the caller chose warns. ``None``, or for a boolean the default + the SDK already behaves as, carries no request and stays quiet, so the + message keeps its signal. + + Args: + consequence: Clause completing "... and has no effect; ", + naming what happened instead (e.g. ``"the session was created + without it"``). + **passed: Parameter name to value. Values of ``None`` are skipped. + """ + for name, value in passed.items(): + if value is None: + continue + logger.warning( + "%s=%r is not supported by this SDK and has no effect; %s.", + name, + value, + consequence, + ) + + +def _is_excluded_base_model(hub_content_name: str) -> bool: + """True if any dash-separated part of the name is in :data:`_EXCLUDED_MODEL_NAMES`. + + Matched on whole parts rather than as a substring, so a name that merely + contains an excluded one (as a longer word ending in it, say) is not caught + with it. This decides both what the listing reports and what ``base_model`` + accepts, so an over-match makes a model unusable rather than just unlisted. + """ + parts = set(hub_content_name.lower().split("-")) + return any(name in parts for name in _EXCLUDED_MODEL_NAMES) + + +def _is_qualified_model_id(value: str) -> bool: + """True if ``value`` is shaped like a model id (``provider/model``). + + The one predicate deciding both what an id is: :func:`resolve_base_model` + routes on it, so a value the listing reports must satisfy it too. Reported + under a value that fails it, a model would be listed and then rejected. + """ + return "/" in value and not value.startswith("arn:") + + +def _hf_id_from_keywords(keywords: List[str]) -> Optional[str]: + """The model id a Hub entry's search keywords advertise, or ``None``. + + The keyword's value is checked, not trusted: it is whatever the Hub + publishes, and only a ``provider/model`` value resolves. + """ + for keyword in keywords: + if keyword.lower().startswith(_HF_ID_KEYWORD_PREFIX): + hf_id = keyword[len(_HF_ID_KEYWORD_PREFIX) :].strip() + return hf_id if hf_id and _is_qualified_model_id(hf_id) else None + return None + + +def _known_model_ids(hub_client: Any) -> Dict[str, str]: + """This client's ``{Hub content name: model id}`` memo, writable in place. + + A client that cannot be referenced weakly gets a throwaway dict, so writes + are discarded rather than shared with another client. + """ + # Membership is the only lookup a WeakKeyDictionary tolerates on a key that + # cannot be referenced weakly; ``get`` raises on one just as ``[]`` does. + try: + if hub_client in _MODEL_ID_BY_HUB_CONTENT: + return _MODEL_ID_BY_HUB_CONTENT[hub_client] + memo: Dict[str, str] = {} + _MODEL_ID_BY_HUB_CONTENT[hub_client] = memo + return memo + except TypeError: + return {} + + +def _public_model_id(hub_content_name: str, keywords: List[str], hub_client: Any) -> str: + """The id a Hub model is published under, and can be passed back as ``base_model``. + + Read from the model's Hub document, the only source that keeps the id's exact + case. Falls back to the list summary's keyword (same id, lowercased), then to + the Hub content name, so every entry is listed under an id that resolves. + + A failed read costs case and not identity where a keyword exists, since + ``base_model`` matches case-insensitively. For a model publishing no keyword + it does change the reported identity, to the Hub content name; that is warned + about, not remembered, and the next call reports the id again. + """ + known = _known_model_ids(hub_client) + if hub_content_name in known: + return known[hub_content_name] + + keyword_id = _hf_id_from_keywords(keywords) + try: + described = _extract_hf_model_id(_describe_hub_content(hub_content_name, hub_client)) + except Exception as e: # noqa: BLE001 - any Hub failure degrades the same way + # Logged: nothing else explains why this one id is spelled unlike the rest. + logger.warning( + "Could not read the published id for Hub model %r (%s); reporting it as " + "%r, which still resolves.", + hub_content_name, + e, + keyword_id or hub_content_name, + ) + return keyword_id or hub_content_name + + model_id = described if described and _is_qualified_model_id(described) else None + model_id = model_id or keyword_id or hub_content_name + known[hub_content_name] = model_id + return model_id + + +def _scan_trainable_hub_models(hub_client: Any) -> Dict[str, str]: + """Return ``{model id: Hub content name}`` for the region's fine-tunable models. + + Queries the SageMaker Public Hub and keeps the models that advertise a + compatible fine-tuning recipe (via their ``HubContentSearchKeywords``), so + the set is discovered live rather than hardcoded and a newly enabled model + appears with no SDK change. + + Args: + hub_client: A boto3 ``sagemaker`` client (its region determines results). + """ + trainable: List["tuple[str, List[str]]"] = [] + next_token = None + while True: + kwargs: dict = {"HubName": _PUBLIC_HUB_NAME, "HubContentType": "Model"} + if next_token: + kwargs["NextToken"] = next_token + resp = hub_client.list_hub_contents(**kwargs) + for item in resp.get("HubContentSummaries", []): + hub_content_name = item.get("HubContentName") + if not hub_content_name: + continue + if _is_excluded_base_model(hub_content_name): + continue + keywords = item.get("HubContentSearchKeywords", []) or [] + if not any(kw.lower().startswith(_TRAINABLE_RECIPE_KEYWORD_PREFIX) for kw in keywords): + continue + trainable.append((hub_content_name, keywords)) + next_token = resp.get("NextToken") + if not next_token: + break + + model_ids = _public_model_ids(trainable, hub_client) + + index: Dict[str, str] = {} + # Walked in page order, so the tie-break stays deterministic however the ids + # were read. + for hub_content_name, _ in trainable: + model_id = model_ids[hub_content_name] + # Two Hub entries publishing one id would otherwise resolve to + # whichever the Hub happened to page first; take the lower name so + # the same call answers the same way twice. Logged because only one + # of the two can be listed, and which one is not the caller's choice. + existing = index.get(model_id) + if existing is not None: + logger.warning( + "Base model id %r is published by more than one model in this " + "region (%s); resolving it to %r. Pass a full Hub model ARN to " + "select the other.", + model_id, + ", ".join(sorted([existing, hub_content_name])), + min(existing, hub_content_name), + ) + if existing is None or hub_content_name < existing: + index[model_id] = hub_content_name + return index + + +def _public_model_ids(trainable: List["tuple[str, List[str]]"], hub_client: Any) -> Dict[str, str]: + """``{Hub content name: model id}`` for each entry, reading unknown ids at once. + + An unknown id costs one Hub call. The calls are independent, so they go out + together and a cold scan costs roughly one round trip rather than one per + model. Fewer than two go on this thread: a pool of one buys nothing. + """ + known = _known_model_ids(hub_client) + resolved = {name: known[name] for name, _ in trainable if name in known} + pending = [(name, kws) for name, kws in trainable if name not in resolved] + + if len(pending) < 2: + resolved.update({name: _public_model_id(name, kws, hub_client) for name, kws in pending}) + return resolved + + from concurrent.futures import ThreadPoolExecutor + + with ThreadPoolExecutor( + max_workers=min(_MAX_HUB_DESCRIBE_WORKERS, len(pending)), + thread_name_prefix="sagemaker-hub-model-id", + ) as pool: + futures = { + name: pool.submit(_public_model_id, name, kws, hub_client) for name, kws in pending + } + resolved.update({name: future.result() for name, future in futures.items()}) + return resolved + + +def _is_hub_index_cached(hub_client: Any) -> bool: + """True if this client's catalogue is already cached, so a miss is a real miss.""" + return hub_client in _HUB_MODEL_INDEX + + +def _hub_model_index(hub_client: Any, *, refresh: bool = False) -> Dict[str, str]: + """``{model id: Hub content name}``, from this client's cache or a fresh scan. + + Args: + hub_client: A boto3 ``sagemaker`` client (it determines the results, so it + is also what the entry is kept under). + refresh: Scan the Hub even when this client is already cached. + """ + # Membership is the only lookup that tolerates a key which cannot be + # referenced weakly; ``get`` would raise on one just as ``[]`` does. + if not refresh and hub_client in _HUB_MODEL_INDEX: + return _HUB_MODEL_INDEX[hub_client] + index = _scan_trainable_hub_models(hub_client) + try: + _HUB_MODEL_INDEX[hub_client] = index + except TypeError: + # A client that cannot be referenced weakly is simply not cached, rather + # than kept alive by this module or shared with another client. + pass + return index + + +def list_base_models_in_region(hub_client: Any) -> List[str]: + """Return the ids of the base models available for fine-tuning in this region. + + Each model is named by its own model id (``provider/model``), spelled as its + provider publishes it wherever the Hub carries one; pass one straight to + :meth:`ServiceClient.create_lora_training_client` as ``base_model``. Results + are region-specific. + + Always scans the Hub rather than answering from the catalogue cache, so a + model enabled since the last call shows up without a new client. + + Args: + hub_client: A boto3 ``sagemaker`` client (its region determines results). + """ + # Sorted case-insensitively: a provider that capitalises its name would + # otherwise sort ahead of every lowercase one, and the list is read by eye. + return sorted(_hub_model_index(hub_client, refresh=True), key=str.lower) + + +def _describe_hub_content(hub_content_name: str, hub_client: Any) -> dict: + """``DescribeHubContent`` for a Hub content name, errors unchanged. + + The raw call, for callers that must tell a lookup failure apart from a model + that simply publishes nothing. :func:`_describe_base_model` is the one that + turns a failure into an actionable message. + """ + return hub_client.describe_hub_content( + HubName=_PUBLIC_HUB_NAME, + HubContentType="Model", + HubContentName=hub_content_name, + ) + + +def _describe_base_model(base_model: str, hub_client: Any) -> dict: + """Look up a base model on the Public Hub and return its content record. + + ``base_model`` may be a Hub content name or a full Hub content ARN. Returns + the ``DescribeHubContent`` response, from which both the versioned ARN and + the model's published id can be read. + + Raises: + ValueError: If the name resolves to no Hub content in this region. + """ + # An ARN already identifies a specific version; pull the content name out of + # it (``...:hub-content/{hub}/Model/{name}[/version]``). Otherwise the value + # is the Hub content name itself. + if base_model.startswith("arn:"): + parts = base_model.split("/") + hub_content_name = parts[3] if len(parts) >= 4 else base_model + else: + hub_content_name = base_model + + try: + return _describe_hub_content(hub_content_name, hub_client) + except Exception as e: + raise ValueError( + f"Could not resolve base model {base_model!r} to a SageMaker Public " + f"Hub model in this region. Pass a model id returned by " + f"list_base_models(), or a full Hub model ARN." + ) from e + + +def _extract_hf_model_id(describe_resp: dict) -> Optional[str]: + """Pull the published ``HuggingfaceId`` out of a ``DescribeHubContent`` response. + + ``HubContentDocument`` is a JSON string; returns the ``HuggingfaceId`` field + from it, or ``None`` if absent/unparseable. + """ + import json + + doc = describe_resp.get("HubContentDocument") + if not doc: + return None + try: + return json.loads(doc).get("HuggingfaceId") or None + except (TypeError, ValueError): + return None + + +def resolve_base_model(base_model: str, hub_client: Any) -> tuple: + """Resolve a base-model identifier to ``(versioned_arn, hf_model_id)`` in one call. + + Accepts three forms of input: + + - A model id (``"provider/model"``), as + :func:`list_base_models_in_region` reports: looked up in the region's + catalogue index, since Hub content names are NOT derivable from ids + mechanically. + - A Hub content name, which is what an ARN and a service-side error name the + model. + - A full Hub content ARN (``"arn:aws:sagemaker:..."``). + + For a Hub content *name*, a single ``DescribeHubContent`` yields both the + versioned ARN and the published id, so callers that need both (session create + + later tokenizer load) avoid describing the model twice. + + For an ARN input the ARN is returned unchanged with ``hf_model_id=None`` (no + Describe). The id is then resolved lazily at tokenizer time via + :func:`resolve_hf_model_id`. + + Returns: + ``(arn, hf_model_id_or_None)``. + + Raises: + ValueError: If a model id matches no supported Hub model, or a Hub + content name resolves to no Hub content in this region. Either + message lists the supported models under both names. + """ + if base_model.startswith("arn:"): + return base_model, None + if _is_qualified_model_id(base_model): + return _resolve_hf_id_to_hub_model(base_model, hub_client) + if _is_excluded_base_model(base_model): + # A Hub content name this client does not support end to end. Refused + # here rather than at session create, and refused on the same rule the + # listing filters on, so a name the listing omits does not resolve. + raise _excluded_base_model_error(base_model) + try: + resp = _describe_base_model(base_model, hub_client) + except ValueError as e: + # A name the Hub does not carry. The Describe alone cannot say what this + # region does offer, so the error lists the ids instead: the usual cause + # is a model's display name, which is neither an id nor a content name. + raise _unsupported_base_model_error( + _hub_model_index(hub_client, refresh=True), base_model + ) from e + return resp["HubContentArn"], _extract_hf_model_id(resp) + + +def _hub_entry_for_model_id(index: Dict[str, str], hf_model_id: str) -> Optional["tuple[str, str]"]: + """``(published id, Hub content name)`` listed under ``hf_model_id``, or ``None``. + + Matched case-insensitively, and the *published* id comes back rather than the + caller's spelling of it: published ids are case-sensitive, so only that + spelling loads a tokenizer. + """ + want = hf_model_id.lower() + for model_id, hub_content_name in index.items(): + if model_id.lower() == want: + return model_id, hub_content_name + return None + + +def _excluded_base_model_error(base_model: str) -> ValueError: + """The error for a base model this client does not support for fine-tuning. + + Lists no catalogue. This input is refused on its name alone, without a Hub + call, and inlining the supported set would turn the cheapest rejection on the + path into a full Hub scan. The message names the one call that answers it. + """ + return ValueError( + f"Base model {base_model!r} is not supported for fine-tuning by this " + f"client. Call list_base_models() for the models available in this region." + ) + + +def _unsupported_base_model_error(index: Dict[str, str], base_model: str) -> ValueError: + """The error for a base model no Hub model in this region publishes. + + Lists the ids and nothing else. Every id here is one the caller can pass + straight back, so the message answers "what should I have typed" without + introducing a second way to name the same model. + """ + # One model per line: a catalogue of any size stays readable, where a single + # joined line grows past anything a terminal or a log viewer wraps well. + if index: + listing = "Supported models:" + "".join( + f"\n {model_id}" for model_id in sorted(index, key=str.lower) + ) + else: + listing = "Supported models: (none)" + return ValueError( + f"Could not resolve base model {base_model!r} to a model supported for " + f"fine-tuning in this region. Pass one of these ids, or a full Hub model " + f"ARN.\n{listing}" + ) + + +def _resolve_hf_id_to_hub_model(hf_model_id: str, hub_client: Any) -> tuple: + """Find the supported Hub model published under ``hf_model_id``. + + A Hub content name is not derivable from a model id by rewriting it, so + the id is looked up in the client's catalogue index, which one Hub scan builds + and every later resolution on that client reuses. Only the name comes from the + cache: it is then described, so the ARN is the model's current version, and a + name the Hub no longer carries is caught rather than trusted. + + Returns: + ``(versioned_arn, published_hf_model_id)``. The published id is returned + rather than the caller's spelling of it, so a case-insensitive match + still yields the canonical id a tokenizer loads by. + + Raises: + ValueError: If no supported model publishes this id. + """ + from_cache = _is_hub_index_cached(hub_client) + index = _hub_model_index(hub_client) + entry = _hub_entry_for_model_id(index, hf_model_id) + if entry is not None: + try: + return _resolved_from_hub_name(entry[1], entry[0], hub_client) + except ValueError: + # The cached name no longer resolves (the entry was renamed or + # withdrawn). Fall through to a fresh scan rather than reporting a + # Hub content name the caller never typed. + pass + elif not from_cache: + # The index above *was* the fresh scan, so a miss is already final and a + # second scan would only ask the same question twice. + raise _unsupported_base_model_error(index, hf_model_id) + + index = _hub_model_index(hub_client, refresh=True) + entry = _hub_entry_for_model_id(index, hf_model_id) + if entry is None: + raise _unsupported_base_model_error(index, hf_model_id) + return _resolved_from_hub_name(entry[1], entry[0], hub_client) + + +def _resolved_from_hub_name(hub_content_name: str, hf_model_id: str, hub_client: Any) -> tuple: + """Describe ``hub_content_name`` and return ``(versioned_arn, hf_model_id)``. + + ``hf_model_id`` is the id the catalogue lists this model under, which is the + id the model publishes. The Hub document is preferred where it carries one, + and neither is the caller's own spelling: a tokenizer loads by exact case. + """ + resp = _describe_base_model(hub_content_name, hub_client) + return resp["HubContentArn"], _extract_hf_model_id(resp) or hf_model_id + + +def resolve_hf_model_id(base_model: str, hub_client: Any) -> str: + """Resolve a base model to its HuggingFace model id for tokenizer loading. + + Reads the ``HuggingfaceId`` published on the model's Public Hub content + record (via ``DescribeHubContent``), never a hand-maintained mapping, + so the tokenizer tracks whatever the hub advertises. A value that is already + a model id (``provider/model``, not an ARN or Hub content name) is returned + as-is so callers can still pass one directly. + + Args: + base_model: A Hub content name or Hub model ARN; or a model id (as + :func:`list_base_models_in_region` reports) to pass through. + hub_client: A boto3 ``sagemaker`` client used for the Hub lookup. + + Raises: + ValueError: If the model has no published HuggingFace id (e.g. a model + whose tokenizer is not publicly available); pass pre-tokenized token + ids in that case. + """ + # A bare HuggingFace id (``org/model``) is not an ARN and not a Hub content + # name (those have no ``/``); pass it straight through. + if _is_qualified_model_id(base_model): + return base_model + + if hub_client is None: + raise ValueError( + f"No tokenizer is available for base model {base_model!r}: a " + f"HuggingFace id could not be looked up. Provide a tokenizer or pass " + f"pre-tokenized token ids instead." + ) + + hf_id = _extract_hf_model_id(_describe_base_model(base_model, hub_client)) + + if not hf_id: + raise ValueError( + f"No public HuggingFace tokenizer is available for base model " + f"{base_model!r}. Pass pre-tokenized token ids instead." + ) + return hf_id + + +def load_tokenizer(base_model: str, hub_client: Any, hf_model_id: Optional[str] = None) -> Any: + """Load a HuggingFace AutoTokenizer for the given base model. + + Args: + base_model: A model id, a Hub content name, or a Hub model ARN. + hub_client: A boto3 ``sagemaker`` client used to look up the model's + published HuggingFace id. Only consulted when ``hf_model_id`` is not + supplied. + hf_model_id: The model's HuggingFace id, if already resolved (e.g. at + session-create time). When given, the Hub lookup is skipped entirely. + + Returns: + A ``transformers.PreTrainedTokenizer`` instance. + + Raises: + ImportError: If ``transformers`` is not installed (install the optional + tokenizer extra: ``pip install sagemaker-train[tokenizer]``). + ValueError: If the model is not supported. + """ + if not hf_model_id and not base_model: + raise ValueError( + "Cannot load tokenizer: no base model or HuggingFace model ID " + "is configured for this session." + ) + + with warnings.catch_warnings(): + warnings.filterwarnings("ignore", message=".*PyTorch.*") + try: + from transformers import AutoTokenizer + except ImportError as e: + raise ImportError( + "Loading a tokenizer requires the 'transformers' package, which " + "is an optional dependency. Install it with " + "`pip install sagemaker-train[tokenizer]`, or pass pre-tokenized " + "token IDs instead of text." + ) from e + + if not hf_model_id: + hf_model_id = resolve_hf_model_id(base_model, hub_client) + logger.info("Loading tokenizer for %r", hf_model_id) + + # Suppress noisy logs and warnings during download + _loggers = ("httpx", "huggingface_hub", "huggingface_hub.utils._http") + saved = {name: logging.getLogger(name).level for name in _loggers} + for name in _loggers: + logging.getLogger(name).setLevel(logging.ERROR) + try: + with warnings.catch_warnings(): + warnings.filterwarnings("ignore") + return AutoTokenizer.from_pretrained(hf_model_id) + finally: + for name, level in saved.items(): + logging.getLogger(name).setLevel(level) + + +def build_lora_config(rank: int, alpha: int) -> dict[str, Any]: + """Build the LoraConfig shape.""" + return {"Rank": rank, "Alpha": alpha} + + +#: Default cap on the number of checkpoints a single list_model_packages call +#: returns. A model package group can hold hundreds of versions; the latest +#: handful is what callers want, so the default is bounded instead of paginating +#: the whole group. +DEFAULT_LIST_MODEL_PACKAGES_LIMIT = 10 + +#: Largest ``MaxResults`` the list APIs accept on one request. A caller asking +#: for more than this pages up to their cap instead, since a larger request is +#: rejected rather than answered. +MAX_RESULTS_PER_PAGE = 100 + +#: How many of the newest packages the "get latest" helpers will describe while +#: looking for a matching kind. The kind is not in the list summary, so each +#: candidate costs one DescribeModelPackage, so this bounds that fan-out. A match +#: can sit well down the list (an RL loop saves sampler weights far more often +#: than state), so an exhausted scan warns instead of returning ``None``, and +#: the caller can retry with a larger ``scan_limit``. +DEFAULT_LATEST_SCAN_LIMIT = 20 + + +def base_model_arn_from_model_package(mp_desc: Dict[str, Any], region: str) -> str: + """Extract the versioned base-model Hub ARN a model package was trained on. + + ``DescribeModelPackage`` publishes the base model as + ``InferenceSpecification.Containers[0].BaseModel`` (``HubContentName`` + + ``HubContentVersion``), not as a top-level field. The version is pinned to + what the weights were trained on; resolving the name via the Hub instead + would silently return the latest version. + + Args: + mp_desc: A ``DescribeModelPackage`` response. + region: Region for the returned ARN. + + Raises: + ValueError: If the response carries no base-model reference. + """ + containers = (mp_desc.get("InferenceSpecification") or {}).get("Containers") or [] + base = (containers[0].get("BaseModel") or {}) if containers else {} + hub_name = base.get("HubContentName") + hub_version = base.get("HubContentVersion") + if not hub_name or not hub_version: + raise ValueError( + "Model package " + f"{mp_desc.get('ModelPackageArn', '')!r} carries no " + "base-model reference (InferenceSpecification.Containers[0].BaseModel)." + ) + # Partition taken from the model package's own ARN rather than assumed: + # ``aws`` is wrong in the China and GovCloud partitions, and every other ARN + # this package builds resolves its partition rather than hardcoding one. + partition = _arn_partition(mp_desc.get("ModelPackageArn")) + return ( + f"arn:{partition}:sagemaker:{region}:aws:hub-content/" + f"{_PUBLIC_HUB_NAME}/Model/{hub_name}/{hub_version}" + ) + + +#: Fields in a full ARN: ``arn:partition:service:region:account:resource``. +_ARN_FIELD_COUNT = 6 + + +def _arn_partition(arn: Optional[str], default: str = "aws") -> str: + """Partition field of ``arn``, or ``default`` when it cannot be read. + + The whole ARN shape is checked before the field is trusted, not just the + presence of a colon: a short or malformed value would otherwise hand back + part of the resource id as the partition. + + Defaults rather than raising: the partition only differs from ``aws`` outside + the commercial regions, and the caller still needs an answer. + """ + parts = (arn or "").split(":") + if len(parts) < _ARN_FIELD_COUNT or parts[0] != "arn" or not parts[1]: + return default + return parts[1] + + +def list_model_package_summaries( + sagemaker_client: Any, + model_package_group_arn: str, + max_results: Optional[int] = DEFAULT_LIST_MODEL_PACKAGES_LIMIT, +) -> List[Dict[str, Any]]: + """Return ``ListModelPackages`` summaries for a group, newest first. + + Auto-paginates only as far as ``max_results``. ``max_results=None`` returns + every version in the group (caller opts into a full scan). The summary + carries arn / version / creation time / description, but NOT the checkpoint + kind (``IsCheckpoint`` lives in the InferenceSpecification, returned only by + ``DescribeModelPackage``). + + Args: + sagemaker_client: A boto3 ``sagemaker`` client. + model_package_group_arn: The session's model package group (name or ARN). + max_results: Cap on summaries returned; None for all. + """ + summaries: list = [] + next_token = None + while True: + kwargs: dict = { + "ModelPackageGroupName": model_package_group_arn, + "SortBy": "CreationTime", + "SortOrder": "Descending", + } + if next_token: + kwargs["NextToken"] = next_token + if max_results is not None: + remaining = max_results - len(summaries) + if remaining <= 0: + break + kwargs["MaxResults"] = min(remaining, MAX_RESULTS_PER_PAGE) + resp = sagemaker_client.list_model_packages(**kwargs) + summaries.extend(resp.get("ModelPackageSummaryList", [])) + next_token = resp.get("NextToken") + if not next_token: + break + if max_results is not None: + return summaries[:max_results] + return summaries + + +def describe_model_package_kind(sagemaker_client: Any, model_package_arn: str): + """Return the :class:`ModelPackageKind` for a model package via Describe. + + ``IsCheckpoint=True`` marks a (save_state) checkpoint; + absent/false marks a (save_weights_for_sampler) model. + + Args: + sagemaker_client: A boto3 ``sagemaker`` client. + model_package_arn: ARN of the model package version to describe. + """ + from sagemaker.train.training_session.types import ModelPackageKind + + resp = sagemaker_client.describe_model_package(ModelPackageName=model_package_arn) + is_checkpoint = False + spec = resp.get("InferenceSpecification") or {} + containers = spec.get("Containers") or [] + if containers and containers[0].get("IsCheckpoint"): + is_checkpoint = True + return ModelPackageKind.CHECKPOINT if is_checkpoint else ModelPackageKind.MODEL + + +#: ``ManagedStorageType`` of a model package whose weights SageMaker stores on +#: your behalf, not in your own S3 location. Describe does not return an +#: S3 URI for one, and deleting the package is what releases its storage. +MANAGED_STORAGE_RESTRICTED = "Restricted" + +#: Max keys ``DeleteObjects`` accepts per call, which is also the max a +#: ``ListObjectsV2`` page returns, so one listed page is one delete batch. +_DELETE_OBJECTS_BATCH = 1000 + + +def split_s3_uri(s3_uri: str) -> "tuple[str, str]": + """Split ``s3://bucket/key/prefix`` into ``("bucket", "key/prefix")``. + + The key is empty when the URI names only a bucket. + + Args: + s3_uri: An ``s3://`` URI. + + Raises: + ValueError: If the URI is not an ``s3://`` URI or names no bucket. + """ + scheme = "s3://" + if not s3_uri.startswith(scheme): + raise ValueError(f"{s3_uri!r} is not an s3:// URI.") + bucket, _, key = s3_uri[len(scheme) :].partition("/") + if not bucket: + raise ValueError(f"{s3_uri!r} names no bucket.") + return bucket, key + + +def describe_model_package_artifacts( + sagemaker_client: Any, model_package_arn: str +) -> Dict[str, Optional[str]]: + """Return where a model package's artifacts live, via Describe. + + Reads ``InferenceSpecification.Containers[0].ModelDataSource.S3DataSource`` + for the location and its type, plus the top-level ``ManagedStorageType``. + ``s3_uri`` is ``None`` when the package exposes no location, which is the + case for one whose storage SageMaker manages + (:data:`MANAGED_STORAGE_RESTRICTED`). + + Args: + sagemaker_client: A boto3 ``sagemaker`` client. + model_package_arn: ARN of the model package version to describe. + """ + resp = sagemaker_client.describe_model_package(ModelPackageName=model_package_arn) + spec = resp.get("InferenceSpecification") or {} + containers = spec.get("Containers") or [] + source = ( + ((containers[0].get("ModelDataSource") or {}).get("S3DataSource") or {}) + if containers + else {} + ) + return { + "s3_uri": source.get("S3Uri") or None, + # "S3Prefix" (everything under the location) or "S3Object" (one object). + # Read from the package, not inferred from whether the URI ends in "/". + "s3_data_type": source.get("S3DataType") or None, + "managed_storage_type": resp.get("ManagedStorageType") or None, + } + + +def delete_s3_objects_under( + s3_client: Any, s3_uri: str, *, s3_data_type: Optional[str] = None +) -> "tuple[int, List[str]]": + """Delete the S3 objects a model package's location covers. + + Returns ``(deleted_count, failed_keys)``. ``DeleteObjects`` reports per-key + failures in the response body, so a partial failure comes back in + ``failed_keys`` instead of raising. + + Each listed page is deleted before the next is requested, so no full key list + is held in memory. The continuation token resumes past keys already returned, + so the deletes do not disturb it. + + Args: + s3_client: A boto3 ``s3`` client. + s3_uri: Location to empty, as carried by the model package. + s3_data_type: ``"S3Object"`` to delete exactly that one key; anything + else (``"S3Prefix"`` or unset) treats the URI as a prefix. + + Raises: + ValueError: If the URI names a bucket with no key under it, where + deleting "its" objects would empty the whole bucket. + """ + bucket, key = split_s3_uri(s3_uri) + if not key: + raise ValueError( + f"{s3_uri!r} names a bucket root, not a location inside it. Refusing " + "to delete: this would empty the whole bucket." + ) + + if s3_data_type == "S3Object": + return _delete_keys(s3_client, bucket, [key]) + + # A prefix is matched on the segment boundary. Without the trailing slash, + # "run-1" would also match the contents of "run-10". + prefix = key if key.endswith("/") else key + "/" + deleted = 0 + failed: List[str] = [] + token: Optional[str] = None + while True: + kwargs: Dict[str, Any] = {"Bucket": bucket, "Prefix": prefix} + if token: + kwargs["ContinuationToken"] = token + page = s3_client.list_objects_v2(**kwargs) + keys = [obj["Key"] for obj in page.get("Contents") or []] + if keys: + page_deleted, page_failed = _delete_keys(s3_client, bucket, keys) + deleted += page_deleted + failed.extend(page_failed) + token = page.get("NextContinuationToken") + if not page.get("IsTruncated") or not token: + break + return deleted, failed + + +def _delete_keys(s3_client: Any, bucket: str, keys: List[str]) -> "tuple[int, List[str]]": + """``DeleteObjects`` these keys, in batches of :data:`_DELETE_OBJECTS_BATCH`. + + Runs quiet, so the response carries only failures and the deleted count is + the batch minus those. + """ + deleted = 0 + failed: List[str] = [] + for start in range(0, len(keys), _DELETE_OBJECTS_BATCH): + batch = keys[start : start + _DELETE_OBJECTS_BATCH] + resp = s3_client.delete_objects( + Bucket=bucket, + Delete={"Objects": [{"Key": k} for k in batch], "Quiet": True}, + ) + batch_failed = [err["Key"] for err in resp.get("Errors") or [] if err.get("Key")] + deleted += len(batch) - len(batch_failed) + failed.extend(batch_failed) + return deleted, failed + + +def _resolve_tokenizer(tokenizer: Any, tokenizer_loader: Any): + """Return a tokenizer, lazily loading one via ``tokenizer_loader`` if needed. + + ``tokenizer`` (already cached) wins; otherwise the zero-arg + ``tokenizer_loader`` (e.g. ``client.get_tokenizer``) is invoked on demand, + so string input works without a manual ``get_tokenizer()`` call. + Raises ValueError if neither yields a tokenizer. + """ + if tokenizer is not None: + return tokenizer + if tokenizer_loader is not None: + return tokenizer_loader() + raise ValueError( + "String input requires a tokenizer, but none is available for this " + "session's base model. Provide pre-tokenized token IDs instead." + ) + + +def prompt_to_model_input( + prompt: Any, + tokenizer: Any = None, + chunk_type: str = "encoded_text", + tokenizer_loader: Any = None, +) -> dict[str, Any]: + """Build a ModelInput shape from a sampling prompt. + + Accepts the same prompt forms as a single training datum: + - A list of token IDs: ``[1, 2, 3]`` + - A string (tokenized via ``tokenizer`` or lazily via ``tokenizer_loader``) + - A pre-built ModelInput dict (passthrough): ``{"chunks": [...]}`` + - A typed :class:`ModelInput` / :class:`EncodedTextChunk`, or any + object with ``chunks`` + + Returns the ``{"chunks": [...]}`` ModelInput shape expected by the sampling + operation. + """ + # Typed inputs: built-in types or compatible external objects. A typed + # input built from text (ModelInput.from_text) carries the text rather than + # chunks, so it is tokenized here too: this is the one place with a + # tokenizer, and serializing it untokenized raises. + if isinstance(prompt, (ModelInput, EncodedTextChunk, ModelInputChunk)): + model_input = ModelInput.from_any(prompt) + if model_input._pending_text is not None: + model_input = model_input._tokenized_with( + _resolve_tokenizer(tokenizer, tokenizer_loader) + ) + return model_input._to_api() + if not isinstance(prompt, (dict, list, str)) and hasattr(prompt, "chunks"): + return ModelInput.from_any(prompt)._to_api() + + # Already a ModelInput dict + if isinstance(prompt, dict) and "chunks" in prompt: + return prompt + + # String: tokenize (lazily loading a tokenizer only when needed) + if isinstance(prompt, str): + token_ids = _resolve_tokenizer(tokenizer, tokenizer_loader).encode(prompt) + # List of token IDs + elif isinstance(prompt, list) and all(isinstance(t, int) for t in prompt): + token_ids = prompt + else: + raise TypeError( + f"Unsupported prompt type: {type(prompt).__name__}. " + "Expected a list of token IDs, a string, or a ModelInput dict." + ) + + return {"chunks": [{"tokens": token_ids, "type": chunk_type}]} + + +def tokenize_to_datum( + token_ids: list[int], + loss_fn_inputs: dict[str, Any] | None = None, + chunk_type: str = "encoded_text", +) -> dict[str, Any]: + """Build a single Datum shape from token IDs. + + ``lossFnInputs`` is only included when ``loss_fn_inputs`` is provided. + The service rejects an auto-generated ``labels`` entry. The accepted keys + depend on the loss function (e.g. ``advantages``, ``weights``, + ``target_tokens``, ``logprobs``, ``ref_log_prob``), so the caller supplies + the inputs appropriate for their loss function. + """ + datum: dict[str, Any] = { + "modelInput": {"chunks": [{"tokens": token_ids, "type": chunk_type}]}, + } + if loss_fn_inputs is not None: + datum["lossFnInputs"] = _serialize_loss_fn_inputs(loss_fn_inputs) + return datum + + +def _batch_element_kind(element: Any) -> str: + """Classify one batch element, for detecting a batch that mixes kinds. + + Only the distinction between the interpretations :func:`batch_to_datum_list` + applies matters here, so an unrecognised element gets its own kind and lands + in the mixed path, where coercing it raises a clear error. + """ + if isinstance(element, (Datum, ModelInput, EncodedTextChunk, ModelInputChunk)): + return "typed" + if isinstance(element, dict): + return "dict" + if isinstance(element, str): + return "text" + if isinstance(element, list) and all(isinstance(t, int) for t in element): + return "tokens" + if not isinstance(element, (dict, list, str, int)) and ( + hasattr(element, "model_input") or hasattr(element, "chunks") + ): + return "typed" + return "unsupported" + + +def _typed_datum_to_api( + element: Any, + tokenizer: Any, + tokenizer_loader: Any, +) -> dict[str, Any]: + """Serialize one typed batch element, tokenizing text it still carries. + + A ``Datum`` or ``ModelInput`` built from text (``ModelInput.from_text``, or a + caller's own object exposing a text prompt) has no chunks to send yet. The + tokenizer lives out here, so this is where that text becomes tokens; + otherwise serializing it raises. + """ + datum = Datum.from_any(element) + if datum._pending_text is not None: + datum = datum._tokenized_with(_resolve_tokenizer(tokenizer, tokenizer_loader)) + return datum._to_api() + + +def _mixed_batch_to_datum_list( + batch: Any, + tokenizer: Any, + loss_fn_inputs: dict[str, Any] | None, + tokenizer_loader: Any, +) -> list[dict[str, Any]]: + """Coerce a batch whose elements are not all the same kind, one at a time. + + Text is the one kind needing outside context (a tokenizer), so it keeps the + dedicated path; everything else goes through :meth:`Datum.from_any`, which + rejects an element it cannot read instead of emitting an empty datum. + """ + tok = None + datums: list[dict[str, Any]] = [] + for element in batch: + if isinstance(element, str): + if tok is None: + tok = _resolve_tokenizer(tokenizer, tokenizer_loader) + datums.append(tokenize_to_datum(tok.encode(element), loss_fn_inputs)) + continue + datum = _typed_datum_to_api(element, tokenizer, tokenizer_loader) + if loss_fn_inputs and "lossFnInputs" not in datum: + datum = {**datum, "lossFnInputs": dict(_serialize_loss_fn_inputs(loss_fn_inputs))} + datums.append(datum) + return datums + + +#: ``lossFnInputs`` entries whose length the service requires to equal the token +#: count of ``modelInput``. Limited to what the service enforces: other entries +#: may legitimately carry their own length. +_PER_TOKEN_LOSS_FN_INPUTS = ( + "target_tokens", + "weights", + "logprobs", + "old_logprobs", +) + + +def model_input_token_count(model_input: Any) -> int | None: + """Token count of a serialized ``modelInput``, or ``None`` if not countable. + + ``None`` for anything that is not the ``{"chunks": [...]}`` shape, so a + length check is skipped rather than guessed at. + """ + if not isinstance(model_input, dict): + return None + chunks = model_input.get("chunks") + if not isinstance(chunks, list): + return None + total = 0 + for chunk in chunks: + if not isinstance(chunk, dict): + return None + tokens = chunk.get("tokens") + if not isinstance(tokens, list): + return None + total += len(tokens) + return total + + +def _wire_model_input_length(datum: dict[str, Any]) -> int | None: + """Token count of a datum's serialized ``modelInput``.""" + return model_input_token_count(datum.get("modelInput")) + + +def _wire_value_length(value: Any) -> int | None: + """Element count of a serialized ``lossFnInputs`` value, or ``None``.""" + if isinstance(value, dict): + data = value.get("data") + return len(data) if isinstance(data, list) else None + if isinstance(value, list): + return len(value) + return None + + +#: The loss inputs each loss function requires on every datum. The field is +#: optional in general, but which entries a given loss needs is not: a request +#: missing one is refused, and the refusal names neither the datum nor the key. +_LOSS_FN_INPUT_KEYS = { + "cross_entropy": ("target_tokens", "weights"), + "importance_sampling": ("target_tokens", "weights", "logprobs", "advantages"), + "ppo": ("target_tokens", "weights", "logprobs", "advantages"), + "cispo": ("target_tokens", "weights", "logprobs", "advantages"), +} + + +def _validate_loss_fn_inputs(datums: list[dict[str, Any]], loss_fn_name: str) -> None: + """Check every datum carries the loss inputs ``loss_fn_name`` requires. + + Raised here rather than left to the request so the message can name the datum + and the keys. An operation that reaches the service without them is refused + after it has been accepted and queued, and that refusal identifies neither. + """ + required = _LOSS_FN_INPUT_KEYS.get(loss_fn_name) + if not required: + return + for index, datum in enumerate(datums): + if not isinstance(datum, dict): + continue + loss_fn_inputs = datum.get("lossFnInputs") + present = loss_fn_inputs if isinstance(loss_fn_inputs, dict) else {} + missing = [key for key in required if key not in present] + if not missing: + continue + raise ValueError( + f"datum[{index}]: loss function {loss_fn_name!r} requires " + f"{', '.join(required)} in its loss inputs, and " + f"{', '.join(missing)} {'is' if len(missing) == 1 else 'are'} missing. " + f"Pass them per datum, or for the whole batch with " + f"loss_fn_inputs={{...}}; each is one value per token of the datum, " + f"the prompt included." + ) + + +def _validate_datum_lengths(datums: list[dict[str, Any]]) -> None: + """Check per-token ``lossFnInputs`` against the datum's token count. + + Every per-token entry must be as long as ``modelInput``, prompt positions + included. Raised locally because the service accepts the mismatch and only + fails the operation later. + """ + for index, datum in enumerate(datums): + if not isinstance(datum, dict): + continue + loss_fn_inputs = datum.get("lossFnInputs") + if not isinstance(loss_fn_inputs, dict) or not loss_fn_inputs: + continue + expected = _wire_model_input_length(datum) + if expected is None: + continue + for key in _PER_TOKEN_LOSS_FN_INPUTS: + if key not in loss_fn_inputs: + continue + actual = _wire_value_length(loss_fn_inputs[key]) + if actual is None or actual == expected: + continue + raise ValueError( + f"datum[{index}]: lossFnInputs[{key!r}] has {actual} values but " + f"modelInput has {expected} tokens. Every per-token entry must " + f"be the same length as modelInput, covering the prompt as well " + f"as the response, e.g. weights=[0.0] * prompt_len + [1.0] * " + f"response_len for a {expected}-token input. Response-only " + f"values are the usual cause of this mismatch." + ) + + +def batch_to_datum_list( + batch: Any, + tokenizer: Any = None, + loss_fn_inputs: dict[str, Any] | None = None, + tokenizer_loader: Any = None, + loss_fn_name: str | None = None, +) -> list[dict[str, Any]]: + """Convert a user-provided batch to a list of Datum dicts. + + Accepts: + - A list of token ID lists: ``[[1,2,3], [4,5,6]]`` + - A list of strings (tokenized via ``tokenizer`` or lazily via + ``tokenizer_loader``): ``["hello", "world"]`` + - A list of pre-built datum dicts (passthrough) + - A list of typed :class:`Datum` / :class:`ModelInput` objects, or + compatible objects with the same attributes + + When ``loss_fn_inputs`` is provided, it is attached to every constructed + datum. For pre-built datum dicts it fills in only those that do not already + carry their own. When omitted, no ``lossFnInputs`` is emitted (the service + rejects an auto-generated ``labels`` entry). + + Args: + loss_fn_name: The loss function the batch will be sent with. Given, every + datum is checked for the keys that loss needs. Omitted, only the + lengths are checked. + + Raises: + ValueError: A datum is missing a loss input ``loss_fn_name`` needs, or a + per-token ``lossFnInputs`` entry is not the same length as the + datum's ``modelInput``. + """ + datums = _coerce_batch_to_datum_list(batch, tokenizer, loss_fn_inputs, tokenizer_loader) + if loss_fn_name is not None: + _validate_loss_fn_inputs(datums, loss_fn_name) + _validate_datum_lengths(datums) + return datums + + +def _coerce_batch_to_datum_list( + batch: Any, + tokenizer: Any = None, + loss_fn_inputs: dict[str, Any] | None = None, + tokenizer_loader: Any = None, +) -> list[dict[str, Any]]: + """Coerce a batch to wire-shape datum dicts, per :func:`batch_to_datum_list`.""" + if not batch: + return [] + + first = batch[0] + + # A batch of bare token IDs is one datum, not a list of elements. + if isinstance(first, int): + # That reading holds only if every element is a token ID. A batch that + # merely starts with one was read as a token list regardless, which put + # the remaining elements inside a token array and sent them. + others = sorted({type(e).__name__ for e in batch if not isinstance(e, int)}) + if others: + raise TypeError( + "Unsupported batch: one starting with a token ID must hold only " + f"token IDs, but this one also holds {others}. Pass a list of token " + "ID lists, a list of strings, or a list of datum dicts." + ) + else: + kinds = {_batch_element_kind(element) for element in batch} + if len(kinds) > 1: + # Mixed kinds: the homogeneous branches below read the whole batch + # through the interpretation chosen for ``first``, which would apply + # the wrong one to every other element. Coerce element by element + # instead so each is read as what it actually is. + return _mixed_batch_to_datum_list(batch, tokenizer, loss_fn_inputs, tokenizer_loader) + + # Typed Datum / ModelInput objects: built-in or compatible external objects. + if isinstance(first, (Datum, ModelInput)) or ( + not isinstance(first, (dict, list, str, int)) + and (hasattr(first, "model_input") or hasattr(first, "chunks")) + ): + datums = [_typed_datum_to_api(d, tokenizer, tokenizer_loader) for d in batch] + if not loss_fn_inputs: + return datums + serialized = _serialize_loss_fn_inputs(loss_fn_inputs) + return [ + d if "lossFnInputs" in d else {**d, "lossFnInputs": dict(serialized)} for d in datums + ] + + # Already a datum dict: fill loss inputs only where absent. Copy rather + # than mutate the caller's dicts, and give each datum its own + # loss_fn_inputs so a later edit to one does not bleed into the others. + if isinstance(first, dict) and all(isinstance(d, dict) and "modelInput" in d for d in batch): + if not loss_fn_inputs: + return batch + serialized = _serialize_loss_fn_inputs(loss_fn_inputs) + return [ + d if "lossFnInputs" in d else {**d, "lossFnInputs": dict(serialized)} for d in batch + ] + + # A datum dict spelled with ``model_input`` or ``prompt`` is read one + # element at a time, the same way a mixed batch is. A dict naming no input + # at all falls through to the error below, which lists what a batch may + # hold; coercing it would instead build a datum with no input at all. + if all(isinstance(d, dict) for d in batch) and any( + "model_input" in d or "prompt" in d for d in batch + ): + return _mixed_batch_to_datum_list(batch, tokenizer, loss_fn_inputs, tokenizer_loader) + + # List of token ID lists + if isinstance(first, list) and all(isinstance(t, int) for t in first): + return [tokenize_to_datum(ids, loss_fn_inputs) for ids in batch] + + # List of strings: tokenize (lazily loading a tokenizer only when needed) + if isinstance(first, str): + tok = _resolve_tokenizer(tokenizer, tokenizer_loader) + return [tokenize_to_datum(tok.encode(text), loss_fn_inputs) for text in batch] + + # Single token ID list (not wrapped) + if isinstance(first, int): + return [tokenize_to_datum(batch, loss_fn_inputs)] + + raise TypeError( + f"Unsupported batch type: {type(first).__name__}. " + "Expected list of token ID lists, list of strings, or list of datum dicts." + ) + + +# --- forward-backward chunking ---------------------------------------------- +# +# A batch too large for one request is split into chunks and submitted as +# consecutive forwardBackward ops. Gradients accumulate until the next +# optim_step, so this is equivalent to sending the whole batch at once. Each +# chunk is capped by datum count and by an estimated byte budget. + +#: Max datums per forward-backward chunk. +MAX_CHUNK_LEN = 1024 +#: Metrics that count what a batch held rather than describe it per datum. A +#: chunked batch splits the count, so the whole-batch value is the total and +#: not an average. Named explicitly: an unrecognized metric is treated as a +#: per-datum quantity, which is the common case but still an assumption. +SUMMED_METRICS = frozenset({"global_token_num"}) + +#: Approximate serialized-byte budget per chunk (leaves headroom under the +#: runtime request-size limit for envelope + signing overhead). +MAX_CHUNK_BYTES = 10_000_000 + + +def _estimate_datum_bytes(datum: Dict[str, Any]) -> int: + """Roughly estimate a datum's serialized size in bytes. + + Counts the token payload (the dominant term: token IDs plus per-token + loss-function inputs such as advantages/weights/logprobs) and adds a small + fixed per-token allowance for JSON structure. Deliberately an estimate: + it only needs to be monotonic in the real size to keep chunks under the + limit, not exact. + """ + tokens = 0 + for chunk in datum.get("modelInput", {}).get("chunks", []): + tokens += len(chunk.get("tokens", [])) + lfi = datum.get("lossFnInputs") or {} + for value in lfi.values(): + if isinstance(value, dict): + data = value.get("data") + if isinstance(data, list): + tokens += len(data) + # ~10 bytes per token/value covers the JSON digits, quoting, and commas. + return tokens * 10 + + +def chunk_datum_list( + datum_list: List[Dict[str, Any]], + *, + max_len: int = MAX_CHUNK_LEN, + max_bytes: int = MAX_CHUNK_BYTES, +) -> List[List[Dict[str, Any]]]: + """Split a datum list into chunks bounded by count and estimated bytes. + + Each chunk holds at most ``max_len`` datums and stays under ``max_bytes`` of + estimated payload, except that a single datum larger than ``max_bytes`` is + kept on its own (it cannot be split further here). Datum order is preserved, + so submitting the chunks in order reproduces the original batch. + """ + if not datum_list: + return [] + + chunks: List[List[Dict[str, Any]]] = [] + current: List[Dict[str, Any]] = [] + current_bytes = 0 + for datum in datum_list: + datum_bytes = _estimate_datum_bytes(datum) + # Close the current chunk if adding this datum would exceed either + # bound (but never emit an empty chunk just because one datum is huge). + if current and (current_bytes + datum_bytes > max_bytes or len(current) >= max_len): + chunks.append(current) + current = [] + current_bytes = 0 + current.append(datum) + current_bytes += datum_bytes + if current: + chunks.append(current) + return chunks + + +def _result_datum_count(result: Any) -> int: + """Datums a chunk's result covers, inferred from its per-datum loss outputs. + + Zero when the chunk carries none, which is the usual case: loss outputs are + only returned when the caller asked for them. Used only as a fallback, since + inferring the count from an optional field silently makes every chunk weigh + the same. + """ + return len(result.loss_fn_outputs or []) + + +def combine_forward_backward_results( + results: List[Any], datum_counts: Optional[List[int]] = None +) -> Any: + """Fold per-chunk forward-backward results into one, in chunk order. + + Scalar metrics are combined so the caller sees the whole-batch step: + ``loss`` and other per-example means are recomputed as a datum-count- + weighted average over the chunks that reported them (each chunk reports + the mean over its own datums), while ``grad_norm``, a property of the + accumulated gradient after all chunks, is taken from the final chunk and + omitted when that chunk does not report it. The counts in + :data:`SUMMED_METRICS` are totalled, since each chunk counts part of the + same batch. Non-scalar or unrecognized metrics fall back to the last + chunk that reported them. + ``loss_fn_outputs`` lists are concatenated in chunk order so per-datum + outputs line up with the batch. + + Args: + results: One result per chunk, in chunk order. + datum_counts: Datums per chunk, in the same order. The caller built the + chunks, so it is the only place the real sizes are known; without + them the weights fall back to the per-datum loss outputs, which are + usually absent and would make an uneven split (1024, 1024, 100) + report a plain mean instead of the weighted one. + + Single-element input returns that result unchanged, so an unchunked call is + unaffected. + """ + if len(results) == 1: + return results[0] + + # Local import avoids a module cycle (types imports nothing from here). + from sagemaker.train.training_session.types import ForwardBackwardOutput + + per_chunk_metrics = [r.metrics for r in results] + if datum_counts is not None and len(datum_counts) == len(results): + weights = [max(count, 1) for count in datum_counts] + else: + weights = [max(_result_datum_count(r), 1) for r in results] + + combined_metrics: Dict[str, float] = {} + # Keyed off every chunk, not just the last: a metric the earlier chunks + # report and the final one omits would otherwise vanish from the result. + for key in dict.fromkeys(k for m in per_chunk_metrics for k in m): + present = [(m[key], w) for m, w in zip(per_chunk_metrics, weights) if key in m] + values = [v for v, _ in present] + if any(not isinstance(v, (int, float)) for v in values): + combined_metrics[key] = values[-1] + elif key == "grad_norm" or key.endswith("/grad_norm"): + # A gradient norm describes the gradient accumulated across all + # chunks, so only the final chunk's is the whole-batch value and an + # average of norms is not one. Reported under a namespaced key, so + # the suffix is matched rather than the whole name. + if key in per_chunk_metrics[-1]: + combined_metrics[key] = per_chunk_metrics[-1][key] + elif key in SUMMED_METRICS: + # The chunks each carry part of one batch, so the total is the + # batch's. Averaging reported a count no chunk measured, smaller + # than a single chunk's own. + combined_metrics[key] = sum(values) + else: + # Weighted over the chunks that reported it, so a metric only some + # chunks carry is still their mean and not diluted by the rest. + key_w = sum(w for _, w in present) + combined_metrics[key] = sum(v * w for v, w in present) / key_w + + # Per-datum loss outputs concatenate in chunk order, so they line up with + # the batch and read exactly as an unchunked pass would. + entries: List[Any] = [] + for r in results: + entries.extend(r.loss_fn_outputs or []) + combined_outputs: Optional[Any] = entries or None + + raw: Dict[str, Any] = {"forwardBackward": {"metrics": combined_metrics}} + if combined_outputs is not None: + raw["forwardBackward"]["lossFnOutputs"] = combined_outputs + return ForwardBackwardOutput._from_api(raw) diff --git a/sagemaker-train/src/sagemaker/train/training_session/exceptions.py b/sagemaker-train/src/sagemaker/train/training_session/exceptions.py new file mode 100644 index 0000000000..7307294012 --- /dev/null +++ b/sagemaker-train/src/sagemaker/train/training_session/exceptions.py @@ -0,0 +1,212 @@ +# 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. +"""Exceptions for the training_session client interface. + +A call that fails on the wire raises the error ``botocore`` models for it: +``ValidationException``, ``ConflictException``, ``ResourceNotFoundException``, +``ThrottlingException``, ``InternalServerException``, reachable both as +attributes of the client (``client.exceptions.ValidationException``) and as +``botocore.exceptions.ClientError``, whose ``response`` carries the service's +own error code and message. The SDK passes those through untouched, so nothing +here has to be caught to read what the service said. + +The exceptions defined here cover what a wire error cannot express: an +asynchronous operation that was accepted and later failed +(:class:`OperationError`), and two conditions the SDK detects on the client side +(:class:`ResultMismatchError`, :class:`ArtifactDeletionError`). A condition with +nothing specific to report, such as an operation asked of a client that is +already closed, raises the base :class:`TrainingSessionError` itself, so catching +that one catches every error defined here. + +Running out of time waiting is the exception to that: a wait that outlives its +``timeout`` raises the built-in ``TimeoutError``, matching what +``concurrent.futures`` and ``asyncio`` raise for the same thing, so an existing +``except TimeoutError`` keeps working. Catch it alongside +:class:`TrainingSessionError` to cover both. Use :func:`is_retryable` to decide +whether a failed call is worth sending again. +""" + +from __future__ import annotations + +from typing import List, Optional + +from botocore.exceptions import ClientError +from botocore.exceptions import ConnectionError as BotoConnectionError +from botocore.exceptions import HTTPClientError + +#: Error codes for a failure that says nothing about the request's validity, so +#: sending it again may well succeed. +_RETRYABLE_ERROR_CODES = frozenset({"ThrottlingException", "InternalServerException"}) + +#: Status codes carrying the same meaning, for a failure returned without an +#: error code to read. +_RETRYABLE_STATUS_CODES = frozenset({429, 500, 502, 503, 504}) + + +def is_retryable(error: BaseException) -> bool: + """Return whether sending the failed call again may succeed. + + True for throttling, an HTTP 5xx, and connection-level failures such as a + read or connect timeout: the request did not arrive, or arrived and could not + be served for a reason unrelated to what it asked for. False for everything + else, including a rejected request, since resending that one gets the same + answer. + + The SDK already retries its own internal calls on this basis. Use it to + decide the same thing for a call of your own. Resolve the future inside the + ``try``: submitting is non-blocking, so the operation's own failure surfaces + where the result is awaited, not where it was queued:: + + future = training_client.forward_backward(batch, loss_fn="cross_entropy") + try: + result = future.result() + except Exception as error: + if not is_retryable(error): + raise + ... + + Args: + error: The exception raised by the failed call. + + Returns: + ``True`` if the failure is transient, ``False`` otherwise. + """ + if isinstance(error, (BotoConnectionError, HTTPClientError)): + # No response, so no status code: the request may never have arrived. + return True + if isinstance(error, ClientError): + response = error.response or {} + code = (response.get("Error") or {}).get("Code") or "" + status = (response.get("ResponseMetadata") or {}).get("HTTPStatusCode") + return code in _RETRYABLE_ERROR_CODES or ( + isinstance(status, int) and status in _RETRYABLE_STATUS_CODES + ) + return False + + +class TrainingSessionError(Exception): + """Base exception for training session errors. + + Every error class in this module derives from this one, and it is also raised + directly where the message says all there is to say: asking a closed client + for another operation, for instance. A wait that runs out of time raises the + built-in ``TimeoutError`` instead, so bound waits are caught the same way as + with any other future. + """ + + +class OperationError(TrainingSessionError): + """Raised when an asynchronous operation completed in a failed state. + + The operation was accepted, so nothing was raised when it was submitted; it + then failed while running, and that surfaces here when the future for it is + resolved. A request the service rejects outright raises ``ClientError`` + instead, at submit time. + + ``operation_id`` identifies the failed operation server-side and is the one + detail to quote when asking for help with it. + """ + + def __init__( + self, + operation_type: str, + message: str, + *, + operation_id: Optional[str] = None, + ): + self.operation_type = operation_type + self.operation_id = operation_id + self.message = message + subject = f"{operation_type} operation {operation_id}" if operation_id else operation_type + super().__init__(f"{subject} failed: {message}") + + def __reduce__(self) -> tuple: + """Rebuild through ``__init__`` so the error survives pickling. + + A training loop that spreads work across processes moves exceptions + between them. Without this the default rebuild calls ``__init__`` with + the single formatted message and raises ``TypeError`` about a missing + argument, replacing the real failure with an unrelated one. + """ + return ( + _rebuild_operation_error, + (type(self), self.operation_type, self.message, self.operation_id), + ) + + +def _rebuild_operation_error( + cls: type, operation_type: str, message: str, operation_id: Optional[str] +) -> "OperationError": + """Reconstruct an :class:`OperationError`; used only by pickling. + + Takes the class so a subclass round-trips as itself, the way the two + siblings do by rebuilding through ``self.__class__``. + """ + return cls(operation_type, message, operation_id=operation_id) + + +class ResultMismatchError(TrainingSessionError): + """Raised when a completed operation's result does not match its request. + + A per-token result must be as long as the sequence sent. A different length + means the result belongs to some other request, so it is raised instead of + returned: the values would otherwise be read as this request's own. + """ + + def __init__(self, operation_type: str, expected: int, received: int): + self.operation_type = operation_type + self.expected = expected + self.received = received + super().__init__( + f"{operation_type} returned {received} values for a {expected}-token " + "sequence. The result does not belong to this request. Issuing these " + "operations one at a time avoids it." + ) + + def __reduce__(self) -> tuple: + """Rebuild through ``__init__``; see :meth:`OperationError.__reduce__`.""" + return (self.__class__, (self.operation_type, self.expected, self.received)) + + +class ArtifactDeletionError(TrainingSessionError): + """Raised when deleting a model package's S3 artifacts did not fully succeed. + + The model package is left in place, so it still points at the objects listed + on ``failed_keys``. Re-run :meth:`ServiceClient.delete_model_package` once + the cause is addressed (most often a missing S3 permission). The message + truncates ``failed_keys``; the attribute is complete. + """ + + #: Keys listed in the message before it truncates. + _KEYS_IN_MESSAGE = 5 + + def __init__(self, model_package_arn: str, s3_uri: str, failed_keys: List[str]): + self.model_package_arn = model_package_arn + self.s3_uri = s3_uri + self.failed_keys = failed_keys + shown = ", ".join(failed_keys[: self._KEYS_IN_MESSAGE]) + if len(failed_keys) > self._KEYS_IN_MESSAGE: + shown += f", ... ({len(failed_keys)} total, see .failed_keys)" + super().__init__( + f"Could not delete every object under {s3_uri} for model package " + f"{model_package_arn}: {shown}. The model package was NOT deleted, so " + "it still points at these objects; retry once the cause is " + "addressed (commonly s3:DeleteObject on that location)." + ) + + def __reduce__(self) -> tuple: + """Rebuild through ``__init__``; see :meth:`OperationError.__reduce__`.""" + return ( + self.__class__, + (self.model_package_arn, self.s3_uri, self.failed_keys), + ) diff --git a/sagemaker-train/src/sagemaker/train/training_session/futures.py b/sagemaker-train/src/sagemaker/train/training_session/futures.py new file mode 100644 index 0000000000..8e07bfe55f --- /dev/null +++ b/sagemaker-train/src/sagemaker/train/training_session/futures.py @@ -0,0 +1,865 @@ +# 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. +"""APIFuture and request-status model for async training/sample operations. + +A single ``APIFuture`` serves training, sampling and checkpointing operations. +It is generic in the type it resolves to: each method returns an +``APIFuture[T]`` whose ``result()`` is a ``T``, the operation's own response +type. ``APIFuture`` is the one future type callers see; the two subclasses below +are internal, so a method that may return either is still annotated +``APIFuture[T]``. + +``timeout`` is the only knob callers pass, and it bounds the whole wait on every +future in this module. A future that resolves in stages +(:class:`PendingAPIFuture` waits for its submit and then polls; +:class:`CombinedFuture` waits for several sub-futures) turns ``timeout`` into a +single deadline at the top of ``result()`` and hands each stage only the time +still left, so ``result(timeout=60)`` returns or raises :class:`TimeoutError` +within 60 seconds regardless of how many stages ran. The default is +``timeout=None`` waits indefinitely; pass an explicit ``timeout`` to bound a +specific wait. + +Every future can be waited on either way. ``result()`` blocks the calling +thread. ``await`` (or :meth:`result_async`) suspends a coroutine instead: no +thread is held while the operation runs, only for the individual status +requests. So ``asyncio.gather`` over hundreds of futures keeps them all in +flight. + +The polling cadence is the SDK's to pick, so it is set at construction rather +than per wait: ``poll_interval``, ``time_fn``, ``sleep_fn``, ``async_sleep_fn`` +and ``rand_fn`` are constructor seams the SDK and its tests use (tests inject a +fake clock and a fixed jitter draw, and run without real waits). Callers receive +futures, never build them. +""" + +from __future__ import annotations + +import asyncio +import concurrent.futures +import random +import threading +import time +from enum import Enum +from typing import ( + Any, + Awaitable, + Callable, + Dict, + Generator, + Generic, + List, + Optional, + Tuple, + TypeVar, +) + +from sagemaker.train.training_session.exceptions import ( + OperationError, + TrainingSessionError, + is_retryable, +) + +#: The names this module contributes to the public API. The future subclasses +#: below are returned, never constructed by callers, so they are reached through +#: :class:`APIFuture` rather than by name; the defaults and helpers around them +#: are wiring. Naming the surface here keeps the generated reference to what a +#: caller can act on. +__all__ = ["APIFuture", "RequestStatus"] + +#: The type an :class:`APIFuture` resolves to. +_T = TypeVar("_T") + +#: Default seconds between poll attempts. +DEFAULT_POLL_INTERVAL = 2 +#: Default wait bound for ``result()``. ``None`` means wait indefinitely, which +#: is the default: a long-running operation that succeeds should not raise +#: ``TimeoutError`` just because it outran an arbitrary clock. Pass an explicit +#: ``timeout`` (seconds) to bound a specific wait. +DEFAULT_TIMEOUT = None +#: Max consecutive transient failures tolerated on a submit or a poll before +#: giving up: about five minutes at the backoff below. Enough to ride out a +#: traffic spike, while an endpoint that is genuinely unavailable still surfaces +#: instead of spinning forever. To bound a wait by elapsed time instead, pass it +#: an explicit ``timeout``. +DEFAULT_MAX_TRANSIENT_POLL_ERRORS = 150 +#: Full-jitter exponential backoff bounds, in seconds. The ceiling doubles per +#: consecutive failure from the base and holds at the max; the delay slept is +#: drawn from below it. +DEFAULT_BACKOFF_BASE_SECONDS = 0.5 +DEFAULT_BACKOFF_MAX_SECONDS = 4.0 + + +#: Marks a future that has not resolved yet. A plain ``None`` cannot: it is a +#: value a ``parse_fn`` may legitimately produce, and treating it as "unresolved" +#: would send the future back to the wire on every ``result()`` call. +_UNSET: Any = object() + + +def _backoff_delay( + attempt: int, + *, + base: float = DEFAULT_BACKOFF_BASE_SECONDS, + cap: float = DEFAULT_BACKOFF_MAX_SECONDS, + rand_fn: Callable[[float, float], float] = random.uniform, +) -> float: + """Seconds to wait before retry ``attempt`` (0-based), with full jitter. + + Backing off eases the request rate a throttled service just refused. Drawing + the delay from below the ceiling, rather than sleeping the ceiling itself, + spreads callers throttled together so they do not all return together. + """ + ceiling = min(cap, base * (2**attempt)) + return rand_fn(0.0, ceiling) + + +def call_with_transient_retry( + fn: Callable[[], Any], + *, + max_attempts: int = DEFAULT_MAX_TRANSIENT_POLL_ERRORS, + backoff_base: float = DEFAULT_BACKOFF_BASE_SECONDS, + backoff_max: float = DEFAULT_BACKOFF_MAX_SECONDS, + sleep_fn: Callable[[float], None] = time.sleep, + rand_fn: Callable[[float, float], float] = random.uniform, +) -> Any: + """Call ``fn``, retrying transient failures up to ``max_attempts``. + + For idempotent runtime calls (e.g. ExecuteTrainingOperation submitted with a + stable clientToken + sequenceNumber). Errors that + :func:`~sagemaker.train.training_session.exceptions.is_retryable` rejects + propagate immediately; the rest wait out a :func:`_backoff_delay` first, so a + throttled submit is not sent straight back at the rate that throttled it. + """ + attempt = 0 + while True: + try: + return fn() + except Exception as error: + if not is_retryable(error): + raise + attempt += 1 + if attempt >= max_attempts: + raise + sleep_fn( + _backoff_delay( + attempt - 1, + base=backoff_base, + cap=backoff_max, + rand_fn=rand_fn, + ) + ) + + +class RequestStatus(str, Enum): + """Lifecycle states of an asynchronous training or sampling operation. + + An operation moves from ``Pending`` to ``InProgress`` to a terminal state: + ``Completed`` (success), ``Failed`` (the operation errored), or ``Skipped`` + (the session failed and abandoned its remaining queued operations). + """ + + PENDING = "Pending" + IN_PROGRESS = "InProgress" + COMPLETED = "Completed" + FAILED = "Failed" + SKIPPED = "Skipped" + + +#: Terminal success state. +_SUCCESS_STATUS = RequestStatus.COMPLETED + +#: Terminal failure states: reaching any of these raises OperationError. +#: (``Skipped`` is terminal: it marks operations abandoned after the session +#: itself failed.) +_TERMINAL_FAILURE_STATUSES = frozenset({RequestStatus.FAILED, RequestStatus.SKIPPED}) + +#: Non-terminal states: keep polling. Named so the three sets can be asserted to +#: partition :class:`RequestStatus` exactly, which is what makes an unrecognized +#: state impossible to poll on forever. +_PENDING_STATUSES = frozenset({RequestStatus.PENDING, RequestStatus.IN_PROGRESS}) + + +def _format_timeout(timeout: Optional[float]) -> str: + """Render a wait bound for a message, tolerating an unbounded one. + + ``None`` means wait indefinitely, and formatting it as a number raises, + which would replace whatever error was being reported with a ``TypeError``. + """ + return "no limit" if timeout is None else f"{timeout:.1f}s" + + +def _deadline_from(timeout: Optional[float], time_fn: Callable[[], float]) -> Optional[float]: + """Absolute monotonic deadline for ``timeout``, or ``None`` for unbounded.""" + return None if timeout is None else time_fn() + timeout + + +def _remaining(deadline: Optional[float], time_fn: Callable[[], float]) -> Optional[float]: + """Seconds left before ``deadline``; ``None`` when unbounded. + + Clamped at zero so an already-passed deadline hands the next stage a + non-negative budget it will immediately time out on, not a negative one + whose meaning varies by API. + """ + if deadline is None: + return None + return max(0.0, deadline - time_fn()) + + +def _coerce_status(raw: Any) -> RequestStatus: + """Map a raw service status value to a ``RequestStatus``. + + Unknown values are treated as pending (``IN_PROGRESS``) instead of raising, + so a new server state never crashes a polling client mid-loop. + """ + try: + return RequestStatus(raw) + except ValueError: + return RequestStatus.IN_PROGRESS + + +class APIFuture(Generic[_T]): + """Future for an async runtime operation (training or sample). + + Call :meth:`result` to block until the operation reaches a terminal state. + The same future type serves both training and sample operations; the caller + supplies a ``poll_fn`` that fetches the latest response and a ``parse_fn`` + that turns a successful response into the public result object. + + Generic in the result type: ``parse_fn`` fixes it, so ``result()`` on an + ``APIFuture[T]`` returns a ``T``. + + Args: + operation_id: Server-assigned id of the operation (for messages). + poll_fn: Zero-arg callable returning the latest operation response dict. + Each call should hit ``Get*OperationResult`` for this operation. + parse_fn: Callable mapping a completed response dict to the result + object returned by :meth:`result`. + status_key: Key under which the response carries its status string. + operation_type: Label used in error messages (e.g. ``"training"``). + poll_interval: Seconds between poll attempts. + time_fn: Monotonic clock, injectable for tests. Defaults to + :func:`time.monotonic`. + sleep_fn: Sleep function used by :meth:`result`, injectable for tests. + Defaults to :func:`time.sleep`. + async_sleep_fn: Awaitable sleep used by :meth:`result_async`, injectable + for tests. Defaults to :func:`asyncio.sleep`. + max_transient_poll_errors: Consecutive transient poll errors (throttling + / 5xx on Get*OperationResult) tolerated before giving up. A + successful poll resets the counter. + rand_fn: Jitter draw for the backoff between failed polls, injectable for + tests. Defaults to :func:`random.uniform`. + """ + + def __init__( + self, + operation_id: str, + poll_fn: Callable[[], Dict[str, Any]], + parse_fn: Callable[[Dict[str, Any]], _T], + *, + status_key: str = "status", + operation_type: str = "operation", + poll_interval: float = DEFAULT_POLL_INTERVAL, + time_fn: Callable[[], float] = time.monotonic, + sleep_fn: Callable[[float], None] = time.sleep, + async_sleep_fn: Callable[[float], Awaitable[None]] = asyncio.sleep, + max_transient_poll_errors: int = DEFAULT_MAX_TRANSIENT_POLL_ERRORS, + rand_fn: Callable[[float, float], float] = random.uniform, + ) -> None: + self._operation_id = operation_id + self._poll_fn = poll_fn + self._parse_fn = parse_fn + self._status_key = status_key + self._operation_type = operation_type + self._poll_interval = poll_interval + self._time_fn = time_fn + self._sleep_fn = sleep_fn + self._async_sleep_fn = async_sleep_fn + self._max_transient_poll_errors = max_transient_poll_errors + self._rand_fn = rand_fn + self._result: Any = _UNSET + #: The terminal failure, once one is known. Memoized beside the + #: result so a resolved future reports itself resolved either way. + self._error: Any = _UNSET + self._lock = threading.Lock() + + @property + def operation_id(self) -> str: + """Return the server-assigned operation id.""" + return self._operation_id + + # Shared by both waits -------------------------------------------------- + # result() and result_async() differ only in how they wait: one sleeps the + # calling thread, the other suspends a coroutine. Everything that decides + # what the wait *means* lives here, so the two cannot drift apart. + + # Concurrent waiters on one future are allowed to poll independently, but + # only one may resolve it: `parse_fn` carries the operation's completion + # side effects, so it must run once. The lock therefore covers classify + + # memoize, and never a poll: holding it across the wire would make a + # bounded result(timeout=...) block on another waiter's unbounded poll, and + # in the async path would occupy the worker threads result_async needs to + # poll at all. + + def _transient_backoff(self, consecutive_transient: int) -> float: + """Seconds to wait after ``consecutive_transient`` failed status reads. + + Not ``poll_interval``: a poll that answered keeps the normal cadence, + while one that failed backs off, since under a spike that cadence is part + of what is being shed. + """ + return _backoff_delay(consecutive_transient - 1, rand_fn=self._rand_fn) + + def _classify_poll(self, resp: Dict[str, Any]) -> Tuple[bool, Any]: + """Read one poll response. + + Returns ``(True, result)`` once the operation has completed and + ``(False, None)`` while it is still running. + + Raises: + OperationError: The response reports a terminal failure state. + """ + status = _coerce_status(resp.get(self._status_key)) + if status == _SUCCESS_STATUS: + return True, self._parse_fn(resp) + if status in _TERMINAL_FAILURE_STATUSES: + raise OperationError( + self._operation_type, + resp.get( + "errorDetails", + f"Operation ended with status: {status.value}", + ), + operation_id=self._operation_id, + ) + return False, None + + def _classify_and_memoize(self, resp: Dict[str, Any]) -> Tuple[bool, Any]: + """Read one poll response and memoize the outcome, at most once. + + Returns what :meth:`_classify_poll` returns. When two waiters poll + concurrently and both see a completed response, the first through here + parses and memoizes; the second is served that same result object and + never calls ``parse_fn`` again. + + A terminal failure is memoized the same way a success is. Both are the + operation's final answer, so a later wait replays it rather than asking + the service again, and the future reports itself resolved either way. + That covers the failures this SDK judges as well as the ones the service + reports: a result that does not match the request it was awaited for is + as final as an operation the service calls failed, and ``parse_fn`` is + where it is caught. + """ + with self._lock: + if self._error is not _UNSET: + raise self._error + if self._result is not _UNSET: + return True, self._result + try: + completed, value = self._classify_poll(resp) + except TrainingSessionError as error: + self._error = error + raise + if completed: + self._result = value + return completed, value + + def _timeout_error(self, timeout: Optional[float]) -> TimeoutError: + """The error both waits raise when ``timeout`` runs out.""" + return TimeoutError( + f"{self._operation_type} operation {self._operation_id} " + f"did not complete within {_format_timeout(timeout)}" + ) + + def result(self, timeout: Optional[float] = DEFAULT_TIMEOUT) -> _T: + """Block until the operation completes and return its parsed result. + + Args: + timeout: Max seconds to wait before raising ``TimeoutError``. + ``None`` waits indefinitely. + + Returns: + The object produced by ``parse_fn`` from the completed response. + + Raises: + OperationError: The operation reached a terminal failure state. + TimeoutError: The operation did not complete within ``timeout``. + Exception: The error botocore raised for a failed poll, either a + transient one that recurred past ``max_transient_poll_errors`` + in a row, or one that was never worth retrying. + """ + if self._error is not _UNSET: + raise self._error + if self._result is not _UNSET: + return self._result + + deadline = _deadline_from(timeout, self._time_fn) + consecutive_transient = 0 + while deadline is None or self._time_fn() < deadline: + try: + resp = self._poll_fn() + except Exception as error: + if not is_retryable(error): + raise + # A transient failure reading the status does not mean the + # operation failed; it is often still running. Keep polling until + # the deadline, but cap consecutive failures so an endpoint that + # is genuinely down surfaces instead of spinning. + consecutive_transient += 1 + if consecutive_transient > self._max_transient_poll_errors: + raise + self._sleep_fn(self._transient_backoff(consecutive_transient)) + continue + + consecutive_transient = 0 + completed, value = self._classify_and_memoize(resp) + if completed: + return value + self._sleep_fn(self._poll_interval) + + raise self._timeout_error(timeout) + + async def result_async(self, timeout: Optional[float] = DEFAULT_TIMEOUT) -> _T: + """Await the operation result without occupying a thread while it runs. + + Equivalent to ``await future``, and the way to await with a ``timeout`` + of your own: ``await future`` cannot take arguments, so it uses the + default. + + Args: + timeout: Max seconds to wait before raising ``TimeoutError``. + ``None`` waits indefinitely. + + Raises: + OperationError: The operation reached a terminal failure state. + TimeoutError: The operation did not complete within ``timeout``. + Exception: The error botocore raised for a failed poll, either a + transient one that recurred past ``max_transient_poll_errors`` + in a row, or one that was never worth retrying. + """ + if self._error is not _UNSET: + raise self._error + if self._result is not _UNSET: + return self._result + + deadline = _deadline_from(timeout, self._time_fn) + consecutive_transient = 0 + while deadline is None or self._time_fn() < deadline: + try: + # The status request itself blocks, so give it a worker thread + # for that one request only, not for the whole operation. + resp = await asyncio.to_thread(self._poll_fn) + except Exception as error: + if not is_retryable(error): + raise + consecutive_transient += 1 + if consecutive_transient > self._max_transient_poll_errors: + raise + await self._async_sleep_fn(self._transient_backoff(consecutive_transient)) + continue + + consecutive_transient = 0 + completed, value = self._classify_and_memoize(resp) + if completed: + return value + await self._async_sleep_fn(self._poll_interval) + + raise self._timeout_error(timeout) + + def __await__(self) -> Generator[Any, None, _T]: + """Make APIFuture awaitable so ``await future`` works in async code. + + Uses the default ``timeout``; call :meth:`result_async` to pass your + own. Every subclass routes ``await`` through its own + :meth:`result_async`, so overriding that one method is enough. + """ + return self.result_async().__await__() + + # The rest of the standard future surface ------------------------------- + # Inherited by both subclasses: each keeps the same ``_result`` slot, so + # asking whether a future has resolved means the same thing for all three. + + def done(self) -> bool: + """Whether this future has resolved, without waiting for it. + + ``True`` once the operation's outcome is known, whether that is a result + or a failure, so a loop that drains futures by checking this terminates + on a failed operation as well as a successful one. A future whose + operation has finished on the service but has not been waited on yet + still reports ``False``: nothing has polled for it. + + A failed *request* is not an outcome and is not recorded here. A call + that could not be made or answered, credentials that expired mid-run + being the clearest case, says nothing about the operation, which may + still be running and may still succeed, so a later wait asks again + rather than replaying an answer that was never about the operation. + """ + return self._result is not _UNSET or self._error is not _UNSET + + def exception(self, timeout: Optional[float] = DEFAULT_TIMEOUT) -> Optional[BaseException]: + """Wait for the operation and return the error it failed with. + + Returns ``None`` if the operation succeeded, so a caller can inspect + several futures without wrapping each in ``try``. + + A failed operation hands back the same error object on every wait, the + way any settled future does, so identity comparisons hold. Re-raising + one object repeatedly does lengthen its traceback each time, which is + worth knowing if a retry loop waits on the same future many times over; + read ``__cause__`` and the message rather than the traceback depth. + + Args: + timeout: Max seconds to wait before raising ``TimeoutError``. + ``None`` waits indefinitely. + + Raises: + TimeoutError: The operation did not complete within ``timeout``. + """ + try: + self.result(timeout=timeout) + except TimeoutError: + raise + except Exception as error: # noqa: BLE001 - reported, not handled + return error + return None + + def cancel(self) -> bool: + """Always ``False``: a submitted operation cannot be called back. + + Present because the rest of this surface follows the standard future + protocol, and a caller looping over futures should not have to special + case this one. The service has the operation by the time a caller holds + the future, so there is nothing to take back; stop the session to end + the work it is doing. + """ + return False + + def cancelled(self) -> bool: + """Always ``False``: see :meth:`cancel`.""" + return False + + +class PendingAPIFuture(APIFuture[_T]): + """An :class:`APIFuture` whose submit round trip is still in flight. + + Returned the moment an operation is queued on the client's background + submitter, before ``Execute*Operation`` has assigned an operation id. + Resolving it waits for the submit first (a submit error surfaces here, + not at the call site), then delegates to the :class:`APIFuture` the + submit produced. ``await`` works unchanged and holds no thread through + either stage. + + Args: + submit_future: ``concurrent.futures.Future`` resolving to the real + :class:`APIFuture` once the submit round trip completes (or + raising the submit error). + operation_type: Label used in error messages (e.g. ``"training"``). + keepalive: The submitting client's ``_SessionHandle``. Held only to keep + it reachable: the automatic session stop fires when the last + reference to the handle goes away, so holding one here guarantees a + session is never stopped while an operation against it is still + outstanding. Dropping the client but keeping its futures stays safe. + time_fn: Monotonic clock used to split ``timeout`` across the two + stages, injectable for tests. Defaults to :func:`time.monotonic`. + """ + + def __init__( + self, + submit_future: Any, + *, + operation_type: str = "operation", + keepalive: Any = None, + time_fn: Callable[[], float] = time.monotonic, + ) -> None: + # No super().__init__: poll_fn/parse_fn/operation_id all arrive with + # the resolved inner future; this subclass only defers to it. + self._submit_future = submit_future + self._operation_type = operation_type + self._result: Any = _UNSET + #: The terminal failure, once one is known. Memoized beside the + #: result so a resolved future reports itself resolved either way. + self._error: Any = _UNSET + self._lock = threading.Lock() + self._time_fn = time_fn + # Reachability only; never called. The submit closure and the inner + # future's poll_fn happen to capture their client today, which has the + # same effect by accident; this makes the guarantee explicit so a later + # refactor of those closures cannot silently remove it. + self._keepalive = keepalive + + @property + def operation_id(self) -> str: + """Wait for the submit to complete, then return the operation id.""" + return self._resolve_submit().operation_id + + def _operation_id_if_known(self) -> str: + """The operation id if the submit already landed, else a placeholder. + + For messages built while a wait is timing out: reading + :attr:`operation_id` there would resume the unbounded wait the caller + asked to bound. + """ + if self._submit_future.done(): + try: + return self._submit_future.result().operation_id + except Exception: # noqa: BLE001 - a failed submit has no id to show + return "" + return "" + + def _resolve_submit(self, timeout: Optional[float] = DEFAULT_TIMEOUT) -> APIFuture[_T]: + """Block until the background submit yields the real future. + + Normalizes the wait timeout to a built-in :class:`TimeoutError`: + ``concurrent.futures.TimeoutError`` is only an alias for it from Python + 3.11 on, and callers should not have to know which stage timed out to + know what to catch. + """ + try: + return self._submit_future.result(timeout) + except concurrent.futures.TimeoutError: + raise TimeoutError( + f"{self._operation_type} operation was not submitted within " + f"{_format_timeout(timeout)}" + ) from None + + def _memoize(self, value: Any) -> Any: + """Record ``value`` as this future's result unless one is already set. + + The inner :class:`APIFuture` already parses exactly once, so concurrent + waiters here arrive with the same object; this only keeps them from + disagreeing about which one is *the* result. + """ + with self._lock: + if self._result is _UNSET: + self._result = value + return self._result + + def _memoize_error(self, error: BaseException) -> None: + """Record a terminal failure so this future reports itself resolved. + + The inner future memoizes its own, but a caller holds *this* one, so it + has to carry the same answer. + """ + with self._lock: + if self._error is _UNSET: + self._error = error + + def result(self, timeout: Optional[float] = DEFAULT_TIMEOUT) -> _T: + """Block until submit + operation complete; return the parsed result. + + ``timeout`` bounds the whole wait: one deadline is taken here and the + poll loop gets only what the submit wait left over. ``None`` waits + indefinitely. + """ + if self._error is not _UNSET: + raise self._error + if self._result is not _UNSET: + return self._result + deadline = _deadline_from(timeout, self._time_fn) + inner = self._resolve_submit(timeout=_remaining(deadline, self._time_fn)) + try: + value = inner.result(timeout=_remaining(deadline, self._time_fn)) + except TimeoutError: + # The poll loop only saw what the submit wait left over; report the + # budget the caller actually passed. + raise self._whole_wait_timeout(inner, timeout) from None + except TrainingSessionError as error: + self._memoize_error(error) + raise + return self._memoize(value) + + async def result_async(self, timeout: Optional[float] = DEFAULT_TIMEOUT) -> _T: + """Await submit + operation without occupying a thread while they run. + + Splits ``timeout`` across the two stages exactly as :meth:`result` does. + Neither stage blocks the event loop: the submit is awaited on it, and + the poll loop it hands back is a coroutine. + """ + if self._error is not _UNSET: + raise self._error + if self._result is not _UNSET: + return self._result + deadline = _deadline_from(timeout, self._time_fn) + inner = await self._resolve_submit_async(timeout=_remaining(deadline, self._time_fn)) + try: + value = await inner.result_async(timeout=_remaining(deadline, self._time_fn)) + except TimeoutError: + raise self._whole_wait_timeout(inner, timeout) from None + except TrainingSessionError as error: + self._memoize_error(error) + raise + return self._memoize(value) + + async def _resolve_submit_async( + self, timeout: Optional[float] = DEFAULT_TIMEOUT + ) -> APIFuture[_T]: + """Await the background submit and return the real future. + + Normalizes the wait timeout to a built-in :class:`TimeoutError` for the + same reason :meth:`_resolve_submit` does: ``asyncio.TimeoutError`` is + only an alias for it from Python 3.11 on. + """ + try: + return await asyncio.wait_for( + asyncio.shield(asyncio.wrap_future(self._submit_future)), timeout + ) + except asyncio.TimeoutError: + raise TimeoutError( + f"{self._operation_type} operation was not submitted within " + f"{_format_timeout(timeout)}" + ) from None + + def _whole_wait_timeout(self, inner: APIFuture[_T], timeout: Optional[float]) -> TimeoutError: + """Report the budget the caller passed, not the slice a stage got.""" + return TimeoutError( + f"{self._operation_type} operation {inner.operation_id} " + f"did not complete within {_format_timeout(timeout)}" + ) + + +class CombinedFuture(APIFuture[_T]): + """A future over several sub-futures whose results combine into one. + + Used when a single logical operation is split into multiple runtime ops + (e.g. a large forward-backward batch chunked to fit the request-size limit). + :meth:`result` blocks on every sub-future in submission order and folds + their results through ``combine_fn``, so callers see the same single result + object they would get from an unchunked call. + + An :class:`APIFuture` subclass so that chunking stays invisible in the public + signature: a chunked call is annotated ``APIFuture[T]`` exactly as an + unchunked one is. :attr:`operation_id` and both waits are overridden; + ``__await__`` is inherited and routes through this :meth:`result_async`. + + The sub-futures were submitted in sequence order, so awaiting them in that + same order matches how the service processes them; the first terminal + failure among them propagates from :meth:`result`. + + Args: + futures: Sub-futures to await, in submission (sequence) order. + combine_fn: Callable folding the ordered list of sub-results into the + single result object returned by :meth:`result`. + time_fn: Monotonic clock used to share ``timeout`` across the + sub-futures, injectable for tests. Defaults to + :func:`time.monotonic`. + """ + + def __init__( + self, + futures: List[APIFuture[_T]], + combine_fn: Callable[[List[_T]], _T], + *, + time_fn: Callable[[], float] = time.monotonic, + ) -> None: + # No super().__init__: there is no single operation to poll here, only + # sub-futures to fold. Same reason PendingAPIFuture skips it. + if not futures: + raise ValueError("CombinedFuture requires at least one sub-future") + self._futures = futures + self._combine_fn = combine_fn + self._result: Any = _UNSET + #: The terminal failure, once one is known. Memoized beside the + #: result so a resolved future reports itself resolved either way. + self._error: Any = _UNSET + self._lock = threading.Lock() + self._time_fn = time_fn + + @property + def operation_id(self) -> str: + """Return the first sub-operation's id (representative, for messages).""" + return self._futures[0].operation_id + + def _operation_id_if_known(self) -> str: + """The first sub-operation's id, or a placeholder if it is not resolved. + + Reading :attr:`operation_id` waits for the submit that produced it, with + no bound. That is fine when a caller asks for it, and wrong inside a + timeout message: the wait the caller asked to bound would resume in the + error path. + """ + try: + return self._futures[0]._operation_id_if_known() + except AttributeError: + return self._futures[0].operation_id + + def result(self, timeout: Optional[float] = DEFAULT_TIMEOUT) -> _T: + """Block until every sub-future completes and return the combined result. + + ``timeout`` bounds the whole wait, not each sub-future: one deadline is + taken here and every sub-future gets only the time still left, so a + chunked call is bounded the same way an unchunked one is. ``None`` waits + indefinitely. + """ + if self._error is not _UNSET: + raise self._error + if self._result is not _UNSET: + return self._result + deadline = _deadline_from(timeout, self._time_fn) + try: + sub_results = [ + f.result(timeout=_remaining(deadline, self._time_fn)) for f in self._futures + ] + except TimeoutError: + raise self._chunked_timeout(timeout) from None + except TrainingSessionError as error: + self._memoize_error(error) + raise + return self._combine_and_memoize(sub_results) + + async def result_async(self, timeout: Optional[float] = DEFAULT_TIMEOUT) -> _T: + """Await every sub-future and return the combined result. + + Awaits them in submission order, the same order :meth:`result` waits in, + so a chunked call reports the same first failure either way. ``timeout`` + bounds the whole wait, not each sub-future. + """ + if self._error is not _UNSET: + raise self._error + if self._result is not _UNSET: + return self._result + deadline = _deadline_from(timeout, self._time_fn) + sub_results = [] + try: + for f in self._futures: + sub_results.append( + await f.result_async(timeout=_remaining(deadline, self._time_fn)) + ) + except TimeoutError: + raise self._chunked_timeout(timeout) from None + except TrainingSessionError as error: + self._memoize_error(error) + raise + return self._combine_and_memoize(sub_results) + + def _memoize_error(self, error: BaseException) -> None: + """Record the first sub-future failure; see the base class.""" + with self._lock: + if self._error is _UNSET: + self._error = error + + def _combine_and_memoize(self, sub_results: List[Any]) -> Any: + """Fold the sub-results into this future's result, at most once. + + ``combine_fn`` folds already-parsed values and touches no wire, so it is + cheap to hold the lock across and gets the same once-only guarantee + ``parse_fn`` has. + """ + with self._lock: + if self._result is _UNSET: + self._result = self._combine_fn(sub_results) + return self._result + + def _chunked_timeout(self, timeout: Optional[float]) -> TimeoutError: + """Each sub-future saw only its slice of the budget; report the whole.""" + return TimeoutError( + f"chunked operation {self._operation_id_if_known()} did not complete " + f"within {_format_timeout(timeout)} " + f"({len(self._futures)} sub-operations)" + ) diff --git a/sagemaker-train/src/sagemaker/train/training_session/sampling_client.py b/sagemaker-train/src/sagemaker/train/training_session/sampling_client.py new file mode 100644 index 0000000000..fc1511d954 --- /dev/null +++ b/sagemaker-train/src/sagemaker/train/training_session/sampling_client.py @@ -0,0 +1,476 @@ +# 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. +"""SamplingClient: sample from a saved checkpoint (or base model). + +Created via :meth:`TrainingClient.create_sampling_client`, +:meth:`TrainingClient.save_weights_and_get_sampling_client`, or +:meth:`ServiceClient.create_sampling_client` (sampling without training), not +directly by callers. Only the last of those creates the backing session, and so +only it stops that session when used as a context manager; the other two share +the TrainingClient's session handle, which is what marks them non-owning. +Turns a prompt + :class:`SamplingParams` into a sampling operation and returns +an :class:`APIFuture` over the +:class:`SampleResponse`. Submits are non-blocking: the wire call runs on a +background submitter thread, and sample ops are order-independent, so many +requests can be queued concurrently in seconds. + +``model_package_arn`` selects which checkpoint to sample from; ``""`` means the +base model. It is sent on the wire as ``modelPackageArn``, or omitted entirely +for the base model. +""" + +from __future__ import annotations + +import logging +from typing import Any, Callable, Dict, List, Optional, TypeVar + +from sagemaker.train.training_session._utils import ( + model_input_token_count, + prompt_to_model_input, +) +from sagemaker.train.training_session.exceptions import ( + ResultMismatchError, + TrainingSessionError, +) +from sagemaker.train.training_session._session_handle import _SessionHandle +from sagemaker.train.training_session._submitter import ( + _CLOSED_MESSAGE, + BackgroundSubmitter, + run_blocking, +) +from sagemaker.train.training_session.futures import ( + APIFuture, + PendingAPIFuture, + call_with_transient_retry, +) +from sagemaker.train.training_session.types import ( + SampleResponse, + SamplingParams, +) + +#: Sampling from weights that were just saved can briefly return a transient +#: error while those weights are still being made servable. Sampling carries no +#: clientToken and is side-effect-free, so the submit is safe to retry. +#: Budget = 150 attempts over the backoff, about 5 min of warmup tolerance. +_SAMPLE_SUBMIT_MAX_ATTEMPTS = 150 + +#: Result type of a submitted sampling op, fixed by the ``parse_fn`` passed to +#: :meth:`SamplingClient._submit_sample_op`. +_T = TypeVar("_T") + +logger = logging.getLogger(__name__) + + +def _parse_sample_result(raw: Dict[str, Any]) -> SampleResponse: + """Extract a :class:`SampleResponse` from a sampling-operation result body. + + The payload is returned under ``operationResult.sample``. + """ + operation_result = raw.get("operationResult") or {} + return SampleResponse._from_api(operation_result.get("sample") or {}) + + +def _make_compute_logprobs_parser( + expected_tokens: Optional[int], +) -> Callable[[Dict[str, Any]], List[Optional[float]]]: + """Build the parser for a compute-logprobs result of a known token count. + + The payload is returned under ``operationResult.computeLogprobs`` (the + compute-logprobs result variant). Returns the bare logprob list, matching + what interactive-training loops expect from ``compute_logprobs``, rather + than a wrapper type. + + The count is bound here so the parser can reject a result belonging to a + different request; see :class:`ResultMismatchError`. ``None`` (a prompt the + service tokenizes) leaves the result unchecked. + """ + + def parse(raw: Dict[str, Any]) -> List[Optional[float]]: + operation_result = raw.get("operationResult") or {} + logprobs = (operation_result.get("computeLogprobs") or {}).get("logprobs", []) + if expected_tokens is not None and len(logprobs) != expected_tokens: + raise ResultMismatchError("computeLogprobs", expected_tokens, len(logprobs)) + return logprobs + + return parse + + +class SamplingClient: + """Samples from a checkpoint; ``model_package_arn=""`` = base model. + + Args: + training_session_name: Name of the backing TrainingSession resource. + runtime: Client for the training-session runtime. + model_package_arn: Model package ARN of the checkpoint to sample from; + ``""`` = base model. + base_model: Base model (HF model ID or SageMaker Hub ARN), used to + resolve a tokenizer via :meth:`get_tokenizer`. + tokenizer: Optional pre-loaded tokenizer for string prompts. + hub_client: Optional SageMaker Hub client, used to look up the + HuggingFace model id behind a Hub ARN when a tokenizer is needed. + Only consulted when neither ``tokenizer`` nor ``hf_model_id`` is set, + and borrowed from ``service_client`` at that point if omitted here. + hf_model_id: Optional HuggingFace model id to load the tokenizer from. + Set it when ``base_model`` is an ARN whose id cannot be looked up, or + to pin the tokenizer to a specific repo. + service_client: The :class:`ServiceClient` that created this client, + used to stop the session when this client owns it. + session_handle: Existing :class:`_SessionHandle` to share, passed by + :meth:`TrainingClient.create_sampling_client` so this client and its + TrainingClient refcount one session. When omitted a fresh handle is + created for ``training_session_name``. + submitter: Existing :class:`BackgroundSubmitter` to share, passed by + :meth:`TrainingClient.create_sampling_client` so refreshing a sampler + does not add an event loop and thread each time. When omitted this + client builds and owns its own. + + Whether this client *owns* its session is read off ``session_handle`` rather + than passed separately: sharing a handle means a TrainingClient is driving + the session, so this client must not stop it; a fresh handle means the + session was created implicitly for this client alone and has no other owner. + Ownership governs :meth:`stop` and ``__exit__`` only; the automatic stop is + the handle's job either way. Submitter ownership is read the same way but + tracked separately: a client can share a handle and still own its submitter. + """ + + def __init__( + self, + training_session_name: str, + runtime: Any, + *, + model_package_arn: str = "", + base_model: str = "", + tokenizer: Any = None, + hub_client: Any = None, + hf_model_id: Optional[str] = None, + service_client: Any = None, + session_handle: Optional[_SessionHandle] = None, + submitter: Optional[BackgroundSubmitter] = None, + ) -> None: + self.training_session_name = training_session_name + self.model_package_arn = model_package_arn + self._runtime = runtime + self._base_model = base_model + self._tokenizer = tokenizer + self._hub_client = hub_client + self._hf_model_id = hf_model_id + self._service_client = service_client + # Ownership IS "did I have to make my own handle?": a shared handle + # means a TrainingClient is driving the same session, so the session + # outlives whichever of the two is collected first and is stopped when + # the last one goes. + self._owns_session = session_handle is None + # Sample ops are order-independent: submits run concurrently on the + # background submitter and return their future immediately. A submitter + # is an event loop on a daemon thread, so sharing one keeps a training + # loop's sampler refresh down to a plain object allocation. + self._owns_submitter = submitter is None + self._submitter = submitter or BackgroundSubmitter( + name=f"sagemaker-sampling-submit-{training_session_name}" + ) + self._handle = session_handle or _SessionHandle(training_session_name, service_client) + if self._owns_submitter: + # A shared submitter is already registered by its owner, and the + # handle keeps every submitter it is given, so registering again + # would pile up references to the one submitter and drain it once + # per entry. + self._handle.register_submitter(self._submitter) + # The TrainingClient keeps submitting on a shared submitter after this + # client is done with it, so retirement is tracked here and not on the + # submitter. + self._closed = False + + def get_tokenizer(self) -> Any: + """Return a HuggingFace tokenizer for the session's base model. + + The runtime accepts only token IDs, so callers with text must tokenize + client-side before :meth:`sample`. The tokenizer is loaded once and + cached (and reused for string prompts passed to :meth:`sample`). + + Returns: + A ``transformers.PreTrainedTokenizer``. + + Raises: + ValueError: If the base model has no tokenizer available (e.g. a + model whose tokenizer is not published to HuggingFace); pass + pre-tokenized token IDs in that case. + ImportError: If ``transformers`` is not installed (install the + optional extra: ``pip install sagemaker-train[tokenizer]``). + """ + if self._tokenizer is None: + from sagemaker.train.training_session._utils import ( + load_tokenizer, + ) + + # The Hub client is only needed to resolve the tokenizer id when no + # hf_model_id was supplied; skip building it otherwise. + hub_client = self._hub_client if self._hf_model_id else self._get_hub_client() + self._tokenizer = load_tokenizer(self._base_model, hub_client, self._hf_model_id) + return self._tokenizer + + def _get_hub_client(self) -> Any: + """Return the Hub client, building it lazily via the service client. + + Usually ``None`` at construction and built here on first use, only when + a base-model ARN's tokenizer id must be resolved and no ``hf_model_id`` + was supplied. Deferring it keeps sampling by token IDs working where + sagemaker-core is not installed. + """ + if self._hub_client is None and self._service_client is not None: + self._hub_client = self._service_client._get_hub_client() + return self._hub_client + + # ------------------------------------------------------------------ + # Sampling + # ------------------------------------------------------------------ + + def sample( + self, + prompt: Any, + num_samples: int = 1, + sampling_params: "Optional[SamplingParams | Dict[str, Any] | Any]" = None, + *, + include_prompt_logprobs: bool = False, + topk_prompt_logprobs: Optional[int] = None, + ) -> APIFuture[SampleResponse]: + """Submit a sampling operation; returns an :class:`APIFuture`. + + Args: + prompt: A list of token IDs, a string (requires a tokenizer), or a + pre-built ModelInput dict. + num_samples: Number of sequences to generate. + sampling_params: Generation parameters. Uses defaults if omitted. + Accepts a :class:`SamplingParams`, a dict of the same fields, + or any object carrying at least one of them, so an existing + config object can be passed as-is. + include_prompt_logprobs: Request per-token logprobs over the prompt. + topk_prompt_logprobs: If set, request the top-k logprobs per prompt + position. + + Returns: + An :class:`APIFuture` resolving to a :class:`SampleResponse`. + Call ``.result()`` on it (or ``await`` it) to get the response. The + submit itself runs on a background thread: this returns immediately, many + requests can be queued concurrently, and a submit error surfaces + when the future is resolved, not here. + """ + # from_any: accept SamplingParams, a dict, or any object with the same attributes. + params = ( + SamplingParams.from_any(sampling_params) + if sampling_params is not None + else SamplingParams() + ) + model_input = prompt_to_model_input( + prompt, tokenizer=self._tokenizer, tokenizer_loader=self.get_tokenizer + ) + + # The sampling operation takes an operationConfig union; we send the + # ``sample`` variant. + sample_config: dict[str, Any] = { + "prompt": model_input, + "numSamples": num_samples, + "samplingParams": params._to_api(), + } + # Select which checkpoint to sample from; omit the field to sample the + # base model. + if self.model_package_arn: + sample_config["modelPackageArn"] = self.model_package_arn + if include_prompt_logprobs: + sample_config["includePromptLogprobs"] = True + if topk_prompt_logprobs is not None: + sample_config["topkPromptLogprobs"] = topk_prompt_logprobs + + return self._submit_sample_op({"sample": sample_config}, _parse_sample_result) + + async def sample_async( + self, + prompt: Any, + num_samples: int = 1, + sampling_params: "Optional[SamplingParams | Dict[str, Any] | Any]" = None, + *, + include_prompt_logprobs: bool = False, + topk_prompt_logprobs: Optional[int] = None, + ) -> SampleResponse: + """Coroutine twin of :meth:`sample`; awaits the response directly. + + ``resp = await sc.sample_async(...)`` resolves to the + :class:`SampleResponse`. Sample ops are order-independent, so unlike + the training ``*_async`` methods there is no sequence position to hold + and therefore no intermediate future to return. + """ + return await self.sample( + prompt, + num_samples=num_samples, + sampling_params=sampling_params, + include_prompt_logprobs=include_prompt_logprobs, + topk_prompt_logprobs=topk_prompt_logprobs, + ) + + def compute_logprobs(self, prompt: Any) -> APIFuture[List[Optional[float]]]: + """Submit a compute-logprobs operation; returns an :class:`APIFuture`. + + Unlike :meth:`sample`, this generates no new tokens; it scores the given + prompt and returns one log probability per prompt token. Useful for + evaluation, reranking, or importance weighting. Non-blocking, like + :meth:`sample`. + + Args: + prompt: A list of token IDs, a string (requires a tokenizer), or a + pre-built ModelInput dict. + + Returns: + An :class:`APIFuture` resolving to a list with one logprob per + prompt token (the first entry is ``None``, since it has no preceding context). + + Raises: + ResultMismatchError: On resolution, when the result's length does + not match the sequence sent. + """ + model_input = prompt_to_model_input( + prompt, tokenizer=self._tokenizer, tokenizer_loader=self.get_tokenizer + ) + + # The sampling operation's operationConfig union, computeLogprobs + # variant. Select the checkpoint the same way as sample: omit + # modelPackageArn to score against the base model. + logprobs_config: dict[str, Any] = {"prompt": model_input} + if self.model_package_arn: + logprobs_config["modelPackageArn"] = self.model_package_arn + return self._submit_sample_op( + {"computeLogprobs": logprobs_config}, + _make_compute_logprobs_parser(model_input_token_count(model_input)), + ) + + async def compute_logprobs_async(self, prompt: Any) -> List[Optional[float]]: + """Coroutine twin of :meth:`compute_logprobs`; awaits the logprobs directly.""" + return await self.compute_logprobs(prompt) + + def _submit_sample_op( + self, + operation_config: Dict[str, Any], + parse_fn: Callable[[Dict[str, Any]], _T], + ) -> APIFuture[_T]: + """Queue a sample-path op; immediately return a future over its result. + + The sampling operation takes no clientToken, unlike the training op: + sampling is side-effect-free, so the service does not dedupe it. The + submit retries through transient unavailability: right after + save_weights_for_sampler the service may still be binding to the new + checkpoint. Safe to retry: no clientToken, side-effect-free. + """ + if self._closed: + raise TrainingSessionError(_CLOSED_MESSAGE) + + def _wire_submit() -> APIFuture[_T]: + resp = call_with_transient_retry( + lambda: self._runtime.execute_sampling_operation( + trainingSessionName=self.training_session_name, + operationConfig=operation_config, + ), + max_attempts=_SAMPLE_SUBMIT_MAX_ATTEMPTS, + ) + operation_id = resp["operationId"] + return APIFuture( + operation_id=operation_id, + poll_fn=lambda: self._runtime.get_sampling_operation_result( + trainingSessionName=self.training_session_name, + operationId=operation_id, + ), + parse_fn=parse_fn, + status_key="status", + operation_type="sample", + ) + + return PendingAPIFuture( + self._submitter.submit_unordered(run_blocking(_wire_submit)), + operation_type="sample", + keepalive=self._handle, + ) + + # ------------------------------------------------------------------ + # Session lifecycle + # ------------------------------------------------------------------ + + def _retire(self) -> None: + """Refuse further work, and release the background submitter if it is ours. + + A shared submitter belongs to the :class:`TrainingClient` that handed + this client out, so closing it would stop that client submitting and is + left to its teardown. This client still stops accepting operations, so + retirement fails fast here either way. + """ + self._closed = True + if self._owns_submitter: + self._submitter.close() + + def stop(self) -> None: + """Stop this client's training session. + + Only meaningful for a client created by + :meth:`ServiceClient.create_sampling_client`, which creates a session + implicitly and therefore owns it. For a client created by + :meth:`TrainingClient.create_sampling_client` the session belongs to the + TrainingClient, so this retires only this client; stop that session + through its :class:`TrainingClient` instead. + + Drains the background submitter first (best effort) when this client owns + it, so every operation already queued reaches the service before the stop + request. The client accepts no new operations afterwards. + """ + self._retire() + if not self._owns_session: + # Shared session: the handle belongs to the TrainingClient driving + # it, so do NOT disarm: that would cancel the automatic stop for a + # session this client has no say over, turning a no-op into a leak. + return + if self._service_client is None: + # Nothing was stopped, so leave the automatic stop armed. + logger.warning( + "SamplingClient owns session %r but has no ServiceClient to " + "stop it with; stop it via ServiceClient.stop_session(...).", + self.training_session_name, + ) + return + try: + self._service_client.stop_session(self.training_session_name) + except Exception: + # Keep the automatic stop armed: it is the only remaining + # client-side attempt at releasing the session. + self._handle.note_explicit_stop_failed() + raise + # Stopped, so a later garbage collection would only repeat the call. + self._handle.disarm() + + def __enter__(self) -> "SamplingClient": + """Return this client, so it can be used as a context manager.""" + return self + + def __exit__(self, exc_type: Any, exc_val: Any, exc_tb: Any) -> None: + """Stop the session on leaving the block, if this client created it.""" + # Drain queued submits so in-flight requests reach the service, then + # stop the session iff this client created it. A client handed out by + # TrainingClient.create_sampling_client shares the TrainingClient's + # session and must not stop it, since that session outlives this block. + if self._owns_session: + try: + self.stop() + except Exception: + logger.warning( + "Failed to stop implicitly-created session %r on exit; it " + "will be stopped when this client is collected, or after " + "its inactivity timeout.", + self.training_session_name, + exc_info=True, + ) + else: + self._retire() diff --git a/sagemaker-train/src/sagemaker/train/training_session/service_client.py b/sagemaker-train/src/sagemaker/train/training_session/service_client.py new file mode 100644 index 0000000000..189774459b --- /dev/null +++ b/sagemaker-train/src/sagemaker/train/training_session/service_client.py @@ -0,0 +1,1313 @@ +# 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. +"""ServiceClient: entry point for creating and managing training sessions. + +Holds session-level defaults (role, MPG, S3 output path) so that child-client +creation is concise. Use :meth:`ServiceClient.create_lora_training_client` to +start a training session and get back a :class:`TrainingClient`, or +:meth:`ServiceClient.create_sampling_client` to sample without training. +""" + +from __future__ import annotations + +import asyncio +import logging +import os +import time +import uuid +from typing import TYPE_CHECKING, Any, Dict, List, Optional + +import boto3 +from botocore.config import Config +from botocore.exceptions import ClientError + +from sagemaker.train.training_session._endpoints import ( + _SAGEMAKER_ENDPOINT_ENV, + _resolve_region, +) + +if TYPE_CHECKING: + from sagemaker.core.helper.session_helper import Session +from sagemaker.train.training_session.exceptions import ArtifactDeletionError +from sagemaker.train.training_session._utils import ( + DEFAULT_LATEST_SCAN_LIMIT, + DEFAULT_LIST_MODEL_PACKAGES_LIMIT, + MANAGED_STORAGE_RESTRICTED, + MAX_RESULTS_PER_PAGE, + base_model_arn_from_model_package, + build_lora_config, + delete_s3_objects_under, + describe_model_package_artifacts, + describe_model_package_kind, + list_base_models_in_region, + list_model_package_summaries, + resolve_base_model, + _warn_unsupported_args, + _RETRY_CONFIG_CONSEQUENCE, +) +from sagemaker.train.training_session.sampling_client import SamplingClient +from sagemaker.train.training_session.training_client import TrainingClient +from sagemaker.train.training_session.types import ( + DEFAULT_LORA_ALPHA, + DEFAULT_LORA_RANK, + ModelPackageDeletion, + ModelPackageInfo, + ModelPackageKind, + SessionStatus, + TrainingSessionSummary, + VpcConfig, +) + +logger = logging.getLogger(__name__) + +_RUNNING_STATUS = "Running" + +_TERMINAL_STATUSES = frozenset({"Failed", "Stopped", "Stopping", "Blocked"}) +#: States from which the service accepts a DeleteTrainingSession call. +_DELETABLE_STATUSES = frozenset({"Stopped", "Failed", "DeletionFailed"}) + +#: States in which a StopTrainingSession call has nothing left to do: the session +#: is already on its way down, already terminal, or being removed. The service +#: rejects the state transition instead of treating the call as a no-op, so +#: :meth:`ServiceClient.stop_session` absorbs it. +_STOP_IS_NOOP_STATUSES = frozenset({"Stopping", "Stopped", "Failed", "Deleting", "DeletionFailed"}) + +#: Name of the model package group used when the caller supplies none. One shared +#: group per account/region, so kind-filtered discovery has to look past packages +#: it did not write; see :data:`DEFAULT_LATEST_SCAN_LIMIT`. +DEFAULT_MODEL_PACKAGE_GROUP_NAME = "sagemaker-default" + +#: Prefix for auto-generated session names, when no ``session_name_prefix`` is set. +DEFAULT_SESSION_NAME_PREFIX = "sagemaker-training-session" + + +def _session_name_of(session: str) -> str: + """Accept either a session name or its ARN and return the name. + + ``list_sessions`` hands back ``TrainingSessionArn`` while the stop/delete + APIs take a name, so callers otherwise have to split the ARN themselves. A + session ARN ends ``.../training-session/``; the name is the final + segment. A bare name has no ``/`` and is returned unchanged. + """ + return session.rsplit("/", 1)[-1] if session.startswith("arn:") else session + + +def _session_and_region(region: Optional[str]) -> "tuple[Session, str]": + """Return the session every call goes through, and the Region it runs in. + + One session serves this client: its clients carry the configuration every + other sagemaker-core caller gets, and it answers for the default output + bucket. It is pinned to the resolved Region so that the Region a session + runs in is the Region its output is written to. + + Args: + region: Region to run in, resolved from the environment when not given. + """ + from sagemaker.core.helper.session_helper import Session as _Session + from sagemaker.train.defaults import TrainDefaults + + resolved = _resolve_region(region) + session = TrainDefaults.get_sagemaker_session( + sagemaker_session=_Session(boto_session=boto3.Session(region_name=resolved)) + ) + return session, resolved + + +def _stop_already_done(error: ClientError) -> bool: + """True when a StopTrainingSession failure only means "nothing left to stop". + + The control plane validates the session state transition and answers a stop + it cannot apply (``Stopping -> Stopping``, ``Stopped -> Stopping``) with a + ``ValidationException`` instead of ignoring it. Reaching that is normal, not a + failure: two teardown paths can legitimately target the same session, an + explicit stop and the automatic stop that runs when the last client + referencing the session is collected, and they can run in either order. + + Matching is on the state named in the message, not the message as a whole, so + it holds if the service rewords the error. + """ + err = error.response.get("Error") or {} + if err.get("Code") != "ValidationException": + return False + message = err.get("Message") or "" + return any( + f"invalid session state transition from {status} to " in message + for status in _STOP_IS_NOOP_STATUSES + ) + + +#: Not-found phrases that mean a model package is already gone. A +#: ``ResourceNotFound`` code says so on its own; ``ValidationException`` also +#: covers real problems (a malformed ARN, a package still referenced), so it +#: counts as absent only when the message says the package is not there. +_ABSENT_MODEL_PACKAGE_PHRASES = ("does not exist", "not found", "cannot be found") + + +def _model_package_absent(error: ClientError) -> bool: + """True when a model package Describe/Delete failed only because it is gone. + + Lets :meth:`ServiceClient.delete_model_package` be idempotent: deleting a + package twice, or racing another caller doing the same, reports "nothing + left to delete" instead of raising. + """ + err = error.response.get("Error") or {} + code = err.get("Code") + if code in ("ResourceNotFound", "ResourceNotFoundException"): + return True + if code != "ValidationException": + return False + message = (err.get("Message") or "").lower() + return any(phrase in message for phrase in _ABSENT_MODEL_PACKAGE_PHRASES) + + +def _model_package_group_exists(error: Exception) -> bool: + """True when creating a model package group failed only because it is there. + + Lets the default group be created on demand without racing another caller + doing the same. ``ResourceInUse`` is the modelled answer, so it is checked + first; the message is only consulted when the code does not settle it. + + Reads the response defensively: the registry client's error classes are its + own, and one may carry the reason in its string alone. + """ + err = getattr(error, "response", None) or {} + if (err.get("Error") or {}).get("Code") in ("ResourceInUse", "ResourceInUseException"): + return True + message = (err.get("Error") or {}).get("Message") or str(error) + return "already exists" in message.lower() + + +def _build_create_session_args( + *, + session_name: str, + base_model_arn: str, + model_package_group_arn: str, + lora_config: Dict[str, Any], + role_arn: str, + s3_path: str, + kms_key_arn: Optional[str], + vpc_config: Optional[VpcConfig], +) -> Dict[str, Any]: + """Build a CreateTrainingSession request. + + Shared by both create paths, so the training and sampling-only requests + cannot drift apart. Model identity goes in ``ModelConfig`` and the output + location in ``OutputConfig``; both are required by the API. + """ + s3_location: Dict[str, Any] = {"Path": s3_path} + if kms_key_arn: + s3_location["KmsKeyArn"] = kms_key_arn + + args: Dict[str, Any] = { + "TrainingSessionName": session_name, + "ModelPackageGroupArn": model_package_group_arn, + "LoraConfig": lora_config, + "RoleArn": role_arn, + "ModelConfig": {"BaseModelArn": base_model_arn}, + "OutputConfig": {"S3Location": s3_location}, + } + if vpc_config is not None: + args["VpcConfig"] = vpc_config._to_api() + return args + + +def _warn_unsupported(**passed: Any) -> None: + """Warn for each accepted-but-unhonoured session-creation keyword argument. + + Thin wrapper over :func:`_warn_unsupported_args` carrying this module's + consequence clause. + """ + _warn_unsupported_args("the session was created without it", **passed) + + +def _warm_hub_imports() -> None: + """Import the sagemaker-core / boto3 modules the Hub client needs, once. + + Called from :meth:`ServiceClient.__init__`, not left to + :meth:`ServiceClient._get_hub_client`, and the reason is subtle: CPython + retains the frame chain of a first-time import, which pins the frame that + triggered it, and therefore whatever that frame returns. Performed on + first use inside ``create_sampling_client`` or ``create_lora_training_client`` + (which call ``_get_hub_client`` to resolve a tokenizer or a base-model + name), it pins the client those factories return, so the client is never + collected and its session is not auto-stopped until interpreter exit. + + Warming here moves the pin onto this call's frame, which holds a + ``ServiceClient`` that lives for the process anyway. Only the first + ``ServiceClient`` in a process pays it; later ones hit the already-imported + fast path and retain nothing. + + Guarded, so sagemaker-core stays an optional dependency, and this keeps + the *package* import cheap, which is what ``_get_hub_client``'s laziness was + protecting. Only constructing a ``ServiceClient`` pays the import cost. + """ + try: + import boto3 # noqa: F401 + import sagemaker.core.helper.session_helper # noqa: F401 + except Exception: # pragma: no cover - optional dependency + logger.debug("sagemaker-core not importable; Hub lookups may fail later") + + +def _extract_describe_summary(resp: dict) -> dict: + """Extract the session info dict from a DescribeTrainingSession response.""" + if "TrainingSessionDetail" in resp: + return resp["TrainingSessionDetail"] + return {k: v for k, v in resp.items() if k != "ResponseMetadata"} + + +def _describe_failure_suffix(summary: dict) -> str: + """Return the reason a session failed, as a message suffix. + + Empty string when the service reported no reason, leaving the status to + speak for itself. + """ + reason = summary.get("ErrorDetails") or summary.get("ErrorMessage") + return f" ({reason})" if reason else "" + + +class ServiceClient: + """Entry point for creating and managing training sessions. + + ``ServiceClient`` holds session-level defaults (``role_arn``, optional + ``model_package_group_arn``, ``s3_output_path``, ``kms_key_arn``, + ``vpc_config``) so that child-client creation needs only the model + identity. + + Anything left unset gets a SageMaker-conventional default: + + * **Session name**: ``{session_name_prefix}-{8 hex}``, prefix defaulting to + :data:`DEFAULT_SESSION_NAME_PREFIX`. ``session_name_prefix`` names sessions + only; it has no effect on the model package group or the S3 path. + * **Model package group**: the account's shared + :data:`DEFAULT_MODEL_PACKAGE_GROUP_NAME`, created on first use. Pass + ``model_package_group_arn`` to keep one experiment's checkpoints in their + own group. + * **S3 output**: ``s3://{sagemaker default bucket}/{session name}/output``, + resolved through the same ``Session.default_bucket()`` the broader + SageMaker PySDK uses, so an existing ``sagemaker_config`` override applies. + + Every client this ``ServiceClient`` creates cleans up after itself: a session + is stopped once the last client referencing it becomes unreachable, so one is + not left Running by a script that exits or a notebook cell that rebinds its + client. The automatic path stops the session and leaves deletion to an + explicit call, so a stopped session stays describable and a run that went + wrong can still be inspected. Prefer a ``with`` block or an explicit + ``.stop()`` where you can: the automatic stop is a backstop, not a substitute + for scoping the session. Durability comes from a saved checkpoint, which + :meth:`create_sampling_client` takes directly. + + Example:: + + from sagemaker.train.training_session import ServiceClient + + svc = ServiceClient( + region="us-west-2", + role_arn="arn:aws:iam::123456789012:role/SageMakerTrainingSessionsRole", + model_package_group_arn="arn:aws:sagemaker:us-west-2:123456789012:model-package-group/my-group", + s3_output_path="s3://my-bucket/output", + ) + # rank and alpha are keyword-only, defaulting to DEFAULT_LORA_RANK + # and DEFAULT_LORA_ALPHA. + tc = svc.create_lora_training_client( + base_model="", + rank=8, + alpha=16, + ) + sp = svc.create_sampling_client(base_model="") + """ + + def __init__( + self, + region: Optional[str] = None, + *, + role_arn: str, + model_package_group_arn: Optional[str] = None, + s3_output_path: Optional[str] = None, + kms_key_arn: Optional[str] = None, + vpc_config: "Optional[VpcConfig | dict[str, Any]]" = None, + session_name_prefix: Optional[str] = None, + max_pool_connections: Optional[int] = None, + ) -> None: + # Imported lazily to keep ``import`` of the package cheap. + from sagemaker.train.training_session._clients import ( + _control_plane_client, + _runtime_client, + ) + + self._sagemaker_session, self._region = _session_and_region(region) + boto_session = self._sagemaker_session.boto_session + self._role_arn = role_arn + self._model_package_group_arn = model_package_group_arn + self._s3_output_path = s3_output_path + self._kms_key_arn = kms_key_arn + self._vpc_config = VpcConfig.from_any(vpc_config) if vpc_config is not None else None + self._session_name_prefix = session_name_prefix + self._sagemaker_client = _control_plane_client( + region=self._region, boto_session=boto_session + ) + runtime_kwargs: Dict[str, Any] = { + "region": self._region, + "boto_session": boto_session, + } + if max_pool_connections is not None: + runtime_kwargs["max_pool_connections"] = max_pool_connections + self._runtime = _runtime_client(**runtime_kwargs) + # Lazily-built ``sagemaker`` clients (via sagemaker-core). The Hub one + # serves model discovery and name -> ARN resolution; the model-registry + # one serves model-package and checkpoint lookups. They are separate + # because the two registries are not always reachable at the same + # endpoint; see the accessors. + self._hub_client: Optional[Any] = None + self._model_registry_client: Optional[Any] = None + # Built on first use by delete_model_package(delete_artifacts=True). + # Nothing else here touches S3, so most callers never build it. + self._s3_client: Optional[Any] = None + # Must happen here, not on first use in a factory; see the docstring. + # Without it, the first client a factory returns is never collected and + # its session is not auto-stopped until interpreter exit. + _warm_hub_imports() + + def _build_core_sagemaker_client(self, endpoint_url: Optional[str]) -> Any: + """Return the ``sagemaker`` client used for Hub and model-registry lookups. + + Normally this is the client on the + :class:`~sagemaker.core.helper.session_helper.Session` this + ``ServiceClient`` holds, so lookups carry the same configuration and + credentials as every other call it makes. + + An ``endpoint_url`` needs its own client, which is built from that same + session so it still signs with those credentials. + """ + # Imported here, not at module scope: this is the expensive import in + # the package, and only these lookups need it. + from sagemaker.core.user_agent import get_user_agent_extra_suffix + + if not endpoint_url: + return self._sagemaker_session.sagemaker_client + return self._sagemaker_session.boto_session.client( + "sagemaker", + endpoint_url=endpoint_url, + config=Config(user_agent_extra=get_user_agent_extra_suffix()), + ) + + def _get_hub_client(self) -> Any: + """Return the client used for SageMaker Public Hub lookups. + + The Public Hub is published per Region, so this client resolves through + the Region's own endpoint: a base model's latest version is whatever the + Region publishes, and reading it anywhere else can name a version the + Region does not serve. + """ + if self._hub_client is None: + self._hub_client = self._build_core_sagemaker_client(None) + return self._hub_client + + def _get_model_registry_client(self) -> Any: + """Return the client used for model-package and checkpoint lookups. + + Model packages are registered against the same endpoint the sessions + are, so this client follows that endpoint rather than the Region default. + Otherwise it would read a different registry than the one the session + wrote to. + """ + endpoint = os.environ.get(_SAGEMAKER_ENDPOINT_ENV) + if not endpoint: + # Both registries answer on the Region's own endpoint, so one client + # serves both. This is the only path a default configuration takes. + return self._get_hub_client() + if self._model_registry_client is None: + self._model_registry_client = self._build_core_sagemaker_client(endpoint) + return self._model_registry_client + + def _get_s3_client(self) -> Any: + """Return a region-pinned ``s3`` client, built on first use. + + Signs with the same credentials as the rest of this client: the + caller's, not the ``role_arn`` the sessions run under. Deleting a model + package's artifacts therefore needs ``s3:ListBucket`` and + ``s3:DeleteObject`` on that location for whoever is calling. + """ + if self._s3_client is None: + self._s3_client = self._sagemaker_session.boto_session.client("s3") + return self._s3_client + + def list_base_models(self) -> List[str]: + """List the base models available for fine-tuning in this client's region. + + Queries the SageMaker Public Hub and returns the models this client can + fine-tune end to end, which is narrower than every model the Hub tags as + fine-tunable. Results are region-specific. + + Each model is named by its own model id (``"provider/model"``), spelled as + its provider publishes it wherever the Hub carries one. Pass one straight to + :meth:`create_lora_training_client` as ``base_model``. The set changes as + models are added, so call this rather than hard-coding an id. + """ + return list_base_models_in_region(self._get_hub_client()) + + # ------------------------------------------------------------------ + # Model package discovery + # ------------------------------------------------------------------ + # These wrap ListModelPackages (newest first). The kind (checkpoint vs model) + # is not in the list summary, so the kind-specific helpers describe + # candidates lazily until one matches. + + def list_model_packages( + self, + *, + kind: Optional[ModelPackageKind] = None, + model_package_group_arn: Optional[str] = None, + max_results: Optional[int] = DEFAULT_LIST_MODEL_PACKAGES_LIMIT, + ) -> List[ModelPackageInfo]: + """List model packages, newest first. + + Each :class:`ModelPackageInfo` carries the ARN, version, creation time, + and description. Pass its ``model_package_arn`` to + :meth:`TrainingClient.load_state` or + :meth:`ServiceClient.create_sampling_client`. + + When ``kind`` is supplied, the method describes each candidate to + resolve its kind and returns only those matching the filter. Without + ``kind``, the ``kind`` field on each result is left unset (stays at one + API call regardless of group size). + + Args: + kind: Optional filter; only return packages of this kind. + model_package_group_arn: The model package group (name or ARN). + When omitted, uses the value set on this :class:`ServiceClient`. + max_results: Cap on packages returned (default 10, newest first). + Pass ``None`` to return every version in the group. + """ + resolved_group = model_package_group_arn or self._model_package_group_arn + if not resolved_group: + raise ValueError("No model_package_group_arn provided and none set on ServiceClient.") + summaries = list_model_package_summaries( + self._get_model_registry_client(), resolved_group, max_results + ) + results = [ModelPackageInfo._from_summary(s) for s in summaries] + if kind is None: + return results + filtered: List[ModelPackageInfo] = [] + for info in results: + info.kind = describe_model_package_kind( + self._get_model_registry_client(), info.model_package_arn + ) + if info.kind == kind: + filtered.append(info) + return filtered + + def get_latest_model_package( + self, + kind: ModelPackageKind, + *, + model_package_group_arn: Optional[str] = None, + scan_limit: int = DEFAULT_LATEST_SCAN_LIMIT, + ) -> Optional[ModelPackageInfo]: + """Return the latest model package of a given kind, or ``None``. + + Lists newest first and describes candidates in order, returning the + first whose kind matches. Use :attr:`ModelPackageKind.CHECKPOINT` for + resumable state (``save_state``) or :attr:`ModelPackageKind.MODEL` for + deployable weights (``save_weights_for_sampler``). + + Args: + kind: The kind of package to search for. + model_package_group_arn: The model package group. When omitted, + uses the value set on this :class:`ServiceClient`. + scan_limit: Most-recent packages to describe while searching + (default :data:`DEFAULT_LATEST_SCAN_LIMIT`). Exhausting it + without a match logs a warning: ``None`` from a truncated search + does not mean no such package exists. + """ + resolved_group = model_package_group_arn or self._model_package_group_arn + if not resolved_group: + raise ValueError("No model_package_group_arn provided and none set on ServiceClient.") + return self._latest_of_kind(resolved_group, kind, scan_limit) + + def _latest_of_kind( + self, model_package_group_arn: str, kind: ModelPackageKind, scan_limit: int + ) -> Optional[ModelPackageInfo]: + """Newest model package matching ``kind``, describing at most ``scan_limit``.""" + summaries = list_model_package_summaries( + self._get_model_registry_client(), + model_package_group_arn, + max_results=scan_limit, + ) + for summary in summaries: + ckpt = ModelPackageInfo._from_summary(summary) + ckpt.kind = describe_model_package_kind( + self._get_model_registry_client(), ckpt.model_package_arn + ) + if ckpt.kind == kind: + return ckpt + if len(summaries) >= scan_limit: + # Exhausted the budget without a match, so a matching package may + # well exist just past it. WARNING, not INFO: the caller is about to + # read None as "nothing to resume from" and start over from the base + # model, and that is not a conclusion to reach from an INFO line. + logger.warning( + "No %s package found in the %d most recent versions of group %s, " + "but the group has more; scan_limit stopped the search, so this " + "is not proof that none exists. Retry with a larger scan_limit, or " + "pass a group dedicated to this run.", + kind.value, + len(summaries), + model_package_group_arn, + ) + elif summaries: + logger.info( + "No %s package among the %d versions of group %s.", + kind.value, + len(summaries), + model_package_group_arn, + ) + return None + + def create_lora_training_client( + self, + base_model: str, + *, + rank: int = DEFAULT_LORA_RANK, + alpha: int = DEFAULT_LORA_ALPHA, + training_session_name: Optional[str] = None, + model_package_group_arn: Optional[str] = None, + s3_output_path: Optional[str] = None, + kms_key_arn: Optional[str] = None, + vpc_config: "Optional[VpcConfig | dict[str, Any]]" = None, + seed: Optional[int] = None, + train_mlp: Optional[bool] = None, + train_attn: Optional[bool] = None, + train_unembed: Optional[bool] = None, + user_metadata: Optional[Dict[str, str]] = None, + ) -> TrainingClient: + """Create a TrainingSession resource and return a TrainingClient. + + Blocks until the session reaches ``Running`` before returning. + + Every parameter after ``base_model`` is keyword-only, deliberately: + ``create_lora_training_client(model, 16)`` does not say what 16 is, and a + bare number at a call site is exactly the kind of thing that gets + mis-read as ``alpha``. Requiring ``rank=16`` costs a caller nothing. + + ``seed``, ``train_mlp``, ``train_attn``, ``train_unembed`` and + ``user_metadata`` are accepted for compatibility with existing training + scripts but are **not honoured**; each one that is set logs a warning + naming it. + + Session-level parameters (``model_package_group_arn``, ``s3_output_path``) + default to the values set on this :class:`ServiceClient`. Pass them here + to override per-session. + + Args: + base_model: A model id from :meth:`list_base_models` + (``"provider/model"``), or a full Hub model ARN to pin a specific + model version. An id is resolved through the Hub to that model's + current version. + rank: LoRA rank (default 32). + alpha: LoRA alpha / scaling factor (default 64). + training_session_name: Name for the new session. When omitted, the + SDK auto-generates ``{session_name_prefix}-{8 hex}``. + model_package_group_arn: Model package group ARN for checkpoints. + Falls back to :class:`ServiceClient`'s default. When neither is + set, the SDK uses the account's shared default group + (:data:`DEFAULT_MODEL_PACKAGE_GROUP_NAME`), creating it on first + use. + s3_output_path: S3 path for outputs. Falls back to + :class:`ServiceClient`'s default. When neither is set, the SDK + uses the SageMaker default bucket convention. + kms_key_arn: ARN of an AWS KMS key used to encrypt the session's + outputs (saved weights and checkpoints) in Amazon S3. Falls + back to :class:`ServiceClient`'s default. When neither is set, + outputs use Amazon S3 managed encryption. + vpc_config: :class:`VpcConfig` (or a dict with ``subnets`` and + ``security_group_ids``) to run the session inside your VPC. + Falls back to :class:`ServiceClient`'s default. + seed: Accepted but not used; setting it logs a warning. LoRA + initialisation is seeded by the service and is not caller-settable. + train_mlp: Accepted but not used; setting it logs a warning. + LoRA is applied to the layer set the service selects. + train_attn: Accepted but not used; setting it logs a warning. + train_unembed: Accepted but not used; setting it logs a warning. + user_metadata: Accepted but not used; setting it logs a warning. + Nothing is recorded against the session. Tag the model package + group instead if you need run metadata to persist. + + Returns: + A :class:`TrainingClient` bound to the running session. + """ + _warn_unsupported( + seed=seed, + train_mlp=train_mlp, + train_attn=train_attn, + train_unembed=train_unembed, + user_metadata=user_metadata, + ) + resolved_mpg = model_package_group_arn or self._model_package_group_arn + resolved_s3 = s3_output_path or self._s3_output_path + resolved_kms = kms_key_arn or self._kms_key_arn + resolved_vpc = ( + VpcConfig.from_any(vpc_config) if vpc_config is not None else self._vpc_config + ) + resolved_name = training_session_name or self._generate_session_name() + + # Resolve the base model: ARN passthrough vs name lookup + if base_model.startswith("arn:"): + resolved_arn, resolved_hf_model_id = base_model, None + else: + resolved_arn, resolved_hf_model_id = resolve_base_model( + base_model, self._get_hub_client() + ) + + if not resolved_mpg: + resolved_mpg = self._auto_create_mpg() + if not resolved_s3: + resolved_s3 = self._default_s3_output_path(resolved_name) + + args = _build_create_session_args( + session_name=resolved_name, + base_model_arn=resolved_arn, + model_package_group_arn=resolved_mpg, + lora_config=build_lora_config(rank, alpha), + role_arn=self._role_arn, + s3_path=resolved_s3, + kms_key_arn=resolved_kms, + vpc_config=resolved_vpc, + ) + + resp = self._sagemaker_client.create_training_session(**args) + logger.info( + "Created training session %r (arn=%s)", + resolved_name, + resp.get("TrainingSessionArn", ""), + ) + self._wait_until_running_or_stop(resolved_name) + return TrainingClient( + training_session_name=resolved_name, + runtime=self._runtime, + base_model=resolved_arn, + hub_client=None, + hf_model_id=resolved_hf_model_id, + model_package_group_arn=resolved_mpg, + service_client=self, + ) + + async def create_lora_training_client_async( + self, base_model: str, **kwargs: Any + ) -> TrainingClient: + """Coroutine twin of :meth:`create_lora_training_client`. + + Session creation waits for the session to reach ``Running``, which can + take minutes on a cold model. Running it on a worker thread keeps that + wait off the event loop, so an ``asyncio`` caller can provision and + prepare data concurrently. + + ``base_model`` is the only positional parameter here too, and forwarding + ``**kwargs`` is what keeps ``rank`` keyword-only through this path rather + than quietly reintroducing a positional form the sync method rejects. + """ + return await asyncio.to_thread(self.create_lora_training_client, base_model, **kwargs) + + def create_training_client_from_state( + self, + model_package_arn: str, + **kwargs: Any, + ) -> TrainingClient: + """Create a session from saved model weights. + + Convenience over :meth:`create_lora_training_client` + ``load_state``: + creates a fresh session (forwarding all the same keyword arguments) and + loads ``model_package_arn`` into it, so the returned client resumes from + that checkpoint. + + This restores **only the model weights, not optimizer state**; the + optimizer resets. To also restore optimizer state (e.g. Adam momentum), + use :meth:`create_training_client_from_state_with_optimizer`. The + checkpoint carries both; which factory you call decides what is + restored from it. + + Args: + model_package_arn: ARN of the model package to load on startup. + **kwargs: Forwarded to :meth:`create_lora_training_client`. When + ``base_model`` is omitted, it is read from the checkpoint's + model package metadata (pinned to the version the weights were + trained on). + """ + return self._create_training_client_from_state( + model_package_arn, load_optimizer=False, kwargs=kwargs + ) + + async def create_training_client_from_state_async( + self, model_package_arn: str, **kwargs: Any + ) -> TrainingClient: + """Coroutine twin of :meth:`create_training_client_from_state`.""" + return await asyncio.to_thread( + self.create_training_client_from_state, + model_package_arn=model_package_arn, + **kwargs, + ) + + def create_training_client_from_state_with_optimizer( + self, model_package_arn: str, **kwargs: Any + ) -> TrainingClient: + """Like :meth:`create_training_client_from_state`, restoring the optimizer.""" + return self._create_training_client_from_state( + model_package_arn, load_optimizer=True, kwargs=kwargs + ) + + async def create_training_client_from_state_with_optimizer_async( + self, model_package_arn: str, **kwargs: Any + ) -> TrainingClient: + """Coroutine twin of :meth:`create_training_client_from_state_with_optimizer`.""" + return await asyncio.to_thread( + self.create_training_client_from_state_with_optimizer, + model_package_arn=model_package_arn, + **kwargs, + ) + + def _create_training_client_from_state( + self, model_package_arn: str, *, load_optimizer: bool, kwargs: Dict[str, Any] + ) -> TrainingClient: + if "base_model" not in kwargs: + + mp_desc = self._get_model_registry_client().describe_model_package( + ModelPackageName=model_package_arn + ) + kwargs = { + **kwargs, + "base_model": base_model_arn_from_model_package(mp_desc, self._region), + } + tc = self.create_lora_training_client(**kwargs) + # Wait for the restore, so the client returned is ready to train. + try: + if load_optimizer: + tc.load_state_with_optimizer(model_package_arn).result() + else: + tc.load_state(model_package_arn).result() + except BaseException: + # No client is being returned, so nothing would own the session the + # line above created; stop it instead of leaving it to the inactivity + # timeout. Never masks the original failure. + try: + tc.stop() + except Exception: + logger.warning( + "Could not stop training session %r after restoring state " + "from %r failed; it will be stopped after its inactivity " + "timeout.", + tc.training_session_name, + model_package_arn, + exc_info=True, + ) + raise + return tc + + def create_sampling_client( + self, + *, + base_model: Optional[str] = None, + model_package_arn: Optional[str] = None, + retry_config: Optional[Any] = None, + ) -> SamplingClient: + """Create a SamplingClient, implicitly creating a training session. + + The SDK creates a training session under the hood and returns a + :class:`SamplingClient` bound to it. The session uses SDK-generated + defaults for MPG and S3 output (or the values on this + :class:`ServiceClient`). + + Because the session is implicit, the returned client *owns* it: use the + client as a context manager (or call :meth:`SamplingClient.stop`) so the + session is stopped at a point you chose:: + + with svc.create_sampling_client(base_model="") as sp: + out = sp.sample("hello").result() + # session stopped here + + If you neither scope nor stop it, the SDK stops the session once the + client becomes unreachable (on garbage collection or at interpreter exit) + instead of leaving it Running until its inactivity timeout stops it. + That is a backstop for interactive work, not a replacement + for the ``with`` block above: it fires at a time the interpreter picks, + and never at all if the process is killed with ``os._exit`` or + ``SIGKILL``. + + Specify exactly one of ``base_model`` or ``model_package_arn``: + + * ``base_model``: sample from the base model directly. + * ``model_package_arn``: sample from saved weights (the SDK resolves + the base model from the package metadata). + + Args: + base_model: A model id from :meth:`list_base_models`, or a full Hub + model ARN. + model_package_arn: ARN of a model package (saved weights) to sample + from. + retry_config: Accepted but not used; setting it logs a warning. + Throttling and 5xx failures are retried on a fixed budget; any + other error surfaces immediately. + + Returns: + A :class:`SamplingClient` bound to the implicit session. + + Raises: + ValueError: If neither or both of ``base_model`` / + ``model_package_arn`` are supplied. + """ + _warn_unsupported_args(_RETRY_CONFIG_CONSEQUENCE, retry_config=retry_config) + if bool(base_model) == bool(model_package_arn): + raise ValueError("Specify exactly one of base_model or model_package_arn.") + + resolved_hf_model_id = None + if model_package_arn: + mp_desc = self._get_model_registry_client().describe_model_package( + ModelPackageName=model_package_arn + ) + resolved_base_arn = base_model_arn_from_model_package(mp_desc, self._region) + checkpoint_ref = model_package_arn + else: + if base_model.startswith("arn:"): + resolved_base_arn = base_model + else: + resolved_base_arn, resolved_hf_model_id = resolve_base_model( + base_model, self._get_hub_client() + ) + checkpoint_ref = "" + + session_name = self._generate_session_name() + resolved_mpg = self._model_package_group_arn or self._auto_create_mpg() + resolved_s3 = self._s3_output_path or self._default_s3_output_path(session_name) + + args = _build_create_session_args( + session_name=session_name, + base_model_arn=resolved_base_arn, + model_package_group_arn=resolved_mpg, + lora_config=build_lora_config(DEFAULT_LORA_RANK, DEFAULT_LORA_ALPHA), + role_arn=self._role_arn, + s3_path=resolved_s3, + kms_key_arn=self._kms_key_arn, + vpc_config=self._vpc_config, + ) + self._sagemaker_client.create_training_session(**args) + logger.info( + "Created implicit sampling session %r for base model %s", + session_name, + resolved_base_arn, + ) + self._wait_until_running_or_stop(session_name) + + return SamplingClient( + training_session_name=session_name, + runtime=self._runtime, + model_package_arn=checkpoint_ref, + base_model=resolved_base_arn, + # Carry the id resolved above so a later get_tokenizer() does not + # describe the Hub again, and leave the Hub client for the client to + # borrow lazily if it ever needs one, as the training path does. + hub_client=None, + hf_model_id=resolved_hf_model_id, + # This session is implicit: no TrainingClient owns it, so no handle + # is passed in and the SamplingClient builds its own. That is exactly + # what marks it the owner, responsible for the stop on __exit__. + service_client=self, + ) + + async def create_sampling_client_async(self, **kwargs: Any) -> SamplingClient: + """Coroutine twin of :meth:`create_sampling_client`. + + Creates the implicit session on a worker thread, so the wait for it to + reach ``Running`` does not block the event loop. + """ + return await asyncio.to_thread(self.create_sampling_client, **kwargs) + + def list_sessions( + self, + *, + name_contains: Optional[str] = None, + status_equals: "Optional[str | SessionStatus]" = None, + max_results: Optional[int] = None, + ) -> "list[TrainingSessionSummary]": + """List training sessions, auto-paginating. + + Args: + name_contains: Only return sessions whose name contains this string. + status_equals: Only return sessions in this status. Accepts a + :class:`SessionStatus` enum or a raw string (e.g. ``"Running"``). + max_results: Cap the total number of summaries returned. When None, + every matching session is returned. Larger than one page, this + pages until the cap is reached. + + Returns: + A list of :class:`TrainingSessionSummary`. + """ + filters: dict[str, Any] = {} + if name_contains is not None: + filters["NameContains"] = name_contains + if status_equals is not None: + filters["StatusEquals"] = getattr(status_equals, "value", status_equals) + + summaries: list[TrainingSessionSummary] = [] + next_token: Optional[str] = None + while True: + args = dict(filters) + if next_token: + args["NextToken"] = next_token + if max_results is not None: + remaining = max_results - len(summaries) + if remaining <= 0: + break + args["MaxResults"] = min(remaining, MAX_RESULTS_PER_PAGE) + resp = self._sagemaker_client.list_training_sessions(**args) + summaries.extend( + TrainingSessionSummary._from_summary(summary) + for summary in resp.get("TrainingSessionSummaries", []) + ) + next_token = resp.get("NextToken") + if not next_token: + break + + if max_results is not None: + return summaries[:max_results] + return summaries + + def stop_session(self, training_session: str) -> None: + """Stop a running training session. + + Transitions the session toward a terminal ``Stopped`` state. + + Idempotent: a session that is already stopping, already terminal, or + being deleted is left as it is and the call returns normally. That makes + stop order-independent, which is what the rest of the lifecycle relies + on: an explicit ``stop_session`` here and the automatic stop that + follows the last client's collection both land on the same session, and + neither has to know whether the other went first. + + Args: + training_session: Session name or ARN. A ``TrainingSessionArn`` from + :meth:`list_sessions` works as-is; the name is taken from the + ARN's last segment. + """ + self._stop_session_if_active(_session_name_of(training_session)) + + def _stop_session_if_active(self, training_session_name: str) -> bool: + """:meth:`stop_session`, reporting whether a stop was actually issued. + + ``False`` means the session was already stopping or terminal, so nothing + happened. The automatic teardown needs to know the difference: telling a + caller "stopped it for you" about a session that was already down is + noise, and warning them they forgot to clean up is simply wrong. + + Kept private; ``stop_session``'s contract is "the session is not + running when this returns", and that holds either way. + """ + try: + self._sagemaker_client.stop_training_session( + TrainingSessionName=training_session_name, + ) + except ClientError as e: + if not _stop_already_done(e): + raise + logger.debug( + "Training session %r is already stopping or terminal; " "nothing to stop.", + training_session_name, + ) + return False + logger.info("Requested stop for training session %r", training_session_name) + return True + + def delete_session(self, training_session: str, *, stop_if_active: bool = False) -> None: + """Delete a training session. + + The service only deletes sessions in a terminal state (``Stopped``, + ``Failed``, or ``DeletionFailed``); deletion is asynchronous: the + session flips to ``Deleting`` and a background worker removes it + shortly after. Pass ``stop_if_active=True`` to stop a still-running + session before deleting it. + + Args: + training_session: Session name or ARN. A ``TrainingSessionArn`` from + :meth:`list_sessions` works as-is; the name is taken from the + ARN's last segment. + stop_if_active: When True, issue a stop and wait for a terminal + state before deleting. + """ + training_session_name = _session_name_of(training_session) + if stop_if_active: + self.stop_session(training_session_name) + self._wait_until_terminal(training_session_name) + + self._sagemaker_client.delete_training_session( + TrainingSessionName=training_session_name, + ) + logger.info("Requested delete for training session %r", training_session_name) + + def delete_model_package( + self, model_package_arn: str, *, delete_artifacts: bool = False + ) -> ModelPackageDeletion: + """Delete a model package, optionally deleting its S3 artifacts too. + + The model package carries the only reference to the S3 location its + weights were written to, so by default this deletes the package and + leaves the objects: they keep costing storage with nothing pointing at + them. Keep the returned :attr:`~ModelPackageDeletion.s3_uri` if you + intend to clean up later. + + With ``delete_artifacts=True`` the objects go first and the package + second, so a failure part-way leaves the package still pointing at + whatever survived and the call can be retried. + + Idempotent: a package that is already gone returns a result with + ``model_package_deleted=False`` instead of raising. + + .. note:: + :meth:`ServiceClient.create_sampling_client` and + :meth:`TrainingClient.load_state` read a package's artifacts, so + once they are deleted do not rely on this package for either. + + Args: + model_package_arn: ARN of the model package version to delete. + delete_artifacts: When True, delete the S3 objects under the + package's location before deleting the package. Objects whose + storage SageMaker manages on your behalf are left to it; + deleting the package is what releases those. + + Returns: + A :class:`ModelPackageDeletion` describing what was deleted and, when + the artifacts were kept, where they still are. + + Raises: + ArtifactDeletionError: If some object could not be deleted. The + model package is left in place, so a retry can find them again. + ValueError: If the package's location names a bucket root, where + deleting "its" artifacts would empty the whole bucket. + + Example:: + + # Delete the package, keep the weights + result = svc.delete_model_package(ckpt.model_package_arn) + print("weights still at", result.s3_uri) + + # Delete both + result = svc.delete_model_package( + ckpt.model_package_arn, delete_artifacts=True + ) + print("removed", result.objects_deleted, "objects") + """ + try: + artifacts = describe_model_package_artifacts( + self._get_model_registry_client(), model_package_arn + ) + except ClientError as e: + if not _model_package_absent(e): + raise + logger.info( + "Model package %r does not exist; nothing to delete.", + model_package_arn, + ) + return ModelPackageDeletion( + model_package_arn=model_package_arn, + model_package_deleted=False, + retained_reason="the model package does not exist", + ) + + s3_uri = artifacts["s3_uri"] + objects_deleted = 0 + retained_reason: Optional[str] = None + if not delete_artifacts: + retained_reason = "delete_artifacts was not set" + elif artifacts["managed_storage_type"] == MANAGED_STORAGE_RESTRICTED: + # Describe returns no location for these and the caller has no + # access to one: deleting the package is what releases the storage. + # A skip, not an error, so the same call works for any kind. + retained_reason = ( + "this model package's artifacts are stored by SageMaker on your " + "behalf, not in your own S3 location; deleting the package " + "releases them" + ) + elif not s3_uri: + retained_reason = "the model package carries no S3 location" + else: + objects_deleted, failed_keys = delete_s3_objects_under( + self._get_s3_client(), + s3_uri, + s3_data_type=artifacts["s3_data_type"], + ) + if failed_keys: + raise ArtifactDeletionError(model_package_arn, s3_uri, failed_keys) + logger.info( + "Deleted %d object(s) under %s for model package %r.", + objects_deleted, + s3_uri, + model_package_arn, + ) + + try: + self._get_model_registry_client().delete_model_package( + ModelPackageName=model_package_arn + ) + deleted = True + except ClientError as e: + if not _model_package_absent(e): + raise + # Another caller deleted it between the describe and here. + deleted = False + + if retained_reason is not None and s3_uri: + # The package carried the only record of this location and is now + # gone, so say it once at INFO even if the result is discarded. + logger.info( + "Model package %r deleted; its artifacts remain at %s (%s).", + model_package_arn, + s3_uri, + retained_reason, + ) + return ModelPackageDeletion( + model_package_arn=model_package_arn, + model_package_deleted=deleted, + s3_uri=s3_uri, + objects_deleted=objects_deleted, + artifacts_retained=retained_reason is not None, + retained_reason=retained_reason, + ) + + # ------------------------------------------------------------------ + # Internal + # ------------------------------------------------------------------ + + def _generate_session_name(self) -> str: + """Auto-generate a unique training session name.""" + + base = self._session_name_prefix or DEFAULT_SESSION_NAME_PREFIX + short_id = uuid.uuid4().hex[:8] + return f"{base}-{short_id}" + + def _auto_create_mpg(self) -> str: + """Return the ARN of the account's default model package group. + + One shared group per account/region, mirroring how the broader SageMaker + PySDK shares one default bucket: a group per session would leave a + never-collected group behind for every session an unconfigured script + creates, and would scatter a run's checkpoints so that + :meth:`get_latest_model_package` could not find them from a later + process. Created on first use; returns the existing ARN thereafter. + """ + group_name = DEFAULT_MODEL_PACKAGE_GROUP_NAME + try: + resp = self._get_model_registry_client().create_model_package_group( + ModelPackageGroupName=group_name, + ModelPackageGroupDescription=( + "Default group for SageMaker training session checkpoints, " + "created by the sagemaker-train SDK." + ), + ) + return resp["ModelPackageGroupArn"] + except self._get_model_registry_client().exceptions.ClientError as e: + if _model_package_group_exists(e): + resp = self._get_model_registry_client().describe_model_package_group( + ModelPackageGroupName=group_name + ) + return resp["ModelPackageGroupArn"] + raise + + def _default_s3_output_path(self, session_name: str) -> str: + """Generate a default S3 output path using the SageMaker default bucket. + + Follows the same convention as the broader SageMaker PySDK: + ``s3://sagemaker-{region}-{account}/{session_name}/output``, and created + if it does not exist yet, exactly like the broader SageMaker PySDK. + """ + bucket = self._sagemaker_session.default_bucket() + return f"s3://{bucket}/{session_name}/output" + + def _wait_until_running_or_stop(self, training_session_name: str) -> None: + """Wait for ``Running``, stopping the session if it never gets there. + + The create call has already returned by this point, so a failed wait + would otherwise leave a session running with no client bound to it: + nothing holds its handle, so no automatic stop is armed and only the + inactivity timeout would release it. Includes ``KeyboardInterrupt``: + interrupting a cold-model wait is the likeliest way to reach this. + + Best effort, and never masks the original failure, which is what tells + the caller why they have no client. + """ + try: + self._wait_until_running(training_session_name) + except BaseException: + try: + self._stop_session_if_active(training_session_name) + except Exception: + logger.warning( + "Could not stop training session %r after it failed to " + "reach Running; it will be stopped after its inactivity " + "timeout.", + training_session_name, + exc_info=True, + ) + raise + + def _wait_until_running( + self, + training_session_name: str, + timeout: int = 600, + poll_interval: int = 10, + ) -> None: + """Block until the session reaches ``Running`` status.""" + deadline = time.monotonic() + timeout + last_status = None + while time.monotonic() < deadline: + resp = self._sagemaker_client.describe_training_session( + TrainingSessionName=training_session_name, + ) + summary = _extract_describe_summary(resp) + status = summary["TrainingSessionStatus"] + if status != last_status: + logger.info("Session %r status: %s", training_session_name, status) + last_status = status + if status == _RUNNING_STATUS: + return + if status in _TERMINAL_STATUSES: + raise RuntimeError( + f"Session {training_session_name!r} entered terminal state: " + f"{status}{_describe_failure_suffix(summary)}" + ) + time.sleep(poll_interval) + + raise TimeoutError( + f"Session {training_session_name!r} did not reach Running within {timeout}s" + ) + + def _wait_until_terminal( + self, + training_session_name: str, + timeout: int = 600, + poll_interval: int = 10, + ) -> None: + """Block until the session reaches a deletable terminal status.""" + deadline = time.monotonic() + timeout + last_status = None + while time.monotonic() < deadline: + resp = self._sagemaker_client.describe_training_session( + TrainingSessionName=training_session_name, + ) + status = _extract_describe_summary(resp)["TrainingSessionStatus"] + if status != last_status: + logger.info("Session %r status: %s", training_session_name, status) + last_status = status + if status in _DELETABLE_STATUSES: + return + time.sleep(poll_interval) + + raise TimeoutError( + f"Session {training_session_name!r} did not reach a terminal state " + f"within {timeout}s" + ) diff --git a/sagemaker-train/src/sagemaker/train/training_session/training_client.py b/sagemaker-train/src/sagemaker/train/training_session/training_client.py new file mode 100644 index 0000000000..cd8a7d00a2 --- /dev/null +++ b/sagemaker-train/src/sagemaker/train/training_session/training_client.py @@ -0,0 +1,1090 @@ +# 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. +"""TrainingClient: drives the training loop over a TrainingSession resource. + +Owns the monotonic ``sequenceNumber`` for the session and turns each primitive +(``forward_backward`` / ``optim_step`` / ``save_weights_for_sampler`` / ``save_state`` / +``load_state``) into an ``ExecuteTrainingOperation`` call, returning a unified +:class:`APIFuture`. Submits are non-blocking: each method stamps its sequence +slot at call time, queues the wire call on a background submitter thread, and +returns its future immediately, so a training loop can queue a large batch of +operations in seconds and resolve results later. Operations can be pipelined: +submit ``optim_step`` before the preceding ``forward_backward`` future +resolves; the sequence numbers stamped at call time preserve order on the +service side. +""" + +from __future__ import annotations + +import asyncio +import functools +import logging +import threading +import uuid +from typing import Any, Callable, Dict, List, Optional, Type, TypeVar + +from sagemaker.train.training_session.exceptions import TrainingSessionError +from sagemaker.train.training_session._utils import ( + batch_to_datum_list, + chunk_datum_list, + combine_forward_backward_results, + _warn_unsupported_args, + _RETRY_CONFIG_CONSEQUENCE, +) +from sagemaker.train.training_session._session_handle import _SessionHandle +from sagemaker.train.training_session._submitter import ( + BackgroundSubmitter, + run_blocking, +) +from sagemaker.train.training_session.futures import ( + APIFuture, + CombinedFuture, + PendingAPIFuture, + call_with_transient_retry, +) +from sagemaker.train.training_session.sampling_client import ( + SamplingClient, + _SAMPLE_SUBMIT_MAX_ATTEMPTS, +) +from sagemaker.train.training_session.types import ( + AdamParams, + ForwardBackwardOutput, + LoadStateResponse, + LossFunction, + ModelPackageInfo, + ModelPackageKind, + OptimStepResponse, + SaveWeightsResponse, + SaveStateResponse, + _OperationResult, +) + +logger = logging.getLogger(__name__) + +#: The first training op after the session flips to ``Running`` can briefly hit +#: transient unavailability while the runtime finishes warming up (the session +#: can report Running just before the runtime path accepts traffic). The submit +#: is idempotent (stable clientToken + sequenceNumber), so it is safe to retry +#: through the warmup window. Budget = 150 attempts over the backoff, about 5 min. +_TRAINING_SUBMIT_MAX_ATTEMPTS = 150 + +#: Consequence clause for the checkpoint-lifetime knobs the save ops accept but +#: do not honour. Named so ``save_state`` and ``save_weights_for_sampler`` cannot +#: drift into telling the caller two different things about the same behaviour. +_SAVE_TTL_CONSEQUENCE = "the checkpoint was saved under the platform's retention policy instead" + +#: Consequence clause for ``name``. Points at the parameter that *does* carry a +#: caller-supplied label, so the warning is actionable, not just a refusal. +_SAVE_NAME_CONSEQUENCE = ( + "a checkpoint is identified by its model package ARN, not a caller-chosen " + "name; pass description=... for a human-readable label" +) + +#: Consequence clause for ``overwrite``. Distinct from the retention clause: this +#: knob is about versioning, and saying "retention policy" here would answer a +#: question the caller did not ask. +_SAVE_OVERWRITE_CONSEQUENCE = ( + "the save registered a new model package version instead, leaving earlier " "versions in place" +) + +#: Result type of a submitted operation, fixed by the ``result_cls`` passed to +#: :meth:`TrainingClient._submit`. Bounded by the base that declares +#: ``_from_api``, which is what ``_submit`` calls to build the result. +_T = TypeVar("_T", bound=_OperationResult) + + +class TrainingClient: + """Drives forward_backward / optim_step / save_weights_for_sampler over a session. + + Instances are created by :meth:`ServiceClient.create_lora_training_client` + and its ``*_from_state`` variants, not directly by callers. + + Args: + training_session_name: Name of the backing TrainingSession resource. + runtime: Client for the training-session runtime. + base_model: Base model (HF model ID or SageMaker Hub ARN), used to + resolve a tokenizer via :meth:`get_tokenizer`. + tokenizer: Optional pre-loaded tokenizer for string batches. + hub_client: Optional SageMaker Hub client, used to look up the + HuggingFace model id behind a Hub ARN when a tokenizer is needed. + Only consulted when neither ``tokenizer`` nor ``hf_model_id`` is set. + hf_model_id: Optional HuggingFace model id to load the tokenizer from. + Set it when ``base_model`` is an ARN whose id cannot be looked up, or + to pin the tokenizer to a specific repo. + model_package_group_arn: The session's model package group, where + save_state / save_weights_for_sampler register checkpoints. Lets + :meth:`list_model_packages` resolve the group without the caller + re-supplying it. + service_client: The :class:`ServiceClient` that created this client, + used to delegate checkpoint listing and to stop the session. + + The backing session is stopped once the last client referencing it becomes + unreachable, including any :class:`SamplingClient` derived from this one, + which shares this client's session handle. The automatic teardown *stops* the + session, leaving it describable for later inspection; deleting one stays an + explicit :meth:`delete` call. + + Releasing the session is handled for you, so a caller can focus on the + training loop. Anything worth keeping goes in a checkpoint: model packages + are durable on their own, so trained weights outlive both the session and + the process, and a new session can be created from one at any time. + """ + + def __init__( + self, + training_session_name: str, + runtime: Any, + *, + base_model: str = "", + tokenizer: Any = None, + hub_client: Any = None, + hf_model_id: Optional[str] = None, + model_package_group_arn: str = "", + service_client: Any = None, + ) -> None: + self.training_session_name = training_session_name + self._runtime = runtime + self._base_model = base_model + self._tokenizer = tokenizer + self._hub_client = hub_client + self._hf_model_id = hf_model_id + self._sequence_counter = 0 + # Reentrant: a chunked forward-backward holds this across its whole chunk + # loop so the chunks get consecutive sequence numbers, and each chunk's + # submit takes it again to stamp its own. + self._sequence_lock = threading.RLock() + # Base methods queue their ExecuteTrainingOperation round trip here and + # return a future immediately; ordered submits keep wire order equal to + # the sequence numbers stamped at call time. + self._submitter = BackgroundSubmitter( + name=f"sagemaker-training-submit-{training_session_name}" + ) + self._model_package_group_arn = model_package_group_arn + self._service_client = service_client + # Latest model package this client saved, per kind. Lets + # get_latest_model_package answer without an API round-trip, and covers the + # window where a just-saved package is not yet describable. Cache only: a + # fresh client starts empty and does not imply a package it never saved. + self._latest_by_kind: Dict[ModelPackageKind, ModelPackageInfo] = {} + self._cache_lock = threading.Lock() + # This session's lifetime, shared with every SamplingClient derived from + # this client. The handle, not this client, is what CPython collects + # last, so the session is stopped when the last of them goes away and + # never while one is still driving it. + self._handle = _SessionHandle(training_session_name, service_client) + self._handle.register_submitter(self._submitter) + + def get_tokenizer(self) -> Any: + """Return a HuggingFace tokenizer for the session's base model. + + The runtime accepts only token IDs, so callers with text must tokenize + client-side: ``tok = tc.get_tokenizer(); ids = tok.encode("hello")``. + The tokenizer is loaded once and cached (and reused for string batches + passed to :meth:`forward_backward`). + + Returns: + A ``transformers.PreTrainedTokenizer``. + + Raises: + ValueError: If the base model has no tokenizer available (e.g. a + model whose tokenizer is not published to HuggingFace); pass + pre-tokenized token IDs in that case. + ImportError: If ``transformers`` is not installed (install the + optional extra: ``pip install sagemaker-train[tokenizer]``). + """ + if self._tokenizer is None: + from sagemaker.train.training_session._utils import ( + load_tokenizer, + ) + + # The Hub client is only needed to resolve the tokenizer id when no + # hf_model_id was supplied; skip building it otherwise. + hub_client = self._hub_client if self._hf_model_id else self._get_hub_client() + self._tokenizer = load_tokenizer(self._base_model, hub_client, self._hf_model_id) + return self._tokenizer + + def _get_hub_client(self) -> Any: + """Return the Hub client, building it lazily via the service client. + + Usually ``None`` at construction and built here on first use, only when + a base-model ARN's tokenizer id must be resolved and no ``hf_model_id`` + was supplied. + """ + if self._hub_client is None and self._service_client is not None: + self._hub_client = self._service_client._get_hub_client() + return self._hub_client + + # ------------------------------------------------------------------ + # Training primitives + # ------------------------------------------------------------------ + + # Each training op comes in two forms: a sync method that queues the submit + # and returns an :class:`APIFuture` immediately, and an ``*_async`` twin, + # a coroutine for async code, where ``fut = await tc.op_async(...)`` + # completes the enqueue and ``await fut`` yields the result. The sync + # methods do the work; the ``*_async`` twins are thin coroutine wrappers. + + def forward_backward( + self, + batch: Any, + loss_fn: "LossFunction | str" = "cross_entropy", + loss_fn_inputs: Optional[Dict[str, Any]] = None, + loss_fn_config: Optional[Dict[str, float]] = None, + ) -> APIFuture[ForwardBackwardOutput]: + """Submit a forward-backward pass; returns an :class:`APIFuture`. + + Non-blocking: the submit runs on a background thread, so this returns + immediately. Resolve the future with ``.result()`` (or ``await`` it) + for the :class:`ForwardBackwardOutput`; a submit error surfaces there, + not at the call site. Ops pipeline naturally: submit the next op + before this future resolves; the sequence numbers stamped at call time + preserve ordering on the service. + + Args: + batch: Training data: a list of token ID lists, a single token ID + list, a list of strings (requires a tokenizer), or raw datum + dicts (advanced). + loss_fn: Loss function name or :class:`LossFunction` enum. + loss_fn_inputs: Loss inputs applied to each datum that does not + already carry its own. Which keys ``loss_fn`` requires is + checked before the request leaves. + loss_fn_config: Optional loss function configuration (e.g. + ``{"clip_low_threshold": 0.2, "clip_high_threshold": 5.0}`` + for CISPO/PPO clipping). + + Raises: + ValueError: A per-token ``lossFnInputs`` entry is not the same length + as its datum's ``modelInput``. Raised before the request is sent, + naming the datum and the key. + + A batch too large to send in one request is split into chunks (bounded + by datum count and estimated size) submitted as consecutive + ``forwardBackward`` ops. The service accumulates gradients across + consecutive forward-backward ops until the next ``optim_step``, so this + is equivalent to sending the whole batch at once. When the batch fits in + one chunk (the common case) the request is unchanged. The returned + future resolves to a single combined :class:`ForwardBackwardOutput`. + """ + loss_fn = LossFunction(loss_fn) + datum_list = batch_to_datum_list( + batch, + tokenizer=self._tokenizer, + loss_fn_inputs=loss_fn_inputs, + tokenizer_loader=self.get_tokenizer, + loss_fn_name=loss_fn.value, + ) + chunks = chunk_datum_list(datum_list) + # 0 chunks (empty batch) or 1 chunk: submit exactly one op. + if len(chunks) <= 1: + return self._submit_forward_backward( + loss_fn.value, + chunks[0] if chunks else datum_list, + loss_fn_config=loss_fn_config, + ) + # Multiple chunks. Gradients accumulate across consecutive + # forward_backward ops, so the chunk ops must carry contiguous, in-order + # sequence numbers: the whole loop holds the sequence lock, not just each + # submit. Without that, an optim_step from another thread can take a + # number in the middle of the batch and apply a partial gradient. + with self._sequence_lock: + futures = [ + self._submit_forward_backward(loss_fn.value, chunk, loss_fn_config=loss_fn_config) + for chunk in chunks + ] + # The chunk sizes are known only here, and the combined loss is a + # datum-count-weighted mean over them. + return CombinedFuture( + futures, + functools.partial( + combine_forward_backward_results, + datum_counts=[len(chunk) for chunk in chunks], + ), + ) + + async def forward_backward_async( + self, + batch: Any, + loss_fn: "LossFunction | str" = "cross_entropy", + loss_fn_inputs: Optional[Dict[str, Any]] = None, + loss_fn_config: Optional[Dict[str, float]] = None, + ) -> APIFuture[ForwardBackwardOutput]: + """Coroutine twin of :meth:`forward_backward` for async code. + + ``fut = await tc.forward_backward_async(...)`` completes the enqueue, so + the op is sequenced before the next line runs, and ``await fut`` + yields the :class:`ForwardBackwardOutput`. + """ + return self.forward_backward(batch, loss_fn, loss_fn_inputs, loss_fn_config=loss_fn_config) + + def _submit_forward_backward( + self, + loss_fn_name: str, + datum_list: List[Dict[str, Any]], + loss_fn_config: Optional[Dict[str, float]] = None, + ) -> APIFuture[ForwardBackwardOutput]: + """Submit one forwardBackward op for an already-built datum list.""" + fb_config: Dict[str, Any] = { + "lossFnName": loss_fn_name, + "datumList": datum_list, + } + if loss_fn_config is not None: + fb_config["lossFnConfig"] = loss_fn_config + op_config = {"forwardBackward": fb_config} + return self._submit(op_config, ForwardBackwardOutput) + + def optim_step( + self, adam_params: "Optional[AdamParams | Dict[str, Any] | Any]" = None + ) -> APIFuture[OptimStepResponse]: + """Submit an optimizer step; returns an :class:`APIFuture`. + + Non-blocking; the future resolves to an :class:`OptimStepResponse` + with ``metrics``. Pipelines with :meth:`forward_backward`; the + sequence numbers stamped at call time preserve ordering on the service. + + Args: + adam_params: Optimizer parameters. Uses defaults if omitted. + Accepts an :class:`AdamParams`, a dict of the same fields, or + any object carrying at least one of them, so an existing + optimizer config object can be passed as-is. + """ + params = AdamParams.from_any(adam_params) if adam_params is not None else AdamParams() + op_config = {"optimizerStep": {"adamParamsConfig": params._to_api()}} + return self._submit(op_config, OptimStepResponse) + + async def optim_step_async( + self, adam_params: "Optional[AdamParams | Dict[str, Any] | Any]" = None + ) -> APIFuture[OptimStepResponse]: + """Coroutine twin of :meth:`optim_step` for async code. + + ``fut = await tc.optim_step_async(...)`` completes the enqueue, + guaranteeing the step is sequenced after any forward-backward already + submitted, and ``await fut`` yields the :class:`OptimStepResponse`. + """ + return self.optim_step(adam_params) + + def save_weights_for_sampler( + self, + description: str = "", + merge_weights: bool = False, + *, + ttl_seconds: Optional[int] = None, + name: Optional[str] = None, + ) -> APIFuture[SaveWeightsResponse]: + """Submit a save-weights op; returns an :class:`APIFuture`. + + Non-blocking; the future resolves to a + :class:`SaveWeightsResponse` carrying ``model_package_arn``. + The op registers a model package version in the session's model package + group. + + When ``merge_weights=True``, the service merges the LoRA adapter into the + base model and writes full, adapter-free weights (``checkpoints/hf_merged/``) + alongside the adapter. Use this when the checkpoint will be deployed as a + single standalone model, e.g. on GPUs without a working LoRA serving path. + + Args: + description: Optional human-readable description stored on the + model package (e.g. ``"weights at step 100"``). + merge_weights: Whether to also write full adapter-free weights for + standalone deployment. + ttl_seconds: Accepted but not used; setting it logs a warning. + Checkpoint lifetime is governed by the platform's retention + policy, not per-save. Delete the model package version when + you no longer need it. + name: Accepted but not used; setting it logs a warning. + Checkpoints are model package *versions* addressed by ARN, so + there is no caller-chosen name to set. Use ``description``. + """ + _warn_unsupported_args(_SAVE_TTL_CONSEQUENCE, ttl_seconds=ttl_seconds) + _warn_unsupported_args(_SAVE_NAME_CONSEQUENCE, name=name) + save_weights_config: Dict[str, Any] = {} + if merge_weights: + save_weights_config["mergeWeights"] = True + if description: + save_weights_config["description"] = description + op_config = {"saveWeights": save_weights_config} + return self._submit( + op_config, + SaveWeightsResponse, + record_fn=self._cache_on_save(ModelPackageKind.MODEL), + ) + + async def save_weights_for_sampler_async( + self, + description: str = "", + merge_weights: bool = False, + *, + ttl_seconds: Optional[int] = None, + name: Optional[str] = None, + ) -> APIFuture[SaveWeightsResponse]: + """Coroutine twin of :meth:`save_weights_for_sampler` for async code.""" + return self.save_weights_for_sampler( + description=description, + merge_weights=merge_weights, + ttl_seconds=ttl_seconds, + name=name, + ) + + def save_state( + self, + description: str = "", + *, + ttl_seconds: Optional[int] = None, + overwrite: bool = False, + name: Optional[str] = None, + ) -> APIFuture[SaveStateResponse]: + """Submit a save of full training state (weights + optimizer). + + Non-blocking; the future resolves to a :class:`SaveStateResponse` + carrying ``model_package_arn``, always a model package version that + :meth:`load_state` can resume from. + + Args: + description: Optional human-readable description stored on the + model package (e.g. ``"step 100 checkpoint"``). + ttl_seconds: Accepted but not used; setting it logs a warning. + Checkpoint lifetime is governed by the platform's retention + policy, not per-save. Delete the model package version when + you no longer need it. + overwrite: Accepted but not used; setting it logs a warning. + Every save registers a new model package *version*, so a save + never overwrites an earlier one, which is what ``False`` (the + default) asks for, so it does not warn. ``True`` warns and is + ignored. + name: Accepted but not used; setting it logs a warning. + Checkpoints are model package *versions* addressed by ARN, so + there is no caller-chosen name to set. Use ``description``. + """ + _warn_unsupported_args(_SAVE_TTL_CONSEQUENCE, ttl_seconds=ttl_seconds) + # None, not False: the default already describes what happens, so only an + # explicit True is a request that goes unhonoured. + _warn_unsupported_args(_SAVE_OVERWRITE_CONSEQUENCE, overwrite=True if overwrite else None) + _warn_unsupported_args(_SAVE_NAME_CONSEQUENCE, name=name) + save_state_config: Dict[str, Any] = {} + if description: + save_state_config["description"] = description + return self._submit( + {"saveState": save_state_config}, + SaveStateResponse, + record_fn=self._cache_on_save(ModelPackageKind.CHECKPOINT), + ) + + async def save_state_async( + self, + description: str = "", + *, + ttl_seconds: Optional[int] = None, + overwrite: bool = False, + name: Optional[str] = None, + ) -> APIFuture[SaveStateResponse]: + """Coroutine twin of :meth:`save_state` for async code.""" + return self.save_state( + description=description, + ttl_seconds=ttl_seconds, + overwrite=overwrite, + name=name, + ) + + def load_state(self, model_package_arn: str) -> APIFuture[LoadStateResponse]: + """Submit a load of a previously saved training state, weights only. + + Restores the LoRA weights from a :meth:`save_state` checkpoint and + leaves the optimizer freshly initialized, so training continues from + the saved weights with the optimizer reset. To also restore optimizer + state (e.g. Adam momentum), use :meth:`load_state_with_optimizer`. + + The checkpoint itself carries both weights and optimizer state; which + method you call decides what is restored from it. + + Non-blocking; the future resolves to a :class:`LoadStateResponse`. + + Args: + model_package_arn: ARN of the model package to load. + """ + return self._submit_load_state(model_package_arn, load_optimizer=False) + + async def load_state_async(self, model_package_arn: str) -> APIFuture[LoadStateResponse]: + """Coroutine twin of :meth:`load_state` for async code.""" + return self.load_state(model_package_arn) + + def load_state_with_optimizer(self, model_package_arn: str) -> APIFuture[LoadStateResponse]: + """Submit a load of saved training state, weights and optimizer. + + Like :meth:`load_state`, but also restores the optimizer state saved by + :meth:`save_state`, so training resumes exactly where it left off + instead of restarting the optimizer from scratch. + + Non-blocking; the future resolves to a :class:`LoadStateResponse`. + + Args: + model_package_arn: ARN of the model package to load. + """ + return self._submit_load_state(model_package_arn, load_optimizer=True) + + async def load_state_with_optimizer_async( + self, model_package_arn: str + ) -> APIFuture[LoadStateResponse]: + """Coroutine twin of :meth:`load_state_with_optimizer`.""" + return self.load_state_with_optimizer(model_package_arn) + + def _submit_load_state( + self, model_package_arn: str, *, load_optimizer: bool + ) -> APIFuture[LoadStateResponse]: + """Shared submit for the two public load-state methods. + + The ``loadOptimizer`` wire field is always sent explicitly; the choice + is never a caller-facing knob; it is fixed by which of the two public + load methods you call. + """ + op_config = { + "loadState": { + "modelPackageArn": model_package_arn, + "loadOptimizer": load_optimizer, + } + } + return self._submit(op_config, LoadStateResponse) + + def forward( + self, + batch: Any, + loss_fn: "LossFunction | str" = "cross_entropy", + loss_fn_inputs: Optional[Dict[str, Any]] = None, + loss_fn_config: Optional[Dict[str, float]] = None, + ) -> APIFuture[ForwardBackwardOutput]: + """Submit a forward-only pass (no gradient); returns an :class:`APIFuture`. + + Same input shape as :meth:`forward_backward`, but runs only the forward + pass, useful for computing loss-fn outputs / logprobs without updating + optimizer state. Non-blocking; the future resolves to a + :class:`ForwardBackwardOutput` with ``metrics`` and ``loss_fn_outputs``. + + Args: + batch: Training data. + loss_fn: Loss function name or :class:`LossFunction` enum. + loss_fn_inputs: Optional loss inputs applied to each datum. + loss_fn_config: Optional loss function configuration (e.g. + ``{"clip_low_threshold": 0.2, "clip_high_threshold": 5.0}``). + + Raises: + ValueError: A per-token ``lossFnInputs`` entry is not the same length + as its datum's ``modelInput``, as for :meth:`forward_backward`. + + A batch too large to send in one request is chunked as it is for + :meth:`forward_backward`, and the returned future resolves to a single + combined :class:`ForwardBackwardOutput`. + """ + loss_fn = LossFunction(loss_fn) + datum_list = batch_to_datum_list( + batch, + tokenizer=self._tokenizer, + loss_fn_inputs=loss_fn_inputs, + tokenizer_loader=self.get_tokenizer, + loss_fn_name=loss_fn.value, + ) + chunks = chunk_datum_list(datum_list) + if len(chunks) <= 1: + return self._submit_forward( + loss_fn.value, + chunks[0] if chunks else datum_list, + loss_fn_config=loss_fn_config, + ) + # A forward pass leaves no gradient to accumulate, so the chunks are + # independent; unlike forward_backward, their order carries no meaning. + futures = [ + self._submit_forward(loss_fn.value, chunk, loss_fn_config=loss_fn_config) + for chunk in chunks + ] + # Weighted by chunk size for the same reason forward_backward is: an + # uneven split would otherwise report a plain mean over the chunks and + # understate whichever one carried the most datums. + return CombinedFuture( + futures, + functools.partial( + combine_forward_backward_results, + datum_counts=[len(chunk) for chunk in chunks], + ), + ) + + def _submit_forward( + self, + loss_fn_name: str, + datum_list: List[Dict[str, Any]], + loss_fn_config: Optional[Dict[str, float]] = None, + ) -> APIFuture[ForwardBackwardOutput]: + """Submit one forward op for an already-built datum list.""" + fwd_config: Dict[str, Any] = { + "lossFnName": loss_fn_name, + "datumList": datum_list, + } + if loss_fn_config is not None: + fwd_config["lossFnConfig"] = loss_fn_config + op_config = {"forward": fwd_config} + return self._submit(op_config, ForwardBackwardOutput) + + async def forward_async( + self, + batch: Any, + loss_fn: "LossFunction | str" = "cross_entropy", + loss_fn_inputs: Optional[Dict[str, Any]] = None, + loss_fn_config: Optional[Dict[str, float]] = None, + ) -> APIFuture[ForwardBackwardOutput]: + """Coroutine twin of :meth:`forward` for async code.""" + return self.forward(batch, loss_fn, loss_fn_inputs, loss_fn_config=loss_fn_config) + + def _load_weights_for_sampling(self, model_package_arn: str) -> LoadStateResponse: + """Point the sampling path at a checkpoint and block until complete. + + Internal method used by :meth:`save_weights_and_get_sampling_client` to + direct the session's sampling path at a saved checkpoint before returning + a :class:`SamplingClient`. + + This operation is routed through the sampling endpoint (not the training + endpoint) and does NOT consume a training sequence number. + """ + return self._submit_load_weights_for_sampling(model_package_arn).result() + + def _submit_load_weights_for_sampling( + self, model_package_arn: str + ) -> APIFuture[LoadStateResponse]: + """Submit a load-weights-for-sampling op; returns an :class:`APIFuture`. + + Routed via the sampling-operations endpoint, so no sequence number + consumed, so it takes the sampling submit budget, not the training one. + Not an ``*_async`` twin: it is the submitting half of + :meth:`_load_weights_for_sampling`, not a coroutine. + """ + op_config = {"loadWeightsForSampling": {"modelPackageArn": model_package_arn}} + resp = call_with_transient_retry( + lambda: self._runtime.execute_sampling_operation( + trainingSessionName=self.training_session_name, + operationConfig=op_config, + ), + max_attempts=_SAMPLE_SUBMIT_MAX_ATTEMPTS, + ) + operation_id = resp["operationId"] + + def parse_fn(r: Dict[str, Any]) -> LoadStateResponse: + return LoadStateResponse._from_api(r.get("operationResult", {})) + + return APIFuture( + operation_id=operation_id, + poll_fn=lambda: self._runtime.get_sampling_operation_result( + trainingSessionName=self.training_session_name, + operationId=operation_id, + ), + parse_fn=parse_fn, + status_key="status", + operation_type="sample", + ) + + # ------------------------------------------------------------------ + # Model package discovery (delegates to ServiceClient with this session's group) + # ------------------------------------------------------------------ + + def list_model_packages( + self, + *, + kind: "Optional[ModelPackageKind]" = None, + max_results: Optional[int] = None, + ) -> List[ModelPackageInfo]: + """List this session's model packages, newest first. + + Convenience over :meth:`ServiceClient.list_model_packages` that supplies + the session's own model package group. + + Args: + kind: Optional filter; only return packages of this kind + (e.g. ``ModelPackageKind.CHECKPOINT`` or ``ModelPackageKind.MODEL``). + max_results: Cap on packages returned. ``None`` uses the + ServiceClient default (10 newest). + """ + kwargs: Dict[str, Any] = {"model_package_group_arn": self._require_group()} + if kind is not None: + kwargs["kind"] = kind + if max_results is not None: + kwargs["max_results"] = max_results + return self._require_service_client().list_model_packages(**kwargs) + + def get_latest_model_package(self, kind: ModelPackageKind) -> Optional[ModelPackageInfo]: + """Return the latest model package of the given kind THIS client saved. + + Answered purely from this client's in-memory save cache, so no API call, + and immune to the brief lag before a just-saved model package is + describable. Returns ``None`` if this client has not saved a package of + that kind in this process. + + A fresh client starts empty: a package from a session this + client did not drive is not implied here. Use + :meth:`ServiceClient.get_latest_model_package` or + :meth:`ServiceClient.list_model_packages` to discover packages in a + model package group. + + Args: + kind: The kind of model package to look up + (e.g. ``ModelPackageKind.CHECKPOINT`` or ``ModelPackageKind.MODEL``). + + Returns: + A :class:`ModelPackageInfo` with ``model_package_arn`` and ``kind``, + or ``None`` if nothing of that kind was saved by this client. + """ + return self._latest_cached(kind) + + # ------------------------------------------------------------------ + # Save cache + # ------------------------------------------------------------------ + + def _cache_on_save(self, kind: ModelPackageKind) -> Callable[[Any], None]: + """Build a record_fn that caches a saved checkpoint of ``kind``. + + Composed into the save_state / save_weights_for_sampler future's ``parse_fn`` (see + :meth:`_submit`), so it runs once the op completes (sync or async). + """ + + def _record(result: Any) -> None: + arn = result.model_package_arn + if not arn: + return + with self._cache_lock: + self._latest_by_kind[kind] = ModelPackageInfo(model_package_arn=arn, kind=kind) + + return _record + + def _latest_cached(self, kind: ModelPackageKind) -> Optional[ModelPackageInfo]: + """Latest cached checkpoint of ``kind``, or ``None`` if none saved.""" + with self._cache_lock: + return self._latest_by_kind.get(kind) + + def _require_group(self) -> str: + if not self._model_package_group_arn: + raise ValueError( + "This client has no model package group associated, so checkpoints " + "cannot be listed. Use ServiceClient.list_model_packages(...) with an " + "explicit model_package_group_arn instead." + ) + return self._model_package_group_arn + + def _require_service_client(self) -> Any: + if self._service_client is None: + raise ValueError( + "This client was not created by a ServiceClient, so checkpoint " + "lookups are unavailable. Use ServiceClient.list_model_packages(...) " + "directly." + ) + return self._service_client + + def _require_lifecycle_client(self) -> Any: + """The ServiceClient used to stop or delete the session. + + Starting and stopping a session are control-plane operations, and this + client holds only the runtime client, so the ServiceClient that created + the session is what carries them out. + """ + if self._service_client is None: + raise ValueError( + "This client was not created by a ServiceClient, so it cannot " + "stop or delete the session. Use ServiceClient.stop_session(name) " + "or ServiceClient.delete_session(name) instead." + ) + return self._service_client + + # ------------------------------------------------------------------ + # Session lifecycle + # ------------------------------------------------------------------ + + def stop(self) -> None: + """Stop this client's training session. + + Convenience over :meth:`ServiceClient.stop_session` for this session, with no + need to repeat the session name. Safe to call on an already-stopped + session. + + Drains the background submitter first (best effort), so every operation + already queued reaches the service before the stop request; the client + accepts no new operations afterwards. + """ + self._submitter.close() + try: + self._require_lifecycle_client().stop_session(self.training_session_name) + except Exception: + # Keep the automatic stop armed: it is the only remaining + # client-side attempt at releasing the session. + self._handle.note_explicit_stop_failed() + raise + # Stopped, so a later garbage collection would only repeat the call. + self._handle.disarm() + + def delete(self, *, stop_if_active: bool = False) -> None: + """Delete this client's training session. + + Convenience over :meth:`ServiceClient.delete_session` for this session. + The service only deletes terminal sessions; pass ``stop_if_active=True`` + to stop and wait for a terminal state first. Drains the background + submitter first; the client accepts no new operations afterwards. + """ + self._submitter.close() + self._require_lifecycle_client().delete_session( + self.training_session_name, stop_if_active=stop_if_active + ) + # Disarmed only once the delete has succeeded, or a running session the + # service refused to delete (the default, stop_if_active=False) was left + # with nothing client-side still trying to release it; detaching is + # one-way, so it held its GPUs until the inactivity timeout. + self._handle.disarm() + + def save_weights_and_get_sampling_client( + self, + description: str = "", + merge_weights: bool = False, + *, + name: Optional[str] = None, + retry_config: Optional[Any] = None, + ) -> "SamplingClient": + """Save weights, load them for sampling, and return a SamplingClient. + + Blocks until ``save_weights_for_sampler`` completes, points the sampling + path at the freshly saved checkpoint, then returns a :class:`SamplingClient` + bound to it. + + Args: + description: Optional human-readable description stored on the + model package (passed through to :meth:`save_weights_for_sampler`). + merge_weights: Whether to also write full adapter-free weights for + standalone deployment (passed through to :meth:`save_weights_for_sampler`). + name: Accepted but not used; setting it logs a warning. + Checkpoints are model package *versions* addressed by ARN, so + there is no caller-chosen name to set. Use ``description``. + retry_config: Accepted but not used; setting it logs a warning. + Throttling and 5xx failures are retried on a fixed budget; any + other error surfaces immediately. + + Returns: + A :class:`SamplingClient` for sampling from the saved checkpoint. + + Raises: + TrainingSessionError: The save completed without a checkpoint ARN. + Raised rather than falling back: an empty ARN means the base + model to :class:`SamplingClient`, so the caller would have + sampled untrained weights believing they were the trained ones. + """ + result = self.save_weights_for_sampler( + description=description, + merge_weights=merge_weights, + name=name, + ).result() + model_package_arn = result.model_package_arn + if not model_package_arn: + raise TrainingSessionError( + "The weights were saved but the service returned no checkpoint " + "ARN, so there is nothing to sample from. Retry the save, or " + "call save_weights_for_sampler() and inspect its result." + ) + # Point the sampling path at this checkpoint before sampling. + self._load_weights_for_sampling(model_package_arn) + # Forwarded, not warned about here: create_sampling_client owns the + # warning, so passing it on keeps exactly one line in the log. + return self.create_sampling_client( + model_package_arn=model_package_arn, retry_config=retry_config + ) + + async def save_weights_and_get_sampling_client_async( + self, + description: str = "", + merge_weights: bool = False, + *, + name: Optional[str] = None, + retry_config: Optional[Any] = None, + ) -> "SamplingClient": + """Coroutine twin of :meth:`save_weights_and_get_sampling_client`. + + The save has to complete before the checkpoint can be sampled, so this + blocks on a worker thread and not on the event loop, letting an + RL loop keep servicing rollouts while the weights are written. + """ + return await asyncio.to_thread( + self.save_weights_and_get_sampling_client, + description=description, + merge_weights=merge_weights, + name=name, + retry_config=retry_config, + ) + + # ------------------------------------------------------------------ + # Sampling client factory + # ------------------------------------------------------------------ + + def create_sampling_client( + self, + model_package_arn: str = "", + *, + retry_config: Optional[Any] = None, + ) -> "SamplingClient": + """Create a SamplingClient; ``model_package_arn=""`` = base model. + + Args: + model_package_arn: Model package ARN of the checkpoint to sample + from. Empty string samples from the base model. + retry_config: Accepted but not used; setting it logs a warning. + Throttling and 5xx failures are retried on a fixed budget; any + other error surfaces immediately. + + Returns: + A :class:`SamplingClient` bound to this session and checkpoint. + """ + _warn_unsupported_args(_RETRY_CONFIG_CONSEQUENCE, retry_config=retry_config) + # The SamplingClient only needs a Hub client to resolve a tokenizer when + # it has neither a cached tokenizer nor an hf_model_id; build one only in + # that case. Best-effort: if sagemaker-core is not installed, skip the + # Hub client: token-ID prompts still work, but string prompts will fail + # later with a clear error. + hub_client = self._hub_client + if hub_client is None and self._tokenizer is None and not self._hf_model_id: + try: + hub_client = self._get_hub_client() + except ImportError: # ModuleNotFoundError is a subclass + pass + return SamplingClient( + training_session_name=self.training_session_name, + runtime=self._runtime, + model_package_arn=model_package_arn, + base_model=self._base_model, + tokenizer=self._tokenizer, + hub_client=hub_client, + hf_model_id=self._hf_model_id, + service_client=self._service_client, + # Sharing the handle refcounts the session across both clients, so it + # survives whichever is collected first and is stopped when the last + # one goes, in either drop order. It also marks the new client as + # NOT owning the session, so leaving a `with` block on it cannot stop + # a session this training loop is still driving. + session_handle=self._handle, + # Shared for the same reason the handle is: a loop that refreshes + # its sampler calls this every step, and a submitter is an event + # loop on a daemon thread. Safe to share because wire calls run on + # the loop's executor, so the loop itself only dispatches, and + # sampling submits unordered, taking no turn from training's. + submitter=self._submitter, + ) + + async def create_sampling_client_async( + self, + model_package_arn: str = "", + *, + retry_config: Optional[Any] = None, + ) -> "SamplingClient": + """Coroutine twin of :meth:`create_sampling_client`. + + Binding to an existing session issues no control-plane call, so this is + cheap; it exists so an ``asyncio`` caller can use one calling convention + throughout without remembering which factories block. + """ + return await asyncio.to_thread( + self.create_sampling_client, + model_package_arn=model_package_arn, + retry_config=retry_config, + ) + + # ------------------------------------------------------------------ + # Internal + # ------------------------------------------------------------------ + + def _submit( + self, + operation_config: Dict[str, Any], + result_cls: Type[_T], + *, + record_fn: Optional[Callable[[_T], None]] = None, + ) -> APIFuture[_T]: + """Queue an operation; immediately return a future over its result. + + The sequence slot is stamped here, atomically with scheduling on the + ordered background submitter, so wire order always equals call order; + the ExecuteTrainingOperation round trip itself happens off-thread. A + submit error (e.g. a ConflictException) surfaces when the returned + future is resolved, not at call time. + + ``record_fn`` is composed into the future's ``parse_fn``, so it runs + exactly once when the op completes, on whichever path resolves the + future. Used to cache saved checkpoints. + """ + client_token = str(uuid.uuid4()) + + def parse_fn(r: Dict[str, Any]) -> _T: + result = result_cls._from_api(r.get("operationResult", {})) + if record_fn is not None: + record_fn(result) + return result + + with self._sequence_lock: + # Every allocated seq is dispatched exactly once: the ordered + # submitter releases wire calls strictly in allocation order, and a + # failed submit still consumed its slot (the service blocks on the + # gap, matching a lost request's semantics). Use + # `current_sequence_number` to inspect without consuming. + self._sequence_counter += 1 + seq = self._sequence_counter + + def _wire_submit() -> APIFuture[_T]: + # Transient 5xx/throttling are retried; the stable clientToken + # + sequenceNumber make the call idempotent, so a retry never + # double-submits or skips a sequence. + resp = call_with_transient_retry( + lambda: self._runtime.execute_training_operation( + trainingSessionName=self.training_session_name, + sequenceNumber=seq, + clientToken=client_token, + operationConfig=operation_config, + ), + max_attempts=_TRAINING_SUBMIT_MAX_ATTEMPTS, + ) + operation_id = resp["operationId"] + return APIFuture( + operation_id=operation_id, + poll_fn=lambda: self._runtime.get_training_operation_result( + trainingSessionName=self.training_session_name, + operationId=operation_id, + ), + parse_fn=parse_fn, + status_key="status", + operation_type="training", + ) + + submit_future = self._submitter.submit_ordered(run_blocking(_wire_submit)) + return PendingAPIFuture(submit_future, operation_type="training", keepalive=self._handle) + + @property + def current_sequence_number(self) -> int: + """The last sequence number assigned to an operation (read-only). + + Safe to read for diagnostics: unlike the internal sequence allocator, + this does NOT advance the counter, so inspecting it never consumes a + sequence the service expects to see dispatched. + """ + return self._sequence_counter + + def __enter__(self) -> "TrainingClient": + """Return this client, so it can be used as a context manager.""" + return self + + def __exit__(self, exc_type: Any, exc_val: Any, exc_tb: Any) -> None: + """Stop the session on leaving the block.""" + # A failed stop must not replace whatever exception is leaving the + # block, but it must not vanish either: the session is still running. + try: + self.stop() + except Exception: + logger.warning( + "Failed to stop session %r on exit; it will be stopped when " + "this client is collected, or after its inactivity timeout.", + self.training_session_name, + exc_info=True, + ) diff --git a/sagemaker-train/src/sagemaker/train/training_session/types.py b/sagemaker-train/src/sagemaker/train/training_session/types.py new file mode 100644 index 0000000000..a7d8f577e6 --- /dev/null +++ b/sagemaker-train/src/sagemaker/train/training_session/types.py @@ -0,0 +1,1470 @@ +# 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. +"""Configuration and result types for the training_session client interface.""" + +from __future__ import annotations + +import logging +from datetime import datetime +from enum import Enum +from typing import Any, Dict, List, Optional, Union + +from pydantic import ( + BaseModel, + ConfigDict, + Field, + PrivateAttr, + field_validator, + model_serializer, + model_validator, +) + +logger = logging.getLogger(__name__) + +# Default LoRA rank for the SDK (the service's native default is 8). +DEFAULT_LORA_RANK = 32 +DEFAULT_LORA_ALPHA = 64 + + +class _ForbidExtraBase(BaseModel): + """Internal shared base for customer-authored types: unknown fields are rejected. + + Rejecting extras catches construction typos (e.g. ``max_token`` for + ``max_tokens``) at object-creation time instead of silently dropping them. + Not exported: an implementation detail, not API surface. + """ + + model_config = ConfigDict(extra="forbid") + + +class _AllowExtraBase(BaseModel): + """Internal shared base for permissive types: unknown fields are kept. + + Used by the types constructed from external objects, where fields we + don't model must not be an error. Not exported: an implementation detail, + not API surface. + """ + + model_config = ConfigDict(extra="allow") + + +def _params_from_attrs(cls: Any, value: Any) -> Any: + """Build an all-optional params type from an arbitrary object's attributes. + + The types this serves (:class:`AdamParams`, :class:`SamplingParams`) have a + default for every field, so there is no required attribute to read and fail + on the way :meth:`VpcConfig.from_any` does. Reading each field with a default + fallback on its own would therefore accept *any* object at all and quietly + substitute the defaults, so an object using different names (``maxTokens``, + ``temp``) would look like it had been honoured while none of its values were + used. + + Sharing at least one field name is what distinguishes a params-like object + from an unrelated one. Fields the object omits still fall back to the + defaults, which is what lets a caller pass a type modelling fewer knobs. + """ + present = [field for field in cls.model_fields if hasattr(value, field)] + if not present: + raise TypeError( + f"cannot read {cls.__name__} fields from {type(value).__name__}: none of " + f"{sorted(cls.model_fields)} are present. Pass a {cls.__name__}, a dict, " + f"or an object using these field names." + ) + return cls(**_drop_unset({field: getattr(value, field) for field in present})) + + +def _drop_unset(fields: Dict[str, Any]) -> Dict[str, Any]: + """Drop ``None`` entries so the field's own default applies instead. + + A params type that models a knob as optional carries ``None`` for one the + caller never set. Forwarding that ``None`` would fail validation on a field + of ours that is required with a default; ``max_tokens`` is the one that + bites, since the service requires it. Treating "not set" as "use the + default" is the same outcome the caller would get by omitting the field. + """ + return {key: val for key, val in fields.items() if val is not None} + + +class SessionStatus(str, Enum): + """Status of a training session. + + Use with :meth:`ServiceClient.list_sessions` to filter by status without + memorising the service's raw PascalCase strings. + """ + + CREATED = "Created" + SCHEDULED = "Scheduled" + RUNNING = "Running" + STOPPING = "Stopping" + STOPPED = "Stopped" + FAILED = "Failed" + BLOCKED = "Blocked" + EXPIRED = "Expired" + DELETING = "Deleting" + DELETION_FAILED = "DeletionFailed" + + +class ModelPackageKind(str, Enum): + """Kind of model package a save operation registers. + + A session registers one of two kinds of model package: + + * ``CHECKPOINT``: a ``save_state`` checkpoint: sharded LoRA weights plus + optimizer/FSDP state. Use it to resume an interrupted training run via + :meth:`TrainingClient.load_state_with_optimizer` (weights + optimizer) or + :meth:`TrainingClient.load_state` (weights only, optimizer resets). + * ``MODEL``: a ``save_weights_for_sampler`` model: the full (unsharded) LoRA + weights. This is the deployable artifact (no optimizer state), so use it + for downstream tasks like inference/sampling or to start a new training + run. + """ + + CHECKPOINT = "checkpoint" + MODEL = "model" + + +class ModelPackageInfo(_ForbidExtraBase): + """Summary of one model package version in the session's group. + + A lightweight view (ARN, version, creation time, kind), distinct from the + full ``sagemaker.core.resources.ModelPackage`` resource. Returned by + :meth:`ServiceClient.list_model_packages` and + :meth:`ServiceClient.get_latest_model_package`. Pass :attr:`model_package_arn` + straight to :meth:`TrainingClient.load_state` / + :meth:`TrainingClient.load_state_with_optimizer` (resume training), or to + :meth:`ServiceClient.create_sampling_client` (inference). + + Args: + model_package_arn: ARN of the model package version. + version: Model package version number (monotonic within the group). + creation_time: When the model package was registered. + description: The model package description set when it was saved + (a customer-provided string, or a platform default). + kind: Whether this is a ``CHECKPOINT`` (save_state) or ``MODEL`` + (save_weights_for_sampler) package. ``None`` when not resolved: the field + is only available from ``DescribeModelPackage``, so + :meth:`ServiceClient.list_model_packages` leaves it unset (it lists + without describing); the ``get_latest_*`` helpers always populate it. + """ + + model_package_arn: str + version: Optional[int] = None + creation_time: Optional[datetime] = None + description: Optional[str] = None + kind: Optional[ModelPackageKind] = None + + @classmethod + def _from_summary(cls, summary: Dict[str, Any]) -> "ModelPackageInfo": + """Build from a ``ListModelPackages`` summary (no ``kind``, which the summary omits).""" + return cls( + model_package_arn=summary["ModelPackageArn"], + version=summary.get("ModelPackageVersion"), + creation_time=summary.get("CreationTime"), + description=summary.get("ModelPackageDescription"), + ) + + +class TrainingSessionSummary(_ForbidExtraBase): + """Summary of one training session, as listed. + + Returned by :meth:`ServiceClient.list_sessions`. Carries the fields the list + reports; anything else about a session comes from describing it. + + Args: + training_session_arn: ARN of the session. + status: Current session status. + creation_time: When the session was created. + last_modified_time: When the session last changed state. + """ + + training_session_arn: str + #: ``str`` is part of the type so a status this version does not know about + #: still reaches the caller rather than failing the whole list call. + status: Optional[Union[SessionStatus, str]] = None + creation_time: Optional[datetime] = None + last_modified_time: Optional[datetime] = None + + @property + def training_session_name(self) -> str: + """The session name, taken from the tail of :attr:`training_session_arn`.""" + return self.training_session_arn.rsplit("/", 1)[-1] + + @classmethod + def _from_summary(cls, summary: Dict[str, Any]) -> "TrainingSessionSummary": + """Build from a ``ListTrainingSessions`` summary. + + An unrecognised status is kept as its raw string rather than dropped, so + a status this version does not know about still reaches the caller. + """ + raw_status = summary.get("TrainingSessionStatus") + try: + status = SessionStatus(raw_status) if raw_status is not None else None + except ValueError: + status = raw_status + return cls( + training_session_arn=summary["TrainingSessionArn"], + status=status, + creation_time=summary.get("CreationTime"), + last_modified_time=summary.get("LastModifiedTime"), + ) + + +class ModelPackageDeletion(_ForbidExtraBase): + """Outcome of :meth:`ServiceClient.delete_model_package`. + + Args: + model_package_arn: ARN of the model package that was deleted. + model_package_deleted: True when this call deleted the package. False + means it was already gone, so there was nothing left to delete. + s3_uri: The S3 location the model package pointed at, or ``None`` when + the package exposes none. When the artifacts were kept, this is the + only remaining record of where they are. + objects_deleted: Number of S3 objects deleted under :attr:`s3_uri`. Zero + when the artifacts were kept, or when the location was already + empty. On a versioned bucket a delete adds a marker instead of + reclaiming storage, so this counts objects, not bytes freed. + artifacts_retained: True when this call did not attempt to delete the + artifacts, either because ``delete_artifacts`` was False or + because they are not in a location you manage. False with + :attr:`objects_deleted` at zero means the location was already + empty. + retained_reason: Human-readable reason the artifacts were kept, or + ``None`` when they were deleted. Diagnostic text for a log; read + :attr:`artifacts_retained` to branch on. + """ + + model_package_arn: str + model_package_deleted: bool + s3_uri: Optional[str] = None + objects_deleted: int = 0 + artifacts_retained: bool = True + retained_reason: Optional[str] = None + + +class VpcConfig(_ForbidExtraBase): + """VPC configuration for a training session. + + When set, the session's compute runs inside your VPC: traffic to your + Amazon S3 output location and to other resources flows through the + subnets and security groups you specify, so make sure they allow the + required access (for example an S3 gateway endpoint). + + Args: + subnets: IDs of the subnets to run the session in. + security_group_ids: IDs of the security groups to apply. + """ + + subnets: List[str] + security_group_ids: List[str] + + @classmethod + def from_any(cls, value: Any) -> "VpcConfig": + """Coerce a VPC-config-like value to VpcConfig. + + Accepts a :class:`VpcConfig` (passthrough), a dict of fields, or any + object with ``subnets`` / ``security_group_ids`` attributes. + """ + if isinstance(value, cls): + return value + if isinstance(value, dict): + return cls(**value) + return cls( + subnets=list(getattr(value, "subnets")), + security_group_ids=list(getattr(value, "security_group_ids")), + ) + + def _to_api(self) -> Dict[str, Any]: + """Convert to the CreateTrainingSession VpcConfig shape.""" + return { + "Subnets": list(self.subnets), + "SecurityGroupIds": list(self.security_group_ids), + } + + +class TensorDtype(str, Enum): + """Data type for tensor values in loss function inputs. + + Maps to the dtypes the runtime accepts in ``TensorData.dtype``. + """ + + FLOAT32 = "float32" + FLOAT16 = "float16" + BFLOAT16 = "bfloat16" + INT32 = "int32" + INT64 = "int64" + + +#: Dtype spellings that name a wire dtype the enum does not answer to. +_DTYPE_ALIASES = {"half": TensorDtype.FLOAT16} + +#: Float dtypes with no wire slot, sent as ``float32``. Narrowing, so it warns. +#: Any name starting with ``float`` also lands here, which covers the widths +#: that keep being added (``float8_e4m3fn``, ``float128``). +_NARROWED_TO_FLOAT32 = frozenset({"double", "float", "longdouble"}) + +#: Integer and boolean dtypes with no wire slot, sent as ``int64``. Widening or +#: exact, so it is silent. Booleans land here rather than on a float dtype: +#: a mask reaches the loss function as the 0/1 integers it is. +_WIDENED_TO_INT64 = frozenset( + { + "bool", + "int8", + "int16", + "uint8", + "uint16", + "uint32", + "uint64", + "int", + "long", + "short", + } +) + + +def _dtype_name(dtype: Any) -> str: + """Reduce a dtype spec to a bare type name. + + Covers all three spellings: ``float32`` (numpy dtype or wire string), + ``torch.float32``, and ```` (a scalar *type*, not a + dtype instance). The last is unwrapped before the name is taken, so the + result is always a bare name the dtype tables can match. + """ + name = str(dtype) + if name.startswith(""): + name = name[len("")] + return name.rsplit(".", 1)[-1] + + +def _dtype_kind(dtype: Any, name: str) -> Optional[str]: + """Classify a dtype as ``"f"`` (float), ``"i"`` (integral), or ``None``. + + Asks the dtype object itself before falling back to its spelling, so a + width the name tables never heard of still classifies correctly: torch + keeps adding float types (``float8_e4m3fn``), which a fixed name list + would reject. + """ + # torch: torch.dtype carries these as plain bools. + is_float = getattr(dtype, "is_floating_point", None) + if isinstance(is_float, bool): + if is_float: + return "f" + # Complex is neither, and must not fall through to the integer branch: + # torch reports is_floating_point False for it. + return None if getattr(dtype, "is_complex", False) else "i" + # numpy: dtype.kind is a one-character code. + kind = getattr(dtype, "kind", None) + if isinstance(kind, str) and kind: + if kind == "f": + return "f" + return "i" if kind in ("i", "u", "b") else None + # A bare spelling, e.g. from a wire dict's "dtype" field. + if name in _NARROWED_TO_FLOAT32 or name.startswith("float"): + return "f" + return "i" if name in _WIDENED_TO_INT64 else None + + +def _flatten_nested(values: List[Any]) -> "tuple[List[Any], List[int]]": + """Fully flatten a rectangular nested list, returning the values and shape. + + Serves the tensor-materialization path, where the rank is whatever the + source tensor had. Raw ``loss_fn_inputs`` lists go through + :func:`_flatten_with_shape` instead, which caps the rank and names the + offending key. + """ + shape: List[int] = [] + probe: Any = values + while isinstance(probe, list): + shape.append(len(probe)) + probe = probe[0] if probe else None + flat: List[Any] = [] + + def walk(node: Any, depth: int) -> None: + if isinstance(node, list): + if depth >= len(shape) or len(node) != shape[depth]: + raise ValueError( + f"tensor data is ragged: expected {shape} but a row at depth " + f"{depth} has {len(node)} entries; rows must all be the same length" + ) + for child in node: + walk(child, depth + 1) + elif depth != len(shape): + raise ValueError( + f"tensor data mixes nesting depths: expected {shape} but found a " + f"value {len(shape) - depth} level(s) above the deepest row" + ) + else: + flat.append(node) + + walk(values, 0) + return flat, shape + + +class TensorData(_ForbidExtraBase): + """A named tensor with dtype and shape metadata. + + Use in ``loss_fn_inputs`` to pass typed tensor data (advantages, weights, + logprobs, etc.) alongside training datums. + + Tensor-like values (torch tensors, numpy arrays, anything with + ``tolist()``) are held as-is and only flattened to a plain list when the + request is serialized or :attr:`data` is first read, so building datums + from large tensors does not copy them up front. + + Args: + data: Flat list of numeric values, or a tensor-like object converted + lazily. + dtype: Tensor element type. + shape: Tensor dimensions (e.g. ``[batch_size, seq_len]``). + """ + + dtype: TensorDtype = TensorDtype.FLOAT32 + shape: Optional[List[int]] = None + + _data: Optional[List[Any]] = PrivateAttr(default=None) + _source: Any = PrivateAttr(default=None) + + def __init__(self, data: Any = None, **kwargs: Any) -> None: + if data is None: + raise ValueError("TensorData requires 'data'") + super().__init__(**kwargs) + if isinstance(data, list): + self._data = data + else: + self._source = data + + @property + def data(self) -> List[Any]: + """Tensor values as a flat list (materialized from the source on first access).""" + if self._data is None: + source = self._source + if hasattr(source, "tolist"): # torch.Tensor / numpy.ndarray + raw = source.tolist() + self._data = raw if isinstance(raw, list) else [raw] + else: + self._data = list(source) + if any(isinstance(v, list) for v in self._data): + # The wire carries a flat ``data`` array alongside ``shape``, so a + # multi-dimensional source has to be flattened. Left nested, it was + # sent as a list of lists and rejected: the runtime types ``data`` as + # an array of numbers. Done once, in place of the nested value. + flat, nested_shape = _flatten_nested(self._data) + self._data = flat + if self.shape is None: + self.shape = nested_shape + return self._data + + def __eq__(self, other: Any) -> bool: + """Compare on values, dtype and shape, whichever form each was built in. + + The values live outside the model fields and are materialized lazily, so + the inherited comparison saw ``_source`` on one side and ``_data`` on the + other and called two equal tensors different until both had been read. + Reading through :attr:`data` normalizes that, and also settles ``shape``. + """ + if not isinstance(other, TensorData): + return NotImplemented + return self.data == other.data and self.dtype == other.dtype and self.shape == other.shape + + # No __hash__ here on purpose. Defining __eq__ already makes this + # unhashable, which is what the mutable values call for and what every other + # model here does. Supplying one that raises instead would report the type as + # hashable and then fail, sending any check for hashability down the wrong + # branch. + + @model_serializer(mode="wrap") + def _serialize_with_data(self, handler: Any) -> Dict[str, Any]: + """Include the tensor values in ``model_dump()``. + + The values are held outside the model fields so a large tensor is not + copied when a datum is built. That kept them out of the dump entirely, + which made a dumped tensor unusable: it round-tripped to a + ``TensorData`` with no values. Materializing here is the first point + where the caller has asked to see them. + """ + dumped = dict(handler(self)) + dumped["data"] = list(self.data) + return dumped + + @field_validator("dtype", mode="before") + @classmethod + def _validate_dtype(cls, dtype: Any) -> Any: + """Normalize whatever was passed as ``dtype`` before pydantic checks it. + + Runs on the field, so direct construction, + ``TensorData(data=t.tolist(), dtype=t.dtype)``, gets the same coercion + as :meth:`from_any` instead of failing pydantic enum validation because + the field itself names only the wire spellings. + """ + return cls._coerce_dtype(dtype) if dtype is not None else dtype + + @classmethod + def _coerce_dtype(cls, dtype: Any) -> TensorDtype: + """Normalize a dtype spec (enum, string, or torch/numpy dtype) to TensorDtype. + + A dtype the wire has no slot for is widened to one it does: a float to + ``float32`` (narrowing ``float64``, which warns), an integer or boolean + to ``int64``. A dtype that is not numeric at all raises rather than + going out as ``float32``, which the service would reject or, worse, + accept with reinterpreted values. + """ + if isinstance(dtype, TensorDtype): + return dtype + name = _dtype_name(dtype) + try: + return TensorDtype(name) + except ValueError: + pass + aliased = _DTYPE_ALIASES.get(name) + if aliased is not None: + return aliased + kind = _dtype_kind(dtype, name) + if kind == "f": + logger.warning( + "dtype %s is not supported on the wire; sending as float32. " + "Cast to float32 yourself to control the rounding.", + name, + ) + return TensorDtype.FLOAT32 + if kind == "i": + return TensorDtype.INT64 + raise ValueError( + f"unsupported tensor dtype {name!r}: expected a float or integer " + f"dtype (one of {', '.join(d.value for d in TensorDtype)}, or a " + f"numpy/torch dtype convertible to one)" + ) + + @classmethod + def from_any(cls, value: Any) -> "TensorData": + """Coerce a tensor-like value to a TensorData. + + Accepts a :class:`TensorData` (passthrough), a wire-format dict, a + torch tensor / numpy array (anything with ``tolist()``), a plain list, + or any object with ``data``/``dtype`` attributes. + """ + if isinstance(value, cls): + return value + if isinstance(value, dict): + return cls( + data=value["data"], + dtype=cls._coerce_dtype(value.get("dtype", TensorDtype.FLOAT32)), + shape=value.get("shape"), + ) + shape = getattr(value, "shape", None) + if shape is not None and not isinstance(shape, list): + shape = list(shape) # torch.Size / numpy tuple + if hasattr(value, "tolist"): # torch.Tensor / numpy.ndarray, converted lazily + raw: Any = value + else: + raw = getattr(value, "data", value) + if not isinstance(raw, list) and not hasattr(raw, "tolist"): + raw = list(raw) + return cls( + data=raw, + dtype=cls._coerce_dtype(getattr(value, "dtype", TensorDtype.FLOAT32)), + shape=shape, + ) + + def tolist(self) -> List[Any]: + """Return the tensor values as a plain list.""" + return self.data + + def to_torch(self) -> Any: + """Convert to a ``torch.Tensor`` (requires torch).""" + import torch + + dtype_map = { + TensorDtype.FLOAT32: torch.float32, + TensorDtype.FLOAT16: torch.float16, + TensorDtype.BFLOAT16: torch.bfloat16, + TensorDtype.INT32: torch.int32, + TensorDtype.INT64: torch.int64, + } + tensor = torch.tensor(self.data, dtype=dtype_map[self.dtype]) + return tensor.reshape(self.shape) if self.shape is not None else tensor + + def to_numpy(self) -> Any: + """Convert to a ``numpy.ndarray`` (requires numpy).""" + import numpy as np + + # numpy has no bfloat16; fall back to float32 for it. + dtype_name = "float32" if self.dtype is TensorDtype.BFLOAT16 else self.dtype.value + arr = np.asarray(self.data, dtype=np.dtype(dtype_name)) + return arr.reshape(self.shape) if self.shape is not None else arr + + def _to_api(self) -> Dict[str, Any]: + """Convert to the runtime TensorData shape.""" + # Read first, not inline below: materializing can fill in ``shape`` from + # the nesting, and that has to happen before ``shape`` is checked. + data = self.data + d: Dict[str, Any] = {"data": data, "dtype": self.dtype.value} + if self.shape is not None: + d["shape"] = self.shape + return d + + +#: Wire dtype for a raw numeric list passed under a known loss-fn input name. +#: A list carries no dtype of its own, and the choice is not inferable from the +#: values: ``[0, 1, 1]`` is a valid weight mask and a valid token list. +_LOSS_FN_INPUT_DTYPES = { + "target_tokens": TensorDtype.INT64, + "weights": TensorDtype.FLOAT32, + "advantages": TensorDtype.FLOAT32, + "logprobs": TensorDtype.FLOAT32, + "old_logprobs": TensorDtype.FLOAT32, + "ref_log_prob": TensorDtype.FLOAT32, + "temperature": TensorDtype.FLOAT32, +} + + +def _flatten_with_shape(key: str, values: List[Any]) -> "tuple[List[Any], List[int]]": + """Flatten a 1-D or rectangular 2-D list, returning the values and shape. + + The wire carries a flat ``data`` array plus a ``shape``, so a nested list has + to be flattened here rather than sent as-is. + """ + rows = [v for v in values if isinstance(v, list)] + if not rows: + return values, [len(values)] + if len(rows) != len(values): + raise ValueError( + f"loss_fn_inputs[{key!r}] mixes numbers and lists; pass either a flat " + f"list of values or a list of equal-length rows" + ) + widths = {len(r) for r in rows} + if len(widths) != 1: + raise ValueError( + f"loss_fn_inputs[{key!r}] is a ragged nested list (row lengths " + f"{sorted(len(r) for r in rows)}); rows must all be the same length" + ) + if any(isinstance(v, list) for r in rows for v in r): + raise ValueError( + f"loss_fn_inputs[{key!r}] nests more than two levels deep; pass a " + f"tensor or a flat list instead" + ) + return [v for r in rows for v in r], [len(rows), widths.pop()] + + +def _loss_fn_input_dtype(key: str, flat: List[Any]) -> TensorDtype: + """Wire dtype for a raw list under ``key``: by name, else inferred.""" + known = _LOSS_FN_INPUT_DTYPES.get(key) + if known is not None: + return known + # bool is a subclass of int, and a mask of booleans is integral. + if flat and all(isinstance(v, (int, bool)) for v in flat): + return TensorDtype.INT64 + return TensorDtype.FLOAT32 + + +def _serialize_loss_fn_inputs(loss_fn_inputs: Dict[str, Any]) -> Dict[str, Any]: + """Serialize ``loss_fn_inputs`` values to wire shapes. + + :class:`TensorData` and tensor-like values (torch tensors, numpy arrays, + objects with ``data``/``dtype`` attributes) become the runtime TensorData + dict. A raw list of numbers becomes one too, taking its dtype from the input + name: ``target_tokens`` is integral, the rest are floating. A dict already + in the runtime TensorData shape passes through unchanged. + """ + out: Dict[str, Any] = {} + for key, value in loss_fn_inputs.items(): + if isinstance(value, TensorData): + out[key] = value._to_api() + elif hasattr(value, "tolist") and hasattr(value, "shape"): + # torch.Tensor / numpy.ndarray, detected without converting + out[key] = TensorData.from_any(value)._to_api() + elif not isinstance(value, dict) and hasattr(value, "data") and hasattr(value, "dtype"): + out[key] = TensorData.from_any(value)._to_api() + elif isinstance(value, list): + # The wire wants a TensorData here, never a bare array. Left to pass + # through, a plain list was rejected before the request went out. + flat, shape = _flatten_with_shape(key, value) + out[key] = TensorData( + data=flat, dtype=_loss_fn_input_dtype(key, flat), shape=shape + )._to_api() + else: + out[key] = value + return out + + +class EncodedTextChunk(_ForbidExtraBase): + """A chunk of token IDs representing encoded text. + + This is the fundamental building block for model inputs: a sequence of + token IDs. Multiple chunks can be combined into a :class:`ModelInputChunk` + list to represent structured prompts. + + Args: + tokens: List of token IDs. + """ + + tokens: List[int] + + @property + def length(self) -> int: + """Number of tokens in this chunk.""" + return len(self.tokens) + + @classmethod + def from_any(cls, value: Any) -> "EncodedTextChunk": + """Coerce a chunk-like value to an EncodedTextChunk. + + Accepts an :class:`EncodedTextChunk` (passthrough), a + :class:`ModelInputChunk` wrapper (or any object with the same attributes), + a wire dict, a plain token ID list, or any object with ``tokens``. + + Raises: + TypeError: If the value carries no token IDs. Rejected instead of + coerced to an empty chunk, which would send a prompt the model + never sees and report success. + """ + if isinstance(value, cls): + return value + if isinstance(value, str): + raise TypeError( + "a chunk holds token IDs, not text; encode the text with the " + f"model's tokenizer first (got {value!r})" + ) + if isinstance(value, dict): + if "chunk" in value: + return cls.from_any(value["chunk"]) + if "tokens" not in value: + raise TypeError(f"chunk dict carries no 'tokens': {value!r}") + return cls(tokens=list(value["tokens"])) + if isinstance(value, list): + return cls(tokens=list(value)) + inner = getattr(value, "chunk", None) # ModelInputChunk wrapper + if inner is not None: + return cls.from_any(inner) + tokens = getattr(value, "tokens", None) + if tokens is None: + raise TypeError( + f"cannot read token IDs from {type(value).__name__}; pass a " + "list of token IDs or an EncodedTextChunk" + ) + return cls(tokens=list(tokens)) + + def _to_api(self) -> Dict[str, Any]: + """Convert to the runtime chunk shape.""" + return {"tokens": self.tokens, "type": "encoded_text"} + + +class ModelInputChunk(_AllowExtraBase): + """A single chunk within a structured model input. + + Wraps an :class:`EncodedTextChunk` (the only chunk type today). This + layer exists so the type system can accommodate future chunk types + (e.g. image embeddings) without breaking existing code. + + Args: + chunk: The encoded text chunk. + """ + + chunk: EncodedTextChunk + + def _to_api(self) -> Dict[str, Any]: + """Convert to the runtime chunk shape.""" + return self.chunk._to_api() + + +class ModelInput(_AllowExtraBase): + """Structured model input composed of one or more chunks. + + Use as the ``prompt`` argument to :meth:`SamplingClient.sample` or to + build typed training datums. Accepts either pre-built chunks or a + simple prompt (token IDs or text string) via the convenience + constructors :meth:`from_ints` and :meth:`from_text`. + + Args: + chunks: List of input chunks. Each chunk is either a + :class:`ModelInputChunk`, an :class:`EncodedTextChunk`, or a + raw dict matching the wire format. + """ + + chunks: Optional[List[Any]] = None + prompt: Optional[Any] = None + + @model_validator(mode="after") + def _reject_both_inputs(self) -> "ModelInput": + """One input or the other, for the reason a Datum takes one or the other. + + The wire shape is built from ``chunks`` when both are set, so a prompt + passed alongside them goes nowhere. An empty prompt carries nothing and + is not a second input. + """ + if self.chunks is not None and self.prompt: + raise ValueError( + "a ModelInput takes either chunks or prompt, not both. Pass " + "chunks for pre-built input, or prompt for token IDs or text." + ) + return self + + @classmethod + def from_ints(cls, tokens: List[int]) -> "ModelInput": + """Build a ModelInput from a flat list of token IDs.""" + return cls(chunks=[EncodedTextChunk(tokens=list(tokens))]) + + @classmethod + def from_text(cls, text: str) -> "ModelInput": + """Build a ModelInput from text. + + The text is tokenized when the input is submitted, using the session's + tokenizer, the same handling a bare string gets. Passing one to a + client that cannot resolve a tokenizer raises there. + """ + return cls(prompt=text) + + @classmethod + def empty(cls) -> "ModelInput": + """Build a ModelInput with no chunks.""" + return cls(chunks=[]) + + @classmethod + def from_any(cls, value: Any) -> "ModelInput": + """Coerce a prompt-like value to a ModelInput. + + Accepts a :class:`ModelInput` (passthrough), an + :class:`EncodedTextChunk`, a wire dict (``chunks``, ``prompt``, or a + single chunk's ``tokens``), a token ID list, or any object with + ``chunks``. + + Raises: + ValueError: A dict naming none of those keys. Such a dict once became + an input with no chunks, which the service accepts, so a training + example was dropped without anything to see. + """ + if isinstance(value, cls): + return value + if isinstance(value, (EncodedTextChunk, ModelInputChunk)): + return cls(chunks=[EncodedTextChunk.from_any(value)]) + if isinstance(value, dict): + if value.get("chunks") is not None and value.get("prompt"): + raise ValueError( + "a ModelInput takes either chunks or prompt, not both. Pass " + "chunks for pre-built input, or prompt for token IDs or text." + ) + if value.get("chunks") is not None: + return cls(chunks=[EncodedTextChunk.from_any(c) for c in value["chunks"]]) + if "prompt" in value: + return cls(prompt=value["prompt"]) + if "tokens" in value: + # A single chunk rather than a whole input. + return cls(chunks=[EncodedTextChunk.from_any(value)]) + raise ValueError( + f"cannot build a model input from a dict with keys {sorted(value)}: " + "pass 'chunks', 'prompt', or a chunk's 'tokens'" + ) + if isinstance(value, list): + return cls.from_ints(value) + raw_chunks = getattr(value, "chunks", None) + if raw_chunks is not None: + return cls(chunks=[EncodedTextChunk.from_any(c) for c in raw_chunks]) + return cls(prompt=value) + + @property + def length(self) -> int: + """Total number of tokens across all chunks. + + Raises: + ValueError: This input has no token IDs to count, for the reasons + :meth:`to_ints` gives. + """ + return len(self.to_ints()) + + def to_ints(self) -> List[int]: + """Flatten all chunks to a single token ID list. + + Raises: + ValueError: This input carries a prompt that is not token IDs, so + it has none to return yet. Returning an empty list instead + would read a full prompt as empty wherever a token budget is + counted. + """ + if self._pending_text is not None: + raise ValueError( + "this ModelInput holds text that has not been tokenized yet, so " + "it has no token IDs. Build it with " + "ModelInput.from_ints(tokenizer.encode(text)) first." + ) + if self.chunks is not None: + out: List[int] = [] + for c in self.chunks: + out.extend(EncodedTextChunk.from_any(c).tokens) + return out + if isinstance(self.prompt, list): + return list(self.prompt) + if self.prompt is not None: + raise ValueError( + f"this ModelInput holds a prompt of type " + f"{type(self.prompt).__name__}, which is not token IDs, so it " + f"has none to return. Pass token IDs or text." + ) + return [] + + def append(self, chunk: Any) -> "ModelInput": + """Return a new ModelInput with ``chunk`` appended (does not mutate). + + An input built from a prompt of token IDs carries those tokens into the + chunks first, so appending extends it rather than replacing it. + + Raises: + ValueError: This input still holds untokenized text. Appending built + the result from ``chunks`` alone, so the text was dropped and the + appended chunk became the whole input. + """ + if self._pending_text is not None: + raise ValueError( + "cannot append to a ModelInput built from text: the text has not " + "been tokenized yet, so there are no chunks to append to. Build it " + "with ModelInput.from_ints(tokenizer.encode(text)) first." + ) + existing = list(self.chunks or []) + if not existing and self.prompt is not None: + # A token-ID prompt is this input's whole content, the same way + # ``_to_api`` reads it, so it has to survive the append. + existing = [EncodedTextChunk.from_any(self.prompt)] + return ModelInput(chunks=existing + [EncodedTextChunk.from_any(chunk)]) + + def append_int(self, token: int) -> "ModelInput": + """Return a new ModelInput with a single-token chunk appended.""" + return self.append(EncodedTextChunk(tokens=[token])) + + @property + def _pending_text(self) -> Optional[str]: + """The text awaiting tokenization, or ``None`` if there is none.""" + return self.prompt if isinstance(self.prompt, str) else None + + def _tokenized_with(self, tokenizer: Any) -> "ModelInput": + """Return an equivalent ModelInput with any pending text tokenized.""" + text = self._pending_text + if text is None: + return self + return ModelInput.from_ints(tokenizer.encode(text)) + + def _to_api(self) -> Dict[str, Any]: + """Convert to the runtime ModelInput shape (``{"chunks": [...]}``). + + If constructed with ``chunks``, serializes them. If constructed with a + raw ``prompt`` of token IDs, wraps it in the chunks format. + + Raises: + ValueError: The input still holds untokenized text. Only a tokenizer + can turn it into chunks, and this type has none. Reaching the + wire, the text was sent as a field the runtime does not define + and the request was rejected. + """ + if self.chunks is not None: + api_chunks = [] + for c in self.chunks: + if isinstance(c, (ModelInputChunk, EncodedTextChunk)): + api_chunks.append(c._to_api()) + elif isinstance(c, dict): + api_chunks.append(c) + else: + api_chunks.append(EncodedTextChunk.from_any(c)._to_api()) + return {"chunks": api_chunks} + if self.prompt is not None: + if isinstance(self.prompt, list): + return {"chunks": [{"tokens": self.prompt, "type": "encoded_text"}]} + if self._pending_text is not None: + raise ValueError( + "this ModelInput holds text that was never tokenized. Submit " + "it through a client method (which tokenizes for you), or " + "build it from token IDs with ModelInput.from_ints(...)." + ) + raise TypeError( + f"cannot build a model input from {type(self.prompt).__name__}: " + f"pass token IDs, text, or pre-built chunks" + ) + return {"chunks": []} + + +class Datum(_AllowExtraBase): + """A single training datum (model input + optional loss function inputs). + + Use as the batch element for :meth:`TrainingClient.forward_backward` and + :meth:`TrainingClient.forward`. + + Can be constructed from: + - A :class:`ModelInput` with optional loss_fn_inputs + - Simple token IDs (via ``prompt``) for convenience + + Args: + model_input: Structured input (chunks-based). Mutually exclusive with + ``prompt``. + prompt: Convenience shorthand for token IDs or text. Converted to a + :class:`ModelInput` internally. + loss_fn_inputs: Per-datum loss function inputs, keyed by name: + ``target_tokens``, ``weights``, ``advantages``, ``logprobs``, + ``ref_log_prob``, ``temperature``. Which are required depends on the + loss function. Values can be :class:`TensorData` objects, tensors, or + raw lists. Per-token entries must be as long as the token sequence. + """ + + # Merges with the base config; clears the pydantic "model_" prefix guard + # so the ``model_input`` field name is allowed. + model_config = ConfigDict(protected_namespaces=()) + + model_input: Optional[ModelInput] = None + prompt: Optional[Any] = None + loss_fn_inputs: Optional[Dict[str, Any]] = None + + @model_validator(mode="after") + def _reject_both_inputs(self) -> "Datum": + """Enforce the documented exclusivity of ``model_input`` and ``prompt``. + + Both set, the wire shape was built from ``model_input`` and the prompt + went nowhere, so a caller who meant to send it trained on something else. + + Raises: + ValueError: Both are set. + """ + # An empty prompt carries no input, so it is not a second one. Objects + # coerced through ``from_any`` commonly declare ``prompt = ""`` as a + # default, and rejecting those would turn a working call into an error + # while dropping nothing. + if self.model_input is not None and self.prompt: + raise ValueError( + "a Datum takes either model_input or prompt, not both. Pass " + "model_input for chunked input, or prompt for token IDs or text." + ) + return self + + @model_validator(mode="before") + @classmethod + def _reject_datum_weight(cls, data: Any) -> Any: + """Reject a ``weight``, pointing at the loss input that carries one. + + A datum weight is applied through ``loss_fn_inputs["weights"]``, one + value per token, not as a scalar on the datum. This type accepts extra + fields, so an unchecked ``weight`` would sit on the object and never be + sent, leaving out the weighting the caller asked for. + """ + if isinstance(data, dict) and "weight" in data: + raise ValueError( + "Datum has no 'weight' field. Weight a datum through " + "loss_fn_inputs={'weights': [...]}, one value per token of " + "model_input, e.g. 0.0 to mask a token out of the loss." + ) + return data + + @classmethod + def from_any(cls, value: Any) -> "Datum": + """Coerce a datum-like value to a Datum. + + Accepts a :class:`Datum` (passthrough), a :class:`ModelInput` or a bare + chunk, a token ID list or text (the ``prompt`` shorthand), a wire dict + (``modelInput``/``lossFnInputs`` or snake_case keys), or any object with + ``model_input`` / ``prompt`` / ``loss_fn_inputs`` attributes. + + Raises ``TypeError`` for a value carrying none of those. Every field on + this type is optional, so an unrecognised value would otherwise become + an empty ``Datum`` serializing to ``{}``, dropping a training example + from the batch with nothing raised. A mixed batch reaches this path: + the element kind is decided by the *first* entry, and later entries + still come through here one at a time. + """ + if isinstance(value, cls): + return value + if isinstance(value, (ModelInput, EncodedTextChunk, ModelInputChunk)): + # ModelInput.from_any wraps a bare chunk; a ModelInput passes through. + return cls(model_input=ModelInput.from_any(value)) + if isinstance(value, dict): + mi = value.get("modelInput", value.get("model_input")) + lfi = value.get("lossFnInputs", value.get("loss_fn_inputs")) + prompt = value.get("prompt") + cls._reject_datum_weight(value) + if mi is None and lfi is None and prompt is None: + raise TypeError( + f"cannot read Datum fields from dict with keys " + f"{sorted(map(str, value))}: none of 'modelInput', 'prompt', " + f"'lossFnInputs' are present (snake_case also accepted)." + ) + if mi is not None: + # ``prompt`` is passed through rather than ignored, so a dict + # naming both is rejected by the same rule the constructor + # applies instead of quietly losing one of them. + return cls( + model_input=ModelInput.from_any(mi), + prompt=prompt, + loss_fn_inputs=lfi or None, + ) + return cls(prompt=prompt, loss_fn_inputs=lfi or None) + if isinstance(value, (str, list, tuple)): + # The prompt shorthand documented on this type: token IDs or text. + return cls(prompt=list(value) if isinstance(value, tuple) else value) + mi = getattr(value, "model_input", None) + # ``is not None``, not ``hasattr``: an object declaring ``chunks = None`` + # is not a ModelInput-like value, and treating it as one would wrap the + # object itself as a prompt instead of reaching the TypeError below. + if mi is None and getattr(value, "chunks", None) is not None: + # A bare ModelInput-like object (ours already handled above). + return cls(model_input=ModelInput.from_any(value)) + lfi = getattr(value, "loss_fn_inputs", None) + prompt = getattr(value, "prompt", None) + if mi is None and lfi is None and prompt is None: + raise TypeError( + f"cannot read Datum fields from {type(value).__name__}: none of " + f"'model_input', 'chunks', 'prompt', 'loss_fn_inputs' are present. " + f"Pass a Datum, a ModelInput, a token ID list, text, or a dict." + ) + # A caller's own datum type that models a weight is caught rather than + # having it silently dropped: there is nowhere on the wire to put one. + if getattr(value, "weight", None) is not None: + cls._reject_datum_weight({"weight": getattr(value, "weight")}) + return cls( + model_input=ModelInput.from_any(mi) if mi is not None else None, + prompt=prompt, + loss_fn_inputs=dict(lfi) if lfi else None, + ) + + @property + def _pending_text(self) -> Optional[str]: + """The text awaiting tokenization, or ``None`` if there is none.""" + if isinstance(self.prompt, str): + return self.prompt + return self.model_input._pending_text if self.model_input is not None else None + + def _tokenized_with(self, tokenizer: Any) -> "Datum": + """Return an equivalent Datum with any pending text tokenized.""" + text = self._pending_text + if text is None: + return self + return Datum( + model_input=ModelInput.from_ints(tokenizer.encode(text)), + loss_fn_inputs=self.loss_fn_inputs, + ) + + def _to_api(self) -> Dict[str, Any]: + """Convert to the runtime Datum shape. + + Raises: + ValueError: The datum still holds untokenized text; see + :meth:`ModelInput._to_api`. + """ + if self.model_input is not None: + d: Dict[str, Any] = {"modelInput": self.model_input._to_api()} + elif self.prompt is not None: + # Routed through ModelInput so token IDs, text and unusable values + # are each handled the one way, here and for a bare model input. + d = {"modelInput": ModelInput.from_any(self.prompt)._to_api()} + else: + d = {} + if self.loss_fn_inputs is not None: + d["lossFnInputs"] = _serialize_loss_fn_inputs(self.loss_fn_inputs) + return d + + +class LossFunction(str, Enum): + """Built-in loss functions for training operations. + + Supported values: cross_entropy, importance_sampling, ppo, cispo. + """ + + CROSS_ENTROPY = "cross_entropy" + IMPORTANCE_SAMPLING = "importance_sampling" + PPO = "ppo" + CISPO = "cispo" + + +class AdamParams(_ForbidExtraBase): + """Optimizer parameters for Adam. + + Args: + learning_rate: Step size for parameter updates. + beta1: Exponential decay rate for the first moment estimates. + beta2: Exponential decay rate for the second moment estimates. + eps: Small constant for numerical stability. + weight_decay: L2 regularization coefficient. + grad_clip_norm: Gradient clipping norm. + lora_plus_lr_ratio: Learning-rate multiplier for LoRA+ B matrices + relative to A matrices. Typical values are 8-32. When set, the + service applies ``learning_rate * lora_plus_lr_ratio`` to B + matrices while A matrices use ``learning_rate``. + """ + + model_config = ConfigDict(validate_assignment=True) + + learning_rate: float = 1e-4 + beta1: float = 0.9 + beta2: float = 0.999 + eps: float = 1e-8 + weight_decay: float = 0.01 + grad_clip_norm: Optional[float] = None + lora_plus_lr_ratio: Optional[float] = None + + @classmethod + def from_any(cls, value: Any) -> "AdamParams": + """Coerce a params-like value to AdamParams. + + Accepts an :class:`AdamParams` (passthrough), a dict of fields, or any + object sharing at least one field name. Raises ``TypeError`` for an + object that shares none, instead of silently returning all defaults. + """ + if isinstance(value, cls): + return value + if isinstance(value, dict): + return cls(**_drop_unset(value)) + return _params_from_attrs(cls, value) + + def _to_api(self) -> Dict[str, Any]: + """Convert to the runtime AdamParamsConfig shape.""" + d: Dict[str, Any] = { + "learningRate": self.learning_rate, + "beta1": self.beta1, + "beta2": self.beta2, + "eps": self.eps, + "weightDecay": self.weight_decay, + } + if self.grad_clip_norm is not None: + d["gradClipNorm"] = self.grad_clip_norm + if self.lora_plus_lr_ratio is not None: + d["loraPlusLrRatio"] = self.lora_plus_lr_ratio + return d + + +class SamplingParams(_ForbidExtraBase): + """Parameters for sampling / inference operations. + + Args: + max_tokens: Maximum number of tokens to generate. The service requires + a minimum of 512. + temperature: Sampling temperature. Must not be negative; 0 is greedy + decoding. The service applies its own upper limit. + top_p: Nucleus sampling probability threshold, between 0 and 1. It is a + cumulative probability, so a percentage such as ``95`` is rejected + here rather than at the service. + top_k: Top-k sampling parameter. ``-1`` disables it. + seed: Random seed for reproducibility. + stop: Stop sequences. + """ + + model_config = ConfigDict(validate_assignment=True) + + # Only the two fractional parameters are bounded here, and only where the + # range is a property of the value rather than a service limit. A whole + # number outside its range is refused before the request is sent, so + # max_tokens needs no guard; its ceiling, like the rest of the service's + # limits, is the service's to raise and not worth a stale copy of. top_k + # takes -1 to mean "off", so a floor would refuse the documented way to + # disable it. + max_tokens: int = 512 + temperature: Optional[float] = Field(default=None, ge=0) + top_p: Optional[float] = Field(default=None, ge=0, le=1) + top_k: Optional[int] = None + seed: Optional[int] = None + stop: Optional[List[str]] = None + + @field_validator("stop", mode="before") + @classmethod + def _coerce_stop(cls, value: Any) -> Any: + """Accept a single stop sequence as a bare string. + + The service models ``stop`` as a list of text sequences. A lone + sequence is naturally written without the enclosing list, so wrap it + without making the caller remember the difference. Token IDs are + rejected with the reason, because the service takes text only. + """ + if isinstance(value, str): + return [value] + if isinstance(value, (list, tuple)) and any(isinstance(v, int) for v in value): + raise ValueError( + "stop takes text sequences, not token IDs; the service matches " + "stop sequences on decoded text. Pass the text form instead." + ) + return value + + @classmethod + def from_any(cls, value: Any) -> "SamplingParams": + """Coerce a params-like value to SamplingParams. + + Accepts a :class:`SamplingParams` (passthrough), a dict of fields, or any + object sharing at least one field name. Raises ``TypeError`` for an + object that shares none, instead of silently returning all defaults. + """ + if isinstance(value, cls): + return value + if isinstance(value, dict): + return cls(**_drop_unset(value)) + return _params_from_attrs(cls, value) + + def _to_api(self) -> Dict[str, Any]: + """Convert to the runtime SamplingParams shape.""" + d: Dict[str, Any] = {"maxTokens": self.max_tokens} + if self.temperature is not None: + d["temperature"] = self.temperature + if self.top_p is not None: + d["topP"] = self.top_p + if self.top_k is not None: + d["topK"] = self.top_k + if self.seed is not None: + d["seed"] = self.seed + if self.stop is not None: + d["stop"] = self.stop + return d + + +class StopReason(str, Enum): + """Reason why text generation stopped.""" + + STOP = "stop" + LENGTH = "length" + + +class SampledSequence(_ForbidExtraBase): + """A single generated sequence from a sampling operation. + + Args: + stop_reason: Why generation stopped. + tokens: Generated token IDs. + logprobs: Per-token log probabilities (if requested). + """ + + stop_reason: StopReason + tokens: List[int] + logprobs: Optional[List[float]] = None + + @classmethod + def _from_api(cls, raw: Dict[str, Any]) -> "SampledSequence": + """Construct from a service response SampledSequence shape.""" + return cls( + stop_reason=StopReason(raw["stopReason"]), + tokens=raw["tokens"], + logprobs=raw.get("logprobs"), + ) + + +class TokenLogprob(_ForbidExtraBase): + """A token ID and its log probability.""" + + token_id: int + logprob: float + + @classmethod + def _from_api(cls, raw: Dict[str, Any]) -> "TokenLogprob": + """Construct from a service response TokenLogprob shape.""" + return cls(token_id=raw["tokenId"], logprob=raw["logprob"]) + + +class _OperationResult(_ForbidExtraBase): + """Base accessor over the ``operationResult`` of a completed operation. + + Wraps the raw ``OperationResult`` shape returned by + ``GetTrainingOperationResult`` and keeps it internally. Each operation + resolves to a dedicated subclass that derives its accessors from the + sub-result key the service populates for that operation, so callers never + probe the union themselves. Remains a :class:`~pydantic.BaseModel` for + consistency with the rest of the surface. + """ + + _raw: Dict[str, Any] = PrivateAttr(default_factory=dict) + + @classmethod + def _from_api(cls, raw: Dict[str, Any]): + """Construct from a service ``GetTrainingOperationResult`` operationResult.""" + obj = cls() + obj._raw = raw + return obj + + def _sub_result(self, *keys: str) -> Dict[str, Any]: + """Return the first populated sub-result among ``keys``, or an empty dict.""" + for key in keys: + if key in self._raw: + return self._raw[key] or {} + return {} + + +class ForwardBackwardOutput(_OperationResult): + """Result of a completed ``forward_backward`` / ``forward`` operation.""" + + @property + def metrics(self) -> Dict[str, float]: + """Return scalar training metrics for the pass. + + Contains the throughput metrics (``promised_flops``, + ``global_token_num``, ``mfu``, ``estimated_flops``) plus scalar + mirrors of the loss diagnostics (``actor/pg_loss``, + ``actor/grad_norm``, ...); the tensor-form diagnostics are in + :attr:`loss_fn_outputs`. + """ + return self._sub_result("forward", "forwardBackward").get("metrics", {}) + + @property + def loss_fn_outputs(self) -> Optional[List[Dict[str, Any]]]: + """Return loss-function tensor outputs from the pass, one entry per datum. + + Entries carry the loss outputs under ``actor/*`` keys + (e.g. ``actor/pg_loss``, ``actor/ppo_kl``, ``actor/pg_clipfrac``, + ``actor/grad_norm``). + """ + return self._sub_result("forward", "forwardBackward").get("lossFnOutputs") + + +class OptimStepResponse(_OperationResult): + """Result of a completed ``optim_step`` operation.""" + + @property + def metrics(self) -> Dict[str, float]: + """Return the optimizer-step metrics, such as the learning rate.""" + return self._sub_result("optimizerStep").get("metrics", {}) + + +class SaveStateResponse(_OperationResult): + """Result of a completed ``save_state`` operation.""" + + @property + def model_package_arn(self) -> Optional[str]: + """Return the ARN of the model package the checkpoint was registered as.""" + return self._sub_result("saveState").get("modelPackageArn") + + +class SaveWeightsResponse(_OperationResult): + """Result of a completed ``save_weights_for_sampler`` operation.""" + + @property + def model_package_arn(self) -> Optional[str]: + """Return the ARN of the model package the weights were registered as.""" + return self._sub_result("saveWeights").get("modelPackageArn") + + +class LoadStateResponse(_OperationResult): + """Result of a completed ``load_state`` / ``load_state_with_optimizer`` operation. + + Successful resolution of the future is the signal that the state was + loaded; the service returns no payload fields for load operations. + """ + + +class SampleResponse(_ForbidExtraBase): + """Result of a completed sampling operation. + + Args: + sequences: Generated sequences. + prompt_logprobs: Log probabilities for prompt tokens (if requested). + topk_prompt_logprobs: Top-k log probabilities per prompt position + (if requested). + """ + + sequences: List[SampledSequence] + prompt_logprobs: Optional[List[float]] = None + topk_prompt_logprobs: Optional[List[List[TokenLogprob]]] = None + + @classmethod + def _from_api(cls, raw: Dict[str, Any]) -> "SampleResponse": + """Construct from a service sampling-operation ``sample`` result.""" + sequences = [SampledSequence._from_api(s) for s in raw.get("sequences") or []] + topk_raw = raw.get("topkPromptLogprobs") + topk = [[TokenLogprob._from_api(t) for t in pos] for pos in topk_raw] if topk_raw else None + return cls( + sequences=sequences, + prompt_logprobs=raw.get("promptLogprobs"), + topk_prompt_logprobs=topk, + ) diff --git a/sagemaker-train/tests/unit/train/training_session/__init__.py b/sagemaker-train/tests/unit/train/training_session/__init__.py new file mode 100644 index 0000000000..e69de29bb2 diff --git a/sagemaker-train/tests/unit/train/training_session/session_lifecycle_grid.py b/sagemaker-train/tests/unit/train/training_session/session_lifecycle_grid.py new file mode 100644 index 0000000000..cea6a55ae1 --- /dev/null +++ b/sagemaker-train/tests/unit/train/training_session/session_lifecycle_grid.py @@ -0,0 +1,185 @@ +# 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/ +"""Does the automatic-cleanup finalizer fire, for each way a client can go away? + +A training client stops its session when it is released, and "released" covers +more shapes than a single test can express. Two of them cannot be settled from +the CPython docs alone, and neither needs a live session: + +* At interpreter exit: is a call still reachable that late, and is + ``sys.is_finalizing()`` already True by then, which would suppress a guarded + implementation? +* Inside a reference cycle, where refcounting alone cannot reclaim the client. + +Run as a child process, one trigger per invocation:: + + python session_lifecycle_grid.py + +The service is a stub whose ``stop_session`` appends to ````, so the +parent can read the result after this process is gone, which is the only way to +observe an exit-time finalizer at all. ``test_session_lifecycle_grid.py`` drives +every trigger and asserts the grid. +""" + +from __future__ import annotations + +import gc +import importlib.util +import os +import signal +import sys + +# Loaded straight from its file rather than imported as +# ``sagemaker.train.training_session._session_handle``, which would execute the +# package __init__ and pull in its third-party dependencies. Cleanup semantics +# depend on nothing but the standard library, so this harness stays runnable on +# any bare interpreter, which is what lets the same grid run across CPython +# versions without building an environment for each. +_MODULE_PATH = os.path.join( + os.path.dirname(os.path.abspath(__file__)), + "..", + "..", + "..", + "..", + "src", + "sagemaker", + "train", + "training_session", + "_session_handle.py", +) +_spec = importlib.util.spec_from_file_location("_session_handle", _MODULE_PATH) +_session_handle = importlib.util.module_from_spec(_spec) +_spec.loader.exec_module(_session_handle) +_SessionHandle = _session_handle._SessionHandle + +SESSION = "grid-session" + + +class FileRecordingService: + """Appends one line per stop, plus what the interpreter state was. + + Writes and flushes immediately: an exit-time finalizer may be racing the + teardown of the io layer, so buffering the result could lose it. + """ + + def __init__(self, path: str) -> None: + self._path = path + + def stop_session(self, training_session_name: str) -> None: + self._stop_session_if_active(training_session_name) + + def _stop_session_if_active(self, training_session_name: str) -> bool: + # The automatic path calls this one, so the stub has to offer it too; + # True means "the session was running and a stop was issued". + with open(self._path, "a") as fh: + fh.write(f"{training_session_name} finalizing={sys.is_finalizing()}\n") + fh.flush() + return True + + +def _make(path: str) -> _SessionHandle: + return _SessionHandle(SESSION, FileRecordingService(path)) + + +# --- triggers --------------------------------------------------------------- +# Each returns normally unless it is meant to kill the process. + + +def t_del_then_collect(path: str) -> None: + handle = _make(path) + del handle + gc.collect() + + +def t_falls_out_of_scope(path: str) -> None: + def _inner() -> None: + _make(path) + + _inner() + gc.collect() + + +def t_rebound(path: str) -> None: + handle = _make(path) + handle = None # noqa: F841 - the point is dropping the last reference + gc.collect() + + +def t_reference_cycle(path: str) -> None: + # Refcounting cannot reclaim this; only the cycle collector can. + class Holder: + pass + + holder = Holder() + holder.handle = _make(path) + holder.self_ref = holder + del holder + gc.collect() + + +def t_explicit_disarm(path: str) -> None: + handle = _make(path) + handle.disarm() + del handle + gc.collect() + + +def t_process_exit(path: str) -> None: + # Still referenced from this frame at return: only interpreter exit can + # release it. This is the `python train.py` case, and the reason the + # exit-time path has to work at all. + handle = _make(path) + globals()["_kept_alive"] = handle + + +def t_sys_exit(path: str) -> None: + globals()["_kept_alive"] = _make(path) + sys.exit(3) + + +def t_uncaught_exception(path: str) -> None: + globals()["_kept_alive"] = _make(path) + raise RuntimeError("deliberate: testing cleanup on an uncaught exception") + + +def t_sigint(path: str) -> None: + globals()["_kept_alive"] = _make(path) + # A parent that ignores SIGINT passes that disposition to its children, and + # some build environments do. Install the interpreter default so the signal + # raises KeyboardInterrupt here rather than being discarded. + signal.signal(signal.SIGINT, signal.default_int_handler) + os.kill(os.getpid(), signal.SIGINT) + + +def t_sigterm(path: str) -> None: + globals()["_kept_alive"] = _make(path) + # Set explicitly for the same reason as the SIGINT case: an inherited + # disposition would otherwise decide what this measures. + signal.signal(signal.SIGTERM, signal.SIG_DFL) + os.kill(os.getpid(), signal.SIGTERM) + + +def t_os_exit(path: str) -> None: + # Documents the gap: os._exit skips finalizers and atexit entirely, so no + # in-process mechanism can cover it. Such a session is reclaimed once it + # passes its inactivity timeout. + globals()["_kept_alive"] = _make(path) + os._exit(0) + + +TRIGGERS = {name[2:]: fn for name, fn in list(globals().items()) if name.startswith("t_")} + + +def main() -> int: + trigger, path = sys.argv[1], sys.argv[2] + TRIGGERS[trigger](path) + return 0 + + +if __name__ == "__main__": + sys.exit(main()) diff --git a/sagemaker-train/tests/unit/train/training_session/test_clients.py b/sagemaker-train/tests/unit/train/training_session/test_clients.py new file mode 100644 index 0000000000..df40ca6cd7 --- /dev/null +++ b/sagemaker-train/tests/unit/train/training_session/test_clients.py @@ -0,0 +1,440 @@ +# 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/ +"""Unit tests for sagemaker.train.training_session._clients. + +Clients are built with dummy credentials and never reach the network, so these +tests assert on what botocore resolved: endpoint, region, retry budget, and the +operations the model exposes. + +Endpoints are asserted on the prepared request rather than on +``client.meta.endpoint_url``. The two can differ: ``meta.endpoint_url`` is filled +in when the client is built, while the endpoint a request is actually sent to is +resolved per request from the service's endpoint rules. The prepared request is +what the caller's configuration has to reach, so that is what is checked here. +""" + +from __future__ import annotations + +import os + +import boto3 +import pytest +from botocore.awsrequest import AWSResponse +from botocore.config import Config +from botocore.exceptions import UnknownServiceError + +from sagemaker.train.training_session._clients import ( + _CONNECT_TIMEOUT, + _READ_TIMEOUT, + _RUNTIME_SERVICE_NAME, + _control_plane_client, + _require_control_plane_operations, + _runtime_client, +) + +#: The runtime model comes from the installed botocore. Tests that build a real +#: runtime client are gated on it rather than skipped silently, so a botocore +#: without the model reports which coverage it cost. +requires_runtime_model = pytest.mark.skipif( + _RUNTIME_SERVICE_NAME not in boto3.Session(region_name="us-west-2").get_available_services(), + reason=f"the installed botocore does not model {_RUNTIME_SERVICE_NAME}", +) + +#: A control-plane and a runtime call, with the arguments each needs to serialize. +_CONTROL_PLANE_CALL = ("describe_training_session", {"TrainingSessionName": "probe"}) +_RUNTIME_CALL = ( + "get_training_operation_result", + {"trainingSessionName": "probe", "operationId": "op"}, +) + + +class _EmptyBody: + """The streaming interface botocore reads a response body through.""" + + def stream(self, **kwargs): + yield b"{}" + + +def _prepared_url(client, call) -> str: + """Return the URL ``client`` prepares for ``call``, sending nothing. + + The response is answered from ``before-send``, so the request is signed and + fully prepared, exactly as it would go on the wire, and then dropped. + """ + operation, params = call + seen = {} + + def _capture(request, **kwargs): + seen["url"] = request.url + return AWSResponse(request.url, 200, {}, _EmptyBody()) + + client.meta.events.register("before-send.*.*", _capture) + getattr(client, operation)(**params) + return seen["url"] + + +_CONTROL_PLANE_OPERATIONS = frozenset( + { + "CreateTrainingSession", + "DescribeTrainingSession", + "ListTrainingSessions", + "StopTrainingSession", + "DeleteTrainingSession", + } +) + +_RUNTIME_OPERATIONS = frozenset( + { + "ExecuteTrainingOperation", + "GetTrainingOperationResult", + "ExecuteSamplingOperation", + "GetSamplingOperationResult", + } +) + + +@pytest.fixture(autouse=True) +def _dummy_credentials(monkeypatch): + """Let clients build without reaching the credential chain.""" + monkeypatch.setenv("AWS_ACCESS_KEY_ID", "test-key") + monkeypatch.setenv("AWS_SECRET_ACCESS_KEY", "test-secret") + monkeypatch.delenv("AWS_PROFILE", raising=False) + + +# --- runtime client --------------------------------------------------------- + + +@requires_runtime_model +def test_runtime_client_exposes_the_operations_the_sdk_calls(): + client = _runtime_client(region="us-west-2") + + assert _RUNTIME_OPERATIONS <= set(client.meta.service_model.operation_names) + + +@requires_runtime_model +def test_runtime_client_targets_the_regional_production_endpoint(): + client = _runtime_client(region="eu-west-1") + + assert _prepared_url(client, _RUNTIME_CALL).startswith( + "https://training-session-runtime.sagemaker.eu-west-1.api.aws/" + ) + assert client.meta.region_name == "eu-west-1" + + +@requires_runtime_model +def test_runtime_client_honours_the_endpoint_override(monkeypatch): + monkeypatch.setenv( + "SAGEMAKER_TRAINING_SESSION_RUNTIME_ENDPOINT", "https://runtime.example.test" + ) + + client = _runtime_client(region="us-west-2") + + assert _prepared_url(client, _RUNTIME_CALL).startswith("https://runtime.example.test/") + + +@requires_runtime_model +def test_runtime_client_does_not_retry(): + """The submit and poll loops own the retry policy; a second one underneath + them would multiply the wait.""" + client = _runtime_client(region="us-west-2") + + assert client.meta.config.retries["total_max_attempts"] == 1 + + +@requires_runtime_model +def test_runtime_client_pool_is_sized_for_sampling_fan_out(): + client = _runtime_client(region="us-west-2") + + assert client.meta.config.max_pool_connections == 512 + + +@requires_runtime_model +def test_runtime_client_pool_size_is_overridable(): + client = _runtime_client(region="us-west-2", max_pool_connections=8) + + assert client.meta.config.max_pool_connections == 8 + + +@requires_runtime_model +def test_the_runtime_client_shortens_the_wait_for_a_connection(): + """The runtime client sends each call once, so this wait is the caller's wait.""" + assert _runtime_client(region="us-west-2").meta.config.connect_timeout == 10 + + +def test_the_control_plane_client_shortens_it_too(): + """Asserted separately: this client needs no runtime model, so it is not gated.""" + assert _control_plane_client(region="us-west-2").meta.config.connect_timeout == 10 + + +def test_shortening_the_connection_wait_is_a_departure_from_the_default(): + """Without which the assertions above would hold whether or not we set it. + + Ten seconds is a decision; sixty is what a client is given when nobody + decides. Pinning the difference is what makes those two tests mean something. + """ + assert Config().connect_timeout == 60 + assert _CONNECT_TIMEOUT == 10 + + +def test_the_read_wait_is_pinned_to_what_a_client_is_given_by_default(): + """The read wait deliberately matches the default, so watch the default. + + Asserting the value on the client cannot show that we set it, because the + number is the same either way. What is worth catching is the default moving + underneath us: on that day this value stops being the status quo and becomes + a change in behaviour, which should be a decision rather than a surprise. + """ + assert _READ_TIMEOUT == Config().read_timeout + assert _control_plane_client(region="us-west-2").meta.config.read_timeout == _READ_TIMEOUT + + +# --- control-plane client --------------------------------------------------- + + +def test_control_plane_client_exposes_the_training_session_operations(): + client = _control_plane_client(region="us-west-2") + + assert _CONTROL_PLANE_OPERATIONS <= set(client.meta.service_model.operation_names) + + +def test_the_create_request_the_sdk_builds_validates_against_the_api_model(): + """Every key in a built CreateTrainingSession request is one the API models. + + The other create tests drive a stub client, which accepts any key it is + handed. A key the API does not model is refused before the request is sent, + so a stub cannot see the difference between a field that works and one that + fails every call -- only the real model can. + """ + from botocore.validate import validate_parameters + + from sagemaker.train.training_session.service_client import _build_create_session_args + + args = _build_create_session_args( + session_name="session", + base_model_arn="arn:aws:sagemaker:us-west-2:aws:hub-content/H/Model/m/1.0.0", + model_package_group_arn="arn:aws:sagemaker:us-west-2:123456789012:model-package-group/g", + lora_config={"Rank": 32, "Alpha": 64}, + role_arn="arn:aws:iam::123456789012:role/R", + s3_path="s3://bucket/prefix/", + kms_key_arn="arn:aws:kms:us-west-2:123456789012:key/k", + vpc_config=None, + ) + + shape = ( + _control_plane_client(region="us-west-2") + .meta.service_model.operation_model("CreateTrainingSession") + .input_shape + ) + # Raises ParamValidationError naming any unmodelled or missing key. + validate_parameters(args, shape) + + +def test_control_plane_client_targets_the_regional_production_endpoint(): + client = _control_plane_client(region="eu-west-1") + + assert _prepared_url(client, _CONTROL_PLANE_CALL).startswith( + "https://api.sagemaker.eu-west-1.amazonaws.com/" + ) + assert client.meta.region_name == "eu-west-1" + + +def test_control_plane_client_honours_the_endpoint_override(monkeypatch): + monkeypatch.setenv("SAGEMAKER_TRAINING_SESSION_ENDPOINT", "https://cp.example.test") + + client = _control_plane_client(region="us-west-2") + + assert _prepared_url(client, _CONTROL_PLANE_CALL).startswith("https://cp.example.test/") + + +def test_control_plane_client_retries_throttling_and_5xx(): + client = _control_plane_client(region="us-west-2") + + assert client.meta.config.retries == {"mode": "standard", "total_max_attempts": 4} + + +# --- caller configuration both clients have to honour ----------------------- + + +def test_control_plane_client_honours_a_fips_configuration(monkeypatch): + """A caller who requires FIPS endpoints must not silently get plain ones. + + ``AWS_USE_FIPS_ENDPOINT`` is read by botocore, not by this package, so this + holds only while the endpoint is left to botocore to resolve. + """ + monkeypatch.setenv("AWS_USE_FIPS_ENDPOINT", "true") + + assert _prepared_url(_control_plane_client(region="eu-west-1"), _CONTROL_PLANE_CALL).startswith( + "https://api-fips.sagemaker.eu-west-1.amazonaws.com/" + ) + + +@requires_runtime_model +def test_runtime_client_honours_a_fips_configuration(monkeypatch): + monkeypatch.setenv("AWS_USE_FIPS_ENDPOINT", "true") + + assert _prepared_url(_runtime_client(region="eu-west-1"), _RUNTIME_CALL).startswith( + "https://training-session-runtime.sagemaker-fips.eu-west-1.api.aws/" + ) + + +def test_control_plane_client_uses_the_regions_own_dns_suffix(): + """A Region outside the ``aws`` partition does not answer on amazonaws.com.""" + assert _prepared_url( + _control_plane_client(region="cn-north-1"), _CONTROL_PLANE_CALL + ).startswith("https://api.sagemaker.cn-north-1.amazonaws.com.cn/") + + +@requires_runtime_model +def test_runtime_client_uses_the_regions_own_dns_suffix(): + assert _prepared_url(_runtime_client(region="cn-north-1"), _RUNTIME_CALL).startswith( + "https://training-session-runtime.sagemaker.cn-north-1.api.amazonwebservices.com.cn/" + ) + + +def test_no_configured_region_is_an_error_not_a_default(monkeypatch): + """An unconfigured environment must not pick a Region for the caller. + + Creating a session in an unasked-for Region starts GPUs there, so this is + raised rather than defaulted. + """ + for name in ("AWS_REGION", "AWS_DEFAULT_REGION", "AWS_PROFILE"): + monkeypatch.delenv(name, raising=False) + monkeypatch.setenv("AWS_CONFIG_FILE", os.devnull) + monkeypatch.setenv("AWS_SHARED_CREDENTIALS_FILE", os.devnull) + + with pytest.raises(ValueError, match="No AWS Region configured"): + _control_plane_client() + with pytest.raises(ValueError, match="No AWS Region configured"): + _runtime_client() + + +def test_older_botocore_fails_with_an_actionable_message(): + """Checked against the operations the model has, not its version string.""" + + class _ModelWithoutTrainingSessions: + operation_names = ["CreateEndpoint", "DescribeEndpoint"] + + class _Client: + class meta: + service_model = _ModelWithoutTrainingSessions() + + with pytest.raises(RuntimeError, match="does not describe the SageMaker training-session"): + _require_control_plane_operations(_Client()) + + +def test_the_control_plane_message_names_no_version_to_upgrade_to(): + """Naming a release only helps once one carries the operations. + + Pointing at a version that does not have them sends the caller through an + upgrade that changes nothing, so the message reports what is missing and + leaves the release out. + """ + + class _ModelWithoutTrainingSessions: + operation_names = ["CreateEndpoint"] + + class _Client: + class meta: + service_model = _ModelWithoutTrainingSessions() + + with pytest.raises(RuntimeError) as caught: + _require_control_plane_operations(_Client()) + + message = str(caught.value) + assert "pip install" not in message + assert ">=" not in message + + +def test_current_botocore_passes_the_check(): + _require_control_plane_operations(_control_plane_client(region="us-west-2")) + + +def test_a_botocore_without_the_runtime_model_fails_with_an_actionable_message(): + """The lookup failure botocore raises names neither the cause nor the fix.""" + + class _SessionWithoutTheRuntime: + def client(self, service_name, **kwargs): + raise UnknownServiceError(service_name=service_name, known_service_names=["sagemaker"]) + + with pytest.raises(RuntimeError, match="does not describe the SageMaker training-session"): + _runtime_client(region="us-west-2", boto_session=_SessionWithoutTheRuntime()) + + +def test_the_supplied_sessions_region_is_used_when_none_is_named(monkeypatch): + """Handing over a session is itself a choice of Region. + + Both factories resolved the Region from a fresh session instead of the one + passed in, so a caller's explicitly pinned session reached another Region. + """ + for name in ("AWS_REGION", "AWS_DEFAULT_REGION"): + monkeypatch.delenv(name, raising=False) + monkeypatch.setenv("AWS_CONFIG_FILE", os.devnull) + monkeypatch.setenv("AWS_SHARED_CREDENTIALS_FILE", os.devnull) + session = boto3.Session(region_name="eu-west-1") + + assert _control_plane_client(boto_session=session).meta.region_name == "eu-west-1" + assert _runtime_client(boto_session=session).meta.region_name == "eu-west-1" + + +def test_the_supplied_session_outranks_an_ambient_aws_region(monkeypatch): + """A session's own Region is not overridden by one named in the environment. + + Managed compute exports AWS_REGION, so without this these clients could end + up in a different Region from every other client built off that session. + """ + monkeypatch.setenv("AWS_REGION", "us-east-1") + monkeypatch.delenv("AWS_DEFAULT_REGION", raising=False) + monkeypatch.setenv("AWS_CONFIG_FILE", os.devnull) + monkeypatch.setenv("AWS_SHARED_CREDENTIALS_FILE", os.devnull) + session = boto3.Session(region_name="eu-west-1") + + assert _control_plane_client(boto_session=session).meta.region_name == "eu-west-1" + assert _runtime_client(boto_session=session).meta.region_name == "eu-west-1" + + +def test_a_session_without_a_region_is_an_error_not_another_profiles(tmp_path, monkeypatch): + """A supplied session answers alone, so an unset Region is not borrowed. + + Resolving on through the standard chain took the Region from whichever + profile that chain lands on, while the credentials came from this one. + """ + config = tmp_path / "config" + config.write_text("[default]\nregion = us-west-2\n\n[profile noregion]\n") + creds = tmp_path / "credentials" + creds.write_text("[noregion]\naws_access_key_id = k\naws_secret_access_key = s\n") + for name in ("AWS_REGION", "AWS_DEFAULT_REGION", "AWS_PROFILE"): + monkeypatch.delenv(name, raising=False) + monkeypatch.setenv("AWS_CONFIG_FILE", str(config)) + monkeypatch.setenv("AWS_SHARED_CREDENTIALS_FILE", str(creds)) + session = boto3.Session(profile_name="noregion") + + with pytest.raises(ValueError, match="No AWS Region configured"): + _control_plane_client(boto_session=session) + with pytest.raises(ValueError, match="No AWS Region configured"): + _runtime_client(boto_session=session) + + +def test_aws_region_still_answers_when_no_session_is_named(monkeypatch): + """boto3 does not read AWS_REGION, so it stays worth reading explicitly.""" + monkeypatch.setenv("AWS_REGION", "us-east-1") + monkeypatch.delenv("AWS_DEFAULT_REGION", raising=False) + monkeypatch.setenv("AWS_CONFIG_FILE", os.devnull) + monkeypatch.setenv("AWS_SHARED_CREDENTIALS_FILE", os.devnull) + + assert _control_plane_client().meta.region_name == "us-east-1" + assert _runtime_client().meta.region_name == "us-east-1" + + +def test_an_explicit_region_still_wins_over_the_sessions(monkeypatch): + for name in ("AWS_REGION", "AWS_DEFAULT_REGION"): + monkeypatch.delenv(name, raising=False) + session = boto3.Session(region_name="eu-west-1") + + client = _control_plane_client(region="us-west-2", boto_session=session) + + assert client.meta.region_name == "us-west-2" diff --git a/sagemaker-train/tests/unit/train/training_session/test_exceptions.py b/sagemaker-train/tests/unit/train/training_session/test_exceptions.py new file mode 100644 index 0000000000..9a7a7f944b --- /dev/null +++ b/sagemaker-train/tests/unit/train/training_session/test_exceptions.py @@ -0,0 +1,172 @@ +# 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/ +"""Unit tests for sagemaker.train.training_session.exceptions. + +``is_retryable`` is the one classification the SDK still makes over a wire +error, and both the submit retry and the poll loops depend on it: a failure is +either transient, and worth sending again, or it is not. +""" + +from __future__ import annotations + +import pickle + +import pytest +from botocore.exceptions import ( + ClientError, + ConnectTimeoutError, + ParamValidationError, + ReadTimeoutError, +) + +from sagemaker.train.training_session.exceptions import ( + ArtifactDeletionError, + OperationError, + ResultMismatchError, + is_retryable, +) + + +def _client_error(code, message="", status=400): + return ClientError( + { + "Error": {"Code": code, "Message": message}, + "ResponseMetadata": {"HTTPStatusCode": status}, + }, + "ExecuteTrainingOperation", + ) + + +# --- is_retryable ----------------------------------------------------------- + + +@pytest.mark.parametrize( + "code,status", + [("ThrottlingException", 429), ("InternalServerException", 500)], +) +def test_modelled_transient_codes_are_retryable(code, status): + assert is_retryable(_client_error(code, "slow down", status=status)) + + +@pytest.mark.parametrize("status", [429, 500, 502, 503, 504]) +def test_unmodelled_transient_statuses_are_retryable(status): + """A gateway ahead of the service can answer without a modelled code.""" + assert is_retryable(_client_error("SomethingUnmodelled", "gateway", status=status)) + + +@pytest.mark.parametrize( + "code,status", + [ + ("ValidationException", 400), + ("ConflictException", 409), + ("ResourceNotFoundException", 404), + ("AccessDeniedException", 403), + ], +) +def test_request_rejections_are_not_retryable(code, status): + """Resending a rejected request gets the same answer.""" + assert not is_retryable(_client_error(code, "no", status=status)) + + +@pytest.mark.parametrize( + "error", + [ + ReadTimeoutError(endpoint_url="https://runtime.test"), + ConnectTimeoutError(endpoint_url="https://runtime.test"), + ], +) +def test_connection_failures_are_retryable(error): + """No response means no status, and the operation may still be running.""" + assert is_retryable(error) + + +def test_local_request_validation_is_not_retryable(): + """botocore checks the request against the model before sending it, and that + verdict does not change on a second try.""" + assert not is_retryable( + ParamValidationError(report="Missing required parameter: operationConfig") + ) + + +def test_error_without_a_response_body_does_not_crash(): + """A malformed error must not fail inside the classifier.""" + assert not is_retryable(ClientError({}, "ExecuteTrainingOperation")) + + +def test_unrelated_exception_is_not_retryable(): + assert not is_retryable(ValueError("not a wire error")) + + +# --- OperationError --------------------------------------------------------- + + +def test_operation_error_carries_the_operation_id(): + """The id is what identifies the failed operation when asking for help.""" + error = OperationError("training", "out of memory", operation_id="op-7") + + assert error.operation_id == "op-7" + assert error.operation_type == "training" + assert "op-7" in str(error) + assert "out of memory" in str(error) + + +def test_operation_error_message_omits_an_absent_id(): + error = OperationError("training", "out of memory") + + assert error.operation_id is None + assert "None" not in str(error) + assert str(error) == "training failed: out of memory" + + +# --- crossing a process boundary --------------------------------------------- +# A training loop that spreads work across processes moves exceptions between +# them, and each of these takes more constructor arguments than the single +# message the default rebuild would pass. Without an explicit rebuild the real +# failure is replaced by a TypeError about a missing argument. + + +@pytest.mark.parametrize( + "error", + [ + OperationError("training", "out of memory", operation_id="op-1"), + OperationError("training", "out of memory"), + ResultMismatchError("computeLogprobs", 3, 2), + ArtifactDeletionError( + "arn:aws:sagemaker:us-west-2:1:model-package/g/1", "s3://bucket/key", ["k1", "k2"] + ), + ], + ids=["operation-with-id", "operation-without-id", "mismatch", "deletion"], +) +def test_every_error_survives_a_pickle_round_trip(error): + restored = pickle.loads(pickle.dumps(error)) + + assert type(restored) is type(error) + assert str(restored) == str(error) + + +def test_a_round_trip_keeps_the_attributes_callers_read(): + """The attributes are the reason to catch these, so they have to survive.""" + original = OperationError("training", "out of memory", operation_id="op-1") + + restored = pickle.loads(pickle.dumps(original)) + + assert restored.operation_type == "training" + assert restored.operation_id == "op-1" + assert restored.message == "out of memory" + + +def test_a_round_trip_keeps_the_full_failed_key_list(): + """The message truncates the keys; the attribute is what carries them all.""" + keys = [f"k{index}" for index in range(12)] + original = ArtifactDeletionError("arn:mp/1", "s3://bucket/prefix/", keys) + + restored = pickle.loads(pickle.dumps(original)) + + assert restored.failed_keys == keys + assert restored.s3_uri == "s3://bucket/prefix/" + assert restored.model_package_arn == "arn:mp/1" diff --git a/sagemaker-train/tests/unit/train/training_session/test_futures.py b/sagemaker-train/tests/unit/train/training_session/test_futures.py new file mode 100644 index 0000000000..529ad9742c --- /dev/null +++ b/sagemaker-train/tests/unit/train/training_session/test_futures.py @@ -0,0 +1,1556 @@ +# 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/ +"""Unit tests for sagemaker.train.training_session.futures. + +All timing is injected — no test performs a real sleep or wall-clock wait. +""" + +from __future__ import annotations + +import concurrent.futures +import threading + +import pytest +from botocore.exceptions import ClientError + +from sagemaker.train.training_session.exceptions import OperationError, ResultMismatchError +from sagemaker.train.training_session.futures import ( + APIFuture, + CombinedFuture, + DEFAULT_BACKOFF_BASE_SECONDS, + DEFAULT_BACKOFF_MAX_SECONDS, + DEFAULT_MAX_TRANSIENT_POLL_ERRORS, + PendingAPIFuture, + RequestStatus, + _backoff_delay, + _coerce_status, + _PENDING_STATUSES, + _remaining, + _SUCCESS_STATUS, + _TERMINAL_FAILURE_STATUSES, + call_with_transient_retry, +) + + +def _max_jitter(low: float, high: float) -> float: + """A jitter draw pinned to the top of the range, so delays are assertable. + + Real draws are random by design; pinning them makes the *ceiling* schedule + observable, which is the part the policy fixes. + """ + return high + + +class FakeClock: + """Deterministic monotonic clock; ``sleep`` advances it, never waits.""" + + def __init__(self) -> None: + self.now = 0.0 + self.sleeps: list[float] = [] + + def time(self) -> float: + return self.now + + def sleep(self, seconds: float) -> None: + self.sleeps.append(seconds) + self.now += seconds + + async def asleep(self, seconds: float) -> None: + """The awaitable seam ``result_async`` sleeps through.""" + self.sleep(seconds) + + +def _make_future( + clock: FakeClock, + responses, + *, + parse_fn=None, + status_key="status", + poll_interval=2, +): + """Build an APIFuture whose poll_fn yields ``responses`` in order.""" + it = iter(responses) + return APIFuture( + operation_id="op-1", + poll_fn=lambda: next(it), + parse_fn=parse_fn or (lambda r: r.get("operationResult", "done")), + status_key=status_key, + operation_type="training", + poll_interval=poll_interval, + time_fn=clock.time, + sleep_fn=clock.sleep, + async_sleep_fn=clock.asleep, + ) + + +def test_the_module_names_only_what_a_caller_can_act_on(): + """The generated reference documents this module by name, so pin the surface. + + Without a declared surface the two future subclasses are published as though + a caller were meant to build one, alongside the polling defaults and the + retry helper the clients wire up for themselves. + """ + import sagemaker.train.training_session.futures as module + + assert module.__all__ == ["APIFuture", "RequestStatus"] + for internal in ("PendingAPIFuture", "CombinedFuture", "call_with_transient_retry"): + assert hasattr(module, internal), f"{internal} moved; this test needs updating" + assert internal not in module.__all__ + + +# --- status model ----------------------------------------------------------- + + +def test_status_partitions_are_complete_and_disjoint(): + all_states = set(RequestStatus) + assert _SUCCESS_STATUS == RequestStatus.COMPLETED + assert {_SUCCESS_STATUS} | _TERMINAL_FAILURE_STATUSES | _PENDING_STATUSES == all_states + assert not (_TERMINAL_FAILURE_STATUSES & _PENDING_STATUSES) + assert _SUCCESS_STATUS not in _TERMINAL_FAILURE_STATUSES + assert _SUCCESS_STATUS not in _PENDING_STATUSES + + +@pytest.mark.parametrize( + "raw,expected", + [ + ("Completed", RequestStatus.COMPLETED), + ("Failed", RequestStatus.FAILED), + ("InProgress", RequestStatus.IN_PROGRESS), + ("Pending", RequestStatus.PENDING), + ("Skipped", RequestStatus.SKIPPED), + ], +) +def test_coerce_status_known(raw, expected): + assert _coerce_status(raw) == expected + + +@pytest.mark.parametrize("raw", ["SomethingNew", None, "", "completed"]) +def test_coerce_status_unknown_is_pending(raw): + # Unknown/None/case-mismatch must not raise — treat as still-running. + assert _coerce_status(raw) == RequestStatus.IN_PROGRESS + + +# --- result() happy path ---------------------------------------------------- + + +def test_result_returns_on_immediate_completion(): + clock = FakeClock() + fut = _make_future(clock, [{"status": "Completed", "operationResult": {"loss": 1.0}}]) + assert fut.result() == {"loss": 1.0} + assert clock.sleeps == [] # completed first poll, never slept + + +def test_result_polls_through_pending_then_completes(): + clock = FakeClock() + fut = _make_future( + clock, + [ + {"status": "Created"}, + {"status": "Dispatched"}, + {"status": "InProgress"}, + {"status": "Completed", "operationResult": "ok"}, + ], + poll_interval=5, + ) + assert fut.result() == "ok" + assert clock.sleeps == [5, 5, 5] # slept once per pending poll + + +def test_result_is_cached_and_does_not_poll_again(): + clock = FakeClock() + calls = {"n": 0} + + def poll(): + calls["n"] += 1 + return {"status": "Completed", "operationResult": "v"} + + fut = APIFuture( + "op", + poll, + lambda r: r["operationResult"], + time_fn=clock.time, + sleep_fn=clock.sleep, + ) + assert fut.result() == "v" + assert fut.result() == "v" + assert calls["n"] == 1 # second result() served from cache + + +def test_parse_fn_side_effects_fire_once_on_success(): + # APIFuture has no on_success callback by design: completion side + # effects (e.g. caching a checkpoint) are composed into parse_fn by the + # submitter. parse_fn must run exactly once, and only on success. + clock = FakeClock() + calls = {"n": 0} + + def parse_fn(resp): + calls["n"] += 1 + return resp["operationResult"] + + fut = _make_future( + clock, + # Two completed polls available, but the second is never reached: + # result() caches after the first success. + [ + {"status": "Completed", "operationResult": "v"}, + {"status": "Completed", "operationResult": "v"}, + ], + parse_fn=parse_fn, + ) + assert fut.result() == "v" + assert fut.result() == "v" # cached; parse_fn not invoked again + assert calls["n"] == 1 + + +def test_parse_fn_not_called_on_failure(): + clock = FakeClock() + calls = {"n": 0} + + def parse_fn(resp): + calls["n"] += 1 + return resp + + fut = _make_future( + clock, + [{"status": "Failed", "errorDetails": "boom"}], + parse_fn=parse_fn, + ) + with pytest.raises(OperationError): + fut.result() + assert calls["n"] == 0 + + +def test_parse_fn_receives_full_response(): + clock = FakeClock() + seen = {} + fut = _make_future( + clock, + [{"status": "Completed", "extra": 7}], + parse_fn=lambda r: seen.setdefault("resp", r), + ) + fut.result() + assert seen["resp"]["extra"] == 7 + + +def test_custom_status_key(): + clock = FakeClock() + fut = _make_future( + clock, + [{"requestStatus": "Completed", "operationResult": "x"}], + status_key="requestStatus", + ) + assert fut.result() == "x" + + +# --- failure + timeout paths ------------------------------------------------ + + +@pytest.mark.parametrize("bad", ["Failed", "Skipped"]) +def test_result_raises_on_terminal_failure(bad): + clock = FakeClock() + fut = _make_future(clock, [{"status": bad, "errorDetails": "boom"}]) + with pytest.raises(OperationError) as ei: + fut.result() + assert "boom" in str(ei.value) + assert ei.value.operation_type == "training" + # The failed operation must be identifiable from the error alone. + assert ei.value.operation_id == "op-1" + + +def test_failure_without_error_details_uses_status_in_message(): + clock = FakeClock() + fut = _make_future(clock, [{"status": "Failed"}]) + with pytest.raises(OperationError) as ei: + fut.result() + assert "Failed" in str(ei.value) + + +def test_result_times_out_without_real_wait(): + clock = FakeClock() + # Always pending → loop must exit on the deadline, not hang. + fut = APIFuture( + "op-timeout", + poll_fn=lambda: {"status": "InProgress"}, + parse_fn=lambda r: r, + poll_interval=3, + time_fn=clock.time, + sleep_fn=clock.sleep, + ) + with pytest.raises(TimeoutError) as ei: + fut.result(timeout=10) + assert "op-timeout" in str(ei.value) + # Advanced only via injected sleeps; bounded by the deadline. + assert clock.now >= 10 + assert all(s == 3 for s in clock.sleeps) + + +# --- transient poll-error resilience ---------------------------------------- + + +def _transient(status_code=503): + """A botocore error the SDK classifies as retryable: what the wire raises.""" + return ClientError( + { + "Error": {"Code": "InternalServerException", "Message": "blip"}, + "ResponseMetadata": {"HTTPStatusCode": status_code}, + }, + "GetTrainingOperationResult", + ) + + +def _make_future_seq( + clock, + items, + *, + max_transient_poll_errors=5, + poll_interval=2, + rand_fn=_max_jitter, +): + """APIFuture whose poll_fn yields dicts or *raises* a retryable error.""" + it = iter(items) + + def poll(): + nxt = next(it) + if isinstance(nxt, Exception): + raise nxt + return nxt + + return APIFuture( + operation_id="op-transient", + poll_fn=poll, + parse_fn=lambda r: r.get("operationResult", "done"), + operation_type="training", + poll_interval=poll_interval, + time_fn=clock.time, + sleep_fn=clock.sleep, + async_sleep_fn=clock.asleep, + max_transient_poll_errors=max_transient_poll_errors, + rand_fn=rand_fn, + ) + + +def test_transient_poll_error_is_retried_then_completes(): + # The op completes server-side; a single transient blip on the status read + # must NOT abort: a transient on the status read is not an op failure. + clock = FakeClock() + fut = _make_future_seq( + clock, + [ + {"status": "InProgress"}, + _transient(), + {"status": "Completed", "operationResult": "ok"}, + ], + ) + assert fut.result(timeout=100) == "ok" + + +def test_transient_poll_errors_reset_counter_on_success(): + clock = FakeClock() + seq = [] + for _ in range(4): + seq.append(_transient()) + seq.append({"status": "InProgress"}) + seq.append({"status": "Completed", "operationResult": "done"}) + fut = _make_future_seq(clock, seq, max_transient_poll_errors=2, poll_interval=1) + assert fut.result(timeout=1000) == "done" + + +def test_too_many_consecutive_transient_poll_errors_raises(): + clock = FakeClock() + fut = _make_future_seq( + clock, + [_transient() for _ in range(10)], + max_transient_poll_errors=3, + poll_interval=1, + ) + with pytest.raises(ClientError): + fut.result(timeout=1000) + + +# --- backoff policy --------------------------------------------------------- + + +def test_backoff_ceiling_doubles_then_clamps(): + ceilings = [_backoff_delay(n, rand_fn=_max_jitter) for n in range(8)] + assert ceilings == [0.5, 1.0, 2.0, 4.0, 4.0, 4.0, 4.0, 4.0] + assert ceilings[0] == DEFAULT_BACKOFF_BASE_SECONDS + assert max(ceilings) == DEFAULT_BACKOFF_MAX_SECONDS + + +def test_the_default_budget_spans_about_five_minutes(): + """Base, cap and attempt count jointly decide the budget; pin the product. + + Changing any one of them moves it without touching the number that names it. + """ + average_wait = sum( + _backoff_delay(n, rand_fn=lambda low, high: high / 2) + for n in range(DEFAULT_MAX_TRANSIENT_POLL_ERRORS - 1) + ) + assert 4 * 60 <= average_wait <= 6 * 60, f"budget is now {round(average_wait)}s" + + +def test_backoff_draws_from_the_whole_range_below_the_ceiling(): + """Full jitter, not equal jitter: the range starts at zero. + + A floor above zero is a wait every caller shares -- the correlation the + jitter exists to remove. + """ + seen = [] + _backoff_delay(2, rand_fn=lambda low, high: seen.append((low, high)) or 0.0) + assert seen == [(0.0, 2.0)] + + +def test_a_failed_poll_backs_off_while_a_healthy_one_keeps_its_cadence(): + """The two waits are different things and must not share a number. + + A poll that answered keeps the configured cadence; a throttled one cannot, + since that cadence is what was refused. + """ + clock = FakeClock() + fut = _make_future_seq( + clock, + [ + {"status": "InProgress"}, # healthy → poll_interval + _transient(), # 1st consecutive failure → 0.5 + _transient(), # 2nd → 1.0 + _transient(), # 3rd → 2.0 + {"status": "Completed", "operationResult": "ok"}, + ], + poll_interval=7, + ) + assert fut.result(timeout=1000) == "ok" + assert clock.sleeps == [7, 0.5, 1.0, 2.0] + + +def test_backoff_restarts_after_a_poll_gets_through(): + """A recovered endpoint must not stay penalised by the previous burst. + + The consecutive counter already resets on success; the delay schedule has to + reset with it, not stay at the ceiling it climbed to. + """ + clock = FakeClock() + fut = _make_future_seq( + clock, + [ + _transient(), # → 0.5 + _transient(), # → 1.0 + {"status": "InProgress"}, # → poll_interval, counter resets + _transient(), # → 0.5 again, not 2.0 + {"status": "Completed", "operationResult": "ok"}, + ], + poll_interval=3, + ) + assert fut.result(timeout=1000) == "ok" + assert clock.sleeps == [0.5, 1.0, 3, 0.5] + + +# --- call_with_transient_retry (submit-path resilience) --------------------- + + +def test_call_with_transient_retry_succeeds_after_blip(): + # The submit call itself can blip, not just the status read that follows it. + clock = FakeClock() + calls = {"n": 0} + + def fn(): + calls["n"] += 1 + if calls["n"] == 1: + raise _transient() + return {"operationId": "op-9"} + + out = call_with_transient_retry(fn, sleep_fn=clock.sleep, rand_fn=_max_jitter) + assert out == {"operationId": "op-9"} + assert calls["n"] == 2 + assert clock.sleeps == [DEFAULT_BACKOFF_BASE_SECONDS] + + +def test_call_with_transient_retry_gives_up_after_max_attempts(): + clock = FakeClock() + + def fn(): + raise _transient() + + with pytest.raises(ClientError): + call_with_transient_retry(fn, max_attempts=3, sleep_fn=clock.sleep, rand_fn=_max_jitter) + # 3 attempts → 2 sleeps between them, and the second waits longer. + assert clock.sleeps == [0.5, 1.0] + + +def test_call_with_transient_retry_propagates_non_transient(): + def fn(): + raise OperationError("request", "bad input") + + with pytest.raises(OperationError): + call_with_transient_retry(fn, sleep_fn=lambda _s: None) + + +def test_call_with_transient_retry_backs_off_instead_of_hammering(): + """The budget must not be spent at one fixed rate. + + A flat interval reoffers the rate the service just refused, so the delays + have to grow and then hold at the cap. + """ + clock = FakeClock() + + def fn(): + raise _transient() + + with pytest.raises(ClientError): + call_with_transient_retry(fn, max_attempts=8, sleep_fn=clock.sleep, rand_fn=_max_jitter) + assert clock.sleeps == [0.5, 1.0, 2.0, 4.0, 4.0, 4.0, 4.0] + assert len(set(clock.sleeps)) > 1, "a flat cadence is the bug being fixed" + + +def test_call_with_transient_retry_jitters_every_wait(): + """Two callers hitting the same throttle must not wait the same amount. + + Drawn per attempt, not once: one draw reused would still put every caller on + the same relative schedule. + """ + draws = [] + + def fn(): + raise _transient() + + def recording_rand(low, high): + draws.append((low, high)) + return low # exercise the bottom of the range; timing is irrelevant here + + with pytest.raises(ClientError): + call_with_transient_retry( + fn, max_attempts=4, sleep_fn=lambda _s: None, rand_fn=recording_rand + ) + # Every wait is a fresh draw from [0, ceiling), never the ceiling itself. + assert draws == [(0.0, 0.5), (0.0, 1.0), (0.0, 2.0)] + + +# --- async surface ---------------------------------------------------------- +# result_async() polls as a coroutine: it suspends between polls instead of +# handing the whole blocking wait to a thread. These pin that it reaches every +# outcome result() does, and that awaiting many at once is not serialized by a +# thread pool. + + +def test_await_and_result_async_inside_event_loop(): + import asyncio + + clock = FakeClock() + + async def go(): + awaited = await _make_future(clock, [{"status": "Completed", "operationResult": "a"}]) + called = await _make_future( + clock, [{"status": "Completed", "operationResult": "b"}] + ).result_async() + return awaited, called + + assert asyncio.run(go()) == ("a", "b") + + +def test_result_async_polls_through_pending_then_completes(): + """The async wait sleeps through pending polls, on the injected clock.""" + import asyncio + + clock = FakeClock() + fut = _make_future( + clock, + [ + {"status": "Pending"}, + {"status": "InProgress"}, + {"status": "Completed", "operationResult": "ok"}, + ], + poll_interval=5, + ) + assert asyncio.run(fut.result_async()) == "ok" + assert clock.sleeps == [5, 5] + + +def test_result_async_is_cached_and_does_not_poll_again(): + import asyncio + + clock = FakeClock() + calls = {"n": 0} + + def poll(): + calls["n"] += 1 + return {"status": "Completed", "operationResult": "v"} + + fut = APIFuture( + "op", + poll, + lambda r: r["operationResult"], + time_fn=clock.time, + sleep_fn=clock.sleep, + async_sleep_fn=clock.asleep, + ) + assert asyncio.run(fut.result_async()) == "v" + assert asyncio.run(fut.result_async()) == "v" + assert calls["n"] == 1 + + +def test_result_async_shares_the_cache_with_result(): + """One future resolved either way must not poll twice or re-parse.""" + import asyncio + + clock = FakeClock() + parses = {"n": 0} + fut = _make_future( + clock, + [{"status": "Completed", "operationResult": "v"}], + parse_fn=lambda r: parses.update(n=parses["n"] + 1) or r["operationResult"], + ) + assert fut.result() == "v" + # A second poll would raise StopIteration; a second parse would count twice. + assert asyncio.run(fut.result_async()) == "v" + assert parses["n"] == 1 + + +@pytest.mark.parametrize("bad", ["Failed", "Skipped"]) +def test_result_async_raises_on_terminal_failure(bad): + import asyncio + + clock = FakeClock() + fut = _make_future(clock, [{"status": bad, "errorDetails": "boom"}]) + with pytest.raises(OperationError) as ei: + asyncio.run(fut.result_async()) + assert "boom" in str(ei.value) + + +def test_result_async_times_out_without_real_wait(): + import asyncio + + clock = FakeClock() + fut = APIFuture( + "op-timeout", + poll_fn=lambda: {"status": "InProgress"}, + parse_fn=lambda r: r, + poll_interval=3, + time_fn=clock.time, + sleep_fn=clock.sleep, + async_sleep_fn=clock.asleep, + ) + with pytest.raises(TimeoutError) as ei: + asyncio.run(fut.result_async(timeout=10)) + assert "op-timeout" in str(ei.value) + assert clock.now >= 10 + + +def test_result_async_retries_transient_then_completes(): + import asyncio + + clock = FakeClock() + fut = _make_future_seq( + clock, + [ + {"status": "InProgress"}, + _transient(), + {"status": "Completed", "operationResult": "ok"}, + ], + ) + assert asyncio.run(fut.result_async(timeout=100)) == "ok" + + +def test_result_async_gives_up_after_too_many_consecutive_transients(): + import asyncio + + clock = FakeClock() + fut = _make_future_seq( + clock, + [_transient() for _ in range(10)], + max_transient_poll_errors=3, + poll_interval=1, + ) + with pytest.raises(ClientError): + asyncio.run(fut.result_async(timeout=1000)) + + +def test_result_async_backs_off_on_the_same_schedule_as_result(): + """The two waits must not drift: throttling is not gentler on a coroutine.""" + import asyncio + + clock = FakeClock() + fut = _make_future_seq( + clock, + [ + {"status": "InProgress"}, + _transient(), + _transient(), + {"status": "Completed", "operationResult": "ok"}, + ], + poll_interval=7, + ) + assert asyncio.run(fut.result_async(timeout=1000)) == "ok" + assert clock.sleeps == [7, 0.5, 1.0] + + +def test_await_does_not_hold_a_worker_thread_for_the_whole_operation(): + """The regression this async rewrite exists for. + + The wait used to be a blocking call handed to the event loop's default + executor, so each pending ``await`` occupied one of its threads until the + operation finished. Awaiting more operations than that pool has threads + admitted them in waves, and starved anything else that needed a thread. + + Here every future needs two polls to finish, so with a genuine coroutine + wait all of them reach their second poll. If the wait were thread-bound, + at most ``max_workers`` could be in flight at once. + """ + import asyncio + + n = 64 # comfortably above min(32, cpu_count + 4) + inflight = 0 + peak = 0 + + def make(): + state = {"calls": 0} + + def poll(): + nonlocal inflight, peak + state["calls"] += 1 + if state["calls"] == 1: + inflight += 1 + peak = max(peak, inflight) + return {"status": "InProgress"} + inflight -= 1 + return {"status": "Completed", "operationResult": "ok"} + + # Real asyncio.sleep(0) so every wait suspends and resumes for real, + # without spending wall-clock time. + return APIFuture( + "op", + poll, + lambda r: r["operationResult"], + poll_interval=0, + async_sleep_fn=asyncio.sleep, + ) + + async def go(): + return await asyncio.gather(*[f.result_async() for f in [make() for _ in range(n)]]) + + assert asyncio.run(go()) == ["ok"] * n + assert peak == n, f"only {peak}/{n} operations were in flight at once" + + +def test_awaiting_operations_does_not_starve_asyncio_to_thread(): + """A pending await must not consume a default-executor thread. + + ``asyncio.to_thread`` is how the ``create_*_async`` factories run, and it + draws on the same pool the old thread-bound wait did. With one pending await + per worker thread, those factories used to have to wait for an operation to + finish before they could start. + + Asserted by ordering rather than by elapsed time: the factory must run while + every operation is still in flight. If a pending await held a thread, the + factory could only start once one of them had released it by completing. + """ + import asyncio + + polls_per_op = 30 + + def make(): + state = {"calls": 0} + + def poll(): + state["calls"] += 1 + if state["calls"] >= polls_per_op: + return {"status": "Completed", "operationResult": "ok"} + return {"status": "InProgress"} + + # Bounded poll count, so the operations finish on their own even when + # the assertion below fails; a failing test must not hang. + return APIFuture("op", poll, lambda r: r["operationResult"], poll_interval=0.01) + + async def go(): + loop = asyncio.get_running_loop() + await loop.run_in_executor(None, lambda: None) # materialize the pool + workers = loop._default_executor._max_workers + + pending = [ + asyncio.ensure_future(f.result_async()) for f in [make() for _ in range(workers)] + ] + await asyncio.sleep(0.02) # let them all reach their poll loop + + factory_work = await asyncio.to_thread(lambda: "factory ran") + completed_early = sum(1 for t in pending if t.done()) + + await asyncio.gather(*pending) + return factory_work, completed_early + + factory_work, completed_early = asyncio.run(go()) + assert factory_work == "factory ran" + assert completed_early == 0, ( + f"{completed_early} operation(s) had to finish before to_thread could " + f"run; the pending awaits are holding worker threads" + ) + + +# --- timeout is a bound on the whole wait ----------------------------------- +# A future that resolves in stages (submit then poll; several chunks) must not +# spend `timeout` per stage: one deadline is taken up front and each stage gets +# only what is left. + + +class FakeSubmit: + """Stands in for the ``concurrent.futures.Future`` a submitter returns. + + ``wait`` seconds of the injected clock are burned before the inner future is + handed back; a shorter allowance times out the way the real one does. + """ + + def __init__(self, clock: FakeClock, inner, *, wait: float = 0.0) -> None: + self._clock = clock + self._inner = inner + self._wait = wait + + def result(self, timeout=None): + if timeout is not None and timeout < self._wait: + self._clock.now += timeout + raise concurrent.futures.TimeoutError() + self._clock.now += self._wait + return self._inner + + +def test_remaining_is_clamped_at_zero_and_unbounded_stays_none(): + clock = FakeClock() + assert _remaining(None, clock.time) is None + assert _remaining(10.0, clock.time) == 10.0 + clock.now = 25.0 + assert _remaining(10.0, clock.time) == 0.0 + + +def test_result_timeout_none_waits_indefinitely(): + clock = FakeClock() + fut = _make_future( + clock, + [{"status": "InProgress"}] * 5 + [{"status": "Completed", "operationResult": "ok"}], + poll_interval=1000, + ) + # Would raise TimeoutError on any finite deadline this small. + assert fut.result(timeout=None) == "ok" + assert clock.now == 5000 + + +def test_pending_future_timeout_covers_submit_and_poll_together(): + clock = FakeClock() + inner = APIFuture( + "op-inner", + poll_fn=lambda: {"status": "InProgress"}, + parse_fn=lambda r: r, + poll_interval=3, + time_fn=clock.time, + sleep_fn=clock.sleep, + ) + # 4s of the 10s budget goes to the submit, so the poll loop gets 6s, not a + # fresh 10s (which would let the whole call run 14s). + pending = PendingAPIFuture( + FakeSubmit(clock, inner, wait=4.0), + operation_type="training", + time_fn=clock.time, + ) + with pytest.raises(TimeoutError) as ei: + pending.result(timeout=10) + assert clock.now == 10 + # The message reports the budget the caller passed, not the 6s slice the + # poll loop got after the submit wait. + assert "within 10.0s" in str(ei.value) + assert "op-inner" in str(ei.value) + + +def test_pending_future_submit_timeout_raises_builtin_timeout_error(): + clock = FakeClock() + pending = PendingAPIFuture( + FakeSubmit(clock, None, wait=30.0), + operation_type="training", + time_fn=clock.time, + ) + # concurrent.futures.TimeoutError is only an alias for the built-in from + # Python 3.11 on; callers get the built-in on every supported version. + with pytest.raises(TimeoutError) as ei: + pending.result(timeout=5) + assert "was not submitted within" in str(ei.value) + + +def test_pending_future_timeout_none_passes_through(): + clock = FakeClock() + inner = _make_future(clock, [{"status": "Completed", "operationResult": "ok"}]) + pending = PendingAPIFuture( + FakeSubmit(clock, inner), operation_type="training", time_fn=clock.time + ) + assert pending.result(timeout=None) == "ok" + + +def test_combined_future_shares_one_deadline_across_sub_futures(): + clock = FakeClock() + first = _make_future( + clock, + [{"status": "InProgress"}, {"status": "Completed", "operationResult": 1}], + poll_interval=5, + ) + second = APIFuture( + "op-2", + poll_fn=lambda: {"status": "InProgress"}, + parse_fn=lambda r: r, + poll_interval=5, + time_fn=clock.time, + sleep_fn=clock.sleep, + ) + combined = CombinedFuture([first, second], sum, time_fn=clock.time) + # first burns 5s of the 10s budget, so second gets 5s, not its own 10s. + with pytest.raises(TimeoutError) as ei: + combined.result(timeout=10) + assert clock.now == 10 + assert "within 10.0s" in str(ei.value) + assert "2 sub-operations" in str(ei.value) + + +def test_combined_future_returns_combined_result_within_budget(): + clock = FakeClock() + subs = [_make_future(clock, [{"status": "Completed", "operationResult": n}]) for n in (1, 2, 3)] + combined = CombinedFuture(subs, sum, time_fn=clock.time) + assert combined.result(timeout=10) == 6 + + +# --- the same staged-timeout contract, on the async waits -------------------- +# Both subclasses now define their own result_async, so the deadline splitting +# has to be pinned there too rather than inherited from the sync path. +# PendingAPIFuture awaits its submit through asyncio.wrap_future, which needs a +# real concurrent.futures.Future; FakeSubmit will not do. + + +def _submitted(inner): + """A real, already-resolved submit future yielding ``inner``.""" + cf: concurrent.futures.Future = concurrent.futures.Future() + cf.set_result(inner) + return cf + + +def test_pending_future_result_async_resolves_submit_then_polls(): + import asyncio + + clock = FakeClock() + inner = _make_future( + clock, + [{"status": "InProgress"}, {"status": "Completed", "operationResult": "ok"}], + poll_interval=5, + ) + pending = PendingAPIFuture(_submitted(inner), operation_type="training", time_fn=clock.time) + assert asyncio.run(pending.result_async()) == "ok" + assert clock.sleeps == [5] + + +def test_pending_future_result_async_reports_the_whole_budget_on_timeout(): + """The poll stage got what the submit left; the message must not say so.""" + import asyncio + + clock = FakeClock() + inner = APIFuture( + "op-inner", + poll_fn=lambda: {"status": "InProgress"}, + parse_fn=lambda r: r, + poll_interval=3, + time_fn=clock.time, + sleep_fn=clock.sleep, + async_sleep_fn=clock.asleep, + ) + pending = PendingAPIFuture(_submitted(inner), operation_type="training", time_fn=clock.time) + with pytest.raises(TimeoutError) as ei: + asyncio.run(pending.result_async(timeout=10)) + assert "within 10.0s" in str(ei.value) + assert "op-inner" in str(ei.value) + + +def test_pending_future_result_async_submit_timeout_is_builtin_timeout_error(): + """asyncio.TimeoutError is only the built-in from 3.11 on; normalize it.""" + import asyncio + + clock = FakeClock() + never = concurrent.futures.Future() # submit never lands + pending = PendingAPIFuture(never, operation_type="training", time_fn=clock.time) + with pytest.raises(TimeoutError) as ei: + asyncio.run(pending.result_async(timeout=0.05)) + assert "was not submitted within" in str(ei.value) + + +def test_pending_future_result_async_is_cached(): + import asyncio + + clock = FakeClock() + inner = _make_future(clock, [{"status": "Completed", "operationResult": "v"}]) + pending = PendingAPIFuture(_submitted(inner), operation_type="training", time_fn=clock.time) + assert asyncio.run(pending.result_async()) == "v" + # A second run would poll again and raise StopIteration if not cached. + assert asyncio.run(pending.result_async()) == "v" + + +def test_combined_future_result_async_is_cached(): + """Re-awaiting must not re-poll the sub-futures or re-run combine_fn.""" + import asyncio + + clock = FakeClock() + combines = {"n": 0} + + def combine(results): + combines["n"] += 1 + return sum(results) + + subs = [_make_future(clock, [{"status": "Completed", "operationResult": n}]) for n in (1, 2, 3)] + combined = CombinedFuture(subs, combine, time_fn=clock.time) + assert asyncio.run(combined.result_async()) == 6 + # A second pass would exhaust each sub-future's response iterator. + assert asyncio.run(combined.result_async()) == 6 + assert combines["n"] == 1 + + +def test_combined_future_result_async_shares_one_deadline(): + import asyncio + + clock = FakeClock() + first = _make_future( + clock, + [{"status": "InProgress"}, {"status": "Completed", "operationResult": 1}], + poll_interval=5, + ) + second = APIFuture( + "op-2", + poll_fn=lambda: {"status": "InProgress"}, + parse_fn=lambda r: r, + poll_interval=5, + time_fn=clock.time, + sleep_fn=clock.sleep, + async_sleep_fn=clock.asleep, + ) + combined = CombinedFuture([first, second], sum, time_fn=clock.time) + # first burns 5s of the 10s budget, so second gets 5s, not its own 10s. + with pytest.raises(TimeoutError) as ei: + asyncio.run(combined.result_async(timeout=10)) + assert clock.now == 10 + assert "within 10.0s" in str(ei.value) + assert "2 sub-operations" in str(ei.value) + + +# --- the async surface CombinedFuture inherits ------------------------------- +# CombinedFuture is an APIFuture subclass so that chunking stays invisible in +# the public signature, and it no longer defines result_async / __await__ of its +# own. These pin that the inherited pair really does route through its combining +# result(); inheritance is easy to break with a stray override. + + +def test_combined_future_is_an_api_future(): + """The public annotation is APIFuture[T]; a chunked call must satisfy it.""" + clock = FakeClock() + subs = [_make_future(clock, [{"status": "Completed", "operationResult": 1}])] + assert isinstance(CombinedFuture(subs, sum, time_fn=clock.time), APIFuture) + + +def test_combined_future_await_and_result_async_combine(): + import asyncio + + clock = FakeClock() + + def _combined(): + subs = [ + _make_future(clock, [{"status": "Completed", "operationResult": n}]) for n in (1, 2, 3) + ] + return CombinedFuture(subs, sum, time_fn=clock.time) + + async def go(): + return await _combined(), await _combined().result_async() + + # Both inherited paths fold the sub-results, rather than returning one of + # them or the list. + assert asyncio.run(go()) == (6, 6) + + +class TestTimeoutErrorPathIsBounded: + """Reporting a timeout must not resume the wait the caller bounded.""" + + def test_a_chunked_timeout_does_not_wait_on_an_unresolved_submit(self): + """The message names the operation only if the submit already landed. + + Reading ``operation_id`` waits for the submit with no bound, so building + the message from it turned a bounded wait into an unbounded one. + """ + never = concurrent.futures.Future() + pending = PendingAPIFuture(never, operation_type="forwardBackward") + combined = CombinedFuture([pending], lambda rs: rs) + + # The defect made this call block forever rather than return something + # wrong, so the assertion has to be that it returns at all. Building the + # message on another thread bounds that: an unbounded wait leaves the + # thread alive, which is what this fails on. + built: list = [] + builder = threading.Thread( + target=lambda: built.append(str(combined._chunked_timeout(30.0))), + daemon=True, + ) + builder.start() + builder.join(timeout=5) + + assert not builder.is_alive(), "building the timeout message waited on the submit" + assert "not submitted" in built[0] + + def test_an_unbounded_timeout_formats_without_raising(self): + """``None`` means no limit; formatting it as a number would raise.""" + never = concurrent.futures.Future() + combined = CombinedFuture( + [PendingAPIFuture(never, operation_type="forwardBackward")], lambda rs: rs + ) + + assert "no limit" in str(combined._chunked_timeout(None)) + + def test_a_resolved_submit_still_names_its_operation(self): + inner = _make_future(FakeClock(), [{"status": "Completed", "operationResult": "v"}]) + landed: concurrent.futures.Future = concurrent.futures.Future() + landed.set_result(inner) + combined = CombinedFuture( + [PendingAPIFuture(landed, operation_type="forwardBackward")], lambda rs: rs + ) + + assert inner.operation_id in str(combined._chunked_timeout(30.0)) + + +# --- the rest of the standard future surface --------------------------------- +# The name promises the protocol callers know from concurrent.futures, so a loop +# over futures should be able to ask these without special casing. + + +class TestStandardFutureSurface: + """``done()``, ``exception()`` and ``cancel()`` on all three classes.""" + + def test_done_is_false_until_the_result_is_kept(self): + clock = FakeClock() + fut = _make_future( + clock, + [{"status": "InProgress"}, {"status": "Completed", "operationResult": "v"}], + ) + + assert not fut.done() + assert fut.result() == "v" + assert fut.done() + + def test_done_stays_false_for_a_future_nobody_waited_on(self): + """A finished operation is not a resolved future: nothing polled it.""" + clock = FakeClock() + fut = _make_future(clock, [{"status": "Completed", "operationResult": "v"}]) + + assert not fut.done() + + def test_exception_returns_the_failure_rather_than_raising_it(self): + clock = FakeClock() + fut = _make_future(clock, [{"status": "Failed", "errorDetails": "boom"}]) + + error = fut.exception() + + assert isinstance(error, OperationError) + assert "boom" in str(error) + + def test_exception_is_none_on_success(self): + clock = FakeClock() + fut = _make_future(clock, [{"status": "Completed", "operationResult": "v"}]) + + assert fut.exception() is None + + def test_exception_propagates_a_timeout_rather_than_returning_it(self): + """A timeout means the question is unanswered, not that the op failed.""" + clock = FakeClock() + fut = _make_future(clock, [{"status": "InProgress"}] * 10, poll_interval=5) + + with pytest.raises(TimeoutError): + fut.exception(timeout=10) + + def test_a_none_result_is_kept_rather_than_re_polled(self): + """``None`` is a value a parse can produce, not a mark of "unresolved". + + Memoizing on truthiness sends the future back to the wire on every look, + and reports it unresolved forever. No other test in this file has a + ``parse_fn`` that returns ``None``, so this is the only thing standing + between the sentinel and a quiet revert. + """ + clock = FakeClock() + polls = {"n": 0} + + def poll_fn(): + polls["n"] += 1 + return {"status": "Completed", "operationResult": None} + + fut = APIFuture( + operation_id="op-1", + poll_fn=poll_fn, + parse_fn=lambda resp: resp["operationResult"], + operation_type="training", + time_fn=clock.time, + sleep_fn=clock.sleep, + async_sleep_fn=clock.asleep, + ) + + assert fut.result() is None + assert fut.done(), "a None result left the future looking unresolved" + assert fut.result() is None + assert polls["n"] == 1, f"the service was polled {polls['n']} times" + + def test_done_becomes_true_for_a_failed_operation(self): + """A failure resolves the future as surely as a result does. + + Otherwise the loop these methods exist to serve never finishes: a drain + that removes futures once ``done()`` spins on a failed one forever. + """ + clock = FakeClock() + fut = _make_future(clock, [{"status": "Failed", "errorDetails": "boom"}]) + + assert not fut.done() + assert isinstance(fut.exception(), OperationError) + assert fut.done() + + def test_a_failure_the_sdk_judges_resolves_the_future_too(self): + """A mismatched result is as terminal as one the service calls failed. + + The check runs while parsing a *completed* response, so this failure + arrives by a different route than a reported one. Recording only the + reported kind would leave ``done()`` meaning one thing for an operation + the service rejected and another for a result that came back wrong, and + would send every later wait back to the service for an answer that + cannot change. + """ + clock = FakeClock() + polls = {"n": 0} + + def poll_fn(): + polls["n"] += 1 + return {"status": "Completed"} + + def parse_fn(_resp): + raise ResultMismatchError("computeLogprobs", 5, 3) + + fut = APIFuture( + poll_fn=poll_fn, + parse_fn=parse_fn, + operation_type="computeLogprobs", + operation_id="op-mismatch", + status_key="status", + time_fn=clock.time, + sleep_fn=clock.sleep, + ) + + with pytest.raises(ResultMismatchError): + fut.result(timeout=5) + assert fut.done(), "a result the SDK rejected left the future unresolved" + + first = fut.exception(timeout=5) + second = fut.exception(timeout=5) + assert first is second + assert polls["n"] == 1, f"the service was asked again: {polls['n']} polls" + + def test_a_failed_request_is_not_mistaken_for_a_failed_operation(self): + """Being unable to ask says nothing about the operation, so nothing is kept. + + Credentials expiring mid-run is the case that matters: the operation may + well be running and may well succeed, so a later wait has to ask again + rather than replay an answer that was never about the operation. + """ + clock = FakeClock() + attempts = {"n": 0} + + def poll_fn(): + attempts["n"] += 1 + if attempts["n"] == 1: + raise ClientError( + {"Error": {"Code": "ExpiredTokenException", "Message": "expired"}}, + "GetTrainingOperationResult", + ) + return {"status": "Completed"} + + fut = APIFuture( + poll_fn=poll_fn, + parse_fn=lambda resp: "parsed", + operation_type="forwardBackward", + operation_id="op-expired", + status_key="status", + time_fn=clock.time, + sleep_fn=clock.sleep, + ) + + with pytest.raises(ClientError): + fut.result(timeout=5) + assert not fut.done(), "a failed request was recorded as the operation's outcome" + + assert fut.result(timeout=5) == "parsed" + assert fut.done() + + def test_a_failure_is_read_once_and_then_replayed(self): + """The terminal answer is kept, so a second look does not hit the wire.""" + clock = FakeClock() + polls = {"n": 0} + + def poll_fn(): + polls["n"] += 1 + return {"status": "Failed", "errorDetails": "boom"} + + fut = APIFuture( + operation_id="op-1", + poll_fn=poll_fn, + parse_fn=lambda resp: resp, + operation_type="training", + time_fn=clock.time, + sleep_fn=clock.sleep, + async_sleep_fn=clock.asleep, + ) + + first = fut.exception() + second = fut.exception() + + assert first is second, "a second look built a different error" + assert polls["n"] == 1, f"the service was polled {polls['n']} times" + with pytest.raises(OperationError) as raised: + fut.result() + assert raised.value is first + + def test_a_drain_loop_over_a_failed_future_terminates(self): + """The canonical shape a caller writes, which the defect made infinite.""" + clock = FakeClock() + pending = [_make_future(clock, [{"status": "Failed", "errorDetails": "boom"}])] + + for _ in range(5): # bounded so a regression fails rather than hangs + for future in list(pending): + if future.done(): + pending.remove(future) + else: + future.exception() + if not pending: + break + + assert pending == [] + + def test_both_subclasses_report_a_failure_as_done(self): + clock = FakeClock() + inner = _make_future(clock, [{"status": "Failed", "errorDetails": "boom"}]) + pending = PendingAPIFuture(_submitted(inner), operation_type="training", time_fn=clock.time) + sub = _make_future(clock, [{"status": "Failed", "errorDetails": "boom"}]) + combined = CombinedFuture([sub], lambda rs: rs, time_fn=clock.time) + + for future in (pending, combined): + assert not future.done() + assert isinstance(future.exception(), OperationError) + assert future.done(), f"{type(future).__name__} did not resolve on failure" + + def test_cancel_reports_that_it_could_not_cancel(self): + clock = FakeClock() + fut = _make_future(clock, [{"status": "Completed", "operationResult": "v"}]) + + assert fut.cancel() is False + assert fut.cancelled() is False + + def test_both_subclasses_inherit_the_surface(self): + clock = FakeClock() + inner = _make_future(clock, [{"status": "Completed", "operationResult": "v"}]) + pending = PendingAPIFuture(_submitted(inner), operation_type="training", time_fn=clock.time) + combined = CombinedFuture([inner], lambda rs: rs, time_fn=clock.time) + + for future in (pending, combined): + assert future.done() is False + assert future.cancel() is False + assert future.cancelled() is False + + assert pending.result() == "v" + assert pending.done() is True + assert pending.exception() is None + + +# --- concurrent waiters on one future ---------------------------------------- +# Completion side effects are composed into parse_fn (see +# test_parse_fn_side_effects_fire_once_on_success), so two waiters resolving the +# same future must not both run it. These tests race real threads; the injected +# clock keeps them from waiting. + + +class TestConcurrentResolution: + """One future, several waiters: resolve once, agree on the result.""" + + def test_two_threads_resolving_one_future_parse_once(self): + clock = FakeClock() + # Both waiters are released holding a completed response, so both reach + # the commit with a result to write. + gate = threading.Barrier(2) + parses = [] + + def poll_fn(): + gate.wait(timeout=5) + return {"status": "Completed", "operationResult": "v"} + + def parse_fn(resp): + parses.append(resp) + return ["v"] # a fresh object per call, so identity is meaningful + + fut = APIFuture( + operation_id="op-1", + poll_fn=poll_fn, + parse_fn=parse_fn, + operation_type="training", + time_fn=clock.time, + sleep_fn=clock.sleep, + async_sleep_fn=clock.asleep, + ) + + with concurrent.futures.ThreadPoolExecutor(max_workers=2) as pool: + waiters = [pool.submit(fut.result) for _ in range(2)] + first, second = (w.result(timeout=5) for w in waiters) + + assert len(parses) == 1, "parse_fn ran per waiter instead of per future" + assert first is second, "waiters were handed different result objects" + assert fut.result() is first + + def test_two_coroutines_awaiting_one_future_parse_once(self): + import asyncio + + clock = FakeClock() + parses = [] + + def parse_fn(resp): + parses.append(resp) + return ["v"] + + # Enough completed responses for both coroutines to poll independently. + fut = _make_future( + clock, + [{"status": "Completed", "operationResult": "v"}] * 2, + parse_fn=parse_fn, + ) + + async def go(): + return await asyncio.gather(fut.result_async(), fut.result_async()) + + first, second = asyncio.run(go()) + + assert len(parses) == 1 + assert first is second + + def test_combine_fn_runs_once_across_concurrent_waiters(self): + clock = FakeClock() + gate = threading.Barrier(2) + combines = [] + + def combine(sub_results): + combines.append(sub_results) + return list(sub_results) + + sub = _make_future(clock, [{"status": "Completed", "operationResult": 1}]) + combined = CombinedFuture([sub], combine, time_fn=clock.time) + + # Release both waiters only once each holds its sub-result, so both + # reach the fold. + settled = sub.result + + def gated_result(**kwargs): + gate.wait(timeout=5) + return settled(**kwargs) + + sub.result = gated_result + + with concurrent.futures.ThreadPoolExecutor(max_workers=2) as pool: + waiters = [pool.submit(combined.result) for _ in range(2)] + first, second = (w.result(timeout=5) for w in waiters) + + assert len(combines) == 1, "combine_fn ran per waiter instead of per future" + assert first is second + + def test_pending_future_waiters_agree_on_one_result(self): + clock = FakeClock() + gate = threading.Barrier(2) + + class MemoizesNothing: + """An inner future returning a fresh object per call. + + A real inner :class:`APIFuture` memoizes, which would hide whether + ``PendingAPIFuture`` commits its own result once; this leaves it as + the only thing that can. + """ + + operation_id = "op-1" + + def result(self, timeout=None): + gate.wait(timeout=5) + return ["v"] + + pending = PendingAPIFuture( + _submitted(MemoizesNothing()), + operation_type="training", + time_fn=clock.time, + ) + + with concurrent.futures.ThreadPoolExecutor(max_workers=2) as pool: + waiters = [pool.submit(pending.result) for _ in range(2)] + first, second = (w.result(timeout=5) for w in waiters) + + assert first is second + assert pending.result() is first + + def test_the_resolve_lock_is_never_held_across_a_poll(self): + """The lock guards classify + memoize only, never the wire call. + + Held across a poll it would make a bounded ``result(timeout=...)`` wait + out another waiter's unbounded poll, and would occupy the worker threads + ``result_async`` needs to poll on at all. + """ + clock = FakeClock() + fut = _make_future( + clock, + [{"status": "InProgress"}, {"status": "Completed", "operationResult": "v"}], + ) + acquirable = [] + wire = fut._poll_fn + + def probing_poll(): + got = fut._lock.acquire(blocking=False) + acquirable.append(got) + if got: + fut._lock.release() + return wire() + + fut._poll_fn = probing_poll + + assert fut.result() == "v" + assert acquirable == [True, True] + + +def test_a_submit_timeout_leaves_the_queued_operation_alone(): + """A timeout waiting on the submit must not cancel the submit. + + The wait used to pass its own cancellation through to the background submit, + so a timed-out caller silently dropped the operation and every later attempt + to resolve the same future raised ``CancelledError``. + """ + import asyncio + + submit = concurrent.futures.Future() + pending = PendingAPIFuture(submit, operation_type="training") + + async def go(): + with pytest.raises(TimeoutError, match="was not submitted within"): + await pending.result_async(timeout=0) + + asyncio.run(go()) + + assert not submit.cancelled() + + +def test_the_future_is_still_resolvable_after_a_submit_timeout(): + """The documented contract: the operation keeps running and can be picked up.""" + import asyncio + + clock = FakeClock() + inner = _make_future(clock, [{"status": "Completed", "operationResult": "ok"}]) + submit = concurrent.futures.Future() + pending = PendingAPIFuture(submit, operation_type="training", time_fn=clock.time) + + async def timed_out(): + with pytest.raises(TimeoutError): + await pending.result_async(timeout=0) + + asyncio.run(timed_out()) + + submit.set_result(inner) + + assert asyncio.run(pending.result_async(timeout=None)) == "ok" diff --git a/sagemaker-train/tests/unit/train/training_session/test_sampling_client.py b/sagemaker-train/tests/unit/train/training_session/test_sampling_client.py new file mode 100644 index 0000000000..fbecc24948 --- /dev/null +++ b/sagemaker-train/tests/unit/train/training_session/test_sampling_client.py @@ -0,0 +1,821 @@ +# 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/ +"""Unit tests for sagemaker.train.training_session.sampling_client. + +The runtime client is faked; no test performs a real sleep or network call. +Covers request shaping, prompt encoding, optional logprob flags, result +parsing, error propagation, and the TrainingClient factory wiring. +""" + +from __future__ import annotations + +import asyncio +import gc +import os +import threading + +import pytest +from botocore.exceptions import ClientError + +from sagemaker.train.training_session._session_handle import _SessionHandle +from sagemaker.train.training_session.exceptions import ( + ResultMismatchError, + TrainingSessionError, +) +from sagemaker.train.training_session.sampling_client import SamplingClient +from sagemaker.train.training_session.training_client import TrainingClient +from sagemaker.train.training_session.types import ( + SampleResponse, + SampledSequence, + SamplingParams, + StopReason, +) + + +class FakeRuntime: + """Records execute_sampling_operation calls and replays canned responses. + + Stands in for the botocore runtime client the sampling client calls + directly, so ``execute_error`` is a botocore error, exactly what the wire + raises. + """ + + def __init__(self, *, execute_error: Exception | None = None, sample_response=None): + self.execute_calls: list[dict] = [] + self.training_calls: list[dict] = [] + self.get_calls: list[dict] = [] + self._execute_error = execute_error + self._op_counter = 0 + self._sample_response = sample_response or { + "sequences": [ + {"stopReason": "stop", "tokens": [10, 11, 12], "logprobs": [-0.1, -0.2, -0.3]} + ] + } + + # Mirrors the current ExecuteSamplingOperation input shape: trainingSessionName + # + the SampleOperationConfig union. Reject unknown params the way the + # runtime model would, so a stray field (e.g. clientToken or + # the now-deprecated flat prompt/numSamples) fails in unit tests too. + _ALLOWED_SAMPLE_PARAMS = frozenset( + { + "trainingSessionName", + "operationConfig", + } + ) + + def execute_sampling_operation(self, **kwargs): + unknown = set(kwargs) - self._ALLOWED_SAMPLE_PARAMS + if unknown: + raise TypeError(f"Unknown parameter(s) in ExecuteSamplingOperation: {unknown}") + self.execute_calls.append(kwargs) + if self._execute_error is not None: + raise self._execute_error + self._op_counter += 1 + return {"operationId": f"sop-{self._op_counter}"} + + def get_sampling_operation_result(self, **kwargs): + self.get_calls.append(kwargs) + # Mirror the union variant of the most recent submit: route result by + # the operation config type of the corresponding execute call. + last_cfg = self.execute_calls[-1]["operationConfig"] if self.execute_calls else {} + if "computeLogprobs" in last_cfg: + # One logprob per prompt token, first entry None, as the real + # contract. A fixed-length reply would trip the SDK's length guard. + sent = sum(len(c["tokens"]) for c in last_cfg["computeLogprobs"]["prompt"]["chunks"]) + logprobs = [None] + [round(-0.1 * (i + 1), 1) for i in range(1, sent)] + return { + "status": "Completed", + "operationResult": {"computeLogprobs": {"logprobs": logprobs}}, + } + if "loadWeightsForSampling" in last_cfg: + return { + "status": "Completed", + "operationResult": {}, + } + return { + "status": "Completed", + "operationResult": {"sample": self._sample_response}, + } + + # training side — only used to drive save_weights_and_get_sampling_client + def execute_training_operation(self, **kwargs): + self.training_calls.append(kwargs) + self._op_counter += 1 + return {"operationId": f"op-{self._op_counter}"} + + def get_training_operation_result(self, **kwargs): + return { + "status": "Completed", + "operationResult": {"saveWeights": {"modelPackageArn": "arn:mp/7"}}, + } + + +def _client(**kwargs) -> tuple[SamplingClient, FakeRuntime]: + rt = FakeRuntime(**kwargs) + return SamplingClient("sess-1", rt), rt + + +def _sample_cfg(call: dict) -> dict: + """Return the ``sample`` variant of a recorded execute_sampling_operation call.""" + return call["operationConfig"]["sample"] + + +# --- request shaping -------------------------------------------------------- + + +def test_sample_shapes_request_from_token_ids(): + sc, rt = _client() + sc.sample([1, 2, 3], num_samples=4, sampling_params=SamplingParams(max_tokens=64)).result() + call = rt.execute_calls[0] + cfg = _sample_cfg(call) + assert call["trainingSessionName"] == "sess-1" + assert cfg["numSamples"] == 4 + assert cfg["prompt"]["chunks"][0]["tokens"] == [1, 2, 3] + assert cfg["samplingParams"]["maxTokens"] == 64 + # ExecuteSamplingOperation has no clientToken field (sampling is not deduped) + assert "clientToken" not in call + # optional flags omitted unless requested + assert "includePromptLogprobs" not in cfg + assert "topkPromptLogprobs" not in cfg + + +def test_sample_defaults_to_one_sample_and_default_params(): + sc, rt = _client() + sc.sample([5, 6]).result() + cfg = _sample_cfg(rt.execute_calls[0]) + assert cfg["numSamples"] == 1 + assert cfg["samplingParams"]["maxTokens"] == SamplingParams().max_tokens + + +def test_sample_passes_through_prebuilt_model_input(): + sc, rt = _client() + mi = {"chunks": [{"tokens": [9, 9], "type": "encoded_text"}]} + sc.sample(mi).result() + assert _sample_cfg(rt.execute_calls[0])["prompt"] is mi + + +def test_sample_string_prompt_requires_tokenizer(): + sc, _ = _client() + with pytest.raises(ValueError, match="tokenizer"): + sc.sample("hello world") + + +def test_sample_string_prompt_uses_tokenizer(): + rt = FakeRuntime() + + class _Tok: + def encode(self, text): + return [ord(c) % 7 for c in text] + + sc = SamplingClient("sess-1", rt, tokenizer=_Tok()) + sc.sample("hi").result() + assert _sample_cfg(rt.execute_calls[0])["prompt"]["chunks"][0]["tokens"] == [ + ord("h") % 7, + ord("i") % 7, + ] + + +def test_sample_tokenizes_a_model_input_built_from_text(): + """``ModelInput.from_text`` must reach the wire as tokens, like a bare string. + + The runtime ModelInput has only a ``chunks`` member, so the typed path has to + tokenize as well; otherwise passing text through the type was strictly + worse than passing it untyped. + """ + from sagemaker.train.training_session.types import ModelInput + + rt = FakeRuntime() + + class _Tok: + def encode(self, text): + return [1, 2, 3] + + sc = SamplingClient("sess-1", rt, tokenizer=_Tok()) + sc.sample(ModelInput.from_text("hi")).result() + cfg = _sample_cfg(rt.execute_calls[0]) + assert cfg["prompt"] == {"chunks": [{"tokens": [1, 2, 3], "type": "encoded_text"}]} + assert "prompt" not in cfg["prompt"] + + +def test_sample_model_input_from_text_requires_tokenizer(): + from sagemaker.train.training_session.types import ModelInput + + sc, _ = _client() + with pytest.raises(ValueError, match="tokenizer"): + sc.sample(ModelInput.from_text("hello world")) + + +def test_sample_rejects_unsupported_prompt_type(): + sc, _ = _client() + with pytest.raises(TypeError, match="Unsupported prompt type"): + sc.sample(12.5) + + +def test_logprob_flags_attached_when_requested(): + sc, rt = _client() + sc.sample([1, 2], include_prompt_logprobs=True, topk_prompt_logprobs=5).result() + cfg = _sample_cfg(rt.execute_calls[0]) + assert cfg["includePromptLogprobs"] is True + assert cfg["topkPromptLogprobs"] == 5 + + +# --- checkpoint selection (modelPackageArn on the wire) --------------------- +# model_package_arn selects the checkpoint, sent on the wire as modelPackageArn. +# Empty => the field is omitted => the base model is sampled. + + +def test_sample_sends_model_package_arn_when_set(): + rt = FakeRuntime() + sc = SamplingClient("sess-1", rt, model_package_arn="arn:mp/42") + sc.sample([1, 2, 3]).result() + assert _sample_cfg(rt.execute_calls[0])["modelPackageArn"] == "arn:mp/42" + + +def test_sample_omits_model_package_arn_for_base_model(): + sc, rt = _client() # default model_package_arn == "" (base model) + sc.sample([1, 2, 3]).result() + # Field must be ABSENT (not an empty string) so the base model is sampled. + assert "modelPackageArn" not in _sample_cfg(rt.execute_calls[0]) + + +def test_sample_async_sends_model_package_arn_when_set(): + rt = FakeRuntime() + sc = SamplingClient("sess-1", rt, model_package_arn="arn:mp/7") + asyncio.run(sc.sample_async([4, 5])) + assert _sample_cfg(rt.execute_calls[0])["modelPackageArn"] == "arn:mp/7" + + +def test_compute_logprobs_sends_model_package_arn_when_set(): + rt = FakeRuntime() + sc = SamplingClient("sess-1", rt, model_package_arn="arn:mp/9") + sc.compute_logprobs([1, 2, 3]).result() + cfg = rt.execute_calls[0]["operationConfig"]["computeLogprobs"] + assert cfg["modelPackageArn"] == "arn:mp/9" + + +def test_compute_logprobs_omits_model_package_arn_for_base_model(): + sc, rt = _client() # base model + sc.compute_logprobs([1, 2, 3]).result() + cfg = rt.execute_calls[0]["operationConfig"]["computeLogprobs"] + assert "modelPackageArn" not in cfg + + +def test_created_sampling_client_sends_its_model_package_arn_on_sample(): + # End-to-end factory wiring: a client created for a specific checkpoint + # actually puts that ARN on the wire (regression guard — the selector used + # to be stored but silently dropped from the request). + rt = FakeRuntime() + tc = TrainingClient("sess-1", rt) + sc = tc.create_sampling_client(model_package_arn="arn:mp/3") + sc.sample([1]).result() + assert _sample_cfg(rt.execute_calls[0])["modelPackageArn"] == "arn:mp/3" + + +def test_compute_logprobs_sends_union_variant_and_parses_result(): + sc, rt = _client() + result = sc.compute_logprobs([1, 2, 3]).result() + # Sent the computeLogprobs union variant with the tokenized prompt, and no + # sample-only fields. + cfg = rt.execute_calls[0]["operationConfig"] + assert set(cfg) == {"computeLogprobs"} + assert cfg["computeLogprobs"]["prompt"]["chunks"][0]["tokens"] == [1, 2, 3] + # Parsed from operationResult.computeLogprobs.logprobs; the first entry is + # None (first token has no preceding context). + assert result == [None, -0.2, -0.3] + + +def test_compute_logprobs_rejects_result_of_the_wrong_length(): + """A per-token result must be as long as the sequence sent. + + A shorter or longer list belongs to a different request, so it is raised + rather than returned; otherwise the caller consumes another request's + values as its own. + """ + + class _WrongLength(FakeRuntime): + def get_sampling_operation_result(self, **kwargs): + return { + "status": "Completed", + "operationResult": {"computeLogprobs": {"logprobs": [None, -0.2]}}, + } + + rt = _WrongLength() + sc = SamplingClient("sess-1", rt) + with pytest.raises(ResultMismatchError, match="2 values for a 3-token sequence"): + sc.compute_logprobs([1, 2, 3]).result() + + +def test_compute_logprobs_accepts_result_matching_the_prompt(): + """The guard does not fire on the correct length (regression guard).""" + sc, _ = _client() + assert len(sc.compute_logprobs([1, 2, 3, 4]).result()) == 4 + + +def test_compute_logprobs_skips_guard_when_count_unknown(): + """A prompt the service tokenizes has no client-side count to check.""" + from sagemaker.train.training_session.sampling_client import ( + _make_compute_logprobs_parser, + ) + + parse = _make_compute_logprobs_parser(None) + raw = {"operationResult": {"computeLogprobs": {"logprobs": [None, -0.2]}}} + assert parse(raw) == [None, -0.2] + + +def test_compute_logprobs_string_prompt_uses_tokenizer(): + rt = FakeRuntime() + + class _Tok: + def encode(self, text): + return [ord(c) % 7 for c in text] + + sc = SamplingClient("sess-1", rt, tokenizer=_Tok()) + sc.compute_logprobs("hi").result() + cfg = rt.execute_calls[0]["operationConfig"]["computeLogprobs"] + assert cfg["prompt"]["chunks"][0]["tokens"] == [ord("h") % 7, ord("i") % 7] + + +def test_each_call_records_separately(): + sc, rt = _client() + f1 = sc.sample([1]) + f2 = sc.sample([2]) + f1.result() + f2.result() + assert len(rt.execute_calls) == 2 + # Sample submits are unordered (order-independent ops), so assert on the + # set of prompts sent rather than arrival order. + sent = sorted(_sample_cfg(c)["prompt"]["chunks"][0]["tokens"][0] for c in rt.execute_calls) + assert sent == [1, 2] + + +# --- result parsing --------------------------------------------------------- + + +def test_result_parses_sample_response(): + sc, _ = _client() + result = sc.sample([1, 2, 3]).result() + assert isinstance(result, SampleResponse) + assert len(result.sequences) == 1 + seq = result.sequences[0] + assert isinstance(seq, SampledSequence) + assert seq.stop_reason == StopReason.STOP + assert seq.tokens == [10, 11, 12] + assert seq.logprobs == [-0.1, -0.2, -0.3] + + +def test_result_parses_prompt_and_topk_logprobs(): + resp = { + "sequences": [{"stopReason": "length", "tokens": [1]}], + "promptLogprobs": [-0.5, -0.6], + "topkPromptLogprobs": [ + [{"tokenId": 1, "logprob": -0.5}, {"tokenId": 2, "logprob": -1.0}], + [{"tokenId": 3, "logprob": -0.6}], + ], + } + sc, _ = _client(sample_response=resp) + result = sc.sample([1, 2]).result() + assert result.prompt_logprobs == [-0.5, -0.6] + assert result.topk_prompt_logprobs[0][0].token_id == 1 + assert result.topk_prompt_logprobs[0][1].logprob == -1.0 + assert result.sequences[0].stop_reason == StopReason.LENGTH + + +def test_poll_uses_sampling_operation_id(): + sc, rt = _client() + sc.sample([1]).result() + assert rt.get_calls[0]["operationId"] == "sop-1" + assert rt.get_calls[0]["trainingSessionName"] == "sess-1" + + +# --- error propagation ------------------------------------------------------ +# The service's own error reaches the caller unchanged, code and message intact. + + +def test_service_error_propagates_unchanged(): + error = ClientError( + { + "Error": {"Code": "ConflictException", "Message": "session is Stopping"}, + "ResponseMetadata": {"HTTPStatusCode": 409}, + }, + "ExecuteSamplingOperation", + ) + sc, _ = _client(execute_error=error) + with pytest.raises(ClientError) as ei: + sc.sample([1, 2]).result() + assert ei.value.response["Error"]["Code"] == "ConflictException" + assert "session is Stopping" in str(ei.value) + + +# --- future contract (sample / sample_async / compute_logprobs) ------------- + + +def test_sample_async_awaits_response_directly(): + sc, rt = _client() + # The async variant is a coroutine: a single await yields the response. + assert asyncio.iscoroutinefunction(sc.sample_async) + resp = asyncio.run(sc.sample_async([7, 8], num_samples=2)) + cfg = _sample_cfg(rt.execute_calls[0]) + assert cfg["numSamples"] == 2 + assert cfg["prompt"]["chunks"][0]["tokens"] == [7, 8] + assert isinstance(resp, SampleResponse) + assert resp.sequences[0].tokens == [10, 11, 12] + + +def test_sample_returns_unresolved_future(): + # Drop-in contract: sample() returns the future itself; callers unwrap + # with .result() (or await). + sc, rt = _client() + fut = sc.sample([7, 8]) + assert hasattr(fut, "result") + assert rt.get_calls == [] + assert isinstance(fut.result(), SampleResponse) + assert len(rt.get_calls) == 1 + + +def test_compute_logprobs_returns_unresolved_future(): + sc, rt = _client() + fut = sc.compute_logprobs([1, 2, 3]) + assert hasattr(fut, "result") + assert rt.get_calls == [] + assert fut.result() == [None, -0.2, -0.3] + + +def test_compute_logprobs_async_awaits_logprobs_directly(): + sc, _ = _client() + assert asyncio.iscoroutinefunction(sc.compute_logprobs_async) + assert asyncio.run(sc.compute_logprobs_async([1, 2, 3])) == [None, -0.2, -0.3] + + +def test_sample_submits_overlap_and_return_before_wire_completes(): + import threading + import time + + class _GatedRuntime(FakeRuntime): + """Blocks every execute until released, like a slow wire.""" + + def __init__(self): + super().__init__() + self.release = threading.Event() + self.started = [] + + def execute_sampling_operation(self, **kwargs): + self.started.append(1) + assert self.release.wait(timeout=5) + return super().execute_sampling_operation(**kwargs) + + rt = _GatedRuntime() + sc = SamplingClient("sess-1", rt) + futs = [sc.sample([i]) for i in range(4)] + # Unordered submits overlap: all four reach the wire before any completes. + deadline = time.monotonic() + 2 + while len(rt.started) < 4 and time.monotonic() < deadline: + time.sleep(0.001) + assert len(rt.started) == 4 + assert rt.execute_calls == [] + rt.release.set() + for f in futs: + assert isinstance(f.result(), SampleResponse) + + +# --- TrainingClient factory wiring ------------------------------------------ + + +def test_create_sampling_client_binds_session_and_ref(): + rt = FakeRuntime() + tc = TrainingClient("sess-1", rt) + sc = tc.create_sampling_client(model_package_arn="arn:mp/3") + assert isinstance(sc, SamplingClient) + assert sc.training_session_name == "sess-1" + assert sc.model_package_arn == "arn:mp/3" + + +def test_create_sampling_client_defaults_to_base_model(): + rt = FakeRuntime() + tc = TrainingClient("sess-1", rt) + assert tc.create_sampling_client().model_package_arn == "" + + +def test_save_weights_and_get_sampling_client_uses_model_package_arn(): + rt = FakeRuntime() + tc = TrainingClient("sess-1", rt) + sc = tc.save_weights_and_get_sampling_client() + assert isinstance(sc, SamplingClient) + assert sc.model_package_arn == "arn:mp/7" # from the canned saveWeights result + # save_weights_for_sampler goes through the training path... + training_ops = [next(iter(c["operationConfig"])) for c in rt.training_calls] + assert training_ops == ["saveWeights"] + # ...but LWFS routes through the sampling-operations path (no seq consumed). + assert len(rt.execute_calls) == 1 + lwfs = rt.execute_calls[0]["operationConfig"]["loadWeightsForSampling"] + assert lwfs["modelPackageArn"] == "arn:mp/7" + + +def test_load_weights_for_sampling_shapes_request(): + rt = FakeRuntime() + tc = TrainingClient("sess-1", rt) + tc._load_weights_for_sampling("arn:mp/9") + assert len(rt.execute_calls) == 1 + cfg = rt.execute_calls[0]["operationConfig"]["loadWeightsForSampling"] + assert cfg == {"modelPackageArn": "arn:mp/9"} + assert len(rt.training_calls) == 0, "LWFS must not use training path" + + +def test_create_sampling_client_propagates_base_model(monkeypatch): + import sagemaker.train.training_session._utils as utils + + sentinel = object() + monkeypatch.setattr( + utils, + "load_tokenizer", + lambda base_model, hub_client, hf_model_id=None: sentinel, + ) + tc = TrainingClient("sess-1", FakeRuntime(), base_model="Qwen/Qwen3-32B") + sc = tc.create_sampling_client(model_package_arn="arn:mp/3") + # SamplingClient can resolve a tokenizer from the inherited base model. + assert sc.get_tokenizer() is sentinel + + +def test_sampling_client_get_tokenizer_unsupported_raises(monkeypatch): + import sagemaker.train.training_session._utils as utils + + def fake_load(base_model, hub_client, hf_model_id=None): + raise ValueError("No public HuggingFace tokenizer is available") + + monkeypatch.setattr(utils, "load_tokenizer", fake_load) + sc = SamplingClient("sess-1", FakeRuntime(), base_model="arn:.../private/1.0") + with pytest.raises(ValueError, match="No public HuggingFace tokenizer"): + sc.get_tokenizer() + + +def test_save_weights_and_get_sampling_client_raises_without_a_checkpoint(): + """A save with no ARN must raise, not hand back a base-model sampler. + + An empty ``model_package_arn`` means the base model to ``SamplingClient``, so + returning one here sampled untrained weights while the caller believed it was + sampling what had just been saved. + """ + rt = FakeRuntime() + + def _no_mp_result(**kwargs): + return {"status": "Completed", "operationResult": {"saveWeights": {}}} + + rt.get_training_operation_result = _no_mp_result + tc = TrainingClient("sess-1", rt) + + with pytest.raises(TrainingSessionError, match="no checkpoint ARN"): + tc.save_weights_and_get_sampling_client() + + # The save itself still went out; nothing was loaded for sampling. + op_types = [next(iter(c["operationConfig"])) for c in rt.training_calls] + assert op_types == ["saveWeights"] + + +# --- session ownership / context-manager lifecycle -------------------------- +# +# Only ServiceClient.create_sampling_client creates the backing session, so only +# that client stops it. A client from TrainingClient.create_sampling_client +# shares the training session and must leave it running. + + +class FakeServiceClient: + """Records stop_session calls; optionally raises to exercise the error path.""" + + def __init__(self, *, stop_error: Exception | None = None): + self.stopped: list[str] = [] + self._stop_error = stop_error + + def _get_hub_client(self): + # TrainingClient.create_sampling_client lazily resolves a Hub client + # through its ServiceClient; tests never need a real one. + return None + + def stop_session(self, training_session_name: str) -> None: + self._stop_session_if_active(training_session_name) + + def _stop_session_if_active(self, training_session_name: str) -> bool: + self.stopped.append(training_session_name) + if self._stop_error is not None: + raise self._stop_error + return True + + +def test_owned_session_is_stopped_on_context_exit(): + svc = FakeServiceClient() + sc = SamplingClient("sess-owned", FakeRuntime(), service_client=svc) + with sc as entered: + assert entered is sc + assert svc.stopped == [], "must not stop while the block is running" + assert svc.stopped == ["sess-owned"] + + +def test_unowned_session_is_not_stopped_on_context_exit(): + # Workflow A: tc.create_sampling_client() shares the TrainingClient's + # session. Leaving the SamplingClient block must not kill it. + svc = FakeServiceClient() + tc = TrainingClient("sess-1", FakeRuntime(), service_client=svc) + sc = tc.create_sampling_client() + assert sc._owns_session is False + with sc: + pass + assert svc.stopped == [] + + +def test_service_client_sampling_client_owns_its_session(): + sc = SamplingClient("sess-x", FakeRuntime(), service_client=FakeServiceClient()) + assert sc._owns_session is True + + +def test_stop_on_owned_session_calls_service_client(): + svc = FakeServiceClient() + sc = SamplingClient("sess-owned", FakeRuntime(), service_client=svc) + sc.stop() + assert svc.stopped == ["sess-owned"] + + +def test_stop_on_unowned_session_is_a_noop(): + # Sharing a handle is what makes it unowned — the same thing + # TrainingClient.create_sampling_client does. + svc = FakeServiceClient() + sc = SamplingClient( + "sess-shared", + FakeRuntime(), + service_client=svc, + session_handle=_SessionHandle("sess-shared", svc), + ) + sc.stop() + assert svc.stopped == [] + + +def test_stop_on_owned_session_without_service_client_warns(caplog): + sc = SamplingClient("sess-orphan", FakeRuntime()) + with caplog.at_level("WARNING"): + sc.stop() + assert "sess-orphan" in caplog.text + + +def test_owned_session_stopped_even_when_block_raises(): + svc = FakeServiceClient() + sc = SamplingClient("sess-owned", FakeRuntime(), service_client=svc) + with pytest.raises(RuntimeError, match="boom"): + with sc: + raise RuntimeError("boom") + # The session is still reclaimed, and __exit__ does not swallow the error. + assert svc.stopped == ["sess-owned"] + + +def test_stop_failure_on_exit_is_logged_not_raised(caplog): + svc = FakeServiceClient(stop_error=RuntimeError("throttled")) + sc = SamplingClient("sess-owned", FakeRuntime(), service_client=svc) + with caplog.at_level("WARNING"): + with sc: + pass + assert "sess-owned" in caplog.text + + +def test_context_exit_drains_submitter_in_both_ownership_modes(): + owned = SamplingClient("s1", FakeRuntime(), service_client=FakeServiceClient()) + with owned: + pass + shared = SamplingClient("s2", FakeRuntime()) + with shared: + pass + # A drained submitter accepts no further work. + for client in (owned, shared): + with pytest.raises(Exception): + client.sample([1, 2, 3]).result() + + +# --- submitter sharing / per-refresh cost ----------------------------------- +# +# A submitter is an event loop on a daemon thread, and the TrainingClient +# factory is called once per sampler refresh. These tests pin the sharing that +# keeps that off the per-step path, and the ownership rules it needs. + + +def _live_fds(): + """Descriptor count for this process; Linux-only, so the caller skips.""" + return len(os.listdir(f"/proc/{os.getpid()}/fd")) + + +def test_derived_sampling_clients_share_the_training_clients_submitter(): + tc = TrainingClient("sess-1", FakeRuntime()) + clients = [tc.create_sampling_client() for _ in range(5)] + assert all(sc._submitter is tc._submitter for sc in clients) + assert all(sc._owns_submitter is False for sc in clients) + + +def test_a_standalone_sampling_client_builds_and_owns_its_own_submitter(): + """The session-owning path has no parent to borrow from.""" + sc = SamplingClient("sess-x", FakeRuntime(), service_client=FakeServiceClient()) + assert sc._owns_submitter is True + assert sc._submitter is not None + sc.stop() + + +def test_refreshing_the_sampler_costs_no_thread_and_no_descriptor_per_refresh(): + """The leak this sharing exists to remove: +1 thread and +3 fds per step. + + An RL loop refreshes its sampler every step, so an unshared submitter made + the cost per step rather than per session; ~20 refreshes exhausted the + default macOS descriptor budget on top of a socket-heavy baseline. + """ + if not os.path.isdir(f"/proc/{os.getpid()}/fd"): + pytest.skip("descriptor accounting via /proc is Linux-only") + + tc = TrainingClient("sess-1", FakeRuntime()) + # Force the parent's loop into existence first, so it is part of the + # baseline and the measurement is purely what the refreshes add. + tc.create_sampling_client().sample([1, 2, 3]).result() + # Settle cyclic garbage first: an earlier test's parked loop collected + # mid-measurement would drop the count and read as a refresh taking one. + gc.collect() + fds_before, threads_before = _live_fds(), threading.active_count() + + clients = [tc.create_sampling_client() for _ in range(20)] + for sc in clients: + sc.sample([1, 2, 3]).result() + + assert threading.active_count() == threads_before, "a refresh started a thread" + assert _live_fds() == fds_before, "a refresh took a descriptor" + + +def test_a_retired_sampling_client_does_not_close_its_parents_submitter(): + """Without the ownership gate, retiring one sampler ends the training run. + + ``stop`` and ``__exit__`` closed the submitter unconditionally, which under + sharing is the parent's. Every later training operation would then fail with + a closed-submitter error. + """ + tc = TrainingClient("sess-1", FakeRuntime()) + retired = tc.create_sampling_client() + + retired.stop() + with tc.create_sampling_client() as also_retired: + assert also_retired is not None + + # The parent is untouched: it still submits, and its submitter is still open. + assert tc._submitter._closed is False + assert tc.create_sampling_client().sample([1, 2, 3]).result() is not None + + +@pytest.mark.parametrize("retire", [lambda sc: sc.stop(), lambda sc: sc.__exit__(*(None,) * 3)]) +def test_a_retired_sampling_client_refuses_work_even_though_it_shares_a_submitter(retire): + """Retirement has to fail fast, not silently keep submitting. + + Leaving the shared submitter open for its owner means closure cannot be read + off the submitter, so the client records it. Without that, a caller who + retires a sampler and keeps the reference would go on reaching the service + against a session it no longer holds. + """ + tc = TrainingClient("sess-1", FakeRuntime()) + retired = tc.create_sampling_client() + + retire(retired) + + for op in (lambda: retired.sample([1, 2, 3]), lambda: retired.compute_logprobs([1, 2, 3])): + with pytest.raises(TrainingSessionError, match="closed and accepts no new"): + op() + # Refusing is this client's own state; the shared submitter stays usable. + assert tc._submitter._closed is False + + +def test_a_shared_submitter_is_registered_once_not_once_per_refresh(): + """The session handle holds submitters strongly for the whole session. + + Registering a shared submitter per refresh would pile up references to the + single submitter and drain it once per entry at teardown. + """ + tc = TrainingClient("sess-1", FakeRuntime()) + for _ in range(10): + tc.create_sampling_client() + assert tc._handle._submitters == [tc._submitter] + + +def test_a_standalone_client_still_registers_its_own_submitter(): + """Skipping registration for a shared submitter must not skip owned ones.""" + sc = SamplingClient("sess-x", FakeRuntime(), service_client=FakeServiceClient()) + assert sc._handle._submitters == [sc._submitter] + + +def test_stopping_the_training_client_closes_its_derived_samplers_too(): + """The shared submitter runs both, so the parent's teardown ends both. + + A derived client used to hold a private submitter and so stayed submittable + after its parent stopped, but the session it samples is gone by then, so the + submits only failed later and further away. Failing at the client with the + closed-client message is the better of the two. + """ + svc = FakeServiceClient() + tc = TrainingClient("sess-1", FakeRuntime(), service_client=svc) + sc = tc.create_sampling_client() + + tc.stop() + + assert svc.stopped == ["sess-1"] + with pytest.raises(TrainingSessionError, match="closed and accepts no new"): + sc.sample([1, 2, 3]) diff --git a/sagemaker-train/tests/unit/train/training_session/test_script_compatibility.py b/sagemaker-train/tests/unit/train/training_session/test_script_compatibility.py new file mode 100644 index 0000000000..f432270b9c --- /dev/null +++ b/sagemaker-train/tests/unit/train/training_session/test_script_compatibility.py @@ -0,0 +1,618 @@ +# 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. +"""Existing training scripts run against this SDK unchanged. + +Two guarantees, both load bearing for dropping such a script in with no edits: + +* **Calling conventions match.** ``rank`` is positional-or-keyword, and the + ``*_async`` twins of every client factory exist, because such code reaches + for the coroutine form far more often than the blocking one. +* **Unhonoured parameters are loud.** Parameters we accept but cannot implement + log a warning naming the parameter. Accepting them silently would let a run + differ from what the caller asked for with nothing in the log to explain it. +""" + +from __future__ import annotations + +import asyncio +import inspect + +import pytest + +from sagemaker.train.training_session import futures +from sagemaker.train.training_session import service_client as sc_mod +from sagemaker.train.training_session.sampling_client import SamplingClient +from sagemaker.train.training_session.service_client import ServiceClient +from sagemaker.train.training_session.training_client import TrainingClient + +# --- calling conventions ---------------------------------------------------- + + +def test_only_base_model_is_positional(): + """`rank` stays keyword-only on purpose. + + ``create_lora_training_client(model, 16)`` does not say what 16 is, and reads + equally well as ``alpha``. Every surveyed call site passes ``rank=`` as a + keyword already (8 of 8), so requiring it costs nothing in drop-in terms. + ``base_model`` *is* positional, because 4 of those 8 pass it that way. + """ + params = inspect.signature(ServiceClient.create_lora_training_client).parameters + assert params["base_model"].kind is inspect.Parameter.POSITIONAL_OR_KEYWORD + for name in ("rank", "alpha", "seed", "user_metadata"): + assert params[name].kind is inspect.Parameter.KEYWORD_ONLY, name + + +@pytest.mark.parametrize( + "cls", [futures.APIFuture, futures.PendingAPIFuture, futures.CombinedFuture] +) +@pytest.mark.parametrize("method", ["result", "result_async"]) +def test_future_wait_takes_only_timeout(cls, method): + """``timeout`` is the one knob on a wait, and it is positional-or-keyword. + + Both ``future.result(30)`` and ``future.result(timeout=30)`` are how existing + scripts bound a wait, so both must bind. A second parameter would also make + ``result(30, 1)`` mean something here that it means nowhere else; the poll + cadence is set when the future is built instead. + """ + params = inspect.signature(getattr(cls, method)).parameters + assert list(params) == ["self", "timeout"] + assert params["timeout"].kind is inspect.Parameter.POSITIONAL_OR_KEYWORD + + +def test_positional_rank_is_rejected(): + """The ambiguous form must fail loudly, not silently bind to something.""" + sig = inspect.signature(ServiceClient.create_lora_training_client) + with pytest.raises(TypeError): + sig.bind(None, "Qwen/Qwen3-8B", 16) + + +def test_async_twin_also_rejects_positional_rank(): + """The coroutine twin must not be a back door to the positional form.""" + sig = inspect.signature(ServiceClient.create_lora_training_client_async) + with pytest.raises(TypeError): + sig.bind(None, "Qwen/Qwen3-8B", 16) + + +@pytest.mark.parametrize( + "cls, name", + [ + (ServiceClient, "create_lora_training_client_async"), + (ServiceClient, "create_training_client_from_state_async"), + (ServiceClient, "create_training_client_from_state_with_optimizer_async"), + (ServiceClient, "create_sampling_client_async"), + (TrainingClient, "save_weights_and_get_sampling_client_async"), + (TrainingClient, "create_sampling_client_async"), + ], +) +def test_async_factory_twins_exist_and_are_coroutines(cls, name): + fn = getattr(cls, name, None) + assert fn is not None, f"{cls.__name__}.{name} is missing" + assert inspect.iscoroutinefunction(fn), f"{cls.__name__}.{name} must be async def" + + +@pytest.mark.parametrize( + "cls, sync_name", + [ + (TrainingClient, "forward_backward"), + (TrainingClient, "optim_step"), + (TrainingClient, "save_state"), + (TrainingClient, "load_state"), + (TrainingClient, "save_weights_for_sampler"), + (TrainingClient, "forward"), + (SamplingClient, "sample"), + (SamplingClient, "compute_logprobs"), + ], +) +def test_every_operation_has_an_async_twin(cls, sync_name): + twin = getattr(cls, f"{sync_name}_async", None) + assert twin is not None, f"{cls.__name__}.{sync_name}_async is missing" + assert inspect.iscoroutinefunction(twin) + + +# --- unhonoured parameters are loud ---------------------------------------- + + +UNSUPPORTED = [ + ("seed", 1234), + ("train_mlp", False), + ("train_attn", False), + ("train_unembed", False), + ("user_metadata", {"run": "exp-1"}), +] + + +@pytest.mark.parametrize("name, value", UNSUPPORTED) +def test_unsupported_parameter_is_accepted_and_warns(name, value, caplog): + """Accepted so the script runs; warned about so the caller is not misled.""" + from sagemaker.train.training_session import service_client as sc_mod + + with caplog.at_level("WARNING"): + sc_mod._warn_unsupported(**{name: value}) + + assert name in caplog.text + assert "not supported" in caplog.text + + +def test_defaults_do_not_warn(caplog): + """Left unset, none of them fire — or the warning becomes background noise.""" + from sagemaker.train.training_session import service_client as sc_mod + + with caplog.at_level("WARNING"): + sc_mod._warn_unsupported( + seed=None, + train_mlp=None, + train_attn=None, + train_unembed=None, + user_metadata=None, + ) + + assert caplog.text == "" + + +@pytest.mark.parametrize( + "args, kwargs", + [ + # supervised training loop — all keyword + ((), {"base_model": "Qwen/Qwen3-8B", "rank": 32, "user_metadata": {"a": "b"}}), + # RL training loop — base_model positional + (("Qwen/Qwen3-8B",), {"rank": 32, "user_metadata": {"a": "b"}}), + # RL recipe loop — no metadata + ((), {"base_model": "Qwen/Qwen3-8B", "rank": 32}), + ], +) +def test_accepts_every_surveyed_call_shape(args, kwargs): + """The exact shapes used across the 8 surveyed creation sites. + + Asserted against the signature rather than by creating a session: this is a + calling-convention guarantee, and binding proves the call would not raise + TypeError. + """ + sig = inspect.signature(ServiceClient.create_lora_training_client) + bound = sig.bind(None, *args, **kwargs) + assert bound.arguments["base_model"] == "Qwen/Qwen3-8B" + assert bound.arguments["rank"] == 32 + + +# --- unhonoured operation parameters are loud ------------------------------- + + +#: ``(method, kwargs, expected_warned_names)`` for the checkpoint-lifetime knobs +#: the save ops accept for drop-in compatibility but cannot honour. +SAVE_UNSUPPORTED = [ + ("save_state", {"ttl_seconds": 3600}, ["ttl_seconds"]), + ("save_state", {"overwrite": True}, ["overwrite"]), + ("save_state", {"ttl_seconds": 60, "overwrite": True}, ["ttl_seconds", "overwrite"]), + ("save_state", {"name": "ckpt-1"}, ["name"]), + ("save_weights_for_sampler", {"ttl_seconds": 3600}, ["ttl_seconds"]), + ("save_weights_for_sampler", {"name": "ckpt-1"}, ["name"]), +] + + +@pytest.mark.parametrize("method, kwargs, warned", SAVE_UNSUPPORTED) +def test_save_op_lifetime_knob_is_accepted_and_warns(method, kwargs, warned, caplog): + """Accepted so the script runs; warned about so the caller is not misled. + + A caller who asked for a 1-hour TTL and got platform retention has a cost + surprise coming, so the discrepancy has to be in the log. + """ + submitted = {} + + class FakeSelf: + def _submit(self, op_config, *args, **kwargs): + submitted["op_config"] = op_config + return "future" + + def _cache_on_save(self, kind): + return None + + with caplog.at_level("WARNING"): + result = getattr(TrainingClient, method)(FakeSelf(), **kwargs) + + # Proceeds: the op is still submitted, and the unhonoured knob is not on the wire. + assert result == "future" + assert not any( + key.lower() in ("ttl", "ttlseconds", "overwrite", "name") + for config in submitted["op_config"].values() + for key in config + ) + for name in warned: + assert name in caplog.text + assert "not supported" in caplog.text + # The name warning has to name the parameter that does work, or the caller + # is told "no" with nowhere to go. + if "name" in warned: + assert "description" in caplog.text + + +@pytest.mark.parametrize("method", ["save_state", "save_weights_for_sampler"]) +def test_save_op_defaults_do_not_warn(method, caplog): + """``overwrite=False`` is what we already do, so it must stay silent. + + Warning on a default would fire on every save in the loop and train the + caller to filter the log, which is how the TTL warning gets missed too. + """ + + class FakeSelf: + def _submit(self, op_config, *args, **kwargs): + return "future" + + def _cache_on_save(self, kind): + return None + + with caplog.at_level("WARNING"): + getattr(TrainingClient, method)(FakeSelf(), description="step 100") + + assert caplog.text == "" + + +@pytest.mark.parametrize( + "method, kwargs", + [ + ("save_state_async", {"ttl_seconds": 3600, "overwrite": True, "name": "c1"}), + ("save_weights_for_sampler_async", {"ttl_seconds": 3600, "name": "c1"}), + ], +) +def test_save_async_twin_forwards_lifetime_knobs(method, kwargs): + """The twin must not drop them, or the warning silently stops firing.""" + seen = {} + sync_name = method[: -len("_async")] + + class FakeSelf: + pass + + setattr(FakeSelf, sync_name, lambda self, **kw: seen.update(kw) or "future") + + assert asyncio.run(getattr(TrainingClient, method)(FakeSelf(), **kwargs)) == "future" + for name, value in kwargs.items(): + assert seen[name] == value + + +# --- the async twins actually run ------------------------------------------- + + +def test_async_twin_delegates_to_the_sync_path(monkeypatch): + """to_thread keeps the blocking wait off the loop but must still deliver.""" + calls = {} + + class FakeSelf: + def create_sampling_client(self, **kwargs): + calls.update(kwargs) + return "sampling-client" + + coro = ServiceClient.create_sampling_client_async( + FakeSelf(), base_model="arn:aws:sagemaker:::hub-content/m" + ) + assert asyncio.run(coro) == "sampling-client" + assert calls == {"base_model": "arn:aws:sagemaker:::hub-content/m"} + + +def test_async_twin_forwards_base_model_and_keywords(): + seen = {} + + class FakeSelf: + def create_lora_training_client(self, base_model, **kwargs): + seen["base_model"] = base_model + seen.update(kwargs) + return "training-client" + + coro = ServiceClient.create_lora_training_client_async( + FakeSelf(), "Qwen/Qwen3-8B", rank=16, user_metadata={"a": "b"} + ) + assert asyncio.run(coro) == "training-client" + assert seen == { + "base_model": "Qwen/Qwen3-8B", + "rank": 16, + "user_metadata": {"a": "b"}, + } + + +@pytest.mark.parametrize( + "cls, method, kwargs", + [ + ( + ServiceClient, + "create_training_client_from_state_async", + {"model_package_arn": "arn:aws:sagemaker:::model-package/g/1", "rank": 8}, + ), + ( + ServiceClient, + "create_training_client_from_state_with_optimizer_async", + {"model_package_arn": "arn:aws:sagemaker:::model-package/g/1", "rank": 8}, + ), + ( + TrainingClient, + "save_weights_and_get_sampling_client_async", + {"description": "step 100", "merge_weights": True}, + ), + ( + TrainingClient, + "create_sampling_client_async", + {"model_package_arn": "arn:aws:sagemaker:::model-package/g/1"}, + ), + ], +) +def test_client_returning_async_twin_forwards_every_keyword(cls, method, kwargs): + """Every keyword must survive the hop onto the worker thread. + + These twins hand the blocking factory to ``asyncio.to_thread``, which + forwards positionally and by keyword. A dropped keyword here would silently + build the client with a default instead of what the caller asked for. + """ + seen = {} + sync_name = method[: -len("_async")] + + class FakeSelf: + pass + + setattr(FakeSelf, sync_name, lambda self, **kw: seen.update(kw) or "client") + + assert asyncio.run(getattr(cls, method)(FakeSelf(), **kwargs)) == "client" + for name, value in kwargs.items(): + assert seen[name] == value + + +def test_async_twin_does_not_swallow_a_keyword_named_func(): + """``to_thread(func, /, ...)`` takes ``func`` positionally-only. + + So a caller keyword that happens to be named ``func`` is forwarded to the + factory rather than rebinding what ``to_thread`` runs. It fails in the + factory, exactly as the synchronous call would. + """ + + class FakeSelf: + def create_lora_training_client(self, base_model, **kwargs): + raise TypeError(f"unexpected keyword: {sorted(kwargs)}") + + coro = ServiceClient.create_lora_training_client_async( + FakeSelf(), "Qwen/Qwen3-8B", func="not-a-callable" + ) + with pytest.raises(TypeError, match=r"unexpected keyword: \['func'\]"): + asyncio.run(coro) + + +# --- auto-generated resource names ------------------------------------------- + + +def test_session_name_prefix_only_prefixes_session_names(): + """The parameter names sessions and nothing else. + + It used to also feed the model package group name, which is why it was + called ``resource_name``. Now that the group has its own shared default, + the prefix has exactly one job and the name says so. + """ + svc = ServiceClient.__new__(ServiceClient) + svc._session_name_prefix = "my-run" + + first = ServiceClient._generate_session_name(svc) + second = ServiceClient._generate_session_name(svc) + + assert first.startswith("my-run-") + assert first != second, "each session needs its own name" + + +def test_default_session_name_prefix_is_used_when_unset(): + svc = ServiceClient.__new__(ServiceClient) + svc._session_name_prefix = None + + assert ServiceClient._generate_session_name(svc).startswith( + f"{sc_mod.DEFAULT_SESSION_NAME_PREFIX}-" + ) + + +def test_default_model_package_group_is_shared_not_per_session(): + """One group per account/region, like the SageMaker default bucket. + + A group per session would leak a group per run and scatter checkpoints so + that a later process could not find the latest one. + """ + created = [] + + class FakeHub: + exceptions = type("E", (), {"ClientError": RuntimeError}) + + def create_model_package_group(self, **kwargs): + created.append(kwargs["ModelPackageGroupName"]) + return {"ModelPackageGroupArn": "arn:group"} + + svc = ServiceClient.__new__(ServiceClient) + svc._get_hub_client = lambda: FakeHub() + + assert ServiceClient._auto_create_mpg(svc) == "arn:group" + assert ServiceClient._auto_create_mpg(svc) == "arn:group" + assert created == [sc_mod.DEFAULT_MODEL_PACKAGE_GROUP_NAME] * 2 + + +def test_default_model_package_group_name_shape(): + """Named per the convention agreed for this SDK. + + ``sagemaker`` prefix and ``default``, and no account id or region: the + group ARN already carries both, so repeating them is noise. Feature-agnostic + for the same reason the default bucket is. + """ + name = sc_mod.DEFAULT_MODEL_PACKAGE_GROUP_NAME + + assert name == "sagemaker-default" + assert not any(ch.isdigit() for ch in name), "no account id" + assert "us-east" not in name and "us-west" not in name, "no region" + + +# --- retry_config is accepted and loud -------------------------------------- + + +def test_service_client_create_sampling_client_warns_on_retry_config(caplog): + """Existing code passes retry_config; it must not TypeError, and must warn.""" + svc = ServiceClient.__new__(ServiceClient) + + with caplog.at_level("WARNING"): + # Reaches the ValueError only after the warning has been emitted, which + # is the point: the warning fires on the way in, not as a side effect of + # a successful session creation. + with pytest.raises(ValueError): + ServiceClient.create_sampling_client(svc, retry_config={"max": 5}) + + assert "retry_config" in caplog.text + assert "not supported" in caplog.text + + +def test_training_client_create_sampling_client_warns_once_on_retry_config(caplog): + """One warning per call, not one per layer that forwards it.""" + + class FakeSelf: + _hub_client = object() + _tokenizer = object() + _hf_model_id = "m" + _runtime = None + _service_client = None + _handle = None + # None is a real value here, not a stand-in: it makes the new client + # build its own submitter, which is what the standalone path does. + _submitter = None + training_session_name = "s" + _base_model = "b" + + with caplog.at_level("WARNING"): + TrainingClient.create_sampling_client( + FakeSelf(), model_package_arn="", retry_config={"max": 5} + ) + + assert caplog.text.count("retry_config") == 1, caplog.text + + +def test_save_weights_and_get_sampling_client_forwards_retry_config(): + """Forwarding matters: dropping it would silently stop the warning.""" + seen = {} + + class FakeFuture: + def result(self): + return type("R", (), {"model_package_arn": "arn:mp/weights-1"})() + + class FakeSelf: + def save_weights_for_sampler(self, **kwargs): + return FakeFuture() + + def _load_weights_for_sampling(self, model_package_arn): + seen["loaded"] = model_package_arn + + def create_sampling_client(self, **kwargs): + seen.update(kwargs) + return "sampling-client" + + result = TrainingClient.save_weights_and_get_sampling_client( + FakeSelf(), retry_config={"max": 5} + ) + + assert result == "sampling-client" + assert seen["retry_config"] == {"max": 5} + + +# --- kind-filtered discovery must not answer None by accident ---------------- + + +def _fake_group(kinds): + """A ServiceClient whose group holds ``kinds`` (newest first).""" + from sagemaker.train.training_session.types import ModelPackageKind + + arns = [f"arn:mp/{i}" for i in range(len(kinds))] + by_arn = dict(zip(arns, kinds)) + + class FakeHub: + def list_model_packages(self, **kwargs): + limit = kwargs.get("MaxResults", len(arns)) + return { + "ModelPackageSummaryList": [ + {"ModelPackageArn": a, "ModelPackageVersion": i} for i, a in enumerate(arns) + ][:limit] + } + + def describe_model_package(self, ModelPackageName): + is_ckpt = by_arn[ModelPackageName] == ModelPackageKind.CHECKPOINT + return {"InferenceSpecification": {"Containers": [{"IsCheckpoint": is_ckpt}]}} + + svc = ServiceClient.__new__(ServiceClient) + svc._get_hub_client = lambda: FakeHub() + svc._model_package_group_arn = "arn:group" + return svc + + +def test_buried_checkpoint_is_found_once_scan_limit_covers_it(): + """The RL-loop shape: weights saved every step, state every N steps. + + The newest CHECKPOINT sits under N-1 MODEL packages, so with N above the + default scan limit it is out of reach until the caller widens the search. + This pins the recovery route the warning points at: a larger scan_limit does + find it, so None from a truncated scan really was a truncation artefact and + not an absent checkpoint. + """ + from sagemaker.train.training_session.types import ModelPackageKind + + kinds = [ModelPackageKind.MODEL] * 50 + [ModelPackageKind.CHECKPOINT] + svc = _fake_group(kinds) + + assert ( + ServiceClient.get_latest_model_package(svc, ModelPackageKind.CHECKPOINT) is None + ), "50 deep is past the default scan limit" + + found = ServiceClient.get_latest_model_package(svc, ModelPackageKind.CHECKPOINT, scan_limit=100) + assert found is not None + assert found.model_package_arn == "arn:mp/50" + + +def test_exhausted_scan_limit_warns_rather_than_implying_none_exist(caplog): + """None + a truncated search is not evidence; say so at WARNING.""" + from sagemaker.train.training_session.types import ModelPackageKind + + svc = _fake_group([ModelPackageKind.MODEL] * 10) + + with caplog.at_level("WARNING"): + found = ServiceClient.get_latest_model_package( + svc, ModelPackageKind.CHECKPOINT, scan_limit=5 + ) + + assert found is None + assert "scan_limit" in caplog.text + + +def test_fully_scanned_group_with_no_match_does_not_warn(caplog): + """Searched the whole group and found nothing, which IS evidence, so quiet.""" + from sagemaker.train.training_session.types import ModelPackageKind + + svc = _fake_group([ModelPackageKind.MODEL] * 3) + + with caplog.at_level("WARNING"): + found = ServiceClient.get_latest_model_package( + svc, ModelPackageKind.CHECKPOINT, scan_limit=20 + ) + + assert found is None + assert caplog.text == "" + + +def test_existing_default_group_is_reused_not_recreated(): + """Second caller must resolve the ARN rather than fail on AlreadyExists.""" + + class AlreadyExists(RuntimeError): + pass + + class FakeHub: + exceptions = type("E", (), {"ClientError": AlreadyExists}) + + def create_model_package_group(self, **kwargs): + raise AlreadyExists("Model Package Group already exists") + + def describe_model_package_group(self, **kwargs): + return {"ModelPackageGroupArn": f"arn:{kwargs['ModelPackageGroupName']}"} + + svc = ServiceClient.__new__(ServiceClient) + svc._get_hub_client = lambda: FakeHub() + + assert ServiceClient._auto_create_mpg(svc) == (f"arn:{sc_mod.DEFAULT_MODEL_PACKAGE_GROUP_NAME}") diff --git a/sagemaker-train/tests/unit/train/training_session/test_service_client.py b/sagemaker-train/tests/unit/train/training_session/test_service_client.py new file mode 100644 index 0000000000..7abd6f5aa9 --- /dev/null +++ b/sagemaker-train/tests/unit/train/training_session/test_service_client.py @@ -0,0 +1,2190 @@ +# 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/ +"""Unit tests for sagemaker.train.training_session.service_client. + +Both the control-plane and runtime clients are faked by patching the factories +in ``_clients``; the Running-wait loop is driven by a patched ``time`` so no +test performs a real sleep. +""" + +from __future__ import annotations + +import json + +import pytest +from botocore.exceptions import ClientError + +from sagemaker.train.training_session import service_client as sc_mod +from sagemaker.train.training_session.service_client import ServiceClient +from sagemaker.train.training_session.training_client import TrainingClient +from sagemaker.train.training_session.types import ModelPackageKind, SessionStatus + + +class FakeSageMakerClient: + """Control-plane fake: scripts describe_training_session status sequence.""" + + def __init__(self, statuses, *, describe_detail=None, list_pages=None, stop_error=None): + self._statuses = list(statuses) + self._stop_error = stop_error + self.create_calls: list[dict] = [] + self.describe_calls: list[dict] = [] + self.stop_calls: list[dict] = [] + self.delete_calls: list[dict] = [] + self.list_calls: list[dict] = [] + self._describe_detail = describe_detail or {} + # list_pages: list of (summaries, next_token) tuples served in order. + self._list_pages = list(list_pages or []) + + def create_training_session(self, **kwargs): + self.create_calls.append(kwargs) + return {"TrainingSessionArn": "arn:aws:sagemaker:::training-session/x"} + + def describe_training_session(self, **kwargs): + self.describe_calls.append(kwargs) + status = self._statuses.pop(0) if len(self._statuses) > 1 else self._statuses[0] + detail = {"TrainingSessionStatus": status, **self._describe_detail} + return {"TrainingSessionDetail": detail} + + def stop_training_session(self, **kwargs): + self.stop_calls.append(kwargs) + if self._stop_error is not None: + raise self._stop_error + return {} + + def delete_training_session(self, **kwargs): + self.delete_calls.append(kwargs) + return {} + + def list_training_sessions(self, **kwargs): + self.list_calls.append(kwargs) + summaries, next_token = self._list_pages.pop(0) + resp = {"TrainingSessionSummaries": summaries} + if next_token: + resp["NextToken"] = next_token + return resp + + +class FakeClock: + def __init__(self): + self.now = 0.0 + self.sleeps: list[float] = [] + + def monotonic(self): + return self.now + + def sleep(self, seconds): + self.sleeps.append(seconds) + self.now += seconds + + +@pytest.fixture +def patch_clients(monkeypatch): + """Patch client factories + time so ServiceClient runs offline and instantly.""" + + clock = FakeClock() + monkeypatch.setattr(sc_mod.time, "monotonic", clock.monotonic) + monkeypatch.setattr(sc_mod.time, "sleep", clock.sleep) + + holder = {} + + def _install(sagemaker_client, runtime=None): + runtime = runtime if runtime is not None else object() + holder["sm"] = sagemaker_client + holder["rt"] = runtime + import sagemaker.train.training_session._clients as clients_mod + + # ServiceClient imports both client builders lazily inside __init__, so + # patch the source module. + monkeypatch.setattr(clients_mod, "_control_plane_client", lambda **kw: sagemaker_client) + monkeypatch.setattr(clients_mod, "_runtime_client", lambda **kw: runtime) + return clock + + return _install, holder + + +def _create_kwargs(**overrides): + base = dict( + training_session_name="sess-1", + model_package_group_arn="arn:mpg/x", + s3_output_path="s3://b/out/", + ) + base.update(overrides) + return base + + +# The new API takes base_model as positional arg; helper returns the +# positional + kwargs for svc.create_lora_training_client(base_model, **kwargs). +_BASE_MODEL_ARN = "arn:hub-content/Model/base" + + +def test_create_waits_for_running_then_returns_training_client(patch_clients): + install, _ = patch_clients + sm = FakeSageMakerClient(["Created", "Created", "Running"]) + clock = install(sm) + + svc = ServiceClient(region="us-west-2", role_arn="arn:aws:iam::123456789012:role/TestRole") + tc = svc.create_lora_training_client(_BASE_MODEL_ARN, **_create_kwargs(rank=16, alpha=32)) + + assert isinstance(tc, TrainingClient) + assert tc.training_session_name == "sess-1" + # CreateTrainingSession got the LoRA config. + cc = sm.create_calls[0] + assert cc["LoraConfig"] == {"Rank": 16, "Alpha": 32} + # Polled until Running; only slept while pending (no real wait). + assert clock.sleeps and all(s == 10 for s in clock.sleeps) + + +def test_create_raises_on_terminal_status(patch_clients): + install, _ = patch_clients + sm = FakeSageMakerClient(["Created", "Failed"]) + install(sm) + + svc = ServiceClient(role_arn="arn:aws:iam::123456789012:role/TestRole") + with pytest.raises(RuntimeError, match="terminal state"): + svc.create_lora_training_client(_BASE_MODEL_ARN, **_create_kwargs()) + + +# --- create_sampling_client ------------------------------------------------- + + +def test_create_sampling_client_with_base_model(patch_clients): + """Creates an implicit session and returns a SamplingClient sampling the + base model.""" + install, holder = patch_clients + sm = FakeSageMakerClient(["Running"]) + runtime = object() + install(sm, runtime=runtime) + + svc = ServiceClient( + role_arn="arn:aws:iam::123456789012:role/TestRole", + model_package_group_arn="arn:mpg/x", + s3_output_path="s3://b/out/", + ) + svc._hub_client = FakeHubClient() + sampler = svc.create_sampling_client(base_model="hubcontent-vlm-model-b") + + assert sampler.model_package_arn == "" + assert sampler._base_model == svc._hub_client._arn + assert sampler._runtime is runtime + # An implicit session was created. + assert len(sm.create_calls) == 1 + assert sm.create_calls[0]["RoleArn"] == "arn:aws:iam::123456789012:role/TestRole" + # Nothing else owns that implicit session, so this client must stop it. + assert sampler._owns_session is True + assert sampler._service_client is svc + + +def test_sampling_client_context_exit_stops_implicit_session(patch_clients): + """`with svc.create_sampling_client(...)` stops the session it created.""" + install, _ = patch_clients + sm = FakeSageMakerClient(["Running"]) + install(sm) + + svc = ServiceClient( + role_arn="arn:aws:iam::123456789012:role/TestRole", + model_package_group_arn="arn:mpg/x", + s3_output_path="s3://b/out/", + ) + svc._hub_client = FakeHubClient() + + with svc.create_sampling_client(base_model="hubcontent-vlm-model-b") as sp: + session_name = sp.training_session_name + assert sm.stop_calls == [] + + assert [c["TrainingSessionName"] for c in sm.stop_calls] == [session_name] + + +def test_create_sampling_client_with_model_package_arn(patch_clients): + """Creates an implicit session and binds the checkpoint ARN.""" + install, _ = patch_clients + sm = FakeSageMakerClient(["Running"]) + install(sm) + + svc = ServiceClient( + region="us-west-2", + role_arn="arn:aws:iam::123456789012:role/TestRole", + model_package_group_arn="arn:mpg/x", + s3_output_path="s3://b/out/", + ) + # Stub hub client: DescribeModelPackage publishes the base model as + # InferenceSpecification.Containers[0].BaseModel (name + pinned version), + # not as a top-level field. + hub = FakeHubClient() + hub.describe_model_package = lambda **kw: { + "ModelPackageArn": "arn:mp/weights-7", + "InferenceSpecification": { + "Containers": [ + { + "IsCheckpoint": False, + "BaseModel": { + "HubContentName": "hubcontent-text-model-a", + "HubContentVersion": "3.15.0", + }, + } + ] + }, + } + svc._hub_client = hub + + sampler = svc.create_sampling_client(model_package_arn="arn:mp/weights-7") + + assert sampler.model_package_arn == "arn:mp/weights-7" + assert sampler._base_model == ( + "arn:aws:sagemaker:us-west-2:aws:hub-content/" + "SageMakerPublicHub/Model/hubcontent-text-model-a/3.15.0" + ) + + +def test_create_sampling_client_rejects_neither_arg(patch_clients): + """Must supply exactly one of base_model or model_package_arn.""" + install, _ = patch_clients + install(FakeSageMakerClient(["Running"])) + + svc = ServiceClient(role_arn="arn:aws:iam::123456789012:role/TestRole") + with pytest.raises(ValueError, match="exactly one"): + svc.create_sampling_client() + + +def test_create_sampling_client_rejects_both_args(patch_clients): + install, _ = patch_clients + install(FakeSageMakerClient(["Running"])) + + svc = ServiceClient(role_arn="arn:aws:iam::123456789012:role/TestRole") + with pytest.raises(ValueError, match="exactly one"): + svc.create_sampling_client(base_model="provider-a/model-a", model_package_arn="arn:mp/x") + + +# --- stop / delete ---------------------------------------------------------- + + +def test_stop_session_calls_control_plane(patch_clients): + install, _ = patch_clients + sm = FakeSageMakerClient(["Running"]) + install(sm) + + ServiceClient(role_arn="arn:aws:iam::123456789012:role/TestRole").stop_session("sess-1") + assert sm.stop_calls == [{"TrainingSessionName": "sess-1"}] + + +def _stop_state_error(from_state): + """The ValidationException the control plane returns for a rejected stop.""" + return ClientError( + { + "Error": { + "Code": "ValidationException", + "Message": ( + f"Cannot stop training session in state {from_state}: invalid " + f"session state transition from {from_state} to Stopping" + ), + }, + "ResponseMetadata": {"HTTPStatusCode": 400}, + }, + "StopTrainingSession", + ) + + +@pytest.mark.parametrize("state", ["Stopping", "Stopped", "Failed", "Deleting", "DeletionFailed"]) +def test_stop_session_is_idempotent_for_states_with_nothing_to_stop(patch_clients, state): + """A stop the service cannot apply is not an error the caller has to handle.""" + install, _ = patch_clients + sm = FakeSageMakerClient(["Running"], stop_error=_stop_state_error(state)) + install(sm) + + svc = ServiceClient(role_arn="arn:aws:iam::123456789012:role/TestRole") + svc.stop_session("sess-1") # must not raise + assert sm.stop_calls == [{"TrainingSessionName": "sess-1"}] + + +def test_stop_session_repeated_calls_both_return(patch_clients): + """Second stop mirrors an explicit stop followed by the automatic one.""" + install, _ = patch_clients + sm = FakeSageMakerClient(["Running"]) + install(sm) + + svc = ServiceClient(role_arn="arn:aws:iam::123456789012:role/TestRole") + svc.stop_session("sess-1") + sm._stop_error = _stop_state_error("Stopping") + svc.stop_session("sess-1") + assert len(sm.stop_calls) == 2 + + +def test_stop_session_reraises_unrelated_validation_error(patch_clients): + """Only the state-transition case is absorbed; real failures still surface.""" + install, _ = patch_clients + sm = FakeSageMakerClient( + ["Running"], + stop_error=ClientError( + {"Error": {"Code": "ValidationException", "Message": "1 validation error"}}, + "StopTrainingSession", + ), + ) + install(sm) + + svc = ServiceClient(role_arn="arn:aws:iam::123456789012:role/TestRole") + with pytest.raises(ClientError): + svc.stop_session("sess-1") + + +def test_stop_session_reraises_resource_not_found(patch_clients): + install, _ = patch_clients + sm = FakeSageMakerClient( + ["Running"], + stop_error=ClientError( + { + "Error": { + "Code": "ResourceNotFoundException", + "Message": "Training session not found: sess-1", + } + }, + "StopTrainingSession", + ), + ) + install(sm) + + svc = ServiceClient(role_arn="arn:aws:iam::123456789012:role/TestRole") + with pytest.raises(ClientError): + svc.stop_session("sess-1") + + +def test_delete_session_calls_delete_directly(patch_clients): + install, _ = patch_clients + sm = FakeSageMakerClient(["Stopped"]) + install(sm) + + ServiceClient(role_arn="arn:aws:iam::123456789012:role/TestRole").delete_session("sess-1") + assert sm.delete_calls == [{"TrainingSessionName": "sess-1"}] + # Without stop_if_active we don't stop or poll. + assert sm.stop_calls == [] + assert sm.describe_calls == [] + + +def test_delete_session_stops_and_waits_when_requested(patch_clients): + install, _ = patch_clients + # Stopping -> Stopped: the wait loop sleeps once, then deletes. + sm = FakeSageMakerClient(["Stopping", "Stopped"]) + install(sm) + + ServiceClient(role_arn="arn:aws:iam::123456789012:role/TestRole").delete_session( + "sess-1", stop_if_active=True + ) + assert sm.stop_calls == [{"TrainingSessionName": "sess-1"}] + assert sm.delete_calls == [{"TrainingSessionName": "sess-1"}] + # Polled at least once before reaching the deletable state. + assert sm.describe_calls + + +def test_stop_session_accepts_arn(patch_clients): + """A TrainingSessionArn from list_sessions passes straight to stop_session.""" + install, _ = patch_clients + sm = FakeSageMakerClient(["Running"]) + install(sm) + + svc = ServiceClient(role_arn="arn:aws:iam::123456789012:role/TestRole") + svc.stop_session("arn:aws:sagemaker:us-west-2:123456789012:training-session/sess-1") + # The ARN's last segment is used as the name. + assert sm.stop_calls == [{"TrainingSessionName": "sess-1"}] + + +def test_delete_session_accepts_arn(patch_clients): + install, _ = patch_clients + sm = FakeSageMakerClient(["Stopped"]) + install(sm) + + svc = ServiceClient(role_arn="arn:aws:iam::123456789012:role/TestRole") + svc.delete_session("arn:aws:sagemaker:us-west-2:123456789012:training-session/sess-1") + assert sm.delete_calls == [{"TrainingSessionName": "sess-1"}] + + +# --- list ------------------------------------------------------------------- + + +def _summary(arn, status="Running"): + return {"TrainingSessionArn": arn, "TrainingSessionStatus": status} + + +def test_list_single_page(patch_clients): + install, _ = patch_clients + sm = FakeSageMakerClient(["Running"], list_pages=[([_summary("arn/a")], None)]) + install(sm) + + out = ServiceClient(role_arn="arn:aws:iam::123456789012:role/TestRole").list_sessions() + assert [s.training_session_arn for s in out] == ["arn/a"] + assert out[0].status is SessionStatus.RUNNING + # No filters set -> empty request body. + assert sm.list_calls == [{}] + + +def test_list_auto_paginates(patch_clients): + install, _ = patch_clients + sm = FakeSageMakerClient( + ["Running"], + list_pages=[ + ([_summary("arn/a")], "tok1"), + ([_summary("arn/b")], None), + ], + ) + install(sm) + + out = ServiceClient(role_arn="arn:aws:iam::123456789012:role/TestRole").list_sessions() + assert [s.training_session_arn for s in out] == ["arn/a", "arn/b"] + # Second call forwarded the NextToken from the first page. + assert sm.list_calls[1]["NextToken"] == "tok1" + + +def test_list_forwards_filters(patch_clients): + install, _ = patch_clients + sm = FakeSageMakerClient(["Running"], list_pages=[([], None)]) + install(sm) + + ServiceClient(role_arn="arn:aws:iam::123456789012:role/TestRole").list_sessions( + name_contains="exp", status_equals="Running" + ) + assert sm.list_calls[0]["NameContains"] == "exp" + assert sm.list_calls[0]["StatusEquals"] == "Running" + + +def test_list_status_equals_enum_sends_wire_value(patch_clients): + """A SessionStatus enum must serialize to its wire value, not str(enum). + + Regression: str() of a str-mixin Enum is "SessionStatus.EXPIRED" on + Python 3.11+, which the service rejects with a ValidationException. + """ + from sagemaker.train.training_session import SessionStatus + + install, _ = patch_clients + sm = FakeSageMakerClient(["Running"], list_pages=[([], None)]) + install(sm) + + ServiceClient(role_arn="arn:aws:iam::123456789012:role/TestRole").list_sessions( + status_equals=SessionStatus.EXPIRED + ) + assert sm.list_calls[0]["StatusEquals"] == "Expired" + + +def test_list_respects_max_results(patch_clients): + install, _ = patch_clients + # First page already returns 2; max_results=1 caps the result and stops paging. + sm = FakeSageMakerClient( + ["Running"], + list_pages=[([_summary("arn/a"), _summary("arn/b")], "tok1")], + ) + install(sm) + + out = ServiceClient(role_arn="arn:aws:iam::123456789012:role/TestRole").list_sessions( + max_results=1 + ) + assert [s.training_session_arn for s in out] == ["arn/a"] + # Only one page fetched despite a NextToken being present. + assert len(sm.list_calls) == 1 + assert sm.list_calls[0]["MaxResults"] == 1 + + +def test_list_holds_max_results_at_the_page_size(patch_clients): + """A cap above one page is served by paging, not by one oversized request. + + The operation rejects a page larger than its own limit, so forwarding the + caller's cap verbatim would fail the call instead of returning that many. + """ + install, _ = patch_clients + sm = FakeSageMakerClient( + ["Running"], + list_pages=[ + ([_summary(f"arn/{index}") for index in range(100)], "tok1"), + ([_summary("arn/100")], None), + ], + ) + install(sm) + + out = ServiceClient(role_arn="arn:aws:iam::123456789012:role/TestRole").list_sessions( + max_results=150 + ) + + assert len(out) == 101 + assert [call["MaxResults"] for call in sm.list_calls] == [100, 50] + + +def test_list_keeps_a_status_this_version_does_not_know(patch_clients): + """An unrecognised status reaches the caller rather than failing the list.""" + install, _ = patch_clients + sm = FakeSageMakerClient( + ["Running"], list_pages=[([_summary("arn/a", status="SomethingNew")], None)] + ) + install(sm) + + out = ServiceClient(role_arn="arn:aws:iam::123456789012:role/TestRole").list_sessions() + + assert out[0].status == "SomethingNew" + + +# --- base_model resolution ------------------------------------------------- +# create_lora_training_client accepts a friendly base_model name (resolved to a +# Hub ARN via DescribeHubContent) OR a full base_model_arn. Exactly one is +# required. The hub client is stubbed so no live AWS call is made. + + +class FakeHubClient: + """Stubs DescribeHubContent + ListHubContents for base-model resolution.""" + + #: Recipe keyword marking a model as fine-tunable (what the SDK filters on). + _TRAINABLE_KEYWORD = "@recipe:finetuning_mtrl_lora" + + def __init__( + self, + arn="arn:aws:sagemaker:us-west-2:aws:hub-content/SageMakerPublicHub/Model/hubcontent-vlm-model-b/9.9.9", + hub_content_names=None, + keywords_by_name=None, + ): + self._arn = arn + # Hub content names "available in this region" for list_hub_contents. + self._names = ( + hub_content_names + if hub_content_names is not None + else ["hubcontent-vlm-model-b", "hubcontent-text-model-a"] + ) + # Search keywords per content name. Default: every listed model carries + # the fine-tuning recipe keyword, i.e. is trainable. + self._keywords_by_name = ( + keywords_by_name + if keywords_by_name is not None + else {n: [self._TRAINABLE_KEYWORD] for n in self._names} + ) + self.describe_calls: list[dict] = [] + self.list_calls: list[dict] = [] + + def describe_hub_content(self, **kwargs): + name = kwargs.get("HubContentName") + self.describe_calls.append(kwargs) + # DescribeHubContent is the authority on existence: an unknown content + # name raises, mirroring the real API's ResourceNotFound. + if name not in self._names: + raise RuntimeError("ResourceNotFound: no such hub content") + # The real Hub doc is a JSON string carrying the published HuggingfaceId, + # one per model — echoed from this model's own keyword so a multi-model + # catalogue publishes several ids rather than collapsing onto one. + return { + "HubContentArn": self._arn, + "HubContentDocument": json.dumps({"HuggingfaceId": self._hf_id(name)}), + } + + def _hf_id(self, name): + for keyword in self._keywords_by_name.get(name, []): + if keyword.startswith("@huggingface-id:"): + return keyword[len("@huggingface-id:") :] + return "Provider-B/Model-B" + + def list_hub_contents(self, **kwargs): + self.list_calls.append(kwargs) + return { + "HubContentSummaries": [ + { + "HubContentName": n, + "HubContentSearchKeywords": self._keywords_by_name.get(n, []), + } + for n in self._names + ] + } + + +def test_get_hub_client_routes_through_sagemaker_core(patch_clients, monkeypatch): + """The Hub / model-package client comes from a region-pinned sagemaker-core + ``Session`` (not a hand-built boto3 client), so the public Hub and model + package calls go through sagemaker-core's client management.""" + install, _ = patch_clients + install(FakeSageMakerClient(["Running"])) + + captured = {} + sentinel_client = object() + + class FakeSession: + def __init__(self, *, boto_session, sagemaker_client=None): + captured["boto_region"] = boto_session.region_name + captured["injected_client"] = sagemaker_client + self.boto_session = boto_session + self.sagemaker_client = sentinel_client + + import sagemaker.core.helper.session_helper as session_mod + + monkeypatch.setattr(session_mod, "Session", FakeSession) + + svc = ServiceClient(region="eu-west-1", role_arn="arn:aws:iam::123456789012:role/TestRole") + client = svc._get_hub_client() + + # Came straight off the core Session, pinned to this ServiceClient's region. + assert client is sentinel_client + assert captured["boto_region"] == "eu-west-1" + # With no endpoint override, core builds the client itself. + assert captured["injected_client"] is None + # Cached: a second call does not build another Session. + assert svc._get_hub_client() is sentinel_client + + +def test_create_with_base_model_name_resolves_arn(patch_clients): + install, _ = patch_clients + sm = FakeSageMakerClient(["Running"]) + install(sm) + + svc = ServiceClient(region="us-west-2", role_arn="arn:aws:iam::123456789012:role/TestRole") + hub = FakeHubClient() + svc._hub_client = hub # stub the lazily-built boto3 client + + kwargs = _create_kwargs() + # base_model is a Hub content name (as returned by list_base_models()). + svc.create_lora_training_client(base_model="hubcontent-vlm-model-b", **kwargs) + + # Resolved via DescribeHubContent with that Hub content name verbatim. + assert hub.describe_calls[0]["HubContentName"] == "hubcontent-vlm-model-b" + assert hub.describe_calls[0]["HubName"] == "SageMakerPublicHub" + # The resolved ARN is what CreateTrainingSession received. + assert sm.create_calls[0]["ModelConfig"]["BaseModelArn"] == hub._arn + + +def test_create_resolves_hf_id_once_and_reuses_for_tokenizer(patch_clients, monkeypatch): + """Create resolves the model id from the same DescribeHubContent used for the + ARN, and get_tokenizer() reuses it instead of describing the model again.""" + install, _ = patch_clients + sm = FakeSageMakerClient(["Running"]) + install(sm) + + svc = ServiceClient(region="us-west-2", role_arn="arn:aws:iam::123456789012:role/TestRole") + hub = FakeHubClient() + svc._hub_client = hub + + seen = {} + import sagemaker.train.training_session._utils as utils + + def fake_load(base_model, hub_client, hf_model_id=None): + seen["hf_model_id"] = hf_model_id + return object() + + monkeypatch.setattr(utils, "load_tokenizer", fake_load) + + kwargs = _create_kwargs() + tc = svc.create_lora_training_client(base_model="hubcontent-vlm-model-b", **kwargs) + tc.get_tokenizer() + + # The create-time describe carried the model id straight into the tokenizer load. + assert seen["hf_model_id"] == "Provider-B/Model-B" + # Exactly one DescribeHubContent across create + get_tokenizer (no second call). + assert len(hub.describe_calls) == 1 + + +# The sampling path takes the same two guarantees as the training path above: it +# must not build the Hub client for an input that needs no Hub lookup, and it +# must carry the id it did resolve rather than make get_tokenizer look it up +# again. + + +def test_create_sampling_client_with_arn_does_not_build_hub_client(patch_clients): + """An ARN input needs no Hub lookup, so none of sagemaker-core is imported. + + Group and output path are supplied so their defaults do not reach for a + client either, leaving the returned client the only thing that could. + """ + install, _ = patch_clients + install(FakeSageMakerClient(["Running"])) + + svc = ServiceClient( + role_arn="arn:aws:iam::123456789012:role/TestRole", + model_package_group_arn="arn:mpg/x", + s3_output_path="s3://b/out/", + ) + built = {"n": 0} + + def fake_build_hub(): + built["n"] += 1 + return FakeHubClient() + + svc._get_hub_client = fake_build_hub + + sampler = svc.create_sampling_client(base_model=_BASE_MODEL_ARN) + + assert built["n"] == 0, "create must not build the Hub client for an ARN input" + assert sampler._hub_client is None + + +def test_create_sampling_client_resolves_hf_id_once_and_reuses_for_tokenizer( + patch_clients, monkeypatch +): + """The name lookup's model id reaches get_tokenizer, so it describes once.""" + install, _ = patch_clients + install(FakeSageMakerClient(["Running"])) + + svc = ServiceClient( + region="us-west-2", + role_arn="arn:aws:iam::123456789012:role/TestRole", + model_package_group_arn="arn:mpg/x", + s3_output_path="s3://b/out/", + ) + hub = FakeHubClient() + svc._hub_client = hub + + seen = {} + import sagemaker.train.training_session._utils as utils + + def fake_load(base_model, hub_client, hf_model_id=None): + seen["hf_model_id"] = hf_model_id + return object() + + monkeypatch.setattr(utils, "load_tokenizer", fake_load) + + sampler = svc.create_sampling_client(base_model="hubcontent-vlm-model-b") + sampler.get_tokenizer() + + assert seen["hf_model_id"] == "Provider-B/Model-B" + assert len(hub.describe_calls) == 1 + + +def test_sampling_client_lazily_builds_hub_client_when_needed(patch_clients, monkeypatch): + """With an ARN input and no hf_model_id the defer is a defer, not a drop.""" + install, _ = patch_clients + install(FakeSageMakerClient(["Running"])) + + svc = ServiceClient( + role_arn="arn:aws:iam::123456789012:role/TestRole", + model_package_group_arn="arn:mpg/x", + s3_output_path="s3://b/out/", + ) + hub = FakeHubClient() + built = {"n": 0} + + def fake_build_hub(): + built["n"] += 1 + return hub + + svc._get_hub_client = fake_build_hub + + seen = {} + import sagemaker.train.training_session._utils as utils + + def fake_load(base_model, hub_client, hf_model_id=None): + seen["hub_client"] = hub_client + return object() + + monkeypatch.setattr(utils, "load_tokenizer", fake_load) + + sampler = svc.create_sampling_client(base_model=_BASE_MODEL_ARN) + assert built["n"] == 0 # not built at create + + sampler.get_tokenizer() + assert built["n"] == 1 # built lazily on first tokenizer use + assert seen["hub_client"] is hub + + +def test_create_from_state_has_no_load_optimizer_param(): + """from_state is weights-only; the optimizer variant is the separate method.""" + import inspect + + sig = inspect.signature(ServiceClient.create_training_client_from_state) + assert "load_optimizer" not in sig.parameters + + +def test_from_state_takes_the_checkpoint_positionally(): + """The checkpoint reads as the subject of the call, so it stays positional. + + Every other argument is forwarded through ``**kwargs`` and has to be named, + which leaves the ARN as the only thing a caller can pass without a keyword. + """ + import inspect + + for method in ( + ServiceClient.create_training_client_from_state, + ServiceClient.create_training_client_from_state_async, + ServiceClient.create_training_client_from_state_with_optimizer, + ServiceClient.create_training_client_from_state_with_optimizer_async, + ): + param = inspect.signature(method).parameters["model_package_arn"] + assert param.kind is inspect.Parameter.POSITIONAL_OR_KEYWORD, method.__name__ + + +class FakeLoadFuture: + """Stand-in for the APIFuture the load methods return. + + Returning a bare ``None`` from a patched ``load_state`` would let the + factory drop the future without resolving it and still pass. This records + the call instead, so the test can assert the restore was waited on. + """ + + def __init__(self, error=None): + self._error = error + self.result_calls = 0 + + def result(self, timeout=None): + self.result_calls += 1 + if self._error is not None: + raise self._error + return None + + +def test_create_from_state_derives_base_model_from_checkpoint(patch_clients, monkeypatch): + """The documented resume recipe passes only model_package_arn — base_model + is read from the checkpoint's model package metadata, pinned to the version + the weights were trained on.""" + install, _ = patch_clients + sm = FakeSageMakerClient(["Running"]) + install(sm) + + svc = ServiceClient( + region="us-west-2", + role_arn="arn:aws:iam::123456789012:role/TestRole", + model_package_group_arn="arn:mpg/x", + s3_output_path="s3://b/out/", + ) + hub = FakeHubClient() + hub.describe_model_package = lambda **kw: { + "ModelPackageArn": "arn:mp/ckpt-3", + "InferenceSpecification": { + "Containers": [ + { + "IsCheckpoint": True, + "BaseModel": { + "HubContentName": "hubcontent-text-model-a", + "HubContentVersion": "3.15.0", + }, + } + ] + }, + } + svc._hub_client = hub + + loads = [] + futures = [] + + def _record(with_optimizer): + def _load(self, arn): + loads.append((arn, with_optimizer)) + fut = FakeLoadFuture() + futures.append(fut) + return fut + + return _load + + # Each factory must route to its matching load method by name; the + # weights-only / with-optimizer choice is not a flag on one call. + monkeypatch.setattr(TrainingClient, "load_state", _record(False)) + monkeypatch.setattr(TrainingClient, "load_state_with_optimizer", _record(True)) + + svc.create_training_client_from_state(model_package_arn="arn:mp/ckpt-3") + svc.create_training_client_from_state_with_optimizer(model_package_arn="arn:mp/ckpt-3") + + expected_base = ( + "arn:aws:sagemaker:us-west-2:aws:hub-content/" + "SageMakerPublicHub/Model/hubcontent-text-model-a/3.15.0" + ) + assert [c["ModelConfig"]["BaseModelArn"] for c in sm.create_calls] == [ + expected_base, + expected_base, + ] + assert loads == [("arn:mp/ckpt-3", False), ("arn:mp/ckpt-3", True)] + # Both factories must block on the restore, not return a client whose + # checkpoint load is still queued. + assert [f.result_calls for f in futures] == [1, 1] + + +def test_create_from_state_explicit_base_model_skips_describe(patch_clients, monkeypatch): + """Passing base_model explicitly must not trigger a DescribeModelPackage.""" + install, _ = patch_clients + sm = FakeSageMakerClient(["Running"]) + install(sm) + + svc = ServiceClient( + region="us-west-2", + role_arn="arn:aws:iam::123456789012:role/TestRole", + model_package_group_arn="arn:mpg/x", + s3_output_path="s3://b/out/", + ) + hub = FakeHubClient() + + def _boom(**kw): + raise AssertionError("describe_model_package must not be called") + + hub.describe_model_package = _boom + svc._hub_client = hub + + monkeypatch.setattr(TrainingClient, "load_state", lambda self, arn: FakeLoadFuture()) + + svc.create_training_client_from_state( + model_package_arn="arn:mp/ckpt-3", base_model=_BASE_MODEL_ARN + ) + assert sm.create_calls[0]["ModelConfig"]["BaseModelArn"] == "arn:hub-content/Model/base" + + +def test_create_from_state_stops_session_when_restore_fails(patch_clients, monkeypatch): + """A failed restore must not leave the just-created session unowned. + + The factory raises instead of returning a client, so nothing holds the + session's handle and no automatic stop is armed for it. + """ + install, _ = patch_clients + sm = FakeSageMakerClient(["Running"]) + install(sm) + + svc = ServiceClient( + region="us-west-2", + role_arn="arn:aws:iam::123456789012:role/TestRole", + model_package_group_arn="arn:mpg/x", + s3_output_path="s3://b/out/", + ) + + boom = RuntimeError("checkpoint unreadable") + monkeypatch.setattr(TrainingClient, "load_state", lambda self, arn: FakeLoadFuture(error=boom)) + + with pytest.raises(RuntimeError, match="checkpoint unreadable"): + svc.create_training_client_from_state( + model_package_arn="arn:mp/ckpt-3", base_model=_BASE_MODEL_ARN + ) + # The original failure propagates, and the session is torn down. + assert len(sm.stop_calls) == 1 + + +def test_create_with_explicit_arn_skips_hub_lookup(patch_clients): + install, _ = patch_clients + sm = FakeSageMakerClient(["Running"]) + install(sm) + + svc = ServiceClient(role_arn="arn:aws:iam::123456789012:role/TestRole") + hub = FakeHubClient() + svc._hub_client = hub + + svc.create_lora_training_client(_BASE_MODEL_ARN, **_create_kwargs()) + assert sm.create_calls[0]["ModelConfig"]["BaseModelArn"] == "arn:hub-content/Model/base" + assert hub.describe_calls == [] # no resolution needed + + +def test_explicit_hf_model_id_skips_tokenizer_describe(patch_clients, monkeypatch): + """An explicit hf_model_id is used for the tokenizer with no Hub describe. + + With an ARN input the model id is normally resolved lazily at get_tokenizer() + time via DescribeHubContent. Passing hf_model_id supplies it up front, so a + caller that already knows it (or lacks Hub read access) never describes. + """ + install, _ = patch_clients + install(FakeSageMakerClient(["Running"])) + + svc = ServiceClient(role_arn="arn:aws:iam::123456789012:role/TestRole") + hub = FakeHubClient() + svc._hub_client = hub + + seen = {} + import sagemaker.train.training_session._utils as utils + + def fake_load(base_model, hub_client, hf_model_id=None): + seen["hf_model_id"] = hf_model_id + return object() + + monkeypatch.setattr(utils, "load_tokenizer", fake_load) + + tc = svc.create_lora_training_client(_BASE_MODEL_ARN, **_create_kwargs()) + tc.get_tokenizer() + + # ARN path: hf_model_id is None at create time (resolved lazily by tokenizer). + assert seen["hf_model_id"] is None + + +def test_create_with_arn_does_not_build_hub_client(patch_clients): + """Creating a session with base_model_arn must NOT build the Hub client. + + The Hub client is built via sagemaker-core, which is not always installed + (e.g. the vendored/Fargate integ env). Building it eagerly at create time + would import sagemaker-core even though the ARN path never needs the Hub, so + the build must stay deferred until a tokenizer lookup actually needs it. + Regression guard for the 1.13.2 eager-import break. + """ + install, _ = patch_clients + install(FakeSageMakerClient(["Running"])) + + svc = ServiceClient(role_arn="arn:aws:iam::123456789012:role/TestRole") + calls = {"n": 0} + original = svc._get_hub_client + + def counting_get_hub_client(): + calls["n"] += 1 + return original() + + svc._get_hub_client = counting_get_hub_client + + tc = svc.create_lora_training_client(_BASE_MODEL_ARN, **_create_kwargs()) + assert calls["n"] == 0, "create must not build the Hub client for an ARN input" + # And the client received no eagerly-built Hub client. + assert tc._hub_client is None + + +def test_tokenizer_lazily_builds_hub_client_when_needed(patch_clients, monkeypatch): + """With an ARN input and no hf_model_id, get_tokenizer() lazily builds the + Hub client (deferred from create) to resolve the model id — proving the defer + is a defer, not a drop.""" + install, _ = patch_clients + install(FakeSageMakerClient(["Running"])) + + svc = ServiceClient(role_arn="arn:aws:iam::123456789012:role/TestRole") + hub = FakeHubClient() + built = {"n": 0} + + def fake_build_hub(): + built["n"] += 1 + return hub + + svc._get_hub_client = fake_build_hub + + seen = {} + import sagemaker.train.training_session._utils as utils + + def fake_load(base_model, hub_client, hf_model_id=None): + seen["hub_client"] = hub_client + return object() + + monkeypatch.setattr(utils, "load_tokenizer", fake_load) + + # ARN input, no hf_model_id -> tokenizer resolution needs the Hub client. + tc = svc.create_lora_training_client(_BASE_MODEL_ARN, **_create_kwargs()) + assert built["n"] == 0 # not built at create + tc.get_tokenizer() + assert built["n"] == 1 # built lazily on first tokenizer use + assert seen["hub_client"] is hub + + +def test_tokenizer_with_hf_id_never_builds_hub_client(patch_clients, monkeypatch): + """get_tokenizer() with a known hf_model_id must NOT build the Hub client. + + The integ path passes base_model_arn + hf_model_id, so no Hub lookup is + ever needed — building the client here would still import sagemaker-core in + the vendored env. Regression guard: the lazy-build must stay gated on a + missing hf_model_id, not fire unconditionally at tokenizer time. + """ + install, _ = patch_clients + install(FakeSageMakerClient(["Running"])) + + svc = ServiceClient(role_arn="arn:aws:iam::123456789012:role/TestRole") + built = {"n": 0} + + def fake_build_hub(): + built["n"] += 1 + return FakeHubClient() + + svc._get_hub_client = fake_build_hub + + import sagemaker.train.training_session._utils as utils + + monkeypatch.setattr(utils, "load_tokenizer", lambda *a, **k: object()) + + tc = svc.create_lora_training_client(_BASE_MODEL_ARN, **_create_kwargs()) + # ARN path: hub client is NOT built at create time (no name resolution needed). + assert built["n"] == 0, "ARN path must not build the Hub client at create time" + # But it IS built when the tokenizer resolves the HF model id lazily. + tc.get_tokenizer() + assert built["n"] == 1, "Tokenizer needs Hub client to resolve model id from ARN" + + +def test_create_rejects_unresolvable_base_model(patch_clients): + install, _ = patch_clients + install(FakeSageMakerClient(["Running"])) + svc = ServiceClient(role_arn="arn:aws:iam::123456789012:role/TestRole") + svc._hub_client = FakeHubClient() + + kwargs = _create_kwargs() + # Not a real Hub content name -> DescribeHubContent fails and we surface a + # clear, actionable error. + with pytest.raises(ValueError, match="Could not resolve base model"): + svc.create_lora_training_client("not-a-real-model", **kwargs) + + +def test_list_base_models_filters_by_recipe_keyword(patch_clients): + # Only models carrying a fine-tuning recipe keyword are trainable; a model + # present in the hub without one must be excluded. Each trainable model is + # reported by the model id its Hub entry advertises. + install, _ = patch_clients + install(FakeSageMakerClient(["Running"])) + svc = ServiceClient(region="us-west-2", role_arn="arn:aws:iam::123456789012:role/TestRole") + svc._hub_client = FakeHubClient( + hub_content_names=[ + "hubcontent-vlm-model-b", # trainable + "hubcontent-vlm-model-c", # trainable + "some-eval-only-model", # present but not fine-tunable + ], + keywords_by_name={ + "hubcontent-vlm-model-b": [ + "@recipe:finetuning_sft_lora", + "@recipe:finetuning_mtrl_lora", + "@huggingface-id:provider-b/model-b", + ], + "hubcontent-vlm-model-c": [ + "@recipe:finetuning_mtrl_lora", + "@huggingface-id:provider-c/model-c", + ], + "some-eval-only-model": ["@recipe:evaluation_mtrlevaluation"], + }, + ) + + models = svc.list_base_models() + # Both fine-tunable models, by model id, sorted; eval-only excluded. + assert models == ["provider-b/model-b", "provider-c/model-c"] + assert svc._hub_client.list_calls # queried the Hub in-region + + +def test_list_base_models_excludes_unsupported_names(patch_clients): + # The recipe keyword alone does not make a model usable through this client. + # A name the client does not support end to end is excluded even when the + # Hub tags it as fine-tunable. + install, _ = patch_clients + install(FakeSageMakerClient(["Running"])) + svc = ServiceClient(region="us-west-2", role_arn="arn:aws:iam::123456789012:role/TestRole") + svc._hub_client = FakeHubClient( + hub_content_names=[ + "hubcontent-vlm-model-b", + # Matches _EXCLUDED_MODEL_NAMES, so excluded despite the recipe. + "hubcontent-nova-model-x", + ], + keywords_by_name={ + "hubcontent-vlm-model-b": [ + "@recipe:finetuning_mtrl_lora", + "@huggingface-id:provider-b/model-b", + ], + "hubcontent-nova-model-x": [ + "@recipe:finetuning_mtrl_lora", + "@huggingface-id:provider-x/model-x", + ], + }, + ) + + assert svc.list_base_models() == ["provider-b/model-b"] + + +# --- checkpoint discovery --------------------------------------------------- +# Checkpoints are model package versions in the session's group. The hub client +# (svc._get_hub_client()) is the boto3 sagemaker client these helpers call, so +# the fake below stubs list_model_packages + describe_model_package. + + +class FakeCheckpointClient: + """Stubs list_model_packages (newest-first, paginated) + describe_model_package.""" + + def __init__(self, summaries, resumable_arns=()): + # summaries: already in the desired (newest-first) order. + self._summaries = summaries + self._resumable = set(resumable_arns) + self.list_calls: list[dict] = [] + self.describe_calls: list[dict] = [] + + def list_model_packages(self, **kwargs): + self.list_calls.append(kwargs) + items = self._summaries + mr = kwargs.get("MaxResults") + if mr is not None: + items = items[:mr] + return {"ModelPackageSummaryList": items} + + def describe_model_package(self, **kwargs): + arn = kwargs["ModelPackageName"] + self.describe_calls.append(kwargs) + is_ckpt = arn in self._resumable + return { + "InferenceSpecification": {"Containers": [{"IsCheckpoint": True} if is_ckpt else {}]} + } + + +def _ckpt_summary(arn, version, desc="d"): + return { + "ModelPackageArn": arn, + "ModelPackageVersion": version, + "ModelPackageDescription": desc, + "CreationTime": "2026-06-20T00:00:00Z", + } + + +def _svc_with_checkpoints(patch_clients, fake): + install, _ = patch_clients + install(FakeSageMakerClient(["Running"])) + svc = ServiceClient(region="us-west-2", role_arn="arn:aws:iam::123456789012:role/TestRole") + svc._hub_client = fake + return svc + + +def test_list_checkpoints_maps_summaries_no_describe(patch_clients): + fake = FakeCheckpointClient( + [_ckpt_summary("g/3", 3, "step 3"), _ckpt_summary("g/2", 2, "step 2")] + ) + svc = _svc_with_checkpoints(patch_clients, fake) + + out = svc.list_model_packages(model_package_group_arn="g") + assert [c.version for c in out] == [3, 2] + assert out[0].model_package_arn == "g/3" + assert out[0].description == "step 3" + # list view does not resolve kind (no describe calls). + assert out[0].kind is None + assert fake.describe_calls == [] + + +def test_list_checkpoints_respects_max_results(patch_clients): + fake = FakeCheckpointClient( + [_ckpt_summary("g/3", 3), _ckpt_summary("g/2", 2), _ckpt_summary("g/1", 1)] + ) + svc = _svc_with_checkpoints(patch_clients, fake) + + out = svc.list_model_packages(model_package_group_arn="g", max_results=2) + assert [c.version for c in out] == [3, 2] + assert fake.list_calls[0]["MaxResults"] == 2 + + +def test_get_latest_resumable_skips_model_checkpoints(patch_clients): + # Newest is a model checkpoint; the resumable one is older — walk past it. + fake = FakeCheckpointClient( + [_ckpt_summary("g/9", 9), _ckpt_summary("g/8", 8), _ckpt_summary("g/7", 7)], + resumable_arns=["g/8"], + ) + svc = _svc_with_checkpoints(patch_clients, fake) + + latest = svc.get_latest_model_package(ModelPackageKind.CHECKPOINT, model_package_group_arn="g") + assert latest.version == 8 + assert latest.kind == ModelPackageKind.CHECKPOINT + # Described g/9 (model, skipped) then g/8 (resumable, match) — stopped there. + assert [c["ModelPackageName"] for c in fake.describe_calls] == ["g/9", "g/8"] + + +def test_get_latest_model_skips_checkpoints(patch_clients): + fake = FakeCheckpointClient( + [_ckpt_summary("g/9", 9), _ckpt_summary("g/8", 8)], resumable_arns=["g/9"] + ) + svc = _svc_with_checkpoints(patch_clients, fake) + + latest = svc.get_latest_model_package(ModelPackageKind.MODEL, model_package_group_arn="g") + assert latest.version == 8 + assert latest.kind == ModelPackageKind.MODEL + + +def test_get_latest_resumable_returns_none_when_no_match_in_scan(patch_clients): + # All model checkpoints, no resumable -> None after scanning. + fake = FakeCheckpointClient( + [_ckpt_summary("g/2", 2), _ckpt_summary("g/1", 1)], resumable_arns=[] + ) + svc = _svc_with_checkpoints(patch_clients, fake) + + assert ( + svc.get_latest_model_package(ModelPackageKind.CHECKPOINT, model_package_group_arn="g") + is None + ) + + +def test_get_latest_resumable_honors_scan_limit(patch_clients): + fake = FakeCheckpointClient( + [_ckpt_summary("g/3", 3), _ckpt_summary("g/2", 2), _ckpt_summary("g/1", 1)], + resumable_arns=["g/1"], + ) + svc = _svc_with_checkpoints(patch_clients, fake) + + # scan_limit=2 only describes the 2 newest; the resumable g/1 is beyond it. + out = svc.get_latest_model_package( + ModelPackageKind.CHECKPOINT, model_package_group_arn="g", scan_limit=2 + ) + assert out is None + assert fake.list_calls[0]["MaxResults"] == 2 + assert len(fake.describe_calls) == 2 + + +def _record_core_clients(monkeypatch): + """Fake out core's ``Session`` and ``boto3.Session``, recording both. + + Returns ``(recorded, override_client)``. ``recorded["client_calls"]`` holds + one entry per pre-built client, and ``recorded["injected"]`` one per core + ``Session``, so a test can tell which accessor passed an endpoint and which + did not. + """ + import boto3 + + recorded = {"client_calls": [], "injected": [], "regions": []} + override_client = object() + default_client = object() + + class _FakeBotoSession: + def __init__(self, region_name=None): + recorded["regions"].append(region_name) + + def client(self, service, endpoint_url=None, config=None): + recorded["client_calls"].append( + { + "service": service, + "endpoint_url": endpoint_url, + "user_agent_extra": getattr(config, "user_agent_extra", None), + } + ) + return override_client + + class _FakeSession: + def __init__(self, *, boto_session, sagemaker_client=None): + self.boto_session = boto_session + recorded["injected"].append(sagemaker_client) + # Core builds its own when handed none, which is the un-overridden + # path; stand in a sentinel so the test can tell them apart. + self.sagemaker_client = ( + sagemaker_client if sagemaker_client is not None else default_client + ) + + # Stub the session's own clients out first: the fake boto3.Session below + # exists to record the lookup clients' arguments, and building the real ones + # through it would only exercise the fake. + monkeypatch.setattr( + "sagemaker.train.training_session._clients._control_plane_client", + lambda **kwargs: object(), + ) + monkeypatch.setattr( + "sagemaker.train.training_session._clients._runtime_client", + lambda **kwargs: object(), + ) + monkeypatch.setattr(boto3, "Session", _FakeBotoSession) + import sagemaker.core.helper.session_helper as session_mod + + monkeypatch.setattr(session_mod, "Session", _FakeSession) + return recorded, override_client, default_client + + +def test_model_registry_client_honors_endpoint_override(monkeypatch): + """An endpoint override reaches the model-registry client. + + The override needs its own client, built off the same session the rest of + this ``ServiceClient`` uses, so it signs with the same credentials and + carries core's user-agent config like any other. + """ + recorded, override_client, _ = _record_core_clients(monkeypatch) + monkeypatch.setenv("SAGEMAKER_TRAINING_SESSION_ENDPOINT", "https://custom.example.test") + + svc = ServiceClient(region="us-west-2", role_arn="arn:aws:iam::123456789012:role/TestRole") + client = svc._get_model_registry_client() + + assert [c["service"] for c in recorded["client_calls"]] == ["sagemaker"] + assert recorded["client_calls"][0]["endpoint_url"] == "https://custom.example.test" + assert recorded["regions"] == ["us-west-2"] + # Built with core's user-agent suffix, not a bare client. + assert recorded["client_calls"][0]["user_agent_extra"] + assert client is override_client + + +def test_hub_client_ignores_endpoint_override(monkeypatch): + """The Hub client stays on the Region's endpoint even under an override. + + The Public Hub is published per Region, so resolving a base model against + anything else would pick a version the Region does not serve. Only the + model-registry client follows the override. + """ + recorded, override_client, default_client = _record_core_clients(monkeypatch) + monkeypatch.setenv("SAGEMAKER_TRAINING_SESSION_ENDPOINT", "https://custom.example.test") + + svc = ServiceClient(region="us-west-2", role_arn="arn:aws:iam::123456789012:role/TestRole") + hub = svc._get_hub_client() + + # No pre-built client at all, so core resolved the Region's own endpoint. + assert recorded["client_calls"] == [] + assert recorded["injected"] == [None] + assert hub is default_client + assert hub is not override_client + # The two accessors are genuinely different clients under an override. + assert svc._get_model_registry_client() is override_client + + +def test_both_lookup_clients_are_one_client_without_an_override(monkeypatch): + """With no override the two registries share a client, as they always did.""" + recorded, _, default_client = _record_core_clients(monkeypatch) + monkeypatch.delenv("SAGEMAKER_TRAINING_SESSION_ENDPOINT", raising=False) + + svc = ServiceClient(region="us-west-2", role_arn="arn:aws:iam::123456789012:role/TestRole") + + assert svc._get_hub_client() is default_client + assert svc._get_model_registry_client() is svc._get_hub_client() + # One core Session, one client: no extra client built for public callers. + assert recorded["injected"] == [None] + assert recorded["client_calls"] == [] + + +# --------------------------------------------------------------------------- +# Failure reason on a session that never reaches Running +# --------------------------------------------------------------------------- + + +def test_terminal_state_error_includes_error_details(patch_clients): + """A failed session reports why, not just the status.""" + install, _ = patch_clients + sm = FakeSageMakerClient( + ["Failed"], describe_detail={"ErrorDetails": "base model not accessible"} + ) + install(sm) + + svc = ServiceClient(region="us-west-2", role_arn="arn:aws:iam::123456789012:role/TestRole") + with pytest.raises(RuntimeError, match="base model not accessible"): + svc.create_lora_training_client(_BASE_MODEL_ARN, **_create_kwargs()) + + +def test_terminal_state_error_falls_back_to_error_message(patch_clients): + """ErrorMessage is the fallback; used only when ErrorDetails is absent.""" + install, _ = patch_clients + sm = FakeSageMakerClient(["Failed"], describe_detail={"ErrorMessage": "quota exceeded"}) + install(sm) + + svc = ServiceClient(region="us-west-2", role_arn="arn:aws:iam::123456789012:role/TestRole") + with pytest.raises(RuntimeError, match="quota exceeded"): + svc.create_lora_training_client(_BASE_MODEL_ARN, **_create_kwargs()) + + +def test_terminal_state_error_without_reason_names_the_status(patch_clients): + """No reason from the service leaves the status to speak for itself.""" + install, _ = patch_clients + sm = FakeSageMakerClient(["Failed"]) + install(sm) + + svc = ServiceClient(region="us-west-2", role_arn="arn:aws:iam::123456789012:role/TestRole") + with pytest.raises(RuntimeError, match="terminal state: Failed$"): + svc.create_lora_training_client(_BASE_MODEL_ARN, **_create_kwargs()) + + +# --------------------------------------------------------------------------- +# Flat and grouped field forms on session creation +# --------------------------------------------------------------------------- + + +def test_create_sends_model_and_output_config(patch_clients): + """Model identity and output location go in ModelConfig and OutputConfig. + + Nothing else carries either value: the flat ``BaseModelArn`` / + ``OutputS3Path`` pair is not part of the request shape, so sending it is + refused before the request leaves. + """ + install, _ = patch_clients + sm = FakeSageMakerClient(["Running"]) + install(sm) + + svc = ServiceClient(region="us-west-2", role_arn="arn:aws:iam::123456789012:role/TestRole") + svc.create_lora_training_client(_BASE_MODEL_ARN, **_create_kwargs()) + + cc = sm.create_calls[0] + assert cc["ModelConfig"] == {"BaseModelArn": _BASE_MODEL_ARN} + assert cc["OutputConfig"]["S3Location"]["Path"] == "s3://b/out/" + assert "BaseModelArn" not in cc + assert "OutputS3Path" not in cc + + +def test_create_sampling_only_session_sends_the_same_shape(patch_clients): + """The implicit sampling-session path builds the same request, not a variant.""" + install, _ = patch_clients + sm = FakeSageMakerClient(["Running"]) + install(sm) + + svc = ServiceClient( + region="us-west-2", + role_arn="arn:aws:iam::123456789012:role/TestRole", + model_package_group_arn="arn:mpg/x", + s3_output_path="s3://b/out/", + ) + svc.create_sampling_client(base_model=_BASE_MODEL_ARN) + + cc = sm.create_calls[0] + assert cc["ModelConfig"] == {"BaseModelArn": _BASE_MODEL_ARN} + assert cc["OutputConfig"]["S3Location"]["Path"] == "s3://b/out/" + assert "BaseModelArn" not in cc + assert "OutputS3Path" not in cc + + +# --------------------------------------------------------------------------- +# KMS / VPC support on session creation +# --------------------------------------------------------------------------- + + +def test_create_with_kms_sends_output_config(patch_clients): + """Per-call kms_key_arn rides on OutputConfig.S3Location.""" + install, _ = patch_clients + sm = FakeSageMakerClient(["Running"]) + install(sm) + + svc = ServiceClient(region="us-west-2", role_arn="arn:aws:iam::123456789012:role/TestRole") + svc.create_lora_training_client( + _BASE_MODEL_ARN, **_create_kwargs(kms_key_arn="arn:aws:kms:us-west-2:123456789012:key/k1") + ) + + cc = sm.create_calls[0] + assert cc["OutputConfig"] == { + "S3Location": { + "Path": "s3://b/out/", + "KmsKeyArn": "arn:aws:kms:us-west-2:123456789012:key/k1", + } + } + + +def test_create_with_vpc_dict_sends_vpc_config(patch_clients): + """A plain dict vpc_config is coerced and sent in the API shape.""" + install, _ = patch_clients + sm = FakeSageMakerClient(["Running"]) + install(sm) + + svc = ServiceClient(region="us-west-2", role_arn="arn:aws:iam::123456789012:role/TestRole") + svc.create_lora_training_client( + _BASE_MODEL_ARN, + **_create_kwargs( + vpc_config={"subnets": ["subnet-1", "subnet-2"], "security_group_ids": ["sg-1"]} + ), + ) + + assert sm.create_calls[0]["VpcConfig"] == { + "Subnets": ["subnet-1", "subnet-2"], + "SecurityGroupIds": ["sg-1"], + } + + +def test_create_without_kms_vpc_omits_only_those_members(patch_clients): + """OutputConfig always goes out; KmsKeyArn and VpcConfig only when set.""" + install, _ = patch_clients + sm = FakeSageMakerClient(["Running"]) + install(sm) + + svc = ServiceClient(region="us-west-2", role_arn="arn:aws:iam::123456789012:role/TestRole") + svc.create_lora_training_client(_BASE_MODEL_ARN, **_create_kwargs()) + + cc = sm.create_calls[0] + assert cc["OutputConfig"] == {"S3Location": {"Path": "s3://b/out/"}} + assert "VpcConfig" not in cc + + +def test_create_kms_vpc_defaults_from_service_client_and_override(patch_clients): + """ServiceClient-level defaults apply; per-call values win.""" + from sagemaker.train.training_session.types import VpcConfig + + install, _ = patch_clients + sm = FakeSageMakerClient(["Running", "Running"]) + install(sm) + + svc = ServiceClient( + region="us-west-2", + role_arn="arn:aws:iam::123456789012:role/TestRole", + kms_key_arn="arn:aws:kms:us-west-2:123456789012:key/default", + vpc_config=VpcConfig(subnets=["subnet-d"], security_group_ids=["sg-d"]), + ) + + svc.create_lora_training_client(_BASE_MODEL_ARN, **_create_kwargs()) + cc = sm.create_calls[0] + assert cc["OutputConfig"]["S3Location"]["KmsKeyArn"].endswith("key/default") + assert cc["VpcConfig"]["Subnets"] == ["subnet-d"] + + svc.create_lora_training_client( + _BASE_MODEL_ARN, + **_create_kwargs( + training_session_name="sess-2", + kms_key_arn="arn:aws:kms:us-west-2:123456789012:key/override", + vpc_config={"subnets": ["subnet-o"], "security_group_ids": ["sg-o"]}, + ), + ) + cc = sm.create_calls[1] + assert cc["OutputConfig"]["S3Location"]["KmsKeyArn"].endswith("key/override") + assert cc["VpcConfig"]["Subnets"] == ["subnet-o"] + + +def test_sampling_client_inherits_kms_and_vpc(patch_clients): + """The implicit sampling session uses the ServiceClient-level values.""" + install, _ = patch_clients + sm = FakeSageMakerClient(["Running"]) + install(sm) + + svc = ServiceClient( + role_arn="arn:aws:iam::123456789012:role/TestRole", + model_package_group_arn="arn:mpg/x", + s3_output_path="s3://b/out/", + kms_key_arn="arn:aws:kms:us-west-2:123456789012:key/k1", + vpc_config={"subnets": ["subnet-1"], "security_group_ids": ["sg-1"]}, + ) + svc._hub_client = FakeHubClient() + svc.create_sampling_client(base_model="hubcontent-vlm-model-b") + + cc = sm.create_calls[0] + assert cc["OutputConfig"]["S3Location"]["KmsKeyArn"].endswith("key/k1") + assert cc["VpcConfig"] == {"Subnets": ["subnet-1"], "SecurityGroupIds": ["sg-1"]} + + +def test_default_s3_output_path_uses_core_default_bucket(patch_clients, monkeypatch): + """The default S3 output path goes through sagemaker-core + ``Session.default_bucket()``, which creates the default bucket on the + caller behalf when missing, same behavior as the broader SageMaker + PySDK (and what the s3:CreateBucket grant in the caller managed policy + exists for).""" + install, _ = patch_clients + install(FakeSageMakerClient(["Running"])) + + calls = {"default_bucket": 0} + + class FakeSession: + def __init__(self, *, boto_session): + self.boto_session = boto_session + self._region = boto_session.region_name + + def default_bucket(self): + calls["default_bucket"] += 1 + return f"sagemaker-{self._region}-123456789012" + + import sagemaker.core.helper.session_helper as session_mod + + monkeypatch.setattr(session_mod, "Session", FakeSession) + + svc = ServiceClient(region="us-west-2", role_arn="arn:aws:iam::123456789012:role/TestRole") + path = svc._default_s3_output_path("my-session") + + assert path == "s3://sagemaker-us-west-2-123456789012/my-session/output" + assert calls["default_bucket"] == 1 + + +# --- delete_model_package --------------------------------------------------- +# Two clients are involved: the hub client (DescribeModelPackage / +# DeleteModelPackage) and an s3 client built only when artifacts are deleted. + + +class FakeModelPackageClient: + """Stubs DescribeModelPackage / DeleteModelPackage for one package.""" + + def __init__( + self, + s3_uri="s3://bucket/group/req-1/model/", + s3_data_type="S3Prefix", + managed_storage_type=None, + describe_error=None, + delete_error=None, + ): + self._s3_uri = s3_uri + self._s3_data_type = s3_data_type + self._managed_storage_type = managed_storage_type + self._describe_error = describe_error + self._delete_error = delete_error + self.describe_calls: list[dict] = [] + self.delete_calls: list[dict] = [] + + def describe_model_package(self, **kwargs): + self.describe_calls.append(kwargs) + if self._describe_error is not None: + raise self._describe_error + source = {} + if self._s3_uri is not None: + source = {"S3Uri": self._s3_uri, "S3DataType": self._s3_data_type} + resp = { + "InferenceSpecification": { + "Containers": [{"ModelDataSource": {"S3DataSource": source}}] + } + } + if self._managed_storage_type is not None: + resp["ManagedStorageType"] = self._managed_storage_type + return resp + + def delete_model_package(self, **kwargs): + self.delete_calls.append(kwargs) + if self._delete_error is not None: + raise self._delete_error + return {} + + +class FakeS3Client: + """Stubs ListObjectsV2 (paged) + DeleteObjects (per-key Errors, not raises).""" + + def __init__(self, pages=None, errors_for=()): + # pages: list of (keys, next_token). One entry means a single page. + self._pages = list(pages if pages is not None else [(["group/req-1/model/w.bin"], None)]) + self._errors_for = set(errors_for) + self.list_calls: list[dict] = [] + self.delete_calls: list[dict] = [] + + def list_objects_v2(self, **kwargs): + self.list_calls.append(kwargs) + keys, next_token = self._pages.pop(0) + resp = {"Contents": [{"Key": k} for k in keys], "IsTruncated": bool(next_token)} + if next_token: + resp["NextContinuationToken"] = next_token + return resp + + def delete_objects(self, **kwargs): + self.delete_calls.append(kwargs) + sent = [o["Key"] for o in kwargs["Delete"]["Objects"]] + errors = [{"Key": k, "Code": "AccessDenied"} for k in sent if k in self._errors_for] + return {"Errors": errors} if errors else {} + + +_MP_ARN = "arn:aws:sagemaker:us-west-2:123456789012:model-package/my-group/7" + + +def _svc_for_mp_delete(patch_clients, monkeypatch, hub, s3=None): + install, _ = patch_clients + install(FakeSageMakerClient(["Running"])) + svc = ServiceClient(region="us-west-2", role_arn="arn:aws:iam::123456789012:role/TestRole") + svc._hub_client = hub + if s3 is not None: + svc._s3_client = s3 + else: + # Any S3 call at all is a bug in the retain path, so fail loudly. + def _boom(*a, **kw): + raise AssertionError("no S3 client should be built when retaining artifacts") + + monkeypatch.setattr(svc, "_get_s3_client", _boom) + return svc + + +def _client_error(code, message="", op="DescribeModelPackage"): + return ClientError({"Error": {"Code": code, "Message": message}}, op) + + +def test_delete_model_package_retains_artifacts_by_default(patch_clients, monkeypatch): + """Default deletes the package only, and reports where the objects still are; + the package was the only record of that location.""" + hub = FakeModelPackageClient() + svc = _svc_for_mp_delete(patch_clients, monkeypatch, hub) + + result = svc.delete_model_package(_MP_ARN) + + assert result.model_package_deleted is True + assert result.s3_uri == "s3://bucket/group/req-1/model/" + assert result.objects_deleted == 0 + assert result.artifacts_retained is True + assert "delete_artifacts" in result.retained_reason + assert hub.delete_calls == [{"ModelPackageName": _MP_ARN}] + + +def test_delete_model_package_deletes_artifacts_before_package(patch_clients, monkeypatch): + """Order matters: the package carries the location, so it goes last. If it + went first, anything left in S3 would be unreachable.""" + hub = FakeModelPackageClient() + s3 = FakeS3Client(pages=[(["group/req-1/model/a", "group/req-1/model/b"], None)]) + svc = _svc_for_mp_delete(patch_clients, monkeypatch, hub, s3) + + order: list[str] = [] + s3_delete = s3.delete_objects + hub_delete = hub.delete_model_package + s3.delete_objects = lambda **kw: (order.append("s3"), s3_delete(**kw))[1] + hub.delete_model_package = lambda **kw: (order.append("mp"), hub_delete(**kw))[1] + + result = svc.delete_model_package(_MP_ARN, delete_artifacts=True) + + assert order == ["s3", "mp"] + assert result.objects_deleted == 2 + assert result.artifacts_retained is False + assert result.retained_reason is None + + +def test_delete_model_package_paginates_and_batches(patch_clients, monkeypatch): + """Each listed page is deleted before the next is requested, so no full key + list is held and the continuation token still resumes correctly.""" + hub = FakeModelPackageClient() + s3 = FakeS3Client(pages=[(["k/1", "k/2"], "tok-1"), (["k/3"], None)]) + svc = _svc_for_mp_delete(patch_clients, monkeypatch, hub, s3) + + result = svc.delete_model_package(_MP_ARN, delete_artifacts=True) + + assert result.objects_deleted == 3 + assert len(s3.delete_calls) == 2 + assert s3.list_calls[1]["ContinuationToken"] == "tok-1" + + +def test_delete_model_package_prefix_gets_trailing_slash(patch_clients, monkeypatch): + """A location without a trailing slash is matched on the segment boundary, + so "run-1" cannot sweep up "run-10".""" + hub = FakeModelPackageClient(s3_uri="s3://bucket/run-1") + s3 = FakeS3Client(pages=[(["run-1/w.bin"], None)]) + svc = _svc_for_mp_delete(patch_clients, monkeypatch, hub, s3) + + svc.delete_model_package(_MP_ARN, delete_artifacts=True) + + assert s3.list_calls[0] == {"Bucket": "bucket", "Prefix": "run-1/"} + + +def test_delete_model_package_s3_object_deletes_exactly_one_key(patch_clients, monkeypatch): + """S3DataType comes from the package, so a single-object location is deleted + as that one key rather than listed as a prefix.""" + hub = FakeModelPackageClient( + s3_uri="s3://bucket/group/req-1/model.tar.gz", s3_data_type="S3Object" + ) + s3 = FakeS3Client() + svc = _svc_for_mp_delete(patch_clients, monkeypatch, hub, s3) + + result = svc.delete_model_package(_MP_ARN, delete_artifacts=True) + + assert s3.list_calls == [] + assert s3.delete_calls[0]["Delete"]["Objects"] == [{"Key": "group/req-1/model.tar.gz"}] + assert result.objects_deleted == 1 + + +def test_delete_model_package_partial_s3_failure_keeps_package(patch_clients, monkeypatch): + """DeleteObjects reports per-key failures in the body rather than raising. + The package is kept so the survivors stay reachable and a retry works.""" + from sagemaker.train.training_session.exceptions import ArtifactDeletionError + + hub = FakeModelPackageClient() + s3 = FakeS3Client(pages=[(["k/1", "k/2"], None)], errors_for=["k/2"]) + svc = _svc_for_mp_delete(patch_clients, monkeypatch, hub, s3) + + with pytest.raises(ArtifactDeletionError) as excinfo: + svc.delete_model_package(_MP_ARN, delete_artifacts=True) + + assert excinfo.value.failed_keys == ["k/2"] + assert excinfo.value.s3_uri == "s3://bucket/group/req-1/model/" + assert hub.delete_calls == [], "the package must survive a partial artifact delete" + + +def test_delete_model_package_refuses_bucket_root(patch_clients, monkeypatch): + """A location naming only a bucket would empty the whole bucket.""" + hub = FakeModelPackageClient(s3_uri="s3://bucket") + s3 = FakeS3Client() + svc = _svc_for_mp_delete(patch_clients, monkeypatch, hub, s3) + + with pytest.raises(ValueError, match="bucket root"): + svc.delete_model_package(_MP_ARN, delete_artifacts=True) + + assert s3.delete_calls == [] + assert hub.delete_calls == [] + + +def test_delete_model_package_skips_service_managed_artifacts(patch_clients, monkeypatch): + """Artifacts SageMaker stores on the caller behalf have no S3 location the + caller can delete, and deleting the package is what releases them. Asking + for delete_artifacts=True is satisfied, not rejected.""" + hub = FakeModelPackageClient(s3_uri=None, managed_storage_type="Restricted") + svc = _svc_for_mp_delete(patch_clients, monkeypatch, hub) + + result = svc.delete_model_package(_MP_ARN, delete_artifacts=True) + + assert result.artifacts_retained is True + assert result.s3_uri is None + assert "SageMaker" in result.retained_reason + assert hub.delete_calls == [{"ModelPackageName": _MP_ARN}] + + +def test_delete_model_package_skips_when_no_s3_location(patch_clients, monkeypatch): + hub = FakeModelPackageClient(s3_uri=None) + svc = _svc_for_mp_delete(patch_clients, monkeypatch, hub) + + result = svc.delete_model_package(_MP_ARN, delete_artifacts=True) + + assert result.artifacts_retained is True + assert "no S3 location" in result.retained_reason + assert hub.delete_calls == [{"ModelPackageName": _MP_ARN}] + + +@pytest.mark.parametrize( + "error", + [ + _client_error("ValidationException", "Model package my-group/7 does not exist."), + _client_error("ResourceNotFound", "Model package not found"), + ], +) +def test_delete_model_package_is_idempotent_when_already_gone(patch_clients, monkeypatch, error): + """Deleting twice reports nothing left to delete rather than raising.""" + hub = FakeModelPackageClient(describe_error=error) + svc = _svc_for_mp_delete(patch_clients, monkeypatch, hub) + + result = svc.delete_model_package(_MP_ARN, delete_artifacts=True) + + assert result.model_package_deleted is False + assert result.s3_uri is None + assert hub.delete_calls == [] + + +def test_delete_model_package_reraises_unrelated_validation_error(patch_clients, monkeypatch): + """A ValidationException that is not "no such package" is a real failure.""" + hub = FakeModelPackageClient( + describe_error=_client_error("ValidationException", "1 validation error detected") + ) + svc = _svc_for_mp_delete(patch_clients, monkeypatch, hub) + + with pytest.raises(ClientError): + svc.delete_model_package(_MP_ARN) + + +def test_delete_model_package_absent_on_delete_after_describe(patch_clients, monkeypatch): + """Another caller deleted it between the describe and the delete.""" + hub = FakeModelPackageClient( + delete_error=_client_error("ValidationException", "does not exist", op="DeleteModelPackage") + ) + svc = _svc_for_mp_delete(patch_clients, monkeypatch, hub) + + result = svc.delete_model_package(_MP_ARN) + + assert result.model_package_deleted is False + assert result.s3_uri == "s3://bucket/group/req-1/model/" + + +def test_delete_model_package_reraises_unrelated_delete_error(patch_clients, monkeypatch): + hub = FakeModelPackageClient( + delete_error=_client_error( + "ValidationException", "Model package is in use", op="DeleteModelPackage" + ) + ) + svc = _svc_for_mp_delete(patch_clients, monkeypatch, hub) + + with pytest.raises(ClientError): + svc.delete_model_package(_MP_ARN) + + +def test_delete_model_package_empty_location_is_not_retained(patch_clients, monkeypatch): + """Nothing to delete is a completed delete, not a retained one.""" + hub = FakeModelPackageClient() + s3 = FakeS3Client(pages=[([], None)]) + svc = _svc_for_mp_delete(patch_clients, monkeypatch, hub, s3) + + result = svc.delete_model_package(_MP_ARN, delete_artifacts=True) + + assert result.objects_deleted == 0 + assert result.artifacts_retained is False + assert s3.delete_calls == [] + + +def test_delete_model_package_no_data_type_is_treated_as_a_prefix(patch_clients, monkeypatch): + """A package carrying a location but no S3DataType is swept as a prefix.""" + hub = FakeModelPackageClient(s3_uri="s3://bucket/group/req-1/model/", s3_data_type=None) + s3 = FakeS3Client(pages=[(["group/req-1/model/w.bin"], None)]) + svc = _svc_for_mp_delete(patch_clients, monkeypatch, hub, s3) + + result = svc.delete_model_package(_MP_ARN, delete_artifacts=True) + + assert s3.list_calls[0] == {"Bucket": "bucket", "Prefix": "group/req-1/model/"} + assert result.objects_deleted == 1 + + +@pytest.mark.parametrize( + "failing_call, error_code", + [ + # The two permissions this needs, each missing in turn. + ("list_objects_v2", "AccessDenied"), + ("delete_objects", "AccessDenied"), + ], +) +def test_delete_model_package_s3_permission_error_keeps_package( + patch_clients, monkeypatch, failing_call, error_code +): + """An S3 call that raises rather than reporting per-key failures propagates + unchanged, and the package survives so the objects stay reachable.""" + hub = FakeModelPackageClient() + s3 = FakeS3Client() + + def _raise(**kwargs): + raise _client_error(error_code, "Access Denied", op=failing_call) + + setattr(s3, failing_call, _raise) + svc = _svc_for_mp_delete(patch_clients, monkeypatch, hub, s3) + + with pytest.raises(ClientError) as excinfo: + svc.delete_model_package(_MP_ARN, delete_artifacts=True) + + assert excinfo.value.response["Error"]["Code"] == error_code + assert hub.delete_calls == [], "the package must survive a failed artifact delete" + + +def test_delete_model_package_logs_where_retained_artifacts_are(patch_clients, monkeypatch, caplog): + """The package was the only record of the location, so the retain path says + it once at INFO, so a caller who discards the result can still find them.""" + hub = FakeModelPackageClient() + svc = _svc_for_mp_delete(patch_clients, monkeypatch, hub) + + with caplog.at_level("INFO", logger=sc_mod.logger.name): + svc.delete_model_package(_MP_ARN) + + assert "s3://bucket/group/req-1/model/" in caplog.text + + +def test_delete_model_package_does_not_log_a_location_it_has_none_for( + patch_clients, monkeypatch, caplog +): + """Nothing was left behind for the caller to find, so there is nothing to + report.""" + hub = FakeModelPackageClient(s3_uri=None, managed_storage_type="Restricted") + svc = _svc_for_mp_delete(patch_clients, monkeypatch, hub) + + with caplog.at_level("INFO", logger=sc_mod.logger.name): + svc.delete_model_package(_MP_ARN, delete_artifacts=True) + + assert "artifacts remain at" not in caplog.text + + +@pytest.mark.parametrize( + "error", + [ + pytest.param(_client_error("ResourceNotFoundException", ""), id="not-found-exception"), + pytest.param( + _client_error("ValidationException", "Model package cannot be found"), + id="cannot-be-found", + ), + # The phrase is matched case-insensitively: the service capitalises these + # messages inconsistently, and a casing change must not turn the + # idempotent path back into a raise. + pytest.param( + _client_error("ValidationException", "Model Package Does Not Exist."), + id="capitalised-message", + ), + ], +) +def test_delete_model_package_absent_code_variants(patch_clients, monkeypatch, error): + hub = FakeModelPackageClient(describe_error=error) + svc = _svc_for_mp_delete(patch_clients, monkeypatch, hub) + + assert svc.delete_model_package(_MP_ARN).model_package_deleted is False + + +@pytest.mark.parametrize( + "error", + [ + pytest.param(_client_error("AccessDeniedException", "not authorized"), id="access-denied"), + pytest.param(_client_error("ThrottlingException", "Rate exceeded"), id="throttled"), + # No Message at all: reading one must not turn a real error into "gone". + pytest.param( + ClientError({"Error": {"Code": "ValidationException"}}, "DescribeModelPackage"), + id="validation-without-message", + ), + # No Error block at all: reading a code out of it must not raise, and an + # unreadable failure is a failure, not an absent package. + pytest.param(ClientError({}, "DescribeModelPackage"), id="no-error-block"), + ], +) +def test_delete_model_package_reraises_real_describe_errors(patch_clients, monkeypatch, error): + """Only "no such package" is swallowed. A permission or throttling failure + must not read as an idempotent no-op: that would report success while the + package is still there.""" + hub = FakeModelPackageClient(describe_error=error) + svc = _svc_for_mp_delete(patch_clients, monkeypatch, hub) + + with pytest.raises(ClientError): + svc.delete_model_package(_MP_ARN) + + +def test_delete_model_package_is_idempotent_across_repeat_calls(patch_clients, monkeypatch): + """The documented sweep loop can run twice: the second pass finds the + package gone and reports it rather than raising.""" + hub = FakeModelPackageClient() + s3 = FakeS3Client(pages=[(["group/req-1/model/w.bin"], None), ([], None)]) + svc = _svc_for_mp_delete(patch_clients, monkeypatch, hub, s3) + + first = svc.delete_model_package(_MP_ARN, delete_artifacts=True) + hub._describe_error = _client_error("ValidationException", "does not exist") + second = svc.delete_model_package(_MP_ARN, delete_artifacts=True) + + assert (first.model_package_deleted, first.objects_deleted) == (True, 1) + assert (second.model_package_deleted, second.objects_deleted) == (False, 0) + assert len(hub.delete_calls) == 1 + + +def test_delete_artifacts_is_keyword_only(): + """``delete_model_package(arn, True)`` would not say what the True deletes, + and it is the destructive half of the call.""" + import inspect + + params = inspect.signature(ServiceClient.delete_model_package).parameters + assert params["delete_artifacts"].kind is inspect.Parameter.KEYWORD_ONLY + assert params["delete_artifacts"].default is False + + +def test_delete_model_package_public_names_are_exported(): + """Both are named in the method's own docs, so both must be importable from + the package root; a caller cannot annotate or catch what is not exported.""" + import sagemaker.train.training_session as pkg + + from sagemaker.train.training_session.exceptions import ArtifactDeletionError + from sagemaker.train.training_session.types import ModelPackageDeletion + + assert pkg.ModelPackageDeletion is ModelPackageDeletion + assert pkg.ArtifactDeletionError is ArtifactDeletionError + assert {"ModelPackageDeletion", "ArtifactDeletionError"} <= set(pkg.__all__) + + +def test_artifact_deletion_error_truncates_its_message_not_its_keys(patch_clients, monkeypatch): + """A whole prefix can fail on one missing permission, so the message shows a + sample and points at the full list rather than printing thousands of keys.""" + from sagemaker.train.training_session.exceptions import ArtifactDeletionError + + keys = [f"group/req-1/model/shard-{i}" for i in range(7)] + hub = FakeModelPackageClient() + s3 = FakeS3Client(pages=[(keys, None)], errors_for=keys) + svc = _svc_for_mp_delete(patch_clients, monkeypatch, hub, s3) + + with pytest.raises(ArtifactDeletionError) as excinfo: + svc.delete_model_package(_MP_ARN, delete_artifacts=True) + + err = excinfo.value + assert err.failed_keys == keys + assert err.model_package_arn == _MP_ARN + assert err.s3_uri == "s3://bucket/group/req-1/model/" + message = str(err) + assert "(7 total, see .failed_keys)" in message + assert "shard-4" in message and "shard-5" not in message + # The recovery step, not just the failure. + assert "s3:DeleteObject" in message and "NOT deleted" in message + + +def test_artifact_deletion_error_lists_a_short_failure_in_full(patch_clients, monkeypatch): + from sagemaker.train.training_session.exceptions import ArtifactDeletionError + + hub = FakeModelPackageClient() + s3 = FakeS3Client(pages=[(["k/1", "k/2"], None)], errors_for=["k/1", "k/2"]) + svc = _svc_for_mp_delete(patch_clients, monkeypatch, hub, s3) + + with pytest.raises(ArtifactDeletionError) as excinfo: + svc.delete_model_package(_MP_ARN, delete_artifacts=True) + + assert "total, see .failed_keys" not in str(excinfo.value) + assert "k/1, k/2" in str(excinfo.value) + + +@pytest.mark.parametrize( + "count, truncated", + [ + pytest.param(5, False, id="exactly-the-sample-size"), + pytest.param(6, True, id="one-over"), + ], +) +def test_artifact_deletion_error_truncates_only_past_the_sample_size(count, truncated): + """The boundary: a failure list the same length as the sample is printed in + full, so the message never says "5 total, see .failed_keys" while already + showing all five.""" + from sagemaker.train.training_session.exceptions import ArtifactDeletionError + + keys = [f"k/{i}" for i in range(count)] + message = str(ArtifactDeletionError(_MP_ARN, "s3://bucket/k/", keys)) + + assert ("see .failed_keys" in message) is truncated + # Either way the sample itself is present. + assert "k/4" in message + + +def test_artifact_deletion_error_is_a_training_session_error(): + """Catchable with the rest of this client's failures.""" + from sagemaker.train.training_session.exceptions import ( + ArtifactDeletionError, + TrainingSessionError, + ) + + assert issubclass(ArtifactDeletionError, TrainingSessionError) + + +def test_get_s3_client_is_region_pinned_and_cached(patch_clients, monkeypatch): + """The s3 client is built lazily, pinned to this ServiceClient region, and + signed with the caller credentials like every other client here.""" + install, _ = patch_clients + install(FakeSageMakerClient(["Running"])) + + captured = {} + sentinel = object() + + class FakeBotoSession: + def __init__(self, region_name=None): + captured["region"] = region_name + self.region_name = region_name + + def client(self, name): + captured["service"] = name + return sentinel + + class FakeSession: + def __init__(self, *, boto_session): + self.boto_session = boto_session + self.sagemaker_client = object() + + import boto3 + import sagemaker.core.helper.session_helper as session_mod + + # Patched before the client exists: its session is built in the constructor, + # so the s3 client comes off that one session rather than a fresh one. + monkeypatch.setattr(boto3, "Session", FakeBotoSession) + monkeypatch.setattr(session_mod, "Session", FakeSession) + + svc = ServiceClient(region="us-west-2", role_arn="arn:aws:iam::123456789012:role/TestRole") + assert svc._s3_client is None + + assert svc._get_s3_client() is sentinel + assert captured == {"region": "us-west-2", "service": "s3"} + assert svc._get_s3_client() is sentinel diff --git a/sagemaker-train/tests/unit/train/training_session/test_session_lifecycle.py b/sagemaker-train/tests/unit/train/training_session/test_session_lifecycle.py new file mode 100644 index 0000000000..b347363a48 --- /dev/null +++ b/sagemaker-train/tests/unit/train/training_session/test_session_lifecycle.py @@ -0,0 +1,586 @@ +# 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/ +"""Unit tests for session lifecycle management (_SessionHandle). + +Nothing here touches the network. Each test drops its client inside a helper +frame and forces a collection, so the assertion is about what the finalizer did +after the last reference went away. + +The negative tests matter more than the positive ones: a spurious stop kills a +session a training loop is still driving, which is worse than the leak this +feature exists to prevent. +""" + +from __future__ import annotations + +import gc +import logging +import warnings +import weakref + +import pytest + +from sagemaker.train.training_session._session_handle import _SessionHandle +from sagemaker.train.training_session.sampling_client import SamplingClient +from sagemaker.train.training_session.training_client import TrainingClient + + +class FakeServiceClient: + """Records control-plane lifecycle calls; never performs one.""" + + def __init__( + self, + *, + stop_error: Exception | None = None, + delete_error: Exception | None = None, + was_active: bool = True, + ): + self.stopped: list[str] = [] + self.deleted: list[str] = [] + self._stop_error = stop_error + self._delete_error = delete_error + # What _stop_session_if_active reports: False means the session was + # already stopping or terminal, so the stop was a no-op. + self._was_active = was_active + + def _get_hub_client(self): + return None + + def stop_session(self, training_session_name: str) -> None: + self._stop_session_if_active(training_session_name) + + def _stop_session_if_active(self, training_session_name: str) -> bool: + self.stopped.append(training_session_name) + if self._stop_error is not None: + raise self._stop_error + return self._was_active + + def delete_session(self, training_session_name: str, **kwargs) -> None: + self.deleted.append(training_session_name) + if self._delete_error is not None: + raise self._delete_error + + +class FakeTransport: + """Minimal stand-in; no operation in this module reaches the wire.""" + + def execute_sampling_operation(self, **kwargs): # pragma: no cover + raise AssertionError("no test here submits an operation") + + def execute_training_operation(self, **kwargs): # pragma: no cover + raise AssertionError("no test here submits an operation") + + +def _collect(): + """Force the finalizers of anything that just became unreachable.""" + gc.collect() + + +# --- the handle in isolation ------------------------------------------------ + + +def test_handle_stops_session_when_last_reference_drops(): + svc = FakeServiceClient() + + def _scope(): + _SessionHandle("sess-1", svc) + + _scope() + _collect() + assert svc.stopped == ["sess-1"] + + +def test_auto_stop_has_no_opt_out(): + # Deliberate: a session's lifetime is not the caller's problem, so there is + # no flag or method that leaves one Running. Durable state is a checkpoint. + handle = _SessionHandle("sess-1", FakeServiceClient()) + assert handle.auto_stop_armed is True + assert not hasattr(handle, "detach") + + +def test_handle_without_service_client_cannot_arm(): + # Nothing to issue the stop through, so the handle must not pretend it can. + handle = _SessionHandle("sess-1", None) + assert handle.auto_stop_armed is False + + +def test_disarm_prevents_the_stop(): + svc = FakeServiceClient() + + def _scope(): + handle = _SessionHandle("sess-1", svc) + handle.disarm() + + _scope() + _collect() + assert svc.stopped == [] + + +def test_disarm_is_idempotent(): + svc = FakeServiceClient() + handle = _SessionHandle("sess-1", svc) + handle.disarm() + handle.disarm() + assert handle.auto_stop_armed is False + + +def test_the_finalizer_does_not_keep_the_handle_alive(): + # The whole design rests on this: a callback closing over the handle would + # make it immortal and the cleanup would never run at all. + svc = FakeServiceClient() + handle = _SessionHandle("sess-1", svc) + import weakref + + ref = weakref.ref(handle) + del handle + _collect() + assert ref() is None + assert svc.stopped == ["sess-1"] + + +def test_auto_stop_drains_registered_submitters_before_stopping(): + order: list[str] = [] + + class RecordingSubmitter: + def close(self, timeout: float = 5.0) -> None: + order.append("drain") + + class RecordingService(FakeServiceClient): + def _stop_session_if_active(self, training_session_name: str) -> bool: + order.append("stop") + return super()._stop_session_if_active(training_session_name) + + svc = RecordingService() + + def _scope(): + handle = _SessionHandle("sess-1", svc) + handle.register_submitter(RecordingSubmitter()) + + _scope() + _collect() + # Queued operations must reach the service before the stop request. + assert order == ["drain", "stop"] + + +def test_stop_failure_during_cleanup_is_swallowed(caplog): + svc = FakeServiceClient(stop_error=RuntimeError("throttled")) + + def _scope(): + _SessionHandle("sess-1", svc) + + with caplog.at_level("WARNING"): + _scope() + _collect() + # A finalizer must never propagate: the sweeper is the remaining backstop. + assert "sess-1" in caplog.text + + +def _not_found_error(code="ResourceNotFound"): + """The ClientError shape the control plane returns for a missing session.""" + from botocore.exceptions import ClientError + + return ClientError( + { + "Error": {"Code": code, "Message": "Training session not found: sess-1"}, + "ResponseMetadata": {"HTTPStatusCode": 400}, + }, + "StopTrainingSession", + ) + + +@pytest.mark.parametrize("code", ["ResourceNotFound", "ResourceNotFoundException"]) +def test_auto_stop_treats_a_deleted_session_as_done(caplog, code): + """The caller deleted it out of band; the goal state is already reached. + + Warning here would tell someone who cleaned up correctly that the session + still needs sweeping, which is the opposite of the truth. + """ + svc = FakeServiceClient(stop_error=_not_found_error(code)) + + def _scope(): + _SessionHandle("sess-1", svc) + + with caplog.at_level("DEBUG"): + _scope() + _collect() + assert "no longer exists" in caplog.text + # Nothing happened, so nothing above DEBUG is emitted. + assert not [r for r in caplog.records if r.levelno > logging.DEBUG] + + +def test_auto_stop_still_warns_on_a_real_failure(caplog): + """Only "it is gone" is absorbed; anything else stays visible.""" + svc = FakeServiceClient(stop_error=_not_found_error(code="AccessDeniedException")) + + def _scope(): + _SessionHandle("sess-1", svc) + + with caplog.at_level("INFO"): + _scope() + _collect() + assert [r for r in caplog.records if r.levelname == "WARNING"] + + +def test_auto_stop_is_silent_when_the_session_was_already_stopping(caplog): + """Nothing to report, so report nothing, not even the ResourceWarning. + + Telling a caller who already stopped the session that they forgot to would + be both noise and untrue. + """ + svc = FakeServiceClient(was_active=False) + + with warnings.catch_warnings(record=True) as caught: + warnings.simplefilter("always") + with caplog.at_level("DEBUG"): + + def _scope(): + _SessionHandle("sess-quiet", svc) + + _scope() + _collect() + + assert svc.stopped == ["sess-quiet"], "it must still attempt the stop" + assert not [w for w in caught if w.category is ResourceWarning] + assert not [r for r in caplog.records if r.levelno > logging.DEBUG] + + +def test_auto_stop_reports_when_it_really_stopped_something(caplog): + """The other half of the contract: a real stop is still announced.""" + svc = FakeServiceClient(was_active=True) + + with warnings.catch_warnings(record=True) as caught: + warnings.simplefilter("always") + with caplog.at_level("INFO"): + + def _scope(): + _SessionHandle("sess-loud", svc) + + _scope() + _collect() + + assert [w for w in caught if w.category is ResourceWarning] + assert "Auto-stopped unreferenced training session" in caplog.text + + +def test_auto_stop_emits_resource_warning(): + svc = FakeServiceClient() + with warnings.catch_warnings(record=True) as caught: + warnings.simplefilter("always") + + def _scope(): + _SessionHandle("sess-noisy", svc) + + _scope() + _collect() + messages = [str(w.message) for w in caught if w.category is ResourceWarning] + assert any("sess-noisy" in m for m in messages) + + +# --- the invariant: stop, never delete -------------------------------------- + + +def test_auto_stop_never_deletes_the_session(): + svc = FakeServiceClient() + + def _scope(): + tc = TrainingClient("sess-1", FakeTransport(), service_client=svc) + return tc.create_sampling_client() + + sp = _scope() + del sp + _collect() + assert svc.stopped == ["sess-1"] + assert svc.deleted == [], "the automatic path must never issue a delete" + + +def test_explicit_delete_disarms_so_no_stop_follows_it(): + # Otherwise a later GC would stop a session name the caller already removed. + svc = FakeServiceClient() + + def _scope(): + tc = TrainingClient("sess-1", FakeTransport(), service_client=svc) + tc.delete() + + _scope() + _collect() + assert svc.deleted == ["sess-1"] + assert svc.stopped == [] + + +def test_failed_delete_leaves_the_automatic_stop_armed(): + """A refused delete must not cost the session its last client-side release. + + ``delete(stop_if_active=False)``, the default, is refused for a session + that is still running. Detaching is one-way, so disarming before the call + left a running session with nothing trying to release it: it held its GPUs + until the inactivity timeout. + """ + svc = FakeServiceClient(delete_error=RuntimeError("session is not terminal")) + + def _scope(): + tc = TrainingClient("sess-1", FakeTransport(), service_client=svc) + try: + tc.delete() + except RuntimeError as exc: + # Cleared deliberately: the traceback holds this frame, and with it + # the client, which would keep the handle reachable past _collect. + assert "not terminal" in str(exc) + exc.__traceback__ = None + else: + raise AssertionError("delete should have propagated the refusal") + + _scope() + _collect() + assert svc.deleted == ["sess-1"] + assert svc.stopped == ["sess-1"], "the automatic stop is the only release left" + + +def test_lifecycle_without_a_service_client_names_the_lifecycle_call(): + """The error has to be about stopping, not about listing checkpoints.""" + tc = TrainingClient("sess-1", FakeTransport()) + for call in (tc.stop, tc.delete): + with pytest.raises(ValueError, match="cannot stop or delete the session"): + call() + + +# --- SamplingClient over an implicit (owned) session ------------------------ + + +def test_owned_sampling_session_is_stopped_when_client_drops(): + svc = FakeServiceClient() + + def _scope(): + SamplingClient( + "sess-implicit", + FakeTransport(), + service_client=svc, + ) + + _scope() + _collect() + assert svc.stopped == ["sess-implicit"] + + +def test_explicit_stop_then_drop_stops_exactly_once(): + svc = FakeServiceClient() + + def _scope(): + sp = SamplingClient( + "sess-implicit", + FakeTransport(), + service_client=svc, + ) + sp.stop() + + _scope() + _collect() + assert svc.stopped == ["sess-implicit"] + + +def test_context_manager_exit_then_drop_stops_exactly_once(): + svc = FakeServiceClient() + + def _scope(): + with SamplingClient( + "sess-implicit", + FakeTransport(), + service_client=svc, + ): + pass + + _scope() + _collect() + assert svc.stopped == ["sess-implicit"] + + +# --- shared session: TrainingClient + derived SamplingClient ---------------- +# +# These four are the regression tests that keep the automatic stop from becoming an +# outage. Dropping either client must leave the session Running while the other +# is still usable, in EITHER drop order. + + +def test_dropping_training_client_keeps_session_for_live_sampling_client(): + svc = FakeServiceClient() + + def _scope(): + tc = TrainingClient("sess-shared", FakeTransport(), service_client=svc) + return tc.create_sampling_client() + + sp = _scope() + _collect() + assert svc.stopped == [], "a SamplingClient is still bound to this session" + # ...and the surviving client stops it when it in turn goes away. + del sp + _collect() + assert svc.stopped == ["sess-shared"] + + +def test_dropping_sampling_client_keeps_session_for_live_training_client(): + svc = FakeServiceClient() + tc = TrainingClient("sess-shared", FakeTransport(), service_client=svc) + + # Takes the client as an argument rather than closing over it, so the only + # reference dropped below is the one this test holds. + def _scope(client): + client.create_sampling_client() + + _scope(tc) + _collect() + assert svc.stopped == [], "the TrainingClient is still driving this session" + del tc + _collect() + assert svc.stopped == ["sess-shared"] + + +def test_session_survives_until_the_last_of_several_sampling_clients_drops(): + svc = FakeServiceClient() + + def _scope(): + tc = TrainingClient("sess-shared", FakeTransport(), service_client=svc) + return tc.create_sampling_client(), tc.create_sampling_client() + + first, second = _scope() + _collect() + assert svc.stopped == [] + del first + _collect() + assert svc.stopped == [], "one sibling is still bound" + del second + _collect() + assert svc.stopped == ["sess-shared"] + + +def test_derived_sampling_client_shares_the_training_clients_handle(): + tc = TrainingClient("sess-shared", FakeTransport(), service_client=FakeServiceClient()) + sp = tc.create_sampling_client() + assert sp._handle is tc._handle + + +def test_noop_stop_on_derived_client_does_not_disarm_the_shared_session(): + # SamplingClient.stop() is documented as a no-op for a shared session. If it + # also disarmed the handle, that no-op would silently turn into a leak. + svc = FakeServiceClient() + + def _scope(): + tc = TrainingClient("sess-shared", FakeTransport(), service_client=svc) + sp = tc.create_sampling_client() + sp.stop() + assert svc.stopped == [] + assert tc._handle.auto_stop_armed is True + + _scope() + _collect() + assert svc.stopped == ["sess-shared"] + + +def test_clients_expose_no_lifecycle_opt_out(): + # The public surface for a session is stop() / delete() and the `with` block. + # Nothing lets a caller leave one Running; see test_auto_stop_has_no_opt_out. + tc = TrainingClient("sess-shared", FakeTransport(), service_client=FakeServiceClient()) + sp = tc.create_sampling_client() + assert not hasattr(tc, "detach") + assert not hasattr(sp, "detach") + + +# --- outstanding operations ------------------------------------------------- + + +def test_future_keeps_the_session_alive_after_its_client_is_dropped(): + # Dropping the client while a future is unresolved must not stop the session + # the operation is running against. + svc = FakeServiceClient() + + def _scope(): + tc = TrainingClient("sess-1", FakeTransport(), service_client=svc) + from sagemaker.train.training_session.futures import PendingAPIFuture + + return PendingAPIFuture(submit_future=None, operation_type="training", keepalive=tc._handle) + + future = _scope() + _collect() + assert svc.stopped == [], "an operation against this session is outstanding" + del future + _collect() + assert svc.stopped == ["sess-1"] + + +# --- ServiceClient factory wiring ------------------------------------------ + + +def test_service_client_factory_arms_auto_stop(monkeypatch): + sp = _sampling_client_from_service_client(monkeypatch) + assert sp._handle.auto_stop_armed is True + + +def test_implicit_sampling_client_owns_its_session(monkeypatch): + # Ownership is derived from "no handle was passed in", not from a flag. + sp = _sampling_client_from_service_client(monkeypatch) + assert sp._owns_session is True + + +def test_first_client_from_a_service_client_is_collectable(monkeypatch): + """Regression: a first-time import must not pin the client a factory returns. + + CPython retains the frame chain of a first-time import. When + ``_get_hub_client`` performed its lazy imports on first use — inside + ``create_sampling_client`` — that pinned the factory's frame and therefore + the client it returned, so the client was never collected and its session + stayed Running until interpreter exit. Only the first call in a process + performs the import, so only the first client was affected. + ``ServiceClient.__init__`` now warms those imports instead. + + Asserted structurally rather than by simulating a cold import, which cannot + be undone within a process: the guarantee is that the factory does not + trigger the import. + """ + import sys + + from sagemaker.train.training_session import service_client as sc_module + + sp = _sampling_client_from_service_client(monkeypatch) + ref = weakref.ref(sp) + del sp + gc.collect() + assert ref() is None, "the client a factory returns must be collectable" + # And the modules whose first-time import caused the pin are already loaded + # by the time any factory runs. + assert hasattr(sc_module, "_warm_hub_imports") + assert "boto3" in sys.modules + + +def _sampling_client_from_service_client(monkeypatch): + """Build a real ServiceClient with the control plane and Hub stubbed out.""" + from sagemaker.train.training_session import service_client as sc_module + + class StubControlPlane: + def __init__(self, **kwargs): + pass + + def create_training_session(self, **kwargs): + return {"TrainingSessionArn": "arn:aws:sagemaker:::training-session/s"} + + monkeypatch.setattr( + "sagemaker.train.training_session._clients._control_plane_client", + lambda **kwargs: StubControlPlane(), + ) + monkeypatch.setattr( + "sagemaker.train.training_session._clients._runtime_client", + lambda **kwargs: object(), + ) + + svc = sc_module.ServiceClient( + region="us-west-2", + role_arn="arn:aws:iam::123456789012:role/R", + model_package_group_arn="arn:aws:sagemaker:us-west-2:123456789012:model-package-group/g", + s3_output_path="s3://bucket/out", + ) + monkeypatch.setattr(svc, "_wait_until_running", lambda *a, **k: None) + monkeypatch.setattr(svc, "_get_hub_client", lambda: None) + + return svc.create_sampling_client(base_model="arn:aws:sagemaker:::hub-content/m") diff --git a/sagemaker-train/tests/unit/train/training_session/test_session_lifecycle_grid.py b/sagemaker-train/tests/unit/train/training_session/test_session_lifecycle_grid.py new file mode 100644 index 0000000000..e650f5dab9 --- /dev/null +++ b/sagemaker-train/tests/unit/train/training_session/test_session_lifecycle_grid.py @@ -0,0 +1,98 @@ +# Copyright Amazon.com, Inc. or its affiliates. All Rights Reserved. +# +# Licensed under the Apache License, Version 2.0 (the "License"). You +# may not use this file except in compliance with the License. A copy of +# the License is located at +# +# http://aws.amazon.com/apache2.0/ +"""Locks in the measured cleanup-trigger grid (see session_lifecycle_grid.py). + +Each case runs ``session_lifecycle_grid.py`` in a child process, because the two +interesting triggers, interpreter exit and signal death, are not observable +from inside the process they happen in. + +The measured results below were identical on CPython 3.9, 3.11, 3.12 and 3.13, +so they are asserted rather than merely recorded. Two of them are the *absence* +of cleanup, and they are the ones worth re-reading if this file ever fails: + +* ``sigterm``: Python installs no SIGTERM handler, so the default disposition + terminates the process outright and no finalizer runs. Anything supervised + (containers, batch schedulers) must stop its own sessions, or install a + handler; the SDK deliberately does not hijack a process-wide signal. +* ``os_exit``: skips finalizers and atexit by design. No in-process mechanism + can cover it; such a session is reclaimed once it passes its inactivity + timeout. +""" + +from __future__ import annotations + +import os +import subprocess +import sys + +import pytest + +_GRID = os.path.join(os.path.dirname(os.path.abspath(__file__)), "session_lifecycle_grid.py") + +# trigger -> (cleanup expected, expected exit code) +CASES = { + "del_then_collect": (True, 0), + "falls_out_of_scope": (True, 0), + "rebound": (True, 0), + "reference_cycle": (True, 0), + "process_exit": (True, 0), + "sys_exit": (True, 3), + "uncaught_exception": (True, 1), + "sigint": (True, -2), + "explicit_disarm": (False, 0), + "sigterm": (False, -15), + "os_exit": (False, 0), +} + + +#: Seconds a grid child gets. Each one starts an interpreter and exits, so this +#: is generous; bounded because the thing under test is interpreter teardown, and +#: a child that wedges there would otherwise hang the run rather than fail it. +_CHILD_TIMEOUT = 60 + + +def _run(trigger: str, tmp_path) -> tuple[bool, int, str]: + out = tmp_path / f"{trigger}.txt" + try: + proc = subprocess.run( + [sys.executable, _GRID, trigger, str(out)], + capture_output=True, + text=True, + timeout=_CHILD_TIMEOUT, + ) + except subprocess.TimeoutExpired: + pytest.fail(f"{trigger}: grid child did not exit within {_CHILD_TIMEOUT}s") + recorded = out.read_text() if out.exists() else "" + return bool(recorded.strip()), proc.returncode, recorded + + +@pytest.mark.parametrize("trigger", sorted(CASES)) +def test_cleanup_trigger_grid(trigger, tmp_path): + expected_cleanup, expected_rc = CASES[trigger] + fired, rc, recorded = _run(trigger, tmp_path) + assert rc == expected_rc, f"{trigger}: unexpected exit code {rc}" + assert fired is expected_cleanup, ( + f"{trigger}: cleanup fired={fired}, expected {expected_cleanup} " + f"(recorded: {recorded!r})" + ) + + +def test_interpreter_is_not_already_finalizing_when_cleanup_runs(tmp_path): + """The exit-time path must not be suppressed by a ``sys.is_finalizing`` guard. + + Measured on 3.9-3.13: at the point ``weakref.finalize`` invokes pending + callbacks during shutdown, ``sys.is_finalizing()`` is still False and a + normal call is reachable. An implementation that skipped cleanup whenever + ``is_finalizing()`` were True would therefore be correct but useless: it + would never skip, while one that skipped on *shutdown* would drop every + ``python train.py`` session. This test pins the observed behaviour so the + reasoning behind _auto_stop's unguarded attempt stays checkable. + """ + fired, rc, recorded = _run("process_exit", tmp_path) + assert fired and rc == 0 + assert "finalizing=False" in recorded, recorded diff --git a/sagemaker-train/tests/unit/train/training_session/test_submitter.py b/sagemaker-train/tests/unit/train/training_session/test_submitter.py new file mode 100644 index 0000000000..f016219369 --- /dev/null +++ b/sagemaker-train/tests/unit/train/training_session/test_submitter.py @@ -0,0 +1,453 @@ +# 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/ +"""Unit tests for sagemaker.train.training_session._submitter. + +No test performs a network call; waits are sub-millisecond polls with +hard deadlines, so the file stays well under a second per test. +""" + +from __future__ import annotations + +import asyncio +import concurrent.futures +import gc +import os +import threading +import time + +import pytest + +from sagemaker.train.training_session._submitter import BackgroundSubmitter +from sagemaker.train.training_session.exceptions import TrainingSessionError + + +def test_ordered_submits_run_in_call_order(): + sub = BackgroundSubmitter() + order = [] + futures = [] + + def make(i): + async def _op(): + await asyncio.sleep(0.001 if i % 2 == 0 else 0) + order.append(i) + return i + + return _op + + for i in range(20): + futures.append(sub.submit_ordered(make(i))) + results = [f.result(timeout=5) for f in futures] + assert results == list(range(20)) + assert order == list(range(20)) + sub.close() + + +def test_submit_returns_without_waiting_for_completion(): + sub = BackgroundSubmitter() + release = threading.Event() + + async def _op(): + while not release.is_set(): + await asyncio.sleep(0.001) + return 42 + + start = time.monotonic() + futures = [sub.submit_ordered(_op) for _ in range(10)] + elapsed = time.monotonic() - start + assert elapsed < 0.05 + assert not futures[0].done() + release.set() + assert [f.result(timeout=5) for f in futures] == [42] * 10 + sub.close() + + +def test_submit_error_surfaces_at_result_time_and_gate_releases(): + sub = BackgroundSubmitter() + + async def _boom(): + raise ValueError("submit failed") + + async def _ok(): + return "fine" + + bad = sub.submit_ordered(_boom) + good = sub.submit_ordered(_ok) + with pytest.raises(ValueError, match="submit failed"): + bad.result(timeout=5) + assert good.result(timeout=5) == "fine" + sub.close() + + +def test_unordered_submits_overlap(): + sub = BackgroundSubmitter() + gate = threading.Event() + started = [] + + async def _op(): + started.append(1) + while not gate.is_set(): + await asyncio.sleep(0.001) + return len(started) + + f1 = sub.submit_unordered(_op) + f2 = sub.submit_unordered(_op) + deadline = time.monotonic() + 2 + while len(started) < 2 and time.monotonic() < deadline: + time.sleep(0.001) + assert len(started) == 2 + gate.set() + f1.result(timeout=5) + f2.result(timeout=5) + sub.close() + + +def test_close_is_idempotent_and_blocks_new_submits(): + sub = BackgroundSubmitter() + + async def _op(): + return 1 + + assert sub.submit_ordered(_op).result(timeout=5) == 1 + sub.close() + sub.close() + with pytest.raises(TrainingSessionError, match="closed and accepts no new operations"): + sub.submit_ordered(_op) + + +def test_close_before_first_submit_is_a_noop(): + sub = BackgroundSubmitter() + sub.close() + with pytest.raises(TrainingSessionError): + sub.submit_unordered(None) + + +def test_a_rejected_schedule_claims_no_ordered_slot(): + """A slot handed out but never scheduled would stall every later submit. + + Ordered submits release their turn in sequence, so a slot claimed by a + coroutine that never runs is never released, and every submit behind it waits + on that turn forever. The slot is therefore claimed only once the coroutine is + on the loop. + """ + sub = BackgroundSubmitter() + + async def _op(): + return 1 + + assert sub.submit_ordered(_op).result(timeout=5) == 1 + claimed = sub._next_seq + + # A closed loop is what makes scheduling fail outright. + sub._loop.call_soon_threadsafe(sub._loop.stop) + sub._thread.join(timeout=5) + sub._loop.close() + with pytest.raises(RuntimeError): + sub.submit_ordered(_op) + + assert sub._next_seq == claimed + + +def test_closing_while_submits_are_in_flight_never_hangs(): + """close() must not be able to strand a submit mid-handoff. + + The slot is claimed and scheduled under the lock close() takes, so close() + cannot land between the two. If it could, the ordered submits behind the + stranded slot would wait on their turn forever: a hang, not an error. + """ + for _ in range(20): + sub = BackgroundSubmitter() + + async def _op(): + return 1 + + futures = [] + errors = [] + + def _submit(): + for _ in range(5): + try: + futures.append(sub.submit_ordered(_op)) + except TrainingSessionError: + # Submitting after close is an error by design. + pass + except Exception as exc: # noqa: BLE001 + errors.append(exc) + + threads = [threading.Thread(target=_submit) for _ in range(4)] + for t in threads: + t.start() + sub.close() + for t in threads: + t.join(timeout=5) + assert not t.is_alive() + + assert errors == [] + for future in futures: + # Resolving or raising are both fine; blocking is the defect. Only + # the timeout is tolerated as a failure signal, so a hang fails here + # instead of being swallowed with everything else. + try: + future.result(timeout=5) + except concurrent.futures.TimeoutError: # pragma: no cover + pytest.fail("a submit never resolved after close(); close() hung it") + except Exception: # noqa: BLE001 - a failed submit is a resolved one + pass + + +def test_a_cancelled_turn_still_releases_the_next_submit(): + """A submit cancelled while waiting its turn must not strand the queue. + + Reachable from the public API: awaiting a future with a timeout cancels the + submit underneath it. Before, the release only ran for a submit that started, + so the slot stayed taken and every later ordered submit waited on it forever. + """ + sub = BackgroundSubmitter(name="cancel-turn") + started, release_first = threading.Event(), threading.Event() + + async def first(): + started.set() + await asyncio.get_running_loop().run_in_executor(None, release_first.wait) + return "first" + + try: + held = sub.submit_ordered(first) + assert started.wait(5) + + waiting = sub.submit_ordered(lambda: asyncio.sleep(0, result="waiting")) + time.sleep(0.2) # let it reach the turn wait + assert waiting.cancel() + + release_first.set() + assert held.result(timeout=5) == "first" + + after = sub.submit_ordered(lambda: asyncio.sleep(0, result="after")) + assert after.result(timeout=5) == "after" + finally: + release_first.set() + sub.close() + + +def test_close_fails_a_submit_it_could_not_wait_out(): + """Stopping the loop freezes an unfinished submit, so it has to be failed. + + Left alone, its future never resolves, and a caller waiting on the result + with the default unbounded timeout waits for the life of the process. + """ + sub = BackgroundSubmitter(name="strand") + block = threading.Event() + + async def never_finishes(): + await asyncio.get_running_loop().run_in_executor(None, block.wait) + return "never" + + try: + stuck = sub.submit_ordered(never_finishes) + time.sleep(0.2) + sub.close(timeout=0.5) + + with pytest.raises(TrainingSessionError, match="still in flight"): + stuck.result(timeout=5) + finally: + block.set() + + +def test_close_tolerates_a_pending_entry_that_already_settled(): + """Losing the race to fail a submit must not abort the shutdown. + + ``close`` fails whatever is still pending once its wait runs out, and the + loop thread can complete that same future in between, which makes setting an + exception on it raise. Letting that escape skips the stop and the join, + leaking the very thread the shutdown exists to retire. + + The window is a genuine race, so it is set up rather than raced for: a + finished future placed in the pending set puts ``close`` in exactly the state + it would reach by losing, without depending on timing to get there. + """ + sub = BackgroundSubmitter(name="close-settled") + started = sub.submit_unordered(lambda: asyncio.sleep(0, result="ok")) + assert started.result(timeout=5) == "ok" + + settled: concurrent.futures.Future = concurrent.futures.Future() + settled.set_result("already done") + with sub._lock: + sub._pending.add(settled) + + sub.close(timeout=0.05) + + assert settled.result(timeout=1) == "already done", "a settled result was replaced" + assert ( + sub._thread is not None and not sub._thread.is_alive() + ), "the loop thread outlived close(), so the stop or the join was skipped" + + +def _await_parked(sub, count, deadline=5.0): + """Poll until ``count`` submits are parked waiting on their turn.""" + end = time.monotonic() + deadline + while time.monotonic() < end: + if len(sub._waiters) >= count: + return True + time.sleep(0.001) + return False + + +def test_a_cancelled_turn_does_not_let_a_later_submit_overtake(): + """Cancelling a queued submit must not hand on a turn it never held. + + The holder is still on the wire, so releasing its successor early puts two + ordered operations on the wire out of call order, which is the one thing + ordered submission exists to prevent. + """ + sub = BackgroundSubmitter(name="cancel-overtake") + order = [] + started, release_first = threading.Event(), threading.Event() + + async def first(): + order.append("1-start") + started.set() + await asyncio.get_running_loop().run_in_executor(None, release_first.wait) + order.append("1-end") + return "first" + + async def second(): + order.append("2-start") + return "second" + + try: + held = sub.submit_ordered(first) + assert started.wait(5) + + queued_second = sub.submit_ordered(second) + queued_third = sub.submit_ordered(lambda: asyncio.sleep(0, result="third")) + assert _await_parked(sub, 2), "the later submits never parked on their turn" + + assert queued_third.cancel() + # Long enough that a wrongly released successor would have run by now, + # short enough to keep the test well under a second. + time.sleep(0.05) + assert order == ["1-start"], f"a later submit overtook the holder: {order}" + + release_first.set() + assert held.result(timeout=5) == "first" + assert queued_second.result(timeout=5) == "second" + assert order == ["1-start", "1-end", "2-start"] + finally: + release_first.set() + sub.close() + + +def test_the_turn_skips_a_run_of_cancelled_sequences(): + """Several cancellations in a row must not leave the turn stuck on any of them. + + The holder advances past every sequence that was cancelled while queued, so a + submit behind a run of them still gets its turn rather than waiting on a + sequence that will never run. + """ + sub = BackgroundSubmitter(name="cancel-run") + started, release_first = threading.Event(), threading.Event() + + async def first(): + started.set() + await asyncio.get_running_loop().run_in_executor(None, release_first.wait) + return "first" + + try: + held = sub.submit_ordered(first) + assert started.wait(5) + + cancelled = [ + sub.submit_ordered(lambda: asyncio.sleep(0, result="skipped")) for _ in range(3) + ] + last = sub.submit_ordered(lambda: asyncio.sleep(0, result="last")) + assert _await_parked(sub, 4), "the later submits never parked on their turn" + + for future in cancelled: + assert future.cancel() + + release_first.set() + assert held.result(timeout=5) == "first" + assert last.result(timeout=5) == "last" + finally: + release_first.set() + sub.close() + + +# --- descriptor release ------------------------------------------------------ + + +def _live_fds(): + """Descriptor count for this process; Linux-only, so the caller skips.""" + return len(os.listdir(f"/proc/{os.getpid()}/fd")) + + +def test_close_returns_the_loops_descriptors_rather_than_waiting_for_gc(): + """Stopping a loop only parks it; the descriptors go back when it closes. + + An open loop holds its readiness poller and its self-pipe pair, three + descriptors, until the loop object is collected. Whoever holds a reference + therefore decides when they come back, and the session handle holds every + submitter it was given for the life of the session, so leaving this to + collection meant a closed submitter kept its descriptors anyway. + """ + if not os.path.isdir(f"/proc/{os.getpid()}/fd"): + pytest.skip("descriptor accounting via /proc is Linux-only") + + # Settle cyclic garbage first: an earlier test's parked loop collected + # mid-measurement would drop the count below the baseline asserted here. + gc.collect() + before = _live_fds() + sub = BackgroundSubmitter(name="fd-release") + sub.submit_unordered(lambda: asyncio.sleep(0, result="ok")).result(timeout=5) + assert _live_fds() > before, "the loop did not come up, so nothing is measured" + + sub.close() + + assert sub._loop.is_closed(), "close() left the loop open" + # Still referenced by ``sub``, so this passes only because close() released it. + assert _live_fds() == before + + +def test_close_stays_idempotent_once_it_has_closed_the_loop(): + """A second close() must not raise on the loop the first one closed. + + Callers rely on this: the session teardown drains submitters that may have + closed themselves, and stop() followed by delete() closes the same one twice. + """ + sub = BackgroundSubmitter(name="double-close") + sub.submit_unordered(lambda: asyncio.sleep(0, result="ok")).result(timeout=5) + sub.close() + sub.close() + sub.close() + + +def test_a_join_that_runs_out_leaves_the_loop_open_for_a_later_close(): + """Closing a still-running loop raises, so the timeout path must not try. + + Losing the descriptors to an exception here would also skip the caller's own + teardown, and a later close() can still finish the job. + """ + sub = BackgroundSubmitter(name="join-timeout") + block = threading.Event() + entered = threading.Event() + + async def wedge_the_loop(): + # Blocks the loop thread itself, not an executor worker, so run_forever + # cannot return and the join is guaranteed to run out. + entered.set() + block.wait() + + try: + sub.submit_unordered(wedge_the_loop) + assert entered.wait(5) + sub.close(timeout=0.1) + + assert sub._thread is not None and sub._thread.is_alive() + assert not sub._loop.is_closed(), "closing a running loop would have raised" + finally: + block.set() + sub._thread.join(timeout=5) + sub.close() diff --git a/sagemaker-train/tests/unit/train/training_session/test_training_client.py b/sagemaker-train/tests/unit/train/training_session/test_training_client.py new file mode 100644 index 0000000000..44651c55ef --- /dev/null +++ b/sagemaker-train/tests/unit/train/training_session/test_training_client.py @@ -0,0 +1,1170 @@ +# 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/ +"""Unit tests for sagemaker.train.training_session.training_client. + +The runtime client is faked; no test performs a real sleep or network call. +Here we assert request shaping, sequence numbering, and error propagation. +""" + +from __future__ import annotations + +import asyncio + +import pytest +from botocore.exceptions import ClientError + +from sagemaker.train.training_session.training_client import TrainingClient +from sagemaker.train.training_session.types import ( + AdamParams, + ForwardBackwardOutput, + LoadStateResponse, + ModelPackageKind, + OptimStepResponse, + SaveWeightsResponse, + SaveStateResponse, +) + + +def _client_error(code, *, status=400, message="boom", op="ExecuteTrainingOperation"): + """A botocore error shaped the way the runtime client raises one.""" + return ClientError( + { + "Error": {"Code": code, "Message": message}, + "ResponseMetadata": {"HTTPStatusCode": status}, + }, + op, + ) + + +class FakeRuntime: + """Records execute_training_operation calls and replays canned responses. + + Stands in for the botocore runtime client the training client calls + directly, so ``execute_error`` is a botocore error, exactly what the wire + raises. + """ + + def __init__(self, *, execute_error: Exception | None = None): + self.execute_calls: list[dict] = [] + self.get_calls: list[dict] = [] + self.sample_execute_calls: list[dict] = [] + self.sample_get_calls: list[dict] = [] + self._execute_error = execute_error + self._op_counter = 0 + + def execute_training_operation(self, **kwargs): + self.execute_calls.append(kwargs) + if self._execute_error is not None: + raise self._execute_error + self._op_counter += 1 + return {"operationId": f"op-{self._op_counter}"} + + def get_training_operation_result(self, **kwargs): + self.get_calls.append(kwargs) + return { + "status": "Completed", + "operationResult": {"optimizerStep": {"metrics": {"loss": 0.25}}}, + } + + def execute_sampling_operation(self, **kwargs): + self.sample_execute_calls.append(kwargs) + self._op_counter += 1 + return {"operationId": f"sample-op-{self._op_counter}"} + + def get_sampling_operation_result(self, **kwargs): + self.sample_get_calls.append(kwargs) + return { + "status": "Completed", + "operationResult": {}, + } + + +def _client(**kwargs) -> tuple[TrainingClient, FakeRuntime]: + rt = FakeRuntime(**kwargs) + return TrainingClient("sess-1", rt), rt + + +def _ce(length, loss_fn="cross_entropy"): + """The loss inputs ``loss_fn`` requires, sized for datums of ``length`` tokens. + + An operation whose datums lack them is refused, so a test about chunking or + sequencing still has to carry them. + """ + inputs = { + "target_tokens": list(range(length)), + "weights": [1.0] * length, + } + if loss_fn in ("ppo", "cispo", "importance_sampling"): + inputs["logprobs"] = [-0.1] * length + inputs["advantages"] = [1.0] * length + return inputs + + +# --- request shaping -------------------------------------------------------- + + +def test_forward_backward_shapes_request_from_token_ids(): + tc, rt = _client() + tc.forward_backward([[1, 2, 3]], loss_fn="cross_entropy", loss_fn_inputs=_ce(3)).result() + call = rt.execute_calls[0] + assert call["trainingSessionName"] == "sess-1" + fb = call["operationConfig"]["forwardBackward"] + assert fb["lossFnName"] == "cross_entropy" + assert fb["datumList"][0]["modelInput"]["chunks"][0]["tokens"] == [1, 2, 3] + + +def test_forward_backward_accepts_enum_and_attaches_loss_fn_inputs(): + from sagemaker.train.training_session.types import LossFunction + + tc, rt = _client() + # The full set ppo reads, each per-token entry as long as the 2-token + # datum: a partial set is refused, so a plumbing check passes a whole one. + inputs = { + "target_tokens": {"data": [6, 0], "dtype": "int64", "shape": [2]}, + "weights": {"data": [0.0, 1.0], "dtype": "float32", "shape": [2]}, + "logprobs": {"data": [-0.1, -0.2], "dtype": "float32", "shape": [2]}, + "advantages": {"data": [1.0, 1.0], "dtype": "float32", "shape": [2]}, + } + tc.forward_backward([[5, 6]], loss_fn=LossFunction.PPO, loss_fn_inputs=inputs).result() + fb = rt.execute_calls[0]["operationConfig"]["forwardBackward"] + assert fb["lossFnName"] == "ppo" + assert fb["datumList"][0]["lossFnInputs"] == inputs + + +def test_optim_step_defaults_and_explicit_params(): + tc, rt = _client() + tc.optim_step().result() + cfg = rt.execute_calls[0]["operationConfig"]["optimizerStep"]["adamParamsConfig"] + assert cfg["learningRate"] == AdamParams().learning_rate + + tc.optim_step(AdamParams(learning_rate=3e-4, grad_clip_norm=1.0)).result() + cfg2 = rt.execute_calls[1]["operationConfig"]["optimizerStep"]["adamParamsConfig"] + assert cfg2["learningRate"] == 3e-4 + assert cfg2["gradClipNorm"] == 1.0 + + +def test_optim_step_lora_plus_lr_ratio(): + tc, rt = _client() + tc.optim_step(AdamParams(learning_rate=1e-4, lora_plus_lr_ratio=16.0)).result() + cfg = rt.execute_calls[0]["operationConfig"]["optimizerStep"]["adamParamsConfig"] + assert cfg["loraPlusLrRatio"] == 16.0 + assert cfg["learningRate"] == 1e-4 + + +def test_optim_step_omits_lora_plus_lr_ratio_when_none(): + tc, rt = _client() + tc.optim_step(AdamParams(learning_rate=1e-4)).result() + cfg = rt.execute_calls[0]["operationConfig"]["optimizerStep"]["adamParamsConfig"] + assert "loraPlusLrRatio" not in cfg + + +def test_save_weights_for_sampler_save_state_load_state_shapes(): + tc, rt = _client() + tc.save_weights_for_sampler().result() + assert rt.execute_calls[0]["operationConfig"]["saveWeights"] == {} + + tc.save_state().result() + assert rt.execute_calls[1]["operationConfig"]["saveState"] == {} + + tc.load_state("arn:aws:sagemaker:::model-package/x").result() + ls = rt.execute_calls[2]["operationConfig"]["loadState"] + assert ls == { + "modelPackageArn": "arn:aws:sagemaker:::model-package/x", + "loadOptimizer": False, + } + + +def test_save_state_sends_description_when_provided(): + tc, rt = _client() + tc.save_state(description="step 100 checkpoint").result() + assert rt.execute_calls[0]["operationConfig"]["saveState"] == { + "description": "step 100 checkpoint" + } + + +def test_save_state_omits_description_by_default(): + tc, rt = _client() + tc.save_state().result() + assert rt.execute_calls[0]["operationConfig"]["saveState"] == {} + + +def test_save_state_async_sends_description(): + tc, rt = _client() + asyncio.run(tc.save_state_async(description="async save at step 50")).result() + assert rt.execute_calls[0]["operationConfig"]["saveState"] == { + "description": "async save at step 50" + } + + +def test_save_weights_for_sampler_sends_description_when_provided(): + tc, rt = _client() + tc.save_weights_for_sampler(description="weights at step 200").result() + assert rt.execute_calls[0]["operationConfig"]["saveWeights"] == { + "description": "weights at step 200", + } + + +def test_save_weights_for_sampler_omits_description_by_default(): + tc, rt = _client() + tc.save_weights_for_sampler().result() + assert rt.execute_calls[0]["operationConfig"]["saveWeights"] == {} + + +def test_save_weights_for_sampler_async_sends_description(): + tc, rt = _client() + asyncio.run(tc.save_weights_for_sampler_async(description="async weights step 50")).result() + assert rt.execute_calls[0]["operationConfig"]["saveWeights"] == { + "description": "async weights step 50", + } + + +def test_save_weights_for_sampler_omits_merge_flag_by_default(): + tc, rt = _client() + tc.save_weights_for_sampler().result() + assert rt.execute_calls[0]["operationConfig"]["saveWeights"] == {} + + +def test_save_weights_for_sampler_sends_merge_flag_when_requested(): + tc, rt = _client() + tc.save_weights_for_sampler(merge_weights=True).result() + assert rt.execute_calls[0]["operationConfig"]["saveWeights"] == { + "mergeWeights": True, + } + asyncio.run(tc.save_weights_for_sampler_async(merge_weights=True)).result() + assert rt.execute_calls[1]["operationConfig"]["saveWeights"] == { + "mergeWeights": True, + } + + +# --- forward-backward chunking ---------------------------------------------- + + +def _token_datum(n_tokens: int): + return {"modelInput": {"chunks": [{"tokens": list(range(n_tokens)), "type": "encoded_text"}]}} + + +def test_forward_backward_single_chunk_sends_one_op_unchanged(): + # A batch that fits in one chunk submits exactly one op, unchanged. + tc, rt = _client() + tc.forward_backward( + [[1, 2, 3], [4, 5, 6]], loss_fn="cross_entropy", loss_fn_inputs=_ce(3) + ).result() + assert len(rt.execute_calls) == 1 + fb = rt.execute_calls[0]["operationConfig"]["forwardBackward"] + assert fb["lossFnName"] == "cross_entropy" + assert len(fb["datumList"]) == 2 + + +def test_forward_backward_splits_batch_over_count_limit(): + # More datums than the chunk-length cap -> multiple ops, order preserved, + # every datum sent exactly once. + from sagemaker.train.training_session import training_client as tc_mod + + tc, rt = _client() + original = tc_mod.chunk_datum_list + + # Force a tiny count cap so a small batch chunks deterministically. + def tc_mod_chunk(dl, **kw): + return original(dl, max_len=2, max_bytes=10**9) + + try: + tc_mod.chunk_datum_list = tc_mod_chunk + tc.forward_backward( + [[1], [2], [3], [4], [5]], loss_fn="ppo", loss_fn_inputs=_ce(1, "ppo") + ).result() + finally: + tc_mod.chunk_datum_list = original + + assert len(rt.execute_calls) == 3 # 2 + 2 + 1 + sent = [ + d["modelInput"]["chunks"][0]["tokens"][0] + for call in rt.execute_calls + for d in call["operationConfig"]["forwardBackward"]["datumList"] + ] + assert sent == [1, 2, 3, 4, 5] # order preserved, no datum dropped/duplicated + # Sequence numbers are contiguous across the chunk ops, consecutive with + # no gap. The service accumulates gradients over these ops and blocks on any + # missing sequence number, so contiguity (not merely strict increase) is the + # property that keeps a chunked batch equivalent to one un-chunked op. + seqs = [c["sequenceNumber"] for c in rt.execute_calls] + assert seqs == list(range(seqs[0], seqs[0] + len(seqs))) + + +def test_forward_splits_batch_over_count_limit(): + # forward chunks too. It carries the same batch as forward_backward, so a + # batch one of them can send and the other cannot is a surprise: the + # forward-only pass failed on a size the training pass handled. + from sagemaker.train.training_session import training_client as tc_mod + + tc, rt = _client() + original = tc_mod.chunk_datum_list + try: + tc_mod.chunk_datum_list = lambda dl, **kw: original(dl, max_len=2, max_bytes=10**9) + tc.forward([[1], [2], [3], [4], [5]], loss_fn="ppo", loss_fn_inputs=_ce(1, "ppo")).result() + finally: + tc_mod.chunk_datum_list = original + + assert len(rt.execute_calls) == 3 # 2 + 2 + 1 + sent = [ + d["modelInput"]["chunks"][0]["tokens"][0] + for call in rt.execute_calls + for d in call["operationConfig"]["forward"]["datumList"] + ] + assert sent == [1, 2, 3, 4, 5] # order preserved, no datum dropped/duplicated + + +def test_forward_single_chunk_sends_one_op_unchanged(): + tc, rt = _client() + tc.forward([[1, 2, 3], [4, 5, 6]], loss_fn="cross_entropy", loss_fn_inputs=_ce(3)).result() + assert len(rt.execute_calls) == 1 + fwd = rt.execute_calls[0]["operationConfig"]["forward"] + assert fwd["lossFnName"] == "cross_entropy" + assert len(fwd["datumList"]) == 2 + + +def test_small_batch_sends_single_op(): + # A batch comfortably under the module chunk bounds is sent as exactly one op. + rt = FakeRuntime() + tc = TrainingClient("sess-1", rt) + tc.forward_backward( + [[1, 2, 3], [4, 5, 6]], loss_fn="cross_entropy", loss_fn_inputs=_ce(3) + ).result() + assert len(rt.execute_calls) == 1 + + +def test_chunk_datum_list_bounds(): + from sagemaker.train.training_session._utils import chunk_datum_list + + # Empty -> no chunks. + assert chunk_datum_list([]) == [] + # Count bound. + data = [_token_datum(1) for _ in range(5)] + assert [len(c) for c in chunk_datum_list(data, max_len=2, max_bytes=10**9)] == [2, 2, 1] + # Byte bound: each datum ~ 10 tokens * 10 = ~100 bytes; cap at 250 -> 2 per chunk. + big = [_token_datum(10) for _ in range(5)] + assert [len(c) for c in chunk_datum_list(big, max_len=1000, max_bytes=250)] == [2, 2, 1] + # A single oversized datum still gets its own chunk (never dropped). + huge = [_token_datum(1000)] + assert chunk_datum_list(huge, max_len=1000, max_bytes=1) == [huge] + + +def test_combine_forward_backward_results_weights_and_concatenates(): + from sagemaker.train.training_session._utils import ( + combine_forward_backward_results, + ) + + # Single result passes through unchanged. + solo = ForwardBackwardOutput._from_api({"forwardBackward": {"metrics": {"loss": 0.5}}}) + assert combine_forward_backward_results([solo]) is solo + + # Chunks carrying no loss outputs still each count once, so the weighted + # mean does not collapse to the last chunk. + r1 = ForwardBackwardOutput._from_api( + {"forwardBackward": {"metrics": {"loss": 1.0, "grad_norm": 9.9}}} + ) + r2 = ForwardBackwardOutput._from_api( + {"forwardBackward": {"metrics": {"loss": 4.0, "grad_norm": 2.0}}} + ) + combined = combine_forward_backward_results([r1, r2]) + assert abs(combined.metrics["loss"] - 2.5) < 1e-9 # (1.0*1 + 4.0*1)/2 + assert combined.metrics["grad_norm"] == 2.0 # last chunk (accumulated grad) + assert combined.loss_fn_outputs is None + + +def test_combine_forward_backward_results_concatenates_per_datum_entries(): + """Per-datum entries concatenate in chunk order, weighting by entry count.""" + from sagemaker.train.training_session._utils import ( + combine_forward_backward_results, + ) + + r1 = ForwardBackwardOutput._from_api( + { + "forwardBackward": { + "metrics": {"loss": 1.0, "grad_norm": 9.9}, + "lossFnOutputs": [{"loss": [0.1]}, {"loss": [0.2]}], + } + } + ) + r2 = ForwardBackwardOutput._from_api( + { + "forwardBackward": { + "metrics": {"loss": 4.0, "grad_norm": 2.0}, + "lossFnOutputs": [{"loss": [0.3]}], + } + } + ) + combined = combine_forward_backward_results([r1, r2]) + # Datum counts come from the entry counts: 2 and 1 -> (1.0*2 + 4.0*1)/3 + assert abs(combined.metrics["loss"] - 2.0) < 1e-9 + assert combined.metrics["grad_norm"] == 2.0 + assert combined.loss_fn_outputs == [{"loss": [0.1]}, {"loss": [0.2]}, {"loss": [0.3]}] + + +# --- base returns future / async twins are coroutines ------------------------ + + +def test_base_ops_return_future_resolving_to_parsed_result(): + tc, _ = _client() + # Base methods return an unresolved future immediately; .result() (or + # await) yields the parsed per-op result. + assert isinstance( + tc.forward_backward([[1]], loss_fn_inputs=_ce(1)).result(), ForwardBackwardOutput + ) + assert isinstance(tc.optim_step().result(), OptimStepResponse) + assert isinstance(tc.save_state().result(), SaveStateResponse) + assert isinstance(tc.forward([[1]], loss_fn_inputs=_ce(1)).result(), ForwardBackwardOutput) + + +def test_async_ops_are_coroutines_yielding_futures(): + tc, _ = _client() + # *_async methods are coroutines: awaiting one completes the enqueue and + # yields the same APIFuture the base method returns. + for coro_fn in ( + tc.forward_backward_async, + tc.optim_step_async, + tc.save_state_async, + tc.forward_async, + tc.save_weights_for_sampler_async, + tc.load_state_async, + ): + assert asyncio.iscoroutinefunction(coro_fn) + for coro, expected_cls in ( + (tc.forward_backward_async([[1]], loss_fn_inputs=_ce(1)), ForwardBackwardOutput), + (tc.optim_step_async(), OptimStepResponse), + (tc.save_state_async(), SaveStateResponse), + (tc.forward_async([[1]], loss_fn_inputs=_ce(1)), ForwardBackwardOutput), + (tc.save_weights_for_sampler_async(), SaveWeightsResponse), + (tc.load_state_async("arn:mp/1"), LoadStateResponse), + ): + fut = asyncio.run(coro) + assert hasattr(fut, "result") + assert isinstance(fut.result(), expected_cls) + # The internal LWFS helper stays sync and returns its future directly. + fut = tc._submit_load_weights_for_sampling("arn:mp/1") + assert isinstance(fut.result(), LoadStateResponse) + + +def test_load_state_with_optimizer_sets_flag(): + tc, rt = _client() + tc.load_state_with_optimizer("arn:mp/7").result() + cfg = rt.execute_calls[0]["operationConfig"]["loadState"] + assert cfg["loadOptimizer"] is True + assert cfg["modelPackageArn"] == "arn:mp/7" + + +def test_load_state_is_weights_only(): + """Bare load_state must NOT restore the optimizer. + + The negative twin of the test above. Without it the two methods can + silently collapse into the same call, which is what happened when + load_state defaulted load_optimizer=True: load_state_with_optimizer + became unreachable and a caller expecting weights-only resume got the + optimizer restored. Engine-side and client-side both split this by + method name, not by a flag. + """ + tc, rt = _client() + tc.load_state("arn:mp/7").result() + cfg = rt.execute_calls[0]["operationConfig"]["loadState"] + assert cfg["loadOptimizer"] is False + assert cfg["modelPackageArn"] == "arn:mp/7" + + +def test_load_state_exposes_no_load_optimizer_knob(): + """Neither load method takes a load_optimizer parameter. + + ``load_state(path)`` / ``load_state_with_optimizer(path)`` carry the + choice in the method name. A flag re-introduces the ambiguity this pair + exists to remove. + """ + import inspect + + for fn in ( + TrainingClient.load_state, + TrainingClient.load_state_async, + TrainingClient.load_state_with_optimizer, + TrainingClient.load_state_with_optimizer_async, + ): + assert "load_optimizer" not in inspect.signature(fn).parameters, fn.__name__ + + +def test_load_weights_for_sampling_uses_sample_endpoint(): + tc, rt = _client() + tc._load_weights_for_sampling("arn:mp/42") + assert len(rt.sample_execute_calls) == 1 + call = rt.sample_execute_calls[0] + assert call["operationConfig"] == {"loadWeightsForSampling": {"modelPackageArn": "arn:mp/42"}} + assert call["trainingSessionName"] == "sess-1" + assert len(rt.execute_calls) == 0, "LWFS must NOT go through training path" + + +def test_load_weights_for_sampling_does_not_consume_sequence(): + tc, rt = _client() + tc.forward_backward([[1]], loss_fn_inputs=_ce(1)).result() + seq_before = tc.current_sequence_number + tc._load_weights_for_sampling("arn:mp/5") + assert ( + tc.current_sequence_number == seq_before + ), "LWFS must not advance the training sequence counter" + tc.optim_step().result() + seqs = [c["sequenceNumber"] for c in rt.execute_calls] + assert seqs == [1, 2], "Training ops should be 1,2 (no gap from LWFS)" + + +# --- sequence numbering ----------------------------------------------------- + + +def test_sequence_numbers_are_monotonic_and_start_at_one(): + tc, rt = _client() + tc.forward_backward([[1]], loss_fn_inputs=_ce(1)).result() + tc.optim_step().result() + tc.save_state().result() + seqs = [c["sequenceNumber"] for c in rt.execute_calls] + assert seqs == [1, 2, 3] + + +def test_each_submit_uses_unique_client_token(): + tc, rt = _client() + tc.optim_step().result() + tc.optim_step().result() + tokens = {c["clientToken"] for c in rt.execute_calls} + assert len(tokens) == 2 + + +# --- error propagation ------------------------------------------------------ +# The service's own error reaches the caller unchanged: the client neither +# renames it nor hides the code and message the service sent. Submits are queued +# on a background thread, so a submit error surfaces when the returned future is +# resolved; tests resolve via .result(). + + +@pytest.mark.parametrize( + "code", + ["ConflictException", "ResourceNotFoundException", "ValidationException"], +) +def test_service_error_propagates_unchanged(code): + tc, _ = _client(execute_error=_client_error(code, message="from the service")) + with pytest.raises(ClientError) as ei: + tc.optim_step().result() + assert ei.value.response["Error"]["Code"] == code + assert "from the service" in str(ei.value) + + +# --- runtime-warmup transient retry on submit -------------------------------- +# The control plane reports Running before the training runtime endpoint accepts +# traffic, so the first submit can 503 at the runtime front door. The submit is +# idempotent (stable clientToken + sequenceNumber), so it retries through the +# warmup window. Sleeps are patched out so the test stays sub-second. + + +class _FlakyRuntime(FakeRuntime): + """Raises a retryable botocore error ``fail_n`` times, then succeeds.""" + + def __init__(self, *, fail_n: int): + super().__init__() + self._fail_n = fail_n + self._attempts = 0 + + def execute_training_operation(self, **kwargs): + self._attempts += 1 + if self._attempts <= self._fail_n: + raise _client_error("InternalServerException", status=503, message="warming up") + return super().execute_training_operation(**kwargs) + + +def _patch_retry_no_sleep(monkeypatch): + """Make _submit's transient retry sleep instantly (no real waits in tests). + + ``call_with_transient_retry`` captures ``time.sleep`` as a default arg at + import time, so patching the time module doesn't reach it — instead replace + the symbol the client calls with a no-sleep partial of the real function. + """ + import functools + import sagemaker.train.training_session.training_client as tc_mod + from sagemaker.train.training_session.futures import call_with_transient_retry + + no_sleep = functools.partial(call_with_transient_retry, sleep_fn=lambda _s: None) + monkeypatch.setattr(tc_mod, "call_with_transient_retry", no_sleep) + + +def test_submit_retries_transient_then_succeeds(monkeypatch): + _patch_retry_no_sleep(monkeypatch) + rt = _FlakyRuntime(fail_n=8) # > the old 5-attempt budget, < the 30 budget + tc = TrainingClient("sess-1", rt) + + fut = tc.optim_step() # must not raise — rides through the warmup 503s + + assert fut.operation_id == "op-1" # blocks until the submit lands + assert rt._attempts == 9 # 8 failures + 1 success + + +def test_submit_gives_up_when_transient_never_clears(monkeypatch): + from sagemaker.train.training_session.training_client import ( + _TRAINING_SUBMIT_MAX_ATTEMPTS, + ) + + _patch_retry_no_sleep(monkeypatch) + rt = _FlakyRuntime(fail_n=10_000) # never clears + tc = TrainingClient("sess-1", rt) + + with pytest.raises(ClientError): + tc.optim_step().result() + assert rt._attempts == _TRAINING_SUBMIT_MAX_ATTEMPTS + + +# --- result wiring ---------------------------------------------------------- + + +def test_future_result_parses_training_result(): + tc, rt = _client() + fut = tc.optim_step() + result = fut.result() + assert isinstance(result, OptimStepResponse) + assert result.metrics == {"loss": 0.25} + # poll hit the right operation id + assert rt.get_calls[0]["operationId"] == fut.operation_id + + +# --- forward (forward-only op) ---------------------------------------------- + + +def test_forward_shapes_request_from_token_ids(): + tc, rt = _client() + tc.forward([[1, 2, 3]], loss_fn="cross_entropy", loss_fn_inputs=_ce(3)).result() + call = rt.execute_calls[0] + assert call["trainingSessionName"] == "sess-1" + fwd = call["operationConfig"]["forward"] + assert fwd["lossFnName"] == "cross_entropy" + assert fwd["datumList"][0]["modelInput"]["chunks"][0]["tokens"] == [1, 2, 3] + # forward-only must NOT emit a forwardBackward config. + assert "forwardBackward" not in call["operationConfig"] + + +def test_forward_accepts_enum_and_attaches_loss_fn_inputs(): + from sagemaker.train.training_session.types import LossFunction + + tc, rt = _client() + # The full set ppo reads, each per-token entry as long as the 2-token + # datum: a partial set is refused, so a plumbing check passes a whole one. + inputs = { + "target_tokens": {"data": [6, 0], "dtype": "int64", "shape": [2]}, + "weights": {"data": [0.0, 1.0], "dtype": "float32", "shape": [2]}, + "logprobs": {"data": [-0.1, -0.2], "dtype": "float32", "shape": [2]}, + "advantages": {"data": [1.0, 1.0], "dtype": "float32", "shape": [2]}, + } + tc.forward([[5, 6]], loss_fn=LossFunction.PPO, loss_fn_inputs=inputs).result() + fwd = rt.execute_calls[0]["operationConfig"]["forward"] + assert fwd["lossFnName"] == "ppo" + assert fwd["datumList"][0]["lossFnInputs"] == inputs + + +# --- loss_fn_config ----------------------------------------------------------- + + +def test_forward_backward_sends_loss_fn_config(): + tc, rt = _client() + cfg = {"clip_low_threshold": 0.2, "clip_high_threshold": 5.0} + tc.forward_backward( + [[1, 2]], loss_fn="cispo", loss_fn_config=cfg, loss_fn_inputs=_ce(2, "cispo") + ).result() + fb = rt.execute_calls[0]["operationConfig"]["forwardBackward"] + assert fb["lossFnConfig"] == cfg + + +def test_forward_backward_omits_loss_fn_config_when_none(): + tc, rt = _client() + tc.forward_backward([[1, 2]], loss_fn_inputs=_ce(2)).result() + fb = rt.execute_calls[0]["operationConfig"]["forwardBackward"] + assert "lossFnConfig" not in fb + + +def test_forward_sends_loss_fn_config(): + tc, rt = _client() + cfg = {"clip_low_threshold": 0.1} + tc.forward([[1, 2]], loss_fn="ppo", loss_fn_config=cfg, loss_fn_inputs=_ce(2, "ppo")).result() + fwd = rt.execute_calls[0]["operationConfig"]["forward"] + assert fwd["lossFnConfig"] == cfg + + +def test_forward_omits_loss_fn_config_when_none(): + tc, rt = _client() + tc.forward([[1, 2]], loss_fn_inputs=_ce(2)).result() + fwd = rt.execute_calls[0]["operationConfig"]["forward"] + assert "lossFnConfig" not in fwd + + +def test_forward_backward_tokenizes_a_model_input_built_from_text(): + """A typed batch element carrying text is tokenized like a bare string.""" + from sagemaker.train.training_session.types import Datum, ModelInput + + class _Tok: + def encode(self, text): + return [4, 5] + + rt = FakeRuntime() + tc = TrainingClient("sess-1", rt, tokenizer=_Tok()) + tc.forward_backward( + [Datum(model_input=ModelInput.from_text("hi"))], loss_fn_inputs=_ce(2) + ).result() + datum = rt.execute_calls[0]["operationConfig"]["forwardBackward"]["datumList"][0] + assert datum["modelInput"] == {"chunks": [{"tokens": [4, 5], "type": "encoded_text"}]} + + +# --- get_tokenizer ----------------------------------------------------------- + + +def test_get_tokenizer_lazy_loads_and_caches(monkeypatch): + import sagemaker.train.training_session._utils as utils + + calls = [] + sentinel = object() + + def fake_load(base_model, hub_client, hf_model_id=None): + calls.append(base_model) + return sentinel + + monkeypatch.setattr(utils, "load_tokenizer", fake_load) + tc = TrainingClient("sess-1", FakeRuntime(), base_model="Qwen/Qwen3-32B") + tok1 = tc.get_tokenizer() + tok2 = tc.get_tokenizer() + assert tok1 is sentinel and tok2 is sentinel + # Loaded once for the right base model, then cached. + assert calls == ["Qwen/Qwen3-32B"] + + +def test_get_tokenizer_unsupported_model_raises(monkeypatch): + import sagemaker.train.training_session._utils as utils + + def fake_load(base_model, hub_client, hf_model_id=None): + raise ValueError("No public HuggingFace tokenizer is available") + + monkeypatch.setattr(utils, "load_tokenizer", fake_load) + tc = TrainingClient("sess-1", FakeRuntime(), base_model="arn:.../private/1.0") + with pytest.raises(ValueError, match="No public HuggingFace tokenizer"): + tc.get_tokenizer() + + +# --- checkpoint discovery ----------------------------------------------------- +# list_model_packages delegates to the owning ServiceClient with the session's own +# group. latest_*_checkpoint are CACHE-ONLY (see the save-driven cache tests +# below) — they never query the group, so a fresh/resumed client returns None. + + +class _RecordingServiceClient: + """Records the kwargs each checkpoint helper is called with.""" + + def __init__(self): + self.calls: list[tuple] = [] + + def list_model_packages(self, **kwargs): + self.calls.append(("list", kwargs)) + return ["ckpt-list"] + + def _stop_session_if_active(self, training_session_name: str) -> bool: + # Reached only when a client is dropped without an explicit stop; the + # automatic path calls this rather than stop_session. + self.calls.append(("stop", training_session_name)) + return True + + +def test_list_checkpoints_forwards_max_results(): + svc = _RecordingServiceClient() + tc = TrainingClient("sess-1", FakeRuntime(), model_package_group_arn="grp", service_client=svc) + + tc.list_model_packages(max_results=5) + kind, kwargs = svc.calls[0] + assert kind == "list" + assert kwargs == {"model_package_group_arn": "grp", "max_results": 5} + + # Default (None) omits max_results so the ServiceClient default applies. + svc.calls.clear() + tc.list_model_packages() + assert svc.calls[0][1] == {"model_package_group_arn": "grp"} + + +def test_list_checkpoints_requires_group(): + # No group on the client -> clear error pointing to ServiceClient. + tc = TrainingClient("sess-1", FakeRuntime(), service_client=_RecordingServiceClient()) + with pytest.raises(ValueError, match="no model package group"): + tc.list_model_packages() + + +def test_list_checkpoints_requires_service_client(): + tc = TrainingClient("sess-1", FakeRuntime(), model_package_group_arn="grp") + with pytest.raises(ValueError, match="not created by a ServiceClient"): + tc.list_model_packages() + + +def test_latest_by_kind_are_none_without_a_save(): + # Cache-only: a client that has not saved returns None and NEVER queries the + # group, even when both a group and a ServiceClient are available. + svc = _RecordingServiceClient() + tc = TrainingClient("sess-1", FakeRuntime(), model_package_group_arn="grp", service_client=svc) + assert tc.get_latest_model_package(ModelPackageKind.CHECKPOINT) is None + assert tc.get_latest_model_package(ModelPackageKind.MODEL) is None + assert svc.calls == [] # no fallback to the group + + +# --- save-driven checkpoint cache ------------------------------------------ + + +class _SaveRuntime(FakeRuntime): + """FakeRuntime whose ops complete carrying a save result with an ARN. + + ``op_kind`` selects whether the completed op reports ``saveState`` or + ``saveWeights`` (so the save response carries model_package_arn). Each + completed op gets a unique ARN so the cache can be checked for recency. + """ + + def __init__(self, op_kind: str = "saveState", **kwargs): + super().__init__(**kwargs) + self._op_kind = op_kind + + def get_training_operation_result(self, **kwargs): + self.get_calls.append(kwargs) + op_id = kwargs["operationId"] + return { + "status": "Completed", + "operationResult": {self._op_kind: {"modelPackageArn": f"arn-{op_id}"}}, + } + + +def test_save_state_caches_checkpoint_no_api_call(): + svc = _RecordingServiceClient() + tc = TrainingClient( + "sess-1", + _SaveRuntime("saveState"), + model_package_group_arn="grp", + service_client=svc, + ) + + tc.save_state().result() + ckpt = tc.get_latest_model_package(ModelPackageKind.CHECKPOINT) + + assert ckpt.model_package_arn == "arn-op-1" + assert ckpt.kind.value == "checkpoint" + # A resumable save is not exposed as a model checkpoint. + assert tc.get_latest_model_package(ModelPackageKind.MODEL) is None + # Cache hit -> ServiceClient was never queried. + assert svc.calls == [] + + +def test_save_weights_for_sampler_caches_model_checkpoint(): + svc = _RecordingServiceClient() + tc = TrainingClient( + "sess-1", + _SaveRuntime("saveWeights"), + model_package_group_arn="grp", + service_client=svc, + ) + + tc.save_weights_for_sampler().result() + ckpt = tc.get_latest_model_package(ModelPackageKind.MODEL) + + assert ckpt.model_package_arn == "arn-op-1" + assert ckpt.kind.value == "model" + assert svc.calls == [] + + +def test_async_save_populates_cache_on_result(): + # The cache fills on completion, which for the async path happens when + # parse_fn runs at .result() time, not at submit time. + svc = _RecordingServiceClient() + tc = TrainingClient( + "sess-1", + _SaveRuntime("saveState"), + model_package_group_arn="grp", + service_client=svc, + ) + + fut = tc.save_state() + # Not resolved yet -> cache-only helper returns None (no group fallback). + assert tc.get_latest_model_package(ModelPackageKind.CHECKPOINT) is None + assert svc.calls == [] + + fut.result() + assert tc.get_latest_model_package(ModelPackageKind.CHECKPOINT).model_package_arn == "arn-op-1" + assert svc.calls == [] # served from cache + + +def test_latest_returns_most_recent_cached_save(): + tc = TrainingClient( + "sess-1", + _SaveRuntime("saveState"), + model_package_group_arn="grp", + service_client=_RecordingServiceClient(), + ) + + tc.save_state().result() # arn-op-1 + tc.save_state().result() # arn-op-2 + assert tc.get_latest_model_package(ModelPackageKind.CHECKPOINT).model_package_arn == "arn-op-2" + + +def test_latest_per_kind_are_tracked_independently(): + # The two kinds occupy separate cache slots: a model save must not evict + # the resumable entry, and each helper returns the latest of its own kind. + tc = TrainingClient( + "sess-1", + _SaveRuntime("saveState"), + model_package_group_arn="grp", + service_client=_RecordingServiceClient(), + ) + + tc.save_state().result() # arn-op-1 -> resumable + # Flip the same fake to report a model save so the op counter keeps + # advancing (a fresh one would reset it to op-1). + tc._runtime._op_kind = "saveWeights" + tc.save_weights_for_sampler().result() # arn-op-2 -> model + + assert tc.get_latest_model_package(ModelPackageKind.CHECKPOINT).model_package_arn == "arn-op-1" + assert tc.get_latest_model_package(ModelPackageKind.MODEL).model_package_arn == "arn-op-2" + + +# --- session lifecycle sugar ----------------------------------------------- + + +def test_stop_delegates_to_service_client(): + calls = [] + + class _Svc(_RecordingServiceClient): + def stop_session(self, name): + calls.append(("stop", name)) + + tc = TrainingClient("sess-1", FakeRuntime(), service_client=_Svc()) + tc.stop() + assert calls == [("stop", "sess-1")] + + +def test_delete_delegates_with_stop_flag(): + calls = [] + + class _Svc(_RecordingServiceClient): + def delete_session(self, name, *, stop_if_active=False): + calls.append(("delete", name, stop_if_active)) + + tc = TrainingClient("sess-1", FakeRuntime(), service_client=_Svc()) + tc.delete(stop_if_active=True) + assert calls == [("delete", "sess-1", True)] + + +def test_stop_requires_service_client(): + tc = TrainingClient("sess-1", FakeRuntime()) + with pytest.raises(ValueError, match="not created by a ServiceClient"): + tc.stop() + + +# --- calling convention ----------------------------------------------------- + + +def test_the_parameters_that_do_nothing_cannot_be_passed_positionally(): + """A parameter that is ignored must not occupy a position. + + Each of these is accepted for compatibility and does nothing. Positionally + they would be indistinguishable from the parameters that do something, so a + value meant for one would land silently in another; requiring the keyword + makes such a call a ``TypeError``. It also leaves them removable later, + since no caller can depend on their position. + """ + import inspect + + from sagemaker.train.training_session.sampling_client import SamplingClient + from sagemaker.train.training_session.service_client import ServiceClient + + inert = { + "seed", + "train_mlp", + "train_attn", + "train_unembed", + "user_metadata", + "retry_config", + "ttl_seconds", + "name", + "overwrite", + } + positional = [] + for cls in (ServiceClient, TrainingClient, SamplingClient): + for method_name, method in vars(cls).items(): + if method_name.startswith("_") or not callable(method): + continue + try: + signature = inspect.signature(method) + except (TypeError, ValueError): # pragma: no cover - not a function + continue + positional += [ + f"{cls.__name__}.{method_name}({parameter})" + for parameter, spec in signature.parameters.items() + if parameter in inert and spec.kind is spec.POSITIONAL_OR_KEYWORD + ] + + assert positional == [], f"inert parameters accepting a position: {positional}" + + +def test_the_parameters_that_do_something_stay_positional(): + """The guard above must not have swept up the real parameters with them.""" + import inspect + + expected = { + "save_state": ["description"], + "save_weights_for_sampler": ["description", "merge_weights"], + "save_weights_and_get_sampling_client": ["description", "merge_weights"], + "create_sampling_client": ["model_package_arn"], + "load_state": ["model_package_arn"], + } + for method_name, wanted in expected.items(): + signature = inspect.signature(getattr(TrainingClient, method_name)) + actual = [ + parameter + for parameter, spec in signature.parameters.items() + if spec.kind is spec.POSITIONAL_OR_KEYWORD and parameter != "self" + ] + assert actual == wanted, f"{method_name} takes {actual}, expected {wanted}" + + +# --- sequence number accessor ---------------------------------------------- + + +def test_current_sequence_number_is_read_only(): + # The public accessor reports the counter WITHOUT advancing it; reading it + # repeatedly must not consume a sequence (the bug that burned seq 6 in e2e). + tc = TrainingClient("sess-1", FakeRuntime()) + assert tc.current_sequence_number == 0 + assert tc.current_sequence_number == 0 # idempotent — no increment + # A real op advances it by exactly one. + tc.optim_step() + assert tc.current_sequence_number == 1 + + +def test_each_submit_advances_sequence_by_one(): + # 1 increment == 1 dispatched request: the seq sent on the wire matches the + # accessor after each op, with no gaps. + tc, rt = _client() + tc.optim_step().result() + tc.optim_step().result() + seqs = [c["sequenceNumber"] for c in rt.execute_calls] + assert seqs == [1, 2] + assert tc.current_sequence_number == 2 + + +# --- non-blocking submit ------------------------------------------------------ + + +def test_base_methods_return_before_submit_completes(): + import threading + import time + + class _GatedRuntime(FakeRuntime): + """Blocks every execute until released, like a slow wire.""" + + def __init__(self): + super().__init__() + self.release = threading.Event() + + def execute_training_operation(self, **kwargs): + assert self.release.wait(timeout=5) + return super().execute_training_operation(**kwargs) + + rt = _GatedRuntime() + tc = TrainingClient("sess-1", rt) + start = time.monotonic() + futs = [tc.optim_step() for _ in range(10)] + elapsed = time.monotonic() - start + # All ten submits queued without a single wire round trip completing. + assert elapsed < 0.5 + assert rt.execute_calls == [] + rt.release.set() + for f in futs: + assert isinstance(f.result(), OptimStepResponse) + # Wire order matches call order despite the background handoff. + assert [c["sequenceNumber"] for c in rt.execute_calls] == list(range(1, 11)) + + +def test_stop_drains_queued_submits_before_stopping(): + from sagemaker.train.training_session.exceptions import TrainingSessionError + + order = [] + + class _Svc(_RecordingServiceClient): + def stop_session(self, name): + order.append("stop") + + class _OrderRuntime(FakeRuntime): + def execute_training_operation(self, **kwargs): + order.append("submit") + return super().execute_training_operation(**kwargs) + + tc = TrainingClient("sess-1", _OrderRuntime(), service_client=_Svc()) + tc.optim_step() + tc.stop() + assert order == ["submit", "stop"] + # The client accepts no new operations after stop. + with pytest.raises(TrainingSessionError, match="closed and accepts no new operations"): + tc.optim_step() + + +# --- the chunk weighting is wired, not just implemented ------------------------ +# The combiner's arithmetic is covered where it lives. What is only observable +# here is whether the client hands it the chunk sizes: without them an uneven +# split reports a plain mean over chunks, understating whichever chunk carried +# the most datums, and nothing raises. + + +class _PerChunkLossRuntime(FakeRuntime): + """Replays a different loss for each forward/forward_backward op, in order.""" + + def __init__(self, losses): + super().__init__() + self._losses = list(losses) + self._served = 0 + + def get_training_operation_result(self, **kwargs): + self.get_calls.append(kwargs) + loss = self._losses[min(self._served, len(self._losses) - 1)] + self._served += 1 + return { + "status": "Completed", + "operationResult": {"forwardBackward": {"metrics": {"loss": loss}}}, + } + + +def _chunked_loss(method_name, losses, batch, loss_fn_inputs): + """Run one chunked call whose chunks report ``losses``; return the combined loss.""" + from sagemaker.train.training_session import training_client as tc_mod + + rt = _PerChunkLossRuntime(losses) + tc = TrainingClient("sess-1", rt) + original = tc_mod.chunk_datum_list + try: + # Two datums per chunk, so a 5-datum batch splits 2 + 2 + 1. + tc_mod.chunk_datum_list = lambda dl, **kw: original(dl, max_len=2, max_bytes=10**9) + result = getattr(tc, method_name)( + batch, loss_fn="cross_entropy", loss_fn_inputs=loss_fn_inputs + ).result() + finally: + tc_mod.chunk_datum_list = original + assert len(rt.execute_calls) == 3, "the batch did not chunk as expected" + return result.metrics["loss"] + + +@pytest.mark.parametrize("method_name", ["forward_backward", "forward"]) +def test_chunked_loss_is_weighted_by_chunk_size(method_name): + """Both chunked paths weight by datum count, so an uneven split is honest. + + Chunks of 2, 2 and 1 reporting 1.0, 1.0 and 11.0 average to + (2 + 2 + 11) / 5 = 3.0 by datum count, against 4.333 as a plain mean over + chunks. The plain mean is what a missing ``datum_counts`` produces. + """ + combined = _chunked_loss( + method_name, + losses=[1.0, 1.0, 11.0], + batch=[[1], [2], [3], [4], [5]], + loss_fn_inputs=_ce(1), + ) + + assert combined == pytest.approx(3.0), ( + f"{method_name} reported {combined}; 4.333 means the chunk sizes never " + f"reached the combiner" + ) diff --git a/sagemaker-train/tests/unit/train/training_session/test_types.py b/sagemaker-train/tests/unit/train/training_session/test_types.py new file mode 100644 index 0000000000..9d2692db99 --- /dev/null +++ b/sagemaker-train/tests/unit/train/training_session/test_types.py @@ -0,0 +1,1511 @@ +# 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/ +"""Unit tests for sagemaker.train.training_session.types.""" + +from __future__ import annotations + +import pytest +from pydantic import BaseModel, ValidationError + +from sagemaker.train.training_session.types import ( + AdamParams, + Datum, + EncodedTextChunk, + ForwardBackwardOutput, + LoadStateResponse, + LossFunction, + ModelInput, + ModelInputChunk, + ModelPackageDeletion, + ModelPackageInfo, + ModelPackageKind, + OptimStepResponse, + SampleResponse, + SampledSequence, + SamplingParams, + SaveWeightsResponse, + SaveStateResponse, + SessionStatus, + StopReason, + TensorData, + TensorDtype, + TokenLogprob, + VpcConfig, + _serialize_loss_fn_inputs, +) + + +class TestLossFunction: + def test_enum_values(self): + assert LossFunction.PPO.value == "ppo" + assert LossFunction.IMPORTANCE_SAMPLING.value == "importance_sampling" + assert LossFunction.CISPO.value == "cispo" + + def test_from_string(self): + assert LossFunction("ppo") == LossFunction.PPO + + def test_invalid_raises(self): + with pytest.raises(ValueError): + LossFunction("invalid") + + +class TestAdamParams: + def test_defaults(self): + p = AdamParams() + assert p.learning_rate == 1e-4 + assert p.beta1 == 0.9 + assert p.beta2 == 0.999 + + def test_to_api_minimal(self): + api = AdamParams()._to_api() + assert api["learningRate"] == 1e-4 + assert "gradClipNorm" not in api + + def test_to_api_with_grad_clip(self): + api = AdamParams(grad_clip_norm=1.0)._to_api() + assert api["gradClipNorm"] == 1.0 + + def test_to_api_with_lora_plus_lr_ratio(self): + api = AdamParams(lora_plus_lr_ratio=16.0)._to_api() + assert api["loraPlusLrRatio"] == 16.0 + + def test_to_api_omits_lora_plus_lr_ratio_when_none(self): + api = AdamParams()._to_api() + assert "loraPlusLrRatio" not in api + + def test_extra_fields_rejected(self): + with pytest.raises(ValidationError): + AdamParams(unknown=1) + + +class TestSamplingParams: + def test_defaults(self): + # Service requires max_tokens >= 512; default must satisfy it. + p = SamplingParams() + assert p.max_tokens == 512 + + def test_to_api(self): + p = SamplingParams(max_tokens=512, temperature=0.7, top_k=50) + api = p._to_api() + assert api == {"maxTokens": 512, "temperature": 0.7, "topK": 50} + + def test_to_api_minimal(self): + api = SamplingParams()._to_api() + assert api == {"maxTokens": 512} + + @pytest.mark.parametrize( + "field,value", + [ + ("temperature", -0.1), + ("top_p", 95), # a percentage, the mistake this guard exists for + ("top_p", 1.5), + ("top_p", -0.1), + ], + ) + def test_a_probability_outside_its_own_range_is_refused(self, field, value): + """These two carry no meaning outside their range, and nothing else stops them. + + The wire model bounds both, but that bound is not applied before the + request is sent, so an out-of-range value would otherwise travel to the + service and fail there, one background submit away from the call that + made it. + """ + with pytest.raises(ValidationError): + SamplingParams(max_tokens=512, **{field: value}) + + @pytest.mark.parametrize( + "field,value", + [("temperature", 0), ("temperature", 2), ("top_p", 0), ("top_p", 1)], + ) + def test_the_range_ends_are_usable(self, field, value): + """Zero temperature is greedy decoding and top_p of 1 keeps everything.""" + p = SamplingParams(max_tokens=512, **{field: value}) + + assert getattr(p, field) == value + + +class TestParamsFromAny: + """The attribute path of from_any on the all-optional params types. + + Both types default every field, so an object that shares no field names + would otherwise be accepted and silently replaced by our defaults. + """ + + class TheirParams: + """A caller's own params class: same field names, no inheritance.""" + + max_tokens = 1024 + temperature = 0.7 + learning_rate = 0.5 + + class Partial: + """Models fewer knobs than we do; the rest must take our defaults.""" + + temperature = 0.3 + beta1 = 0.5 + + class Unrelated: + """camelCase / misspelt: shares no field name with either type.""" + + maxTokens = 1024 + temp = 0.7 + lr = 0.5 + + @pytest.mark.parametrize("cls", [SamplingParams, AdamParams]) + def test_passthrough(self, cls): + p = cls() + assert cls.from_any(p) is p + + def test_sampling_params_from_duck_typed_object(self): + p = SamplingParams.from_any(self.TheirParams()) + assert p.max_tokens == 1024 + assert p.temperature == 0.7 + + def test_adam_params_from_duck_typed_object(self): + p = AdamParams.from_any(self.TheirParams()) + assert p.learning_rate == 0.5 + + def test_partial_object_keeps_our_defaults_for_the_rest(self): + assert SamplingParams.from_any(self.Partial()).temperature == 0.3 + assert SamplingParams.from_any(self.Partial()).max_tokens == 512 + assert AdamParams.from_any(self.Partial()).beta1 == 0.5 + assert AdamParams.from_any(self.Partial()).learning_rate == 1e-4 + + @pytest.mark.parametrize("cls", [SamplingParams, AdamParams]) + def test_object_sharing_no_field_names_is_rejected(self, cls): + # Regression: this used to return all-defaults, so a caller's values + # were dropped without a word. + with pytest.raises(TypeError, match="none of"): + cls.from_any(self.Unrelated()) + + @pytest.mark.parametrize("cls", [SamplingParams, AdamParams]) + def test_dict_with_an_unknown_key_is_rejected(self, cls): + with pytest.raises(ValidationError): + cls.from_any({"definitely_not_a_field": 1}) + + +class TestOperationResults: + def test_forward_metrics_and_loss_fn_outputs(self): + r = ForwardBackwardOutput._from_api( + {"forward": {"metrics": {"loss": 1.5}, "lossFnOutputs": [{"logprobs": [0.1]}]}} + ) + assert r.metrics == {"loss": 1.5} + assert r.loss_fn_outputs == [{"logprobs": [0.1]}] + + def test_forward_backward_metrics(self): + r = ForwardBackwardOutput._from_api({"forwardBackward": {"metrics": {"grad_norm": 0.2}}}) + assert r.metrics == {"grad_norm": 0.2} + + +class TestLossFnOutputsWireShapes: + """``lossFnOutputs`` is a per-datum list (``TensorDataMapList``). + + The field is either a list, an empty list, or absent; a mapping is not one + of the shapes the response can carry. + """ + + def _fb(self, body): + return ForwardBackwardOutput._from_api({"forwardBackward": body}) + + def test_per_datum_entries(self): + r = self._fb({"lossFnOutputs": [{"actor/pg_loss": [1.0]}, {"actor/pg_loss": [2.0]}]}) + assert r.loss_fn_outputs == [{"actor/pg_loss": [1.0]}, {"actor/pg_loss": [2.0]}] + + def test_absent_is_none_not_empty(self): + """Absent stays distinguishable from an empty result. + + An operation that produced no per-datum outputs omits the field + altogether, which is not the same answer as an empty list. + """ + r = self._fb({"metrics": {"grad_norm": 0.2}}) + assert r.loss_fn_outputs is None + + def test_empty_list_preserved(self): + r = self._fb({"lossFnOutputs": []}) + assert r.loss_fn_outputs == [] + + def test_optimizer_step_metrics(self): + r = OptimStepResponse._from_api({"optimizerStep": {"metrics": {"lr": 1e-4}}}) + assert r.metrics == {"lr": 1e-4} + + def test_model_package_arn_from_save_ops(self): + assert ( + SaveStateResponse._from_api( + {"saveState": {"modelPackageArn": "arn:x"}} + ).model_package_arn + == "arn:x" + ) + assert ( + SaveWeightsResponse._from_api( + {"saveWeights": {"modelPackageArn": "arn:y"}} + ).model_package_arn + == "arn:y" + ) + + def test_empty_result_defaults(self): + r = ForwardBackwardOutput._from_api({}) + assert r.metrics == {} + assert r.loss_fn_outputs is None + assert OptimStepResponse._from_api({}).metrics == {} + assert SaveStateResponse._from_api({}).model_package_arn is None + assert SaveWeightsResponse._from_api({}).model_package_arn is None + assert LoadStateResponse._from_api({}) is not None + + def test_is_basemodel_raw_not_public(self): + raw = {"saveWeights": {"modelPackageArn": "arn:z"}} + r = SaveWeightsResponse._from_api(raw) + assert isinstance(r, BaseModel) + assert r.model_package_arn == "arn:z" + assert "raw" not in type(r).model_fields + + def test_extra_fields_rejected(self): + with pytest.raises(ValidationError): + ForwardBackwardOutput(unexpected=1) + + +class _CompatTensor: + """External-object stand-in with TensorData-shaped attributes.""" + + def __init__(self, data, dtype="float32"): + self.data = data + self.dtype = dtype + + +class _CompatChunk: + """External-object stand-in with EncodedTextChunk-shaped attributes.""" + + def __init__(self, tokens): + self.tokens = tokens + + +class _CompatModelInput: + """External-object stand-in with ModelInput-shaped attributes.""" + + def __init__(self, chunks): + self.chunks = chunks + + +class _CompatDatum: + """External-object stand-in with Datum-shaped attributes.""" + + def __init__(self, model_input, loss_fn_inputs=None): + self.model_input = model_input + self.loss_fn_inputs = loss_fn_inputs or {} + + +class _CompatCompressedTensor: + """External tensor whose ``data`` holds a compressed encoding. + + Stands in for the tensor libraries that store a sparse encoding in ``data`` + and expand it in ``tolist()``, with ``shape`` describing the *dense* tensor. + Reading ``data`` here sends the non-zero values alone against a dense shape, + a mask silently missing most of its entries, so ``tolist()`` has to win. + """ + + def __init__(self, dense_rows, dtype="float32"): + self._dense_rows = dense_rows + self.shape = [len(dense_rows), len(dense_rows[0])] + self.dtype = dtype + # The compressed payload: non-zero values only. + self.data = [v for row in dense_rows for v in row if v] + + def tolist(self): + return [list(row) for row in self._dense_rows] + + +class TestTensorData: + def test_to_api(self): + api = TensorData(data=[1.0, 2.0], dtype=TensorDtype.FLOAT32)._to_api() + assert api == {"data": [1.0, 2.0], "dtype": "float32"} + + def test_to_api_with_shape(self): + api = TensorData(data=[1, 2, 3, 4], dtype=TensorDtype.INT64, shape=[2, 2])._to_api() + assert api["shape"] == [2, 2] + assert api["dtype"] == "int64" + + def test_from_any_passthrough(self): + td = TensorData(data=[1.0]) + assert TensorData.from_any(td) is td + + def test_from_any_dict(self): + td = TensorData.from_any({"data": [1, 2], "dtype": "int32", "shape": [2]}) + assert td.data == [1, 2] + assert td.dtype == TensorDtype.INT32 + assert td.shape == [2] + + def test_from_any_plain_list(self): + td = TensorData.from_any([0.5, 0.7]) + assert td.data == [0.5, 0.7] + assert td.dtype == TensorDtype.FLOAT32 + + def test_from_any_compatible(self): + td = TensorData.from_any(_CompatTensor([1.0, 2.0], dtype="bfloat16")) + assert td.data == [1.0, 2.0] + assert td.dtype == TensorDtype.BFLOAT16 + + def test_from_any_torch_style_dtype_string(self): + td = TensorData.from_any(_CompatTensor([1], dtype="torch.int64")) + assert td.dtype == TensorDtype.INT64 + + def test_from_any_non_numeric_dtype_raises(self): + """A dtype the wire cannot carry is refused, not sent as float32. + + Reinterpreting complex values as float32 produced a request the service + either rejected or trained on with the wrong numbers. + """ + with pytest.raises(ValueError, match="unsupported tensor dtype 'complex128'"): + TensorData.from_any(_CompatTensor([1.0], dtype="complex128")) + + def test_from_any_float64_narrows_to_float32(self): + td = TensorData.from_any(_CompatTensor([1.0], dtype="torch.float64")) + assert td.dtype == TensorDtype.FLOAT32 + + def test_from_any_bool_widens_to_int64(self): + """A boolean mask reaches the loss function as 0/1 integers, not floats.""" + td = TensorData.from_any(_CompatTensor([1], dtype="torch.bool")) + assert td.dtype == TensorDtype.INT64 + + def test_direct_construction_coerces_a_library_dtype(self): + """``TensorData(data=..., dtype=t.dtype)`` is the obvious way to write it. + + The coercion used to live only on ``from_any``, so passing a tensor + library's dtype object, or its plain name, failed enum validation + with no hint that a conversion existed. + """ + np = pytest.importorskip("numpy") + assert TensorData(data=[1.0], dtype=np.dtype("float64")).dtype == TensorDtype.FLOAT32 + assert TensorData(data=[1], dtype=np.dtype("bool")).dtype == TensorDtype.INT64 + assert TensorData(data=[1.0], dtype="float64").dtype == TensorDtype.FLOAT32 + + def test_two_dimensional_source_is_flattened_with_shape(self): + """The wire carries a flat ``data`` array plus ``shape``, never nested rows. + + A 2-D tensor reached the wire as a list of lists, which the runtime types + as an array of numbers and rejects. + """ + np = pytest.importorskip("numpy") + api = TensorData.from_any(np.arange(6, dtype=np.float32).reshape(2, 3))._to_api() + assert api == {"data": [0.0, 1.0, 2.0, 3.0, 4.0, 5.0], "dtype": "float32", "shape": [2, 3]} + + def test_nested_data_without_a_shape_infers_one(self): + api = TensorData(data=[[1.0, 2.0], [3.0, 4.0]])._to_api() + assert api == {"data": [1.0, 2.0, 3.0, 4.0], "dtype": "float32", "shape": [2, 2]} + + def test_ragged_nested_data_raises(self): + with pytest.raises(ValueError, match="ragged"): + TensorData(data=[[1.0, 2.0], [3.0]])._to_api() + + def test_unknown_float_width_narrows_to_float32(self): + """Float widths keep being added; a name table silently rejects each new one.""" + assert TensorData._coerce_dtype("float8_e4m3fn") == TensorDtype.FLOAT32 + assert TensorData._coerce_dtype("float128") == TensorDtype.FLOAT32 + + def test_half_is_float16_not_float32(self): + """``half`` names a dtype the wire carries, so it must not be narrowed.""" + assert TensorData._coerce_dtype("half") == TensorDtype.FLOAT16 + + def test_every_dtype_spelling_agrees(self): + """A scalar type, a dtype instance and a plain name must reach one answer. + + A scalar type stringifies as ````; unwrapping that + badly left ``int32'>``, which matched nothing, so a dtype the wire + carries was refused, and ``float32`` warned about narrowing to itself. + """ + np = pytest.importorskip("numpy") + for spelling in (np.int32, np.dtype("int32"), "int32"): + assert TensorData._coerce_dtype(spelling) == TensorDtype.INT32 + for spelling in (np.float32, np.dtype("float32"), "float32"): + assert TensorData._coerce_dtype(spelling) == TensorDtype.FLOAT32 + with pytest.raises(ValueError, match="'complex128'"): + TensorData._coerce_dtype(np.complex128) + + def test_tolist(self): + assert TensorData(data=[1, 2]).tolist() == [1, 2] + + def test_to_numpy(self): + np = pytest.importorskip("numpy") + arr = TensorData(data=[1.0, 2.0], dtype=TensorDtype.FLOAT32).to_numpy() + assert arr.dtype == np.float32 + assert arr.tolist() == [1.0, 2.0] + + def test_to_numpy_respects_shape(self): + pytest.importorskip("numpy") + arr = TensorData(data=[1, 2, 3, 4], dtype=TensorDtype.INT32, shape=[2, 2]).to_numpy() + assert arr.shape == (2, 2) + + +class TestSerializeLossFnInputs: + def test_tensor_data_serialized(self): + out = _serialize_loss_fn_inputs({"advantages": TensorData(data=[0.1, 0.2])}) + assert out["advantages"] == {"data": [0.1, 0.2], "dtype": "float32"} + + def test_compatible_tensor_serialized(self): + out = _serialize_loss_fn_inputs({"weights": _CompatTensor([1.0], dtype="float16")}) + assert out["weights"] == {"data": [1.0], "dtype": "float16"} + + def test_raw_list_becomes_tensor_data(self): + """The wire wants a TensorData here; a bare list was rejected locally. + + dtype comes from the input name, which a list cannot carry: ``[1, 2, 3]`` + is a valid token list and a valid weight mask. + """ + out = _serialize_loss_fn_inputs({"target_tokens": [1, 2, 3], "weights": [1, 0, 1]}) + assert out["target_tokens"] == {"data": [1, 2, 3], "dtype": "int64", "shape": [3]} + assert out["weights"] == {"data": [1, 0, 1], "dtype": "float32", "shape": [3]} + + def test_raw_list_under_unknown_key_infers_dtype(self): + out = _serialize_loss_fn_inputs({"custom": [1, 2], "custom_f": [1.5]}) + assert out["custom"]["dtype"] == "int64" + assert out["custom_f"]["dtype"] == "float32" + + def test_nested_list_is_flattened_with_shape(self): + out = _serialize_loss_fn_inputs({"weights": [[1.0, 0.0], [0.5, 1.0]]}) + assert out["weights"] == { + "data": [1.0, 0.0, 0.5, 1.0], + "dtype": "float32", + "shape": [2, 2], + } + + def test_ragged_nested_list_raises(self): + with pytest.raises(ValueError, match="ragged nested list"): + _serialize_loss_fn_inputs({"weights": [[1.0, 0.0], [1.0]]}) + + def test_non_list_values_pass_through(self): + """A dict is the wire shape already, or an advanced escape hatch.""" + raw = {"config": {"k": "v"}} + assert _serialize_loss_fn_inputs(raw) == raw + + def test_numpy_array_serialized(self): + np = pytest.importorskip("numpy") + out = _serialize_loss_fn_inputs({"advantages": np.array([0.5, 0.6], dtype=np.float32)}) + assert out["advantages"]["data"] == pytest.approx([0.5, 0.6]) + assert out["advantages"]["dtype"] == "float32" + + def test_two_dimensional_array_is_flattened(self): + np = pytest.importorskip("numpy") + out = _serialize_loss_fn_inputs({"weights": np.ones((2, 3), dtype=np.float32)}) + assert out["weights"] == {"data": [1.0] * 6, "dtype": "float32", "shape": [2, 3]} + + +class TestExternalTensorObjects: + """Values built by another tensor library must serialize to the same wire shape. + + A caller porting an existing training loop passes that library's objects + straight through, so these paths carry real data and cannot be allowed to + reinterpret it. + """ + + def test_expanded_values_win_over_a_compressed_data_attribute(self): + """``tolist()`` is the authority on the values, not ``data``. + + For a compressed encoding the two disagree: ``data`` holds only the + non-zero values while ``shape`` describes the dense tensor. Reading + ``data`` sent a mask with most of its entries missing, against a shape + claiming they were there: accepted, and trained on the wrong numbers. + """ + weights = _CompatCompressedTensor([[1.0, 0.0, 0.0], [0.0, 2.0, 0.0]]) + assert weights.data == [1.0, 2.0] # the compressed payload + out = _serialize_loss_fn_inputs({"weights": weights}) + assert out["weights"] == { + "data": [1.0, 0.0, 0.0, 0.0, 2.0, 0.0], + "dtype": "float32", + "shape": [2, 3], + } + + def test_external_datum_of_external_tensors_round_trips(self): + datum = _CompatDatum( + model_input=_CompatModelInput([_CompatChunk([10, 11, 12])]), + loss_fn_inputs={ + "target_tokens": _CompatTensor([11, 12, 13], dtype="int64"), + "weights": _CompatCompressedTensor([[0.0, 1.0, 1.0]]), + }, + ) + assert Datum.from_any(datum)._to_api() == { + "modelInput": {"chunks": [{"tokens": [10, 11, 12], "type": "encoded_text"}]}, + "lossFnInputs": { + "target_tokens": {"data": [11, 12, 13], "dtype": "int64"}, + "weights": {"data": [0.0, 1.0, 1.0], "dtype": "float32", "shape": [1, 3]}, + }, + } + + def test_raw_lists_and_external_tensors_agree_on_the_wire(self): + """The same values reach the wire identically whichever form they arrive in.""" + np = pytest.importorskip("numpy") + rows = [[0.0, 1.0], [1.0, 0.0]] + from_list = _serialize_loss_fn_inputs({"weights": rows}) + from_array = _serialize_loss_fn_inputs({"weights": np.array(rows, dtype=np.float32)}) + assert from_list == from_array + + +class TestEncodedTextChunk: + def test_length(self): + assert EncodedTextChunk(tokens=[1, 2, 3]).length == 3 + + def test_from_any_passthrough(self): + c = EncodedTextChunk(tokens=[1]) + assert EncodedTextChunk.from_any(c) is c + + def test_from_any_dict(self): + assert EncodedTextChunk.from_any({"tokens": [1, 2]}).tokens == [1, 2] + + def test_from_any_list(self): + assert EncodedTextChunk.from_any([3, 4]).tokens == [3, 4] + + def test_from_any_model_input_chunk_wrapper(self): + wrapper = ModelInputChunk(chunk=EncodedTextChunk(tokens=[5])) + assert EncodedTextChunk.from_any(wrapper).tokens == [5] + + def test_from_any_compatible(self): + assert EncodedTextChunk.from_any(_CompatChunk([7, 8])).tokens == [7, 8] + + +class TestModelInput: + def test_from_ints(self): + mi = ModelInput.from_ints([1, 2, 3]) + assert mi._to_api() == {"chunks": [{"tokens": [1, 2, 3], "type": "encoded_text"}]} + + def test_empty(self): + mi = ModelInput.empty() + assert mi.length == 0 + assert mi._to_api() == {"chunks": []} + + def test_to_ints_flattens_chunks(self): + mi = ModelInput(chunks=[EncodedTextChunk(tokens=[1, 2]), EncodedTextChunk(tokens=[3])]) + assert mi.to_ints() == [1, 2, 3] + + def test_length(self): + assert ModelInput.from_ints([1, 2, 3]).length == 3 + + def test_append_returns_new_input(self): + mi = ModelInput.from_ints([1]) + mi2 = mi.append([2, 3]) + assert mi.to_ints() == [1] + assert mi2.to_ints() == [1, 2, 3] + + def test_append_int(self): + assert ModelInput.from_ints([1]).append_int(9).to_ints() == [1, 9] + + def test_from_any_passthrough(self): + mi = ModelInput.from_ints([1]) + assert ModelInput.from_any(mi) is mi + + def test_from_any_token_list(self): + assert ModelInput.from_any([1, 2]).to_ints() == [1, 2] + + def test_from_any_wire_dict(self): + mi = ModelInput.from_any({"chunks": [{"tokens": [1], "type": "encoded_text"}]}) + assert mi.to_ints() == [1] + + def test_from_any_compatible(self): + compat = _CompatModelInput(chunks=[_CompatChunk([1, 2]), _CompatChunk([3])]) + assert ModelInput.from_any(compat).to_ints() == [1, 2, 3] + + def test_to_api_coerces_compatible_chunk(self): + # A compatible chunk passed directly (bypassing from_any) must still be + # serialized to the wire format rather than sent as-is. + mi = ModelInput(chunks=[_CompatChunk([4, 5])]) + assert mi._to_api() == {"chunks": [{"tokens": [4, 5], "type": "encoded_text"}]} + + +class TestDatum: + def test_to_api_with_model_input(self): + d = Datum(model_input=ModelInput.from_ints([1, 2])) + assert d._to_api() == { + "modelInput": {"chunks": [{"tokens": [1, 2], "type": "encoded_text"}]} + } + + def test_to_api_serializes_tensor_loss_inputs(self): + d = Datum( + model_input=ModelInput.from_ints([1]), + loss_fn_inputs={"advantages": TensorData(data=[0.5])}, + ) + assert d._to_api()["lossFnInputs"]["advantages"] == {"data": [0.5], "dtype": "float32"} + + def test_from_any_passthrough(self): + d = Datum(model_input=ModelInput.from_ints([1])) + assert Datum.from_any(d) is d + + def test_from_any_model_input(self): + d = Datum.from_any(ModelInput.from_ints([1, 2])) + assert d.model_input.to_ints() == [1, 2] + + def test_from_any_wire_dict(self): + d = Datum.from_any( + {"modelInput": {"chunks": [{"tokens": [1]}]}, "lossFnInputs": {"w": [1.0]}} + ) + assert d.model_input.to_ints() == [1] + assert d.loss_fn_inputs == {"w": [1.0]} + + def test_from_any_compatible(self): + compat = _CompatDatum( + model_input=_CompatModelInput([_CompatChunk([1, 2])]), + loss_fn_inputs={"advantages": _CompatTensor([0.1])}, + ) + d = Datum.from_any(compat) + api = d._to_api() + assert api["modelInput"]["chunks"][0]["tokens"] == [1, 2] + assert api["lossFnInputs"]["advantages"] == {"data": [0.1], "dtype": "float32"} + + def test_from_any_empty_loss_inputs_omitted(self): + compat = _CompatDatum(model_input=_CompatModelInput([_CompatChunk([1])])) + assert "lossFnInputs" not in Datum.from_any(compat)._to_api() + + +class TestDatumFromAnyNeverSilentlyEmpty: + """Every field on Datum is optional, so an unhandled value used to coerce to + an empty Datum and serialize to ``{}``: a training example dropped from + the batch with nothing raised. + """ + + def test_token_id_list_is_the_prompt_shorthand(self): + assert Datum.from_any([4, 5, 6])._to_api() == { + "modelInput": {"chunks": [{"tokens": [4, 5, 6], "type": "encoded_text"}]} + } + + def test_text_is_the_prompt_shorthand(self): + """Text is preserved for the client to tokenize, not serialized as-is. + + The runtime Datum has no text member, so serializing text raises here and + the tokenizing happens at the client boundary instead. + """ + d = Datum.from_any("hello") + assert d._pending_text == "hello" + with pytest.raises(ValueError, match="never tokenized"): + d._to_api() + + def test_pending_text_tokenizes_at_the_boundary(self): + class _Tok: + def encode(self, text): + return [11, 12] + + assert Datum.from_any("hello")._tokenized_with(_Tok())._to_api() == { + "modelInput": {"chunks": [{"tokens": [11, 12], "type": "encoded_text"}]} + } + + def test_tuple_of_token_ids(self): + assert Datum.from_any((4, 5))._to_api() == { + "modelInput": {"chunks": [{"tokens": [4, 5], "type": "encoded_text"}]} + } + + def test_bare_chunk_is_wrapped(self): + assert Datum.from_any(EncodedTextChunk(tokens=[7, 8]))._to_api() == { + "modelInput": {"chunks": [{"tokens": [7, 8], "type": "encoded_text"}]} + } + + def test_duck_typed_weight_is_rejected(self): + """A scalar datum weight has no wire member, so it is refused not dropped. + + Accepted and ignored, the weighting the caller asked for was silently + absent from training. + """ + + class TheirDatum: + model_input = ModelInput.from_ints([1, 2]) + weight = 2.5 + + with pytest.raises(ValueError, match="no 'weight' field"): + Datum.from_any(TheirDatum()) + + def test_duck_typed_prompt_is_honoured(self): + class TheirDatum: + prompt = "hi" + + assert Datum.from_any(TheirDatum())._pending_text == "hi" + + def test_unrelated_object_is_rejected(self): + class Unrelated: + foo = 1 + + with pytest.raises(TypeError, match="none of"): + Datum.from_any(Unrelated()) + + def test_object_declaring_empty_fields_is_rejected(self): + """Declaring the attributes is not the same as carrying values. + + A ``chunks`` attribute set to ``None`` must not read as a + ModelInput-like object: that would wrap the object itself as a prompt + and put ``{'prompt': }`` on the wire instead of raising. + """ + + class AllNone: + chunks = None + model_input = None + prompt = None + loss_fn_inputs = None + + with pytest.raises(TypeError, match="none of"): + Datum.from_any(AllNone()) + + def test_object_with_populated_chunks_still_accepted(self): + """The guard must not reject the ModelInput-like objects it exists for.""" + + class TheirModelInput: + chunks = [EncodedTextChunk(tokens=[1, 2])] + + assert Datum.from_any(TheirModelInput())._to_api() == { + "modelInput": {"chunks": [{"tokens": [1, 2], "type": "encoded_text"}]} + } + + def test_unrecognised_dict_is_rejected(self): + with pytest.raises(TypeError, match="none of"): + Datum.from_any({"foo": 1}) + + def test_dict_weight_is_rejected(self): + with pytest.raises(ValueError, match="no 'weight' field"): + Datum.from_any({"modelInput": {"chunks": [{"tokens": [1]}]}, "weight": 2.0}) + + def test_constructor_weight_is_rejected(self): + with pytest.raises(ValueError, match="no 'weight' field"): + Datum(model_input=ModelInput.from_ints([1]), weight=2.0) + + def test_prompt_only_dict(self): + assert Datum.from_any({"prompt": "hi"})._pending_text == "hi" + + def test_loss_inputs_only_dict_is_allowed(self): + """An explicit wire dict is the advanced escape hatch; only a dict with + no recognised key at all is a mistake. + """ + assert Datum.from_any({"lossFnInputs": {"w": [1.0]}})._to_api() == { + "lossFnInputs": {"w": {"data": [1.0], "dtype": "float32", "shape": [1]}} + } + + @pytest.mark.parametrize( + "batch", + [ + [Datum(model_input=ModelInput.from_ints([1, 2, 3])), [4, 5, 6]], + [[4, 5, 6], Datum(model_input=ModelInput.from_ints([1, 2, 3]))], + [ + {"modelInput": {"chunks": [{"tokens": [4, 5, 6]}]}}, + Datum(model_input=ModelInput.from_ints([1, 2, 3])), + ], + ], + ids=["typed-first", "tokens-first", "dict-first"], + ) + def test_mixed_batch_keeps_every_element(self, batch): + """A batch mixing kinds is coerced element by element, in either order. + + The homogeneous fast paths read the whole batch through the one + interpretation chosen for ``batch[0]``, which silently emptied later + elements one way round and serialized a raw Datum object the other. + """ + from sagemaker.train.training_session._utils import ( + batch_to_datum_list, + ) + + datums = batch_to_datum_list(batch) + assert len(datums) == 2 + assert all(d.get("modelInput", {}).get("chunks") for d in datums) + assert sorted(d["modelInput"]["chunks"][0]["tokens"] for d in datums) == [ + [1, 2, 3], + [4, 5, 6], + ] + + def test_mixed_batch_still_tokenizes_text(self): + """Text is the one kind needing outside context, so it keeps the + tokenizer path rather than passing through as a raw prompt. + """ + from sagemaker.train.training_session._utils import ( + batch_to_datum_list, + ) + + class FakeTokenizer: + def encode(self, text): + return [len(text)] + + datums = batch_to_datum_list( + [Datum(model_input=ModelInput.from_ints([1, 2])), "hello"], + tokenizer=FakeTokenizer(), + ) + assert datums[1]["modelInput"]["chunks"][0]["tokens"] == [5] + + +class TestFromAnyParams: + def test_adam_params_passthrough(self): + p = AdamParams(learning_rate=3e-4) + assert AdamParams.from_any(p) is p + + def test_adam_params_from_dict(self): + assert AdamParams.from_any({"learning_rate": 2e-4}).learning_rate == 2e-4 + + def test_adam_params_compatible(self): + class _CompatAdam: + learning_rate = 5e-5 + beta1 = 0.8 + + p = AdamParams.from_any(_CompatAdam()) + assert p.learning_rate == 5e-5 + assert p.beta1 == 0.8 + assert p.beta2 == 0.999 # default filled in + + def test_sampling_params_from_dict(self): + assert SamplingParams.from_any({"max_tokens": 1024}).max_tokens == 1024 + + def test_sampling_params_compatible(self): + class _CompatSampling: + max_tokens = 700 + temperature = 0.9 + + p = SamplingParams.from_any(_CompatSampling()) + assert p.max_tokens == 700 + assert p.temperature == 0.9 + + +class TestParamsUnsetFieldsUseOurDefaults: + """An external params object leaves unset knobs as ``None``. + + Forwarding that ``None`` would fail validation on a field of ours that is + required with a default, so "not set" has to mean "use the default": the + same outcome as omitting the field. + """ + + def test_sampling_unset_max_tokens_falls_back(self): + class _Unset: + max_tokens = None + temperature = 0.7 + + p = SamplingParams.from_any(_Unset()) + assert p.max_tokens == 512 # our default, not a validation error + assert p.temperature == 0.7 + + def test_sampling_unset_max_tokens_in_dict(self): + p = SamplingParams.from_any({"max_tokens": None, "temperature": 0.7}) + assert p.max_tokens == 512 + assert p.temperature == 0.7 + + def test_adam_unset_fields_fall_back(self): + class _Unset: + learning_rate = 5e-5 + beta2 = None + + p = AdamParams.from_any(_Unset()) + assert p.learning_rate == 5e-5 + assert p.beta2 == 0.999 # ours, not None + + def test_explicit_sentinels_still_pass_through(self): + """A deliberate value is honoured, including the service's own defaults.""" + + class _Explicit: + max_tokens = 512 + temperature = 1.0 + top_p = 1.0 + top_k = -1 # the service's documented "no limit" + + api = SamplingParams.from_any(_Explicit())._to_api() + assert api == { + "maxTokens": 512, + "temperature": 1.0, + "topP": 1.0, + "topK": -1, + } + + +class TestSamplingParamsStopCoercion: + def test_bare_string_is_wrapped(self): + assert SamplingParams(max_tokens=512, stop="END").stop == ["END"] + + def test_bare_string_via_from_any(self): + p = SamplingParams.from_any({"max_tokens": 512, "stop": "END"}) + assert p._to_api()["stop"] == ["END"] + + def test_list_of_strings_unchanged(self): + assert SamplingParams(max_tokens=512, stop=["a", "b"]).stop == ["a", "b"] + + def test_token_ids_rejected_with_reason(self): + """The service matches stop sequences on decoded text, so IDs cannot work.""" + with pytest.raises(ValidationError, match="not token IDs"): + SamplingParams(max_tokens=512, stop=[1, 2]) + + +class _CountingTensor: + """Tensor-like object that counts tolist() calls (torch/numpy stand-in).""" + + def __init__(self, values, dtype="float32"): + self._values = values + self.dtype = dtype + self.shape = (len(values),) + self.tolist_calls = 0 + + def tolist(self): + self.tolist_calls += 1 + return list(self._values) + + +class TestTensorDataLazyMaterialization: + def test_from_any_defers_tolist(self): + src = _CountingTensor([1.0, 2.0]) + td = TensorData.from_any(src) + assert src.tolist_calls == 0 + assert td.shape == [2] + assert td.dtype == TensorDtype.FLOAT32 + + def test_data_access_materializes_once(self): + src = _CountingTensor([1.0, 2.0]) + td = TensorData.from_any(src) + assert td.data == [1.0, 2.0] + assert td.data == [1.0, 2.0] + assert src.tolist_calls == 1 + + def test_to_api_materializes_lazily(self): + src = _CountingTensor([1, 2, 3], dtype="int32") + td = TensorData.from_any(src) + assert src.tolist_calls == 0 + api = td._to_api() + assert src.tolist_calls == 1 + assert api == {"data": [1, 2, 3], "dtype": "int32", "shape": [3]} + + def test_serialize_loss_fn_inputs_converts_once(self): + src = _CountingTensor([0.1, 0.2]) + out = _serialize_loss_fn_inputs({"advantages": src}) + assert out["advantages"]["data"] == [0.1, 0.2] + assert src.tolist_calls == 1 + + def test_data_is_required(self): + with pytest.raises(ValueError): + TensorData(dtype=TensorDtype.FLOAT32) + + def test_extra_fields_still_rejected(self): + with pytest.raises(ValidationError): + TensorData(data=[1.0], not_a_field=1) + + +class TestSharedBaseConfigInheritance: + """The shared bases must preserve each type's original ConfigDict behavior.""" + + def test_adam_params_still_forbids_extras(self): + with pytest.raises(ValidationError): + AdamParams(learning_rte=1e-4) + + def test_adam_params_still_validates_assignment(self): + p = AdamParams() + with pytest.raises(ValidationError): + p.learning_rate = "not-a-float" + + def test_sampling_params_still_forbids_extras(self): + with pytest.raises(ValidationError): + SamplingParams(max_token=100) + + def test_datum_still_allows_extras(self): + d = Datum(prompt=[1, 2], custom_field="kept") + assert d.custom_field == "kept" + + def test_datum_model_input_name_allowed(self): + d = Datum(model_input=ModelInput.from_ints([1, 2])) + assert d.model_input.to_ints() == [1, 2] + + +class TestStatusAndKindEnums: + """The public enums must keep the service's raw wire values.""" + + def test_session_status_values(self): + assert SessionStatus.RUNNING.value == "Running" + assert SessionStatus("Stopped") is SessionStatus.STOPPED + assert {s.value for s in SessionStatus} == { + "Created", + "Scheduled", + "Running", + "Stopping", + "Stopped", + "Failed", + "Blocked", + "Expired", + "Deleting", + "DeletionFailed", + } + + def test_model_package_kind_values(self): + assert ModelPackageKind.CHECKPOINT.value == "checkpoint" + assert ModelPackageKind.MODEL.value == "model" + assert ModelPackageKind("model") is ModelPackageKind.MODEL + + def test_stop_reason_values(self): + assert StopReason("length") is StopReason.LENGTH + assert StopReason("stop") is StopReason.STOP + + def test_tensor_dtype_values(self): + assert {d.value for d in TensorDtype} == { + "float32", + "float16", + "bfloat16", + "int32", + "int64", + } + + +class TestModelPackageInfo: + def test_from_summary(self): + info = ModelPackageInfo._from_summary( + { + "ModelPackageArn": "arn:aws:sagemaker:us-west-2:123:model-package/g/3", + "ModelPackageVersion": 3, + "ModelPackageDescription": "after epoch 2", + } + ) + assert info.model_package_arn == "arn:aws:sagemaker:us-west-2:123:model-package/g/3" + assert info.version == 3 + assert info.description == "after epoch 2" + # ListModelPackages summaries carry no kind — it stays unresolved. + assert info.kind is None + + def test_kind_accepts_enum(self): + info = ModelPackageInfo(model_package_arn="arn:x", kind=ModelPackageKind.CHECKPOINT) + assert info.kind is ModelPackageKind.CHECKPOINT + + def test_extra_fields_rejected(self): + with pytest.raises(ValidationError): + ModelPackageInfo(model_package_arn="arn:x", not_a_field=1) + + +class TestModelPackageDeletion: + def test_defaults_describe_a_retained_location(self): + """The default shape is the common one: package gone, objects kept.""" + result = ModelPackageDeletion(model_package_arn="arn:x", model_package_deleted=True) + assert result.s3_uri is None + assert result.objects_deleted == 0 + assert result.artifacts_retained is True + assert result.retained_reason is None + + def test_extra_fields_rejected(self): + with pytest.raises(ValidationError): + ModelPackageDeletion( + model_package_arn="arn:x", model_package_deleted=True, not_a_field=1 + ) + + +class TestTokenLogprob: + def test_from_api(self): + tl = TokenLogprob._from_api({"tokenId": 25, "logprob": -4.22}) + assert tl.token_id == 25 + assert tl.logprob == -4.22 + + def test_extra_fields_rejected(self): + with pytest.raises(ValidationError): + TokenLogprob(token_id=1, logprob=-0.5, not_a_field=1) + + +class TestSampledSequence: + def test_from_api(self): + seq = SampledSequence._from_api( + {"stopReason": "length", "tokens": [4, 5, 6], "logprobs": [-0.7, -0.5, -0.3]} + ) + assert seq.stop_reason is StopReason.LENGTH + assert seq.tokens == [4, 5, 6] + assert seq.logprobs == [-0.7, -0.5, -0.3] + + def test_from_api_without_logprobs(self): + seq = SampledSequence._from_api({"stopReason": "stop", "tokens": [1]}) + assert seq.logprobs is None + + def test_from_api_ignores_token_usage(self): + """tokenUsage is on the wire but deliberately not surfaced. + + The service returns it per sequence; the SDK drops it rather than + committing to public surface it would then have to keep. + """ + seq = SampledSequence._from_api( + { + "stopReason": "stop", + "tokens": [1, 2], + "tokenUsage": {"prefill": 7, "decode": 2}, + } + ) + assert not hasattr(seq, "token_usage") + + +class TestSampleResponseFromApi: + def test_full_response(self): + # Wire shape: topkPromptLogprobs is a list (per prompt position) of + # TokenLogprob objects, converted server-side from [tokenId, logprob] + # pairs. + r = SampleResponse._from_api( + { + "sequences": [{"stopReason": "length", "tokens": [4, 5]}], + "promptLogprobs": [-8.25, -11.30], + "topkPromptLogprobs": [ + [{"tokenId": 25, "logprob": -4.22}, {"tokenId": 654, "logprob": -4.47}] + ], + } + ) + assert len(r.sequences) == 1 + assert r.sequences[0].tokens == [4, 5] + assert r.prompt_logprobs == [-8.25, -11.30] + assert r.topk_prompt_logprobs[0][0].token_id == 25 + assert r.topk_prompt_logprobs[0][1].logprob == -4.47 + + def test_minimal_response(self): + r = SampleResponse._from_api({"sequences": []}) + assert r.sequences == [] + assert r.prompt_logprobs is None + assert r.topk_prompt_logprobs is None + + def test_stop_reason_is_parsed(self): + r = SampleResponse._from_api( + {"sequences": [{"stopReason": "stop", "tokens": [10, 20, 30]}]} + ) + assert r.sequences[0].tokens == [10, 20, 30] + assert r.sequences[0].stop_reason.value == "stop" + + def test_empty_response_no_field(self): + """``sequences`` absent — returns empty list rather than raising.""" + r = SampleResponse._from_api({}) + assert r.sequences == [] + + +class TestVpcConfig: + def test_to_api_shape(self): + vc = VpcConfig(subnets=["subnet-1", "subnet-2"], security_group_ids=["sg-1"]) + assert vc._to_api() == { + "Subnets": ["subnet-1", "subnet-2"], + "SecurityGroupIds": ["sg-1"], + } + + def test_from_any_passthrough(self): + vc = VpcConfig(subnets=["subnet-1"], security_group_ids=["sg-1"]) + assert VpcConfig.from_any(vc) is vc + + def test_from_any_dict(self): + vc = VpcConfig.from_any({"subnets": ["subnet-1"], "security_group_ids": ["sg-1", "sg-2"]}) + assert vc.subnets == ["subnet-1"] + assert vc.security_group_ids == ["sg-1", "sg-2"] + + def test_from_any_compatible(self): + class _Compat: + subnets = ("subnet-1",) + security_group_ids = ("sg-1",) + + vc = VpcConfig.from_any(_Compat()) + assert vc.subnets == ["subnet-1"] + assert vc.security_group_ids == ["sg-1"] + + def test_rejects_unknown_fields(self): + with pytest.raises(ValidationError): + VpcConfig(subnets=["s"], security_group_ids=["g"], vpc_id="vpc-1") + + +class TestResponseTypeNaming: + """The per-operation response-type names ARE the canonical public names.""" + + def test_legacy_names_removed(self): + import sagemaker.train.training_session as ts + import sagemaker.train.training_session.types as types_mod + + for legacy in ("TrainingResult", "SampleResult", "ComputeLogprobsResult"): + assert not hasattr(ts, legacy) + assert not hasattr(types_mod, legacy) + + def test_response_types_exported_from_package_root(self): + import sagemaker.train.training_session as ts + + assert ts.ForwardBackwardOutput is ForwardBackwardOutput + assert ts.OptimStepResponse is OptimStepResponse + assert ts.SaveStateResponse is SaveStateResponse + assert ts.SaveWeightsResponse is SaveWeightsResponse + assert ts.LoadStateResponse is LoadStateResponse + assert ts.SampleResponse is SampleResponse + + def test_each_save_type_reads_only_its_own_operation_key(self): + """The name of each response type is the operation it reads. + + Crossing the two keys must yield nothing rather than the other's ARN. + """ + assert ( + SaveStateResponse._from_api( + {"saveWeights": {"modelPackageArn": "arn:wrong"}} + ).model_package_arn + is None + ) + assert ( + SaveWeightsResponse._from_api( + {"saveState": {"modelPackageArn": "arn:wrong"}} + ).model_package_arn + is None + ) + + def test_retired_weights_era_names_are_gone(self): + """The old names must not linger as aliases: one name per operation.""" + import sagemaker.train.training_session as ts + import sagemaker.train.training_session.types as types_mod + + for retired in ("SaveWeightsForSamplerResponse", "LoadWeightsResponse"): + assert not hasattr(ts, retired) + assert not hasattr(types_mod, retired) + + +class TestModelInputFromAnyRejectsUnknownDicts: + """An unrecognized dict must not become an input with no chunks. + + That shape is valid on the wire, so a dropped training example produced no + error anywhere: the operation succeeded and simply trained on nothing. + """ + + def test_an_unrecognized_dict_raises(self): + with pytest.raises(ValueError, match="cannot build a model input"): + ModelInput.from_any({"totally": "unknown"}) + + def test_the_message_names_the_keys_it_accepts(self): + with pytest.raises(ValueError, match="chunks"): + ModelInput.from_any({"totally": "unknown"}) + + def test_a_bare_chunk_is_read_as_one_chunk(self): + """``{"tokens": [...]}`` is a chunk, and used to silently become empty.""" + assert ModelInput.from_any({"tokens": [1, 2, 3]}).to_ints() == [1, 2, 3] + + def test_the_recognized_shapes_still_work(self): + assert ModelInput.from_any({"chunks": [{"tokens": [4, 5]}]}).to_ints() == [4, 5] + assert ModelInput.from_any({"prompt": "hello"}).prompt == "hello" + + +class TestTensorDataRoundTrip: + """Values are held outside the model fields, but a dump must still carry them.""" + + def test_model_dump_includes_the_data(self): + tensor = TensorData([1.0, 2.0], dtype=TensorDtype.FLOAT32, shape=[2]) + + assert tensor.model_dump()["data"] == [1.0, 2.0] + + def test_a_dump_round_trips(self): + tensor = TensorData([1, 2, 3], dtype=TensorDtype.INT64, shape=[3]) + + restored = TensorData.model_validate(tensor.model_dump()) + + assert restored.data == [1, 2, 3] + assert restored.dtype is TensorDtype.INT64 + assert restored.shape == [3] + + def test_a_tensor_like_source_is_not_materialized_by_construction(self): + """The laziness the private storage exists for has to survive.""" + + class Lazy: + def __init__(self): + self.reads = 0 + + def tolist(self): + self.reads += 1 + return [7, 8] + + source = Lazy() + tensor = TensorData(source, dtype=TensorDtype.INT64, shape=[2]) + assert source.reads == 0 + + assert tensor.data == [7, 8] + assert source.reads == 1 + + +class TestDatumRejectsBothInputs: + """``model_input`` and ``prompt`` are documented as exclusive; enforce it. + + With both set the wire shape came from ``model_input`` and the prompt went + nowhere, so a caller who meant to send it trained on something else. + """ + + def test_both_set_raises(self): + with pytest.raises(ValidationError, match="not both"): + Datum(model_input=ModelInput.from_ints([1, 2]), prompt=[3, 4]) + + def test_either_alone_is_accepted(self): + assert Datum(model_input=ModelInput.from_ints([1, 2])).prompt is None + assert Datum(prompt=[3, 4]).model_input is None + + +class TestModelInputAppendRejectsPendingText: + """Appending to a text-built input dropped the text and kept the chunk.""" + + def test_appending_to_text_raises(self): + with pytest.raises(ValueError, match="has not been tokenized"): + ModelInput.from_text("hello").append(EncodedTextChunk(tokens=[1])) + + def test_appending_to_token_input_still_works(self): + appended = ModelInput.from_ints([1, 2]).append(EncodedTextChunk(tokens=[3])) + + assert appended.to_ints() == [1, 2, 3] + + def test_appending_keeps_a_token_id_prompt(self): + """A prompt of token IDs is content, not pending text, so it survives. + + Only text needs a tokenizer; token IDs are already the input's chunks as + far as the wire is concerned, and dropping them loses training tokens + with nothing raised. + """ + appended = ModelInput(prompt=[1, 2, 3]).append(EncodedTextChunk(tokens=[9])) + + assert appended.to_ints() == [1, 2, 3, 9] + assert appended._to_api() == { + "chunks": [ + {"tokens": [1, 2, 3], "type": "encoded_text"}, + {"tokens": [9], "type": "encoded_text"}, + ] + } + + +class TestDatumFromAnyAppliesTheSameExclusivity: + """``from_any`` is the path every batch element takes, so it must agree. + + The constructor rejected both inputs while this path kept ``model_input`` + and dropped ``prompt``, which is the same silent loss one layer down. + """ + + def test_a_dict_naming_both_is_rejected(self): + with pytest.raises(ValidationError, match="not both"): + Datum.from_any({"model_input": {"chunks": [{"tokens": [1, 2]}]}, "prompt": "hi"}) + + def test_an_empty_prompt_is_not_a_second_input(self): + """``prompt=""`` carries nothing, so it is not a conflict to reject. + + Objects coerced through this path commonly default their prompt to an + empty string, and failing those would break a call that loses nothing. + """ + datum = Datum.from_any({"model_input": {"chunks": [{"tokens": [1, 2]}]}, "prompt": ""}) + + assert datum.model_input is not None + assert datum.model_input.to_ints() == [1, 2] + + def test_an_object_with_an_empty_prompt_and_model_input_is_accepted(self): + class Carrier: + model_input = ModelInput.from_ints([4, 5]) + prompt = "" + loss_fn_inputs = None + + datum = Datum.from_any(Carrier()) + + assert datum.model_input is not None + assert datum.model_input.to_ints() == [4, 5] + + +class TestTensorDataFollowsTheHashProtocol: + """Mutable values mean unhashable, and that has to be reported honestly. + + Supplying a ``__hash__`` that raises registers the type as hashable, so a + caller branching on ``isinstance(x, Hashable)`` takes the hashable path and + then crashes. + """ + + def test_it_reports_itself_unhashable(self): + import collections.abc + + assert not isinstance(TensorData([1, 2]), collections.abc.Hashable) + + def test_hashing_it_raises_type_error(self): + with pytest.raises(TypeError, match="unhashable"): + hash(TensorData([1, 2])) + + +class TestTensorDataEqualityIgnoresAccessOrder: + """Values are materialized lazily, and equality must not depend on that.""" + + def test_a_tensor_like_source_equals_the_same_list(self): + class Lazy: + def tolist(self): + return [1.0, 2.0] + + from_source = TensorData(Lazy(), dtype=TensorDtype.FLOAT32, shape=[2]) + from_list = TensorData([1.0, 2.0], dtype=TensorDtype.FLOAT32, shape=[2]) + + assert from_source == from_list + + def test_different_values_are_still_unequal(self): + assert TensorData([1, 2], dtype=TensorDtype.INT64, shape=[2]) != TensorData( + [9, 9], dtype=TensorDtype.INT64, shape=[2] + ) + + def test_a_different_dtype_is_unequal(self): + assert TensorData([1, 2], dtype=TensorDtype.INT64, shape=[2]) != TensorData( + [1, 2], dtype=TensorDtype.INT32, shape=[2] + ) + + +class TestOldLogprobsIsPinnedLikeLogprobs: + """Both logprob inputs are floats, so neither may be inferred from its values. + + ``old_logprobs`` was absent from the dtype table, so a batch whose values + happened to be whole numbers went out as int64 while ``logprobs`` stayed + float32 -- the same quantity described two ways in one request. + """ + + def test_whole_numbers_still_serialize_as_float32(self): + datum = Datum( + model_input=ModelInput.from_ints([1, 2, 3]), + loss_fn_inputs={"old_logprobs": [0, -1, -2], "logprobs": [0, -1, -2]}, + ) + + sent = datum._to_api()["lossFnInputs"] + + assert sent["old_logprobs"]["dtype"] == "float32" + assert sent["logprobs"]["dtype"] == "float32" + + +class TestModelInputFromAnyRejectsBothInputs: + """``from_any`` must answer chunks-vs-prompt the way the constructor does. + + It took ``chunks`` and dropped ``prompt``, so a dict carrying an empty chunk + list and a real prompt built an empty input that the service accepts, + dropping the example with nothing raised. + """ + + def test_both_keys_raises(self): + with pytest.raises(ValueError, match="not both"): + ModelInput.from_any({"chunks": [{"tokens": [1]}], "prompt": [2, 3]}) + + def test_an_empty_chunk_list_with_a_prompt_also_raises(self): + with pytest.raises(ValueError, match="not both"): + ModelInput.from_any({"chunks": [], "prompt": [2, 3]}) + + def test_a_dumped_chunk_input_still_round_trips(self): + """``model_dump()`` carries ``prompt=None``, which is not a second input.""" + dumped = ModelInput.from_ints([1, 2]).model_dump() + + assert ModelInput.from_any(dumped).to_ints() == [1, 2] + + def test_either_key_alone_is_accepted(self): + assert ModelInput.from_any({"chunks": [{"tokens": [1]}]}).to_ints() == [1] + assert ModelInput.from_any({"prompt": [2, 3]}).to_ints() == [2, 3] + + def test_a_text_built_dump_still_round_trips(self): + """``from_text`` dumps carry ``chunks=None``, which is not a second input. + + A dump names every field, so the exclusivity rule has to read the values + rather than the keys. + """ + dumped = ModelInput.from_text("hello").model_dump() + + assert dumped["chunks"] is None + assert ModelInput.from_any(dumped)._pending_text == "hello" + + def test_a_null_chunks_key_alongside_a_prompt_is_accepted(self): + assert ModelInput.from_any({"chunks": None, "prompt": [1, 2]}).to_ints() == [1, 2] + + +class TestTextInputHasNoTokenIdsYet: + """``to_ints``/``length`` reported a full prompt as empty. + + ``_to_api`` already refuses a text-built input; these two returned ``[]`` and + ``0``, so a token budget counted the prompt as absent instead of failing. + """ + + def test_to_ints_raises(self): + with pytest.raises(ValueError, match="has not been tokenized"): + ModelInput.from_text("hello world").to_ints() + + def test_length_raises(self): + with pytest.raises(ValueError, match="has not been tokenized"): + _ = ModelInput.from_text("hello world").length + + def test_a_prompt_that_is_not_token_ids_also_raises(self): + """``_to_api`` already refuses these, so reading them as empty disagreed. + + Any prompt that is not a token ID list has none to return, not zero of + them, and a token budget counted the difference. + """ + with pytest.raises(ValueError, match="not token IDs"): + ModelInput(prompt=object()).to_ints() + + def test_a_token_input_is_unaffected(self): + assert ModelInput.from_ints([1, 2, 3]).to_ints() == [1, 2, 3] + assert ModelInput.from_ints([1, 2, 3]).length == 3 + + def test_an_empty_input_is_still_empty_rather_than_an_error(self): + assert ModelInput.empty().to_ints() == [] + assert ModelInput.empty().length == 0 diff --git a/sagemaker-train/tests/unit/train/training_session/test_utils.py b/sagemaker-train/tests/unit/train/training_session/test_utils.py new file mode 100644 index 0000000000..bd4871aab2 --- /dev/null +++ b/sagemaker-train/tests/unit/train/training_session/test_utils.py @@ -0,0 +1,1745 @@ +# 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/ +"""Unit tests for sagemaker.train.training_session._utils helpers.""" + +from __future__ import annotations + +import json +import logging +import re +import weakref +from unittest.mock import MagicMock + +import pytest + +from sagemaker.train.training_session._utils import ( + batch_to_datum_list, + combine_forward_backward_results, + _arn_partition, + build_lora_config, + describe_model_package_kind, + list_base_models_in_region, + list_model_package_summaries, + prompt_to_model_input, + resolve_base_model, + resolve_hf_model_id, + tokenize_to_datum, +) +from sagemaker.train.training_session.types import ( + Datum, + EncodedTextChunk, + ModelInput, + ModelPackageKind, + TensorData, +) + + +class TestBuildLoraConfig: + def test_basic(self): + assert build_lora_config(32, 64) == {"Rank": 32, "Alpha": 64} + + def test_same_rank_alpha(self): + assert build_lora_config(16, 16) == {"Rank": 16, "Alpha": 16} + + +class _FakeHub: + """Stubs describe_hub_content, returning a HubContentDocument with (or + without) a HuggingfaceId so tokenizer resolution can be exercised offline.""" + + def __init__(self, huggingface_id=None, raises=False, arn=None): + self._hf_id = huggingface_id + self._raises = raises + self._arn = ( + arn or "arn:aws:sagemaker:us-west-2:aws:hub-content/SageMakerPublicHub/Model/stub/1.0" + ) + + def describe_hub_content(self, **kwargs): + if self._raises: + raise RuntimeError("ResourceNotFound") + doc = {} if self._hf_id is None else {"HuggingfaceId": self._hf_id} + return {"HubContentArn": self._arn, "HubContentDocument": json.dumps(doc)} + + +_MODEL_B_ARN = ( + "arn:aws:sagemaker:us-west-2:aws:hub-content/SageMakerPublicHub/Model/" + "hubcontent-vlm-model-b/1.0.0" +) +_MODEL_A_ARN = ( + "arn:aws:sagemaker:us-west-2:aws:hub-content/SageMakerPublicHub/Model/" + "hubcontent-text-model-a/3.15.0" +) +_MODEL_C_ARN = ( + "arn:aws:sagemaker:us-west-2:aws:hub-content/SageMakerPublicHub/Model/" + "hubcontent-text-model-c/1.2.0" +) + + +class _FakeCatalogHub: + """Stubs list_hub_contents + describe_hub_content over a model catalog. + + Mirrors the real Hub shape: content names are NOT derivable from model ids + (e.g. ``Provider-B/Model-B`` lives at ``hubcontent-vlm-model-b``), so + model-id resolution must go through the published model id. + + Args: + models: ``{hub_content_name: (arn, hf_id_or_None)}``. + advertise_hf_id: Put ``@huggingface-id:`` on the list summary, as the + real Hub does. False leaves only the recipe keyword, so the + Describe fallback is what has to supply the id. + region: Reported as ``meta.region_name``, as a real client does. The + catalogue is cached per client rather than per Region, so this + changes nothing about caching; None leaves the client region-less. + keywords_by_name: Full keyword override per content name, for models + that are not fine-tunable at all. + """ + + _RECIPE_KW = "@recipe:finetuning_mtrl_lora" + + def __init__(self, models, *, advertise_hf_id=False, region=None, keywords_by_name=None): + self._models = models + self._advertise_hf_id = advertise_hf_id + self._keywords_by_name = keywords_by_name or {} + self.list_calls = 0 + self.describe_calls: list[str] = [] + if region is not None: + self.meta = MagicMock(region_name=region) + + def _keywords(self, name): + if name in self._keywords_by_name: + return self._keywords_by_name[name] + keywords = [self._RECIPE_KW] + hf_id = self._models[name][1] + if self._advertise_hf_id and hf_id: + keywords.append(f"@huggingface-id:{hf_id.lower()}") + return keywords + + def list_hub_contents(self, **kwargs): + self.list_calls += 1 + return { + "HubContentSummaries": [ + {"HubContentName": n, "HubContentSearchKeywords": self._keywords(n)} + for n in self._models + ] + } + + def describe_hub_content(self, **kwargs): + name = kwargs["HubContentName"] + self.describe_calls.append(name) + if name not in self._models: + raise RuntimeError("ResourceNotFound") + arn, hf_id = self._models[name] + doc = {} if hf_id is None else {"HuggingfaceId": hf_id} + return {"HubContentArn": arn, "HubContentDocument": json.dumps(doc)} + + +class _SlottedHub: + """A Hub client that cannot be referenced weakly, so it cannot be cached. + + ``__slots__`` on a class deriving straight from ``object`` is what drops + ``__weakref__``; a subclass of a normal class would inherit it back. + """ + + __slots__ = ("list_calls", "describe_calls") + + def __init__(self): + self.list_calls = 0 + self.describe_calls = [] + + def list_hub_contents(self, **kwargs): + self.list_calls += 1 + return { + "HubContentSummaries": [ + { + "HubContentName": "hubcontent-text-model-a", + "HubContentSearchKeywords": [ + "@recipe:finetuning_mtrl_lora", + "@huggingface-id:provider-a/model-a", + ], + } + ] + } + + def describe_hub_content(self, **kwargs): + self.describe_calls.append(kwargs["HubContentName"]) + return { + "HubContentArn": _MODEL_A_ARN, + "HubContentDocument": json.dumps({"HuggingfaceId": "provider-a/model-a"}), + } + + +class TestListBaseModels: + """The listing reports each model by the id its provider gave it.""" + + _MODELS = { + "hubcontent-vlm-model-b": (_MODEL_B_ARN, "Provider-B/Model-B"), + "hubcontent-text-model-a": (_MODEL_A_ARN, "provider-a/model-a"), + } + + def test_reports_the_id_as_its_provider_capitalises_it(self): + # The keyword is free but lowercased; only the document has the exact + # case, which is the spelling a customer copies. + hub = _FakeCatalogHub(self._MODELS, advertise_hf_id=True) + assert list_base_models_in_region(hub) == [ + "provider-a/model-a", + "Provider-B/Model-B", + ] + assert sorted(hub.describe_calls) == sorted(self._MODELS) + + def test_reads_each_id_once_per_client(self): + # Re-listing pays for the catalogue, not for every id in it again. + hub = _FakeCatalogHub(self._MODELS, advertise_hf_id=True, region="us-west-2") + list_base_models_in_region(hub) + hub.describe_calls.clear() + + assert list_base_models_in_region(hub) == [ + "provider-a/model-a", + "Provider-B/Model-B", + ] + assert hub.describe_calls == [] + + def test_reads_only_the_ids_it_does_not_already_know(self): + # A model enabled since the last call has its id read; the rest do not. + # dict(): the fake holds what it is given, and this one grows. + hub = _FakeCatalogHub(dict(self._MODELS), advertise_hf_id=True, region="us-west-2") + list_base_models_in_region(hub) + hub.describe_calls.clear() + hub._models["hubcontent-text-model-c"] = (_MODEL_C_ARN, "Provider-C/Model-C") + + assert "Provider-C/Model-C" in list_base_models_in_region(hub) + assert hub.describe_calls == ["hubcontent-text-model-c"] + + def test_a_client_that_cannot_be_referenced_weakly_still_lists(self): + # Nothing is memoised for it, so it re-reads rather than sharing ids. + hub = _SlottedHub() + assert list_base_models_in_region(hub) == ["provider-a/model-a"] + assert list_base_models_in_region(hub) == ["provider-a/model-a"] + assert hub.describe_calls == ["hubcontent-text-model-a"] * 2 + + def test_reports_the_document_id_with_no_keyword_to_fall_back_on(self): + # No keyword means nothing to degrade to: the document path stands alone. + hub = _FakeCatalogHub(self._MODELS) + assert list_base_models_in_region(hub) == [ + "provider-a/model-a", + "Provider-B/Model-B", + ] + assert sorted(hub.describe_calls) == sorted(self._MODELS) + + def test_falls_back_to_hub_content_name_without_any_hf_id(self): + # A model publishing no model id anywhere is still listed, under the one + # identifier it does have, which create_lora_training_client accepts. + hub = _FakeCatalogHub({"some-private-model": (_MODEL_B_ARN, None)}) + assert list_base_models_in_region(hub) == ["some-private-model"] + + def test_excludes_models_without_a_finetuning_recipe(self): + hub = _FakeCatalogHub( + { + "hubcontent-text-model-a": (_MODEL_A_ARN, "provider-a/model-a"), + "some-eval-only-model": (_MODEL_B_ARN, "org/eval-only"), + }, + advertise_hf_id=True, + keywords_by_name={"some-eval-only-model": ["@recipe:evaluation_mtrlevaluation"]}, + ) + assert list_base_models_in_region(hub) == ["provider-a/model-a"] + + def test_paginates(self): + hub = _FakeCatalogHub(self._MODELS, advertise_hf_id=True) + pages = [ + { + "HubContentSummaries": [ + { + "HubContentName": "hubcontent-text-model-a", + "HubContentSearchKeywords": [ + "@recipe:finetuning_mtrl_lora", + "@huggingface-id:provider-a/model-a", + ], + } + ], + "NextToken": "tok", + }, + { + "HubContentSummaries": [ + { + "HubContentName": "hubcontent-vlm-model-b", + "HubContentSearchKeywords": [ + "@recipe:finetuning_mtrl_lora", + "@huggingface-id:provider-b/model-b", + ], + } + ] + }, + ] + hub.list_hub_contents = MagicMock(side_effect=pages) + assert list_base_models_in_region(hub) == [ + "provider-a/model-a", + "Provider-B/Model-B", + ] + assert hub.list_hub_contents.call_args_list[1][1]["NextToken"] == "tok" + + def test_rescans_the_hub_on_every_call(self): + # A model enabled since the last call must show up without restarting + # the process, so the listing never answers from the cache. + hub = _FakeCatalogHub(self._MODELS, advertise_hf_id=True, region="us-west-2") + list_base_models_in_region(hub) + list_base_models_in_region(hub) + assert hub.list_calls == 2 + + def test_sorts_case_insensitively(self): + # A provider that capitalises its name would otherwise sort ahead of + # every lowercase one, and this list is read by eye. + hub = _FakeCatalogHub(self._MODELS) + assert list_base_models_in_region(hub) == [ + "provider-a/model-a", + "Provider-B/Model-B", + ] + + def test_ignores_a_keyword_that_is_not_a_model_id(self): + # The keyword value is whatever the Hub publishes. A single-segment value + # is not routed to the id path, so listing it would name it as supported + # in the very error that rejects it; the content name is listed instead. + hub = _FakeCatalogHub({"hubcontent-single": (_MODEL_A_ARN, None)}) + hub._keywords_by_name = { + "hubcontent-single": ["@recipe:finetuning_mtrl_lora", "@huggingface-id:singleseg"] + } + + listed = list_base_models_in_region(hub) + + assert listed == ["hubcontent-single"] + # And the whole point: every id the listing reports resolves. + assert resolve_base_model(listed[0], hub)[0] == _MODEL_A_ARN + + def test_a_failed_id_read_costs_case_and_never_the_model(self): + # Falls back to the keyword's copy of the same id: lowercase, still + # every model, still resolves. + hub = _FakeCatalogHub(self._MODELS, advertise_hf_id=True, region="us-west-2") + hub.describe_hub_content = MagicMock( + side_effect=RuntimeError("ThrottlingException: Rate exceeded") + ) + + listed = list_base_models_in_region(hub) + + assert listed == ["provider-a/model-a", "provider-b/model-b"] + del hub.describe_hub_content # drops the stub, restoring the fake's own + assert resolve_base_model("provider-b/model-b", hub)[0] == _MODEL_B_ARN + + def test_a_failed_id_read_is_retried_rather_than_remembered(self): + # Memoising it would pin the model to lowercase for the client's life. + hub = _FakeCatalogHub(self._MODELS, advertise_hf_id=True, region="us-west-2") + hub.describe_hub_content = MagicMock( + side_effect=RuntimeError("ThrottlingException: Rate exceeded") + ) + list_base_models_in_region(hub) + del hub.describe_hub_content # drops the stub, restoring the fake's own + + assert list_base_models_in_region(hub) == [ + "provider-a/model-a", + "Provider-B/Model-B", + ] + + def test_a_failed_id_read_with_no_keyword_falls_back_to_the_content_name(self): + # Nothing else names it, and that name resolves, so it is still listed. + hub = _FakeCatalogHub({"hubcontent-text-model-a": (_MODEL_A_ARN, "Provider-A/Model-A")}) + hub.describe_hub_content = MagicMock(side_effect=RuntimeError("AccessDeniedException")) + + assert list_base_models_in_region(hub) == ["hubcontent-text-model-a"] + + def test_duplicate_id_resolves_deterministically_and_warns(self, caplog): + # Two entries publishing one id: only one can be listed, and which one is + # not the caller's choice, so the drop is logged rather than silent. + hub = _FakeCatalogHub( + { + "hubcontent-model-v2": (_MODEL_B_ARN, "provider-a/model-a"), + "hubcontent-model-v1": (_MODEL_A_ARN, "provider-a/model-a"), + }, + advertise_hf_id=True, + region="us-west-2", + ) + with caplog.at_level(logging.WARNING): + assert list_base_models_in_region(hub) == ["provider-a/model-a"] + + assert "published by more than one model in this region" in caplog.text + assert "hubcontent-model-v1" in caplog.text + assert "hubcontent-model-v2" in caplog.text + # Deterministic across pagination order, which is what the tie-break buys. + assert resolve_base_model("provider-a/model-a", hub)[0] == _MODEL_A_ARN + + +class TestResolveBaseModel: + """Tests the 3 resolution paths: ARN, Hub content name, model id.""" + + def _hub(self, **kwargs): + return _FakeCatalogHub( + { + "hubcontent-vlm-model-b": (_MODEL_B_ARN, "Provider-B/Model-B"), + "hubcontent-text-model-a": (_MODEL_A_ARN, "provider-a/model-a"), + }, + **kwargs, + ) + + def test_arn_passthrough(self): + result_arn, hf_id = resolve_base_model(_MODEL_B_ARN, self._hub()) + assert result_arn == _MODEL_B_ARN + assert hf_id is None + + def test_hub_content_name(self): + result_arn, hf_id = resolve_base_model("hubcontent-vlm-model-b", self._hub()) + assert result_arn == _MODEL_B_ARN + assert hf_id == "Provider-B/Model-B" + + def test_huggingface_id_matched_via_published_id(self): + # Hub content name is NOT slash-to-dash of the model id — resolution must + # match the HuggingfaceId published on each supported model. + result_arn, hf_id = resolve_base_model("Provider-B/Model-B", self._hub()) + assert result_arn == _MODEL_B_ARN + assert hf_id == "Provider-B/Model-B" + + def test_huggingface_id_case_insensitive(self): + result_arn, hf_id = resolve_base_model("provider-b/model-b", self._hub()) + assert result_arn == _MODEL_B_ARN + assert hf_id == "Provider-B/Model-B" # canonical published id returned + + def test_case_insensitive_match_returns_the_catalogued_id_not_the_caller_s(self): + # published ids are case-sensitive, so only the published spelling + # loads a tokenizer. When the keyword carries the id and the Hub document + # omits it, the caller's spelling must still not be what comes back. + class _KeywordOnlyHub(_FakeCatalogHub): + """Advertises the canonical id; its Hub document carries none.""" + + def _keywords(self, name): + return ["@recipe:finetuning_mtrl_lora", "@huggingface-id:Provider-A/Model-A"] + + def describe_hub_content(self, **kwargs): + self.describe_calls.append(kwargs["HubContentName"]) + return {"HubContentArn": _MODEL_A_ARN, "HubContentDocument": "{}"} + + hub = _KeywordOnlyHub( + {"hubcontent-text-model-a": (_MODEL_A_ARN, "Provider-A/Model-A")}, + region="us-west-2", + ) + + _, hf_id = resolve_base_model("provider-a/MODEL-a", hub) + + assert hf_id == "Provider-A/Model-A" + + def test_excluded_hub_content_name_does_not_resolve(self): + # The listing filters this model out, so resolving it by Hub content name + # would let a name the listing never shows start a session anyway. + hub = _FakeCatalogHub( + { + "hubcontent-text-model-a": (_MODEL_A_ARN, "provider-a/model-a"), + "amazon-nova-model-x": (_MODEL_B_ARN, "amazon/nova-model-x"), + }, + advertise_hf_id=True, + region="us-west-2", + ) + with pytest.raises(ValueError, match="not supported for fine-tuning"): + resolve_base_model("amazon-nova-model-x", hub) + # Refused on the name alone: the cheapest rejection on this path must not + # pay for a Hub scan just to format its message. + assert hub.describe_calls == [] + assert hub.list_calls == 0 + + def test_a_name_merely_ending_in_an_excluded_one_still_resolves(self): + # The exclusion matches whole dash-separated parts. Matched as a bare + # substring, a longer word ending in an excluded name would be caught + # with it, and the model would be neither listed nor accepted. + hub = _FakeCatalogHub( + {"vendor-supernova-lite": (_MODEL_A_ARN, "vendor/supernova-lite")}, + advertise_hf_id=True, + region="us-west-2", + ) + assert list_base_models_in_region(hub) == ["vendor/supernova-lite"] + assert resolve_base_model("vendor-supernova-lite", hub)[0] == _MODEL_A_ARN + + def test_huggingface_id_unsupported_raises_with_supported_list(self): + # re.DOTALL: the supported set is listed one model per line. + with pytest.raises( + ValueError, match=re.compile("provider-z/model-z.*Supported models", re.DOTALL) + ): + resolve_base_model("provider-z/model-z", self._hub()) + + def test_unsupported_error_lists_one_model_per_line(self): + # A joined single line grows past what a terminal or log viewer wraps. + with pytest.raises(ValueError) as excinfo: + resolve_base_model("provider-z/model-z", self._hub(advertise_hf_id=True)) + lines = str(excinfo.value).splitlines() + assert lines[1] == "Supported models:" + assert lines[2:] == [ + " provider-a/model-a", + " Provider-B/Model-B", + ] + + def test_unsupported_error_lists_ids_only(self): + # Every value listed is one the caller can pass straight back. Naming the + # models a second way as well would answer "what should I have typed" + # with two answers, one of which the listing never shows. + with pytest.raises(ValueError) as excinfo: + resolve_base_model("provider-z/model-z", self._hub(advertise_hf_id=True)) + message = str(excinfo.value) + assert "provider-a/model-a" in message + assert "hubcontent" not in message + + def test_hub_name_not_found_raises(self): + with pytest.raises(ValueError, match="Could not resolve"): + resolve_base_model("nonexistent-model", self._hub()) + + def test_hub_name_not_found_lists_supported_models(self): + # A display name is neither an id nor a content name, and is the likeliest + # mistake here, so the message says what this region does offer. + with pytest.raises(ValueError, match="Supported models"): + resolve_base_model("model-a", self._hub(advertise_hf_id=True)) + + def test_hf_id_resolution_reuses_the_cached_catalogue(self): + # The first resolution lists the region's models to find the Hub content + # name behind the id; later ones only describe that name. + hub = self._hub(advertise_hf_id=True, region="us-west-2") + resolve_base_model("provider-a/model-a", hub) + hub.describe_calls.clear() + + resolve_base_model("provider-b/model-b", hub) + + assert hub.list_calls == 1 + # The scan's reads are behind us; this describes only the name it needs. + assert hub.describe_calls == ["hubcontent-vlm-model-b"] + + def test_one_client_never_answers_out_of_anothers_catalogue(self): + # A catalogue is kept under the client that produced it, so it cannot + # answer for a client reading a different Hub. That covers every axis at + # once: Region, endpoint, and credentials all change the answer, and the + # client is the only thing that carries all three. + # + # Asserted on the call counts, NOT the error text: were the catalogue + # shared, the second lookup would describe the first's Hub content name, + # fail, self-heal into a rescan, and raise a message matching "Could not + # resolve" just as a legitimate miss does. The describe tells them apart. + first = self._hub(advertise_hf_id=True, region="us-west-2") + resolve_base_model("provider-b/model-b", first) + second = _FakeCatalogHub( + {"hubcontent-text-model-a": (_MODEL_A_ARN, "provider-a/model-a")}, + advertise_hf_id=True, + region="us-east-1", + ) + with pytest.raises(ValueError, match="Could not resolve"): + resolve_base_model("provider-b/model-b", second) + # Never reached for the other client's Hub content name (it does describe + # its own, to read that model's id). + assert "hubcontent-vlm-model-b" not in second.describe_calls + # And one scan: its own miss is final, not retried against a cache. + assert second.list_calls == 1 + + def test_two_clients_reading_one_region_still_hold_separate_catalogues(self): + # Same Region, but a different endpoint or different credentials means a + # different Hub. Neither is reliably readable off a client, so neither is + # trusted to build a key from. + first = self._hub(advertise_hf_id=True, region="us-west-2") + resolve_base_model("provider-b/model-b", first) + same_region = _FakeCatalogHub( + {"hubcontent-text-model-a": (_MODEL_A_ARN, "provider-a/model-a")}, + advertise_hf_id=True, + region="us-west-2", + ) + with pytest.raises(ValueError, match="Could not resolve"): + resolve_base_model("provider-b/model-b", same_region) + assert "hubcontent-vlm-model-b" not in same_region.describe_calls + + def test_a_client_reporting_no_region_is_still_cached(self): + # Caching does not depend on reading a Region off the client, so a client + # that reports none is not penalised with a rescan per resolution. + hub = self._hub(advertise_hf_id=True) + resolve_base_model("provider-a/model-a", hub) + resolve_base_model("provider-b/model-b", hub) + assert hub.list_calls == 1 + + def test_a_client_that_cannot_be_referenced_weakly_still_resolves(self): + # Not cacheable, so it rescans rather than being kept alive by this module + # or sharing another client's catalogue. It must still resolve. + # + # __slots__ on a class deriving straight from object is what removes + # weakref support; on a subclass it would be inherited from the base and + # this would quietly test the ordinary cached path instead. + hub = _SlottedHub() + with pytest.raises(TypeError): + weakref.ref(hub) + + assert resolve_base_model("provider-a/model-a", hub)[0] == _MODEL_A_ARN + assert resolve_base_model("provider-a/model-a", hub)[0] == _MODEL_A_ARN + assert hub.list_calls == 2 # uncached, so it rescans rather than sharing + + def test_warm_cache_still_resolves_a_newly_added_model(self): + # The self-heal must not only produce good errors: a model added after + # the cache warmed has to actually resolve. + hub = self._hub(advertise_hf_id=True, region="us-west-2") + resolve_base_model("provider-a/model-a", hub) + hub._models["hubcontent-text-model-c"] = (_MODEL_C_ARN, "provider-c/model-c") + + result_arn, hf_id = resolve_base_model("provider-c/model-c", hub) + + assert result_arn == _MODEL_C_ARN + assert hf_id == "provider-c/model-c" + + def test_unknown_id_on_a_cold_cache_scans_once(self): + # The miss came from a scan that just ran, so asking again would only put + # the same question to the Hub twice. + hub = self._hub(advertise_hf_id=True, region="us-west-2") + with pytest.raises(ValueError, match="Could not resolve"): + resolve_base_model("provider-z/model-z", hub) + assert hub.list_calls == 1 + + def test_stale_cached_name_is_rescanned(self): + # A Hub content name that has since been withdrawn must not surface as a + # name the caller never typed; the catalogue is rescanned instead. + hub = self._hub(advertise_hf_id=True, region="us-west-2") + resolve_base_model("provider-a/model-a", hub) + del hub._models["hubcontent-text-model-a"] + with pytest.raises(ValueError, match="Could not resolve base model 'provider-a/model-a'"): + resolve_base_model("provider-a/model-a", hub) + + +class TestResolveHfModelId: + def test_direct_hf_model_id_passthrough(self): + # A bare model id needs no Hub lookup. + assert resolve_hf_model_id("provider-c/model-c", _FakeHub()) == "provider-c/model-c" + + def test_hub_arn_reads_huggingface_id_from_doc(self): + arn = ( + "arn:aws:sagemaker:us-west-2:aws:hub-content/SageMakerPublicHub/Model/" + "hubcontent-text-model-a/3.15.0" + ) + hub = _FakeHub(huggingface_id="provider-a/model-a") + assert resolve_hf_model_id(arn, hub) == "provider-a/model-a" + + def test_model_without_huggingface_id_raises(self): + # A model whose Hub record carries no HuggingfaceId (tokenizer not public) + # raises a clear, actionable error. + with pytest.raises(ValueError, match="No public HuggingFace tokenizer"): + resolve_hf_model_id("some-private-model", _FakeHub(huggingface_id=None)) + + +class TestTokenizeToDatum: + def test_basic_omits_loss_fn_inputs(self): + # Without explicit loss_fn_inputs, no lossFnInputs key is emitted — + # the service rejects an auto-generated `labels` entry. + d = tokenize_to_datum([1, 2, 3]) + assert d["modelInput"]["chunks"] == [{"tokens": [1, 2, 3], "type": "encoded_text"}] + assert "lossFnInputs" not in d + + def test_with_loss_fn_inputs(self): + d = tokenize_to_datum([1], loss_fn_inputs={"advantages": [0.0]}) + assert d["lossFnInputs"] == { + "advantages": {"data": [0.0], "dtype": "float32", "shape": [1]} + } + + +class TestBatchToDatumList: + def test_empty(self): + assert batch_to_datum_list([]) == [] + + def test_token_id_lists(self): + result = batch_to_datum_list([[1, 2], [3, 4]]) + assert len(result) == 2 + assert result[0]["modelInput"]["chunks"][0]["tokens"] == [1, 2] + assert all("lossFnInputs" not in d for d in result) + + def test_single_sequence(self): + result = batch_to_datum_list([10, 20, 30]) + assert len(result) == 1 + assert result[0]["modelInput"]["chunks"][0]["tokens"] == [10, 20, 30] + + def test_datum_passthrough(self): + datum = {"modelInput": {"chunks": [{"tokens": [1], "type": "encoded_text"}]}} + result = batch_to_datum_list([datum]) + assert result == [datum] + + def test_datum_loss_inputs_do_not_mutate_caller(self): + # Filling loss_fn_inputs must not write into the caller's dicts, and + # each datum must get its own copy (no shared-reference aliasing). + d1 = {"modelInput": {"chunks": [{"tokens": [1], "type": "encoded_text"}]}} + d2 = {"modelInput": {"chunks": [{"tokens": [2], "type": "encoded_text"}]}} + batch = [d1, d2] + loss = {"labels": [0]} + + result = batch_to_datum_list(batch, loss_fn_inputs=loss) + + # Caller's dicts untouched. + assert "lossFnInputs" not in d1 and "lossFnInputs" not in d2 + # Each result carries the loss inputs, in its own dict (not aliased). + serialized = {"labels": {"data": [0], "dtype": "int64", "shape": [1]}} + assert result[0]["lossFnInputs"] == serialized + assert result[1]["lossFnInputs"] == serialized + assert result[0]["lossFnInputs"] is not result[1]["lossFnInputs"] + assert result[0]["lossFnInputs"] is not loss + # Adding a key to one datum's loss inputs must not bleed into the + # other datum or the caller's source dict. + result[0]["lossFnInputs"]["extra"] = 1 + assert "extra" not in result[1]["lossFnInputs"] + assert "extra" not in loss + + def test_datum_preserves_existing_loss_inputs(self): + # A datum that already carries lossFnInputs is left as-is. + own = {"labels": [7]} + datum = { + "modelInput": {"chunks": [{"tokens": [1], "type": "encoded_text"}]}, + "lossFnInputs": own, + } + result = batch_to_datum_list([datum], loss_fn_inputs={"labels": [0]}) + assert result[0]["lossFnInputs"] is own + + def test_strings_with_tokenizer(self): + tokenizer = MagicMock() + tokenizer.encode.side_effect = [[1, 2], [3, 4]] + result = batch_to_datum_list(["hello", "world"], tokenizer=tokenizer) + assert len(result) == 2 + assert result[0]["modelInput"]["chunks"][0]["tokens"] == [1, 2] + + def test_strings_without_tokenizer_raises(self): + with pytest.raises(ValueError, match="tokenizer"): + batch_to_datum_list(["hello"]) + + def test_strings_auto_load_tokenizer_via_loader(self): + # No tokenizer passed, but a loader is available — it's invoked lazily + # so plain-text batches work without a manual get_tokenizer() call. + loaded = MagicMock() + loaded.encode.side_effect = [[7, 8]] + calls = [] + + def loader(): + calls.append(1) + return loaded + + result = batch_to_datum_list(["hi"], tokenizer_loader=loader) + assert result[0]["modelInput"]["chunks"][0]["tokens"] == [7, 8] + assert calls == [1] # loader invoked exactly once, on demand + + def test_token_ids_never_invoke_loader(self): + # Token-ID input must NOT trigger tokenizer loading (no HF download). + def loader(): + raise AssertionError("tokenizer_loader must not be called for token IDs") + + result = batch_to_datum_list([[1, 2]], tokenizer_loader=loader) + assert result[0]["modelInput"]["chunks"][0]["tokens"] == [1, 2] + + def test_unsupported_type_raises(self): + with pytest.raises(TypeError, match="Unsupported"): + batch_to_datum_list([3.14]) + + +def _per_token(n: int) -> dict: + """Per-token loss inputs for an ``n``-token datum.""" + half = n // 2 + return { + "target_tokens": list(range(1, n + 1)), + "weights": [0.0] * half + [1.0] * (n - half), + "logprobs": [-0.1] * n, + "advantages": [1.0] * n, + } + + +_TOKENS10 = list(range(100, 110)) + + +def _wire_datum(tokens: list[int], loss_fn_inputs: dict | None = None) -> dict: + datum = {"modelInput": {"chunks": [{"tokens": tokens, "type": "encoded_text"}]}} + if loss_fn_inputs is not None: + datum["lossFnInputs"] = loss_fn_inputs + return datum + + +class TestPerTokenLengthValidation: + """Per-token lossFnInputs entries must match the datum's token count. + + A mismatch is accepted by the request path and only fails once the + operation runs, so these cover raising locally instead. + """ + + def test_shorter_than_model_input_raises(self): + batch = [Datum(model_input=ModelInput.from_ints(_TOKENS10), loss_fn_inputs=_per_token(5))] + with pytest.raises(ValueError, match=r"datum\[0\].*'target_tokens'.*5 values.*10 tokens"): + batch_to_datum_list(batch) + + def test_longer_than_model_input_raises(self): + """Too many values is the likelier mistake: response-length, not full-sequence. + + Per-token entries cover the whole sequence, so a caller sizing them to + the response alone lands short, but one that concatenates prompt and + response twice, or reuses a longer datum's inputs, lands long. Both + directions have to raise locally. + """ + batch = [Datum(model_input=ModelInput.from_ints(_TOKENS10), loss_fn_inputs=_per_token(14))] + with pytest.raises(ValueError, match=r"datum\[0\].*'target_tokens'.*14 values.*10 tokens"): + batch_to_datum_list(batch) + + def test_batch_level_longer_than_model_input_raises(self): + with pytest.raises(ValueError, match=r"lossFnInputs\['target_tokens'\]"): + batch_to_datum_list([_TOKENS10], loss_fn_inputs=_per_token(14)) + + def test_prebuilt_wire_dict_longer_checked(self): + with pytest.raises(ValueError, match=r"datum\[0\]"): + batch_to_datum_list([_wire_datum(_TOKENS10, _per_token(14))]) + + def test_matched_lengths_pass(self): + batch = [Datum(model_input=ModelInput.from_ints(_TOKENS10), loss_fn_inputs=_per_token(10))] + result = batch_to_datum_list(batch) + assert len(result[0]["lossFnInputs"]["weights"]["data"]) == 10 + + def test_batch_level_loss_fn_inputs_checked(self): + # loss_fn_inputs passed alongside the batch is merged in after each + # datum is built, so it has to be checked at the end of the batch path. + with pytest.raises(ValueError, match=r"lossFnInputs\['target_tokens'\]"): + batch_to_datum_list([_TOKENS10], loss_fn_inputs=_per_token(5)) + + def test_batch_level_loss_fn_inputs_matched_passes(self): + result = batch_to_datum_list([_TOKENS10], loss_fn_inputs=_per_token(10)) + assert len(result[0]["lossFnInputs"]["target_tokens"]["data"]) == 10 + + def test_prebuilt_wire_dict_checked(self): + with pytest.raises(ValueError, match=r"datum\[0\]"): + batch_to_datum_list([_wire_datum(_TOKENS10, _per_token(5))]) + + def test_offending_datum_index_reported(self): + batch = [ + _wire_datum(_TOKENS10, _per_token(10)), + _wire_datum(_TOKENS10, _per_token(10)), + _wire_datum(_TOKENS10, _per_token(4)), + ] + with pytest.raises(ValueError, match=r"datum\[2\]"): + batch_to_datum_list(batch) + + def test_mixed_batch_checked(self): + batch = [ + Datum(model_input=ModelInput.from_ints(_TOKENS10), loss_fn_inputs=_per_token(10)), + _wire_datum(_TOKENS10, _per_token(4)), + ] + with pytest.raises(ValueError, match=r"datum\[1\]"): + batch_to_datum_list(batch) + + def test_tensor_data_values_checked(self): + batch = [ + Datum( + model_input=ModelInput.from_ints(_TOKENS10), + loss_fn_inputs={ + "weights": TensorData(data=[1.0] * 5, shape=[5]), + "target_tokens": TensorData(data=list(range(5)), shape=[5]), + }, + ) + ] + with pytest.raises(ValueError, match=r"5 values.*10 tokens"): + batch_to_datum_list(batch) + + def test_tensor_data_matched_passes(self): + batch = [ + Datum( + model_input=ModelInput.from_ints(_TOKENS10), + loss_fn_inputs={ + "weights": TensorData(data=[1.0] * 10, shape=[10]), + "target_tokens": TensorData(data=list(range(10)), shape=[10]), + }, + ) + ] + assert len(batch_to_datum_list(batch)) == 1 + + def test_multi_chunk_length_summed(self): + split = {"modelInput": {"chunks": [{"tokens": _TOKENS10[:4]}, {"tokens": _TOKENS10[4:]}]}} + ok = {**split, "lossFnInputs": _per_token(10)} + assert len(batch_to_datum_list([ok])) == 1 + bad = {**split, "lossFnInputs": _per_token(9)} + with pytest.raises(ValueError, match=r"9 values.*10 tokens"): + batch_to_datum_list([bad]) + + # --- entries that are not per-token must not be rejected --- + + def test_scalar_knob_not_checked(self): + # A scalar sits alongside per-token entries with its own length. + batch = [ + Datum( + model_input=ModelInput.from_ints(_TOKENS10), + loss_fn_inputs={**_per_token(10), "temperature": [0.7]}, + ) + ] + assert len(batch_to_datum_list(batch)) == 1 + + def test_unrecognised_key_not_checked(self): + # Only known per-token keys are checked, so a key this version does not + # know about is passed through for the service to interpret. + batch = [ + Datum( + model_input=ModelInput.from_ints(_TOKENS10), + loss_fn_inputs={**_per_token(10), "some_future_knob": [1.0, 2.0]}, + ) + ] + assert len(batch_to_datum_list(batch)) == 1 + + def test_entries_the_service_does_not_enforce_are_not_checked(self): + # ``advantages`` is reshaped rather than length-checked by the service, + # so a request carrying a short one is still accepted and must not be + # rejected here. Keeps this check from being stricter than the service. + batch = [ + Datum( + model_input=ModelInput.from_ints(_TOKENS10), + loss_fn_inputs={"advantages": [0.5], "ref_log_prob": [0.1]}, + ) + ] + assert len(batch_to_datum_list(batch)) == 1 + + def test_text_prompt_not_checked(self): + # Token count is not known until the text is tokenized. + batch = [{"modelInput": {"prompt": "hello world"}, "lossFnInputs": _per_token(5)}] + assert len(batch_to_datum_list(batch)) == 1 + + def test_no_loss_fn_inputs_passes(self): + assert len(batch_to_datum_list([_TOKENS10])) == 1 + + def test_labels_only_not_checked(self): + # ``labels`` is not a per-token entry, so it carries its own length. + assert len(batch_to_datum_list([[1]], loss_fn_inputs={"labels": [0]})) == 1 + + +class _FakeMP: + """Stubs list_model_packages (paginated) + describe_model_package.""" + + def __init__(self, pages=None, describe=None): + # pages: list of (summaries, next_token) served in order. + self._pages = list(pages or []) + self._describe = describe or {} + self.list_calls: list[dict] = [] + self.describe_calls: list[dict] = [] + + def list_model_packages(self, **kwargs): + self.list_calls.append(kwargs) + summaries, next_token = self._pages.pop(0) + resp = {"ModelPackageSummaryList": summaries} + if next_token: + resp["NextToken"] = next_token + return resp + + def describe_model_package(self, **kwargs): + self.describe_calls.append(kwargs) + return self._describe[kwargs["ModelPackageName"]] + + +def _summary(arn, version, desc="d"): + return { + "ModelPackageArn": arn, + "ModelPackageVersion": version, + "ModelPackageDescription": desc, + "CreationTime": "2026-06-20T00:00:00Z", + } + + +class TestListCheckpointSummaries: + def test_sorts_descending_and_caps(self): + client = _FakeMP(pages=[([_summary("a/3", 3), _summary("a/2", 2)], "tok")]) + out = list_model_package_summaries(client, "grp", max_results=2) + assert [s["ModelPackageVersion"] for s in out] == [3, 2] + call = client.list_calls[0] + assert call["ModelPackageGroupName"] == "grp" + assert call["SortBy"] == "CreationTime" + assert call["SortOrder"] == "Descending" + assert call["MaxResults"] == 2 + # Capped at 2 -> stopped paging despite NextToken. + assert len(client.list_calls) == 1 + + def test_auto_paginates_until_limit(self): + client = _FakeMP(pages=[([_summary("a/3", 3)], "tok"), ([_summary("a/2", 2)], None)]) + out = list_model_package_summaries(client, "grp", max_results=10) + assert [s["ModelPackageVersion"] for s in out] == [3, 2] + assert client.list_calls[1]["NextToken"] == "tok" + + def test_none_returns_all(self): + client = _FakeMP(pages=[([_summary("a/2", 2)], "tok"), ([_summary("a/1", 1)], None)]) + out = list_model_package_summaries(client, "grp", max_results=None) + assert len(out) == 2 + # No MaxResults sent when unbounded. + assert "MaxResults" not in client.list_calls[0] + + def test_cap_above_one_page_pages_instead_of_asking_for_it(self): + """A cap larger than the page size is served by paging. + + ListModelPackages accepts at most MAX_RESULTS_PER_PAGE per request and + rejects anything above it, so sending the caller's cap verbatim would + fail the call rather than return that many. + """ + first = [_summary(f"a/{index}", index) for index in range(100)] + client = _FakeMP(pages=[(first, "tok"), ([_summary("a/100", 100)], None)]) + + out = list_model_package_summaries(client, "grp", max_results=150) + + assert len(out) == 101 + assert [call["MaxResults"] for call in client.list_calls] == [100, 50] + + +class TestDescribeModelPackageKind: + def test_is_checkpoint_true_is_checkpoint_kind(self): + client = _FakeMP( + describe={"a/1": {"InferenceSpecification": {"Containers": [{"IsCheckpoint": True}]}}} + ) + assert describe_model_package_kind(client, "a/1") == ModelPackageKind.CHECKPOINT + + def test_is_checkpoint_absent_is_model(self): + client = _FakeMP(describe={"a/1": {"InferenceSpecification": {"Containers": [{}]}}}) + assert describe_model_package_kind(client, "a/1") == ModelPackageKind.MODEL + + def test_no_inference_spec_is_model(self): + client = _FakeMP(describe={"a/1": {}}) + assert describe_model_package_kind(client, "a/1") == ModelPackageKind.MODEL + + +class _CompatChunk: + def __init__(self, tokens): + self.tokens = tokens + + +class _CompatModelInput: + def __init__(self, chunks): + self.chunks = chunks + + +class _CompatDatum: + def __init__(self, model_input, loss_fn_inputs=None): + self.model_input = model_input + self.loss_fn_inputs = loss_fn_inputs or {} + + +class TestTypedInputWiring: + """Typed objects (ours and attribute-compatible external ones) flow through the input paths.""" + + def test_batch_of_datum_objects(self): + batch = [Datum(model_input=ModelInput.from_ints([1, 2]))] + out = batch_to_datum_list(batch) + assert out == [{"modelInput": {"chunks": [{"tokens": [1, 2], "type": "encoded_text"}]}}] + + def test_batch_of_model_input_objects(self): + out = batch_to_datum_list([ModelInput.from_ints([3])]) + assert out[0]["modelInput"]["chunks"][0]["tokens"] == [3] + + def test_batch_of_compatible_datums(self): + batch = [ + _CompatDatum( + model_input=_CompatModelInput([_CompatChunk([1, 2])]), + loss_fn_inputs={"advantages": TensorData(data=[0.5])}, + ) + ] + out = batch_to_datum_list(batch) + assert out[0]["modelInput"]["chunks"][0]["tokens"] == [1, 2] + assert out[0]["lossFnInputs"]["advantages"] == {"data": [0.5], "dtype": "float32"} + + def test_typed_batch_fills_shared_loss_inputs_only_where_absent(self): + batch = [ + Datum(model_input=ModelInput.from_ints([1]), loss_fn_inputs={"w": [1.0]}), + Datum(model_input=ModelInput.from_ints([2])), + ] + out = batch_to_datum_list(batch, loss_fn_inputs={"w": TensorData(data=[9.0])}) + assert out[0]["lossFnInputs"] == {"w": {"data": [1.0], "dtype": "float32", "shape": [1]}} + assert out[1]["lossFnInputs"] == {"w": {"data": [9.0], "dtype": "float32"}} + + def test_typed_batch_shared_loss_inputs_not_aliased(self): + batch = [ + Datum(model_input=ModelInput.from_ints([1])), + Datum(model_input=ModelInput.from_ints([2])), + ] + out = batch_to_datum_list(batch, loss_fn_inputs={"w": [1.0]}) + out[0]["lossFnInputs"]["extra"] = True + assert "extra" not in out[1]["lossFnInputs"] + + def test_tokenize_to_datum_serializes_tensor_data(self): + datum = tokenize_to_datum([1], loss_fn_inputs={"advantages": TensorData(data=[0.1])}) + assert datum["lossFnInputs"]["advantages"] == {"data": [0.1], "dtype": "float32"} + + def test_dict_batch_shared_tensor_loss_inputs_serialized(self): + batch = [{"modelInput": {"chunks": [{"tokens": [1], "type": "encoded_text"}]}}] + out = batch_to_datum_list(batch, loss_fn_inputs={"adv": TensorData(data=[0.2])}) + assert out[0]["lossFnInputs"]["adv"] == {"data": [0.2], "dtype": "float32"} + + def test_prompt_model_input_object(self): + out = prompt_to_model_input(ModelInput.from_ints([1, 2])) + assert out == {"chunks": [{"tokens": [1, 2], "type": "encoded_text"}]} + + def test_prompt_encoded_text_chunk(self): + out = prompt_to_model_input(EncodedTextChunk(tokens=[5])) + assert out == {"chunks": [{"tokens": [5], "type": "encoded_text"}]} + + def test_prompt_compatible_model_input(self): + compat = _CompatModelInput([_CompatChunk([7, 8])]) + out = prompt_to_model_input(compat) + assert out == {"chunks": [{"tokens": [7, 8], "type": "encoded_text"}]} + + def test_prompt_dict_still_passes_through(self): + raw = {"chunks": [{"tokens": [9], "type": "encoded_text"}]} + assert prompt_to_model_input(raw) is raw + + +class TestBaseModelArnFromModelPackage: + def _desc(self, base): + return { + "ModelPackageArn": "arn:mp/weights-7", + "InferenceSpecification": {"Containers": [{"IsCheckpoint": False, "BaseModel": base}]}, + } + + def test_builds_versioned_hub_arn(self): + from sagemaker.train.training_session._utils import ( + base_model_arn_from_model_package, + ) + + arn = base_model_arn_from_model_package( + self._desc( + {"HubContentName": "hubcontent-text-model-a", "HubContentVersion": "3.15.0"} + ), + region="us-west-2", + ) + assert arn == ( + "arn:aws:sagemaker:us-west-2:aws:hub-content/" + "SageMakerPublicHub/Model/hubcontent-text-model-a/3.15.0" + ) + + def test_missing_base_model_raises(self): + from sagemaker.train.training_session._utils import ( + base_model_arn_from_model_package, + ) + + with pytest.raises(ValueError, match="no base-model reference"): + base_model_arn_from_model_package(self._desc({}), region="us-west-2") + + def test_no_containers_raises(self): + from sagemaker.train.training_session._utils import ( + base_model_arn_from_model_package, + ) + + with pytest.raises(ValueError, match="no base-model reference"): + base_model_arn_from_model_package({"InferenceSpecification": {}}, region="us-west-2") + + +class TestSplitS3Uri: + def test_bucket_and_key(self): + from sagemaker.train.training_session._utils import split_s3_uri + + assert split_s3_uri("s3://b/a/b/c/") == ("b", "a/b/c/") + + def test_bucket_only_gives_empty_key(self): + from sagemaker.train.training_session._utils import split_s3_uri + + assert split_s3_uri("s3://b") == ("b", "") + assert split_s3_uri("s3://b/") == ("b", "") + + def test_non_s3_scheme_raises(self): + from sagemaker.train.training_session._utils import split_s3_uri + + with pytest.raises(ValueError, match="not an s3:// URI"): + split_s3_uri("https://bucket.s3.amazonaws.com/key") + + def test_no_bucket_raises(self): + from sagemaker.train.training_session._utils import split_s3_uri + + with pytest.raises(ValueError, match="names no bucket"): + split_s3_uri("s3:///key") + + +class TestDeleteS3ObjectsUnder: + """``DeleteObjects`` caps at 1000 keys per call, so a larger page is split.""" + + class _S3: + def __init__(self, keys): + self._keys = keys + self.delete_calls: list[dict] = [] + + def list_objects_v2(self, **kwargs): + # One oversized page, to prove batching is not left to the lister. + return {"Contents": [{"Key": k} for k in self._keys], "IsTruncated": False} + + def delete_objects(self, **kwargs): + self.delete_calls.append(kwargs) + return {} + + def test_batches_at_the_delete_objects_limit(self): + from sagemaker.train.training_session._utils import ( + delete_s3_objects_under, + ) + + s3 = self._S3([f"p/{i}" for i in range(2500)]) + deleted, failed = delete_s3_objects_under(s3, "s3://b/p/") + + assert (deleted, failed) == (2500, []) + sizes = [len(c["Delete"]["Objects"]) for c in s3.delete_calls] + assert sizes == [1000, 1000, 500] + + def test_runs_quiet_so_only_failures_come_back(self): + from sagemaker.train.training_session._utils import ( + delete_s3_objects_under, + ) + + s3 = self._S3(["p/1"]) + delete_s3_objects_under(s3, "s3://b/p/") + + assert s3.delete_calls[0]["Delete"]["Quiet"] is True + + +class TestDescribeModelPackageArtifacts: + """Describe carries the location in a nested optional structure, so every + level of it can be absent.""" + + @staticmethod + def _client(response): + class _SageMaker: + def __init__(self): + self.calls: list[dict] = [] + + def describe_model_package(self, **kwargs): + self.calls.append(kwargs) + return response + + return _SageMaker() + + def _describe(self, response, arn="arn:aws:sagemaker:us-west-2:123:model-package/g/1"): + from sagemaker.train.training_session._utils import ( + describe_model_package_artifacts, + ) + + client = self._client(response) + return describe_model_package_artifacts(client, arn), client + + def test_reads_location_type_and_managed_storage(self): + artifacts, client = self._describe( + { + "ManagedStorageType": "None", + "InferenceSpecification": { + "Containers": [ + { + "ModelDataSource": { + "S3DataSource": { + "S3Uri": "s3://b/g/req-1/model/", + "S3DataType": "S3Prefix", + } + } + } + ] + }, + } + ) + + assert artifacts == { + "s3_uri": "s3://b/g/req-1/model/", + "s3_data_type": "S3Prefix", + "managed_storage_type": "None", + } + assert client.calls == [ + {"ModelPackageName": "arn:aws:sagemaker:us-west-2:123:model-package/g/1"} + ] + + @pytest.mark.parametrize( + "response", + [ + pytest.param({}, id="no-inference-specification"), + pytest.param({"InferenceSpecification": {}}, id="no-containers-key"), + pytest.param({"InferenceSpecification": {"Containers": []}}, id="empty-containers"), + pytest.param( + {"InferenceSpecification": {"Containers": [{}]}}, id="no-model-data-source" + ), + pytest.param( + {"InferenceSpecification": {"Containers": [{"ModelDataSource": {}}]}}, + id="no-s3-data-source", + ), + pytest.param( + { + "InferenceSpecification": { + "Containers": [{"ModelDataSource": {"S3DataSource": {"S3Uri": ""}}}] + } + }, + id="empty-s3-uri", + ), + pytest.param( + # The legacy container shape. Nothing this SDK saves uses it, + # and reporting no location keeps such a package's artifacts. + {"InferenceSpecification": {"Containers": [{"ModelDataUrl": "s3://b/k"}]}}, + id="legacy-model-data-url", + ), + ], + ) + def test_absent_at_any_level_reads_as_no_location(self, response): + artifacts, _ = self._describe(response) + + assert artifacts == { + "s3_uri": None, + "s3_data_type": None, + "managed_storage_type": None, + } + + def test_managed_storage_without_a_location(self): + from sagemaker.train.training_session._utils import ( + MANAGED_STORAGE_RESTRICTED, + ) + + artifacts, _ = self._describe({"ManagedStorageType": MANAGED_STORAGE_RESTRICTED}) + + assert artifacts["managed_storage_type"] == MANAGED_STORAGE_RESTRICTED + assert artifacts["s3_uri"] is None + + +class TestDeleteS3ObjectsUnderEdges: + class _S3: + """Serves scripted ``(keys, next_token)`` pages, optionally reporting + per-key failures the way ``DeleteObjects`` does: in the body.""" + + def __init__(self, pages, errors_for=(), truncated_without_token=False): + self._pages = list(pages) + self._errors_for = set(errors_for) + self._truncated_without_token = truncated_without_token + self.list_calls: list[dict] = [] + self.delete_calls: list[dict] = [] + + def list_objects_v2(self, **kwargs): + self.list_calls.append(kwargs) + keys, next_token = self._pages.pop(0) + resp = {"IsTruncated": bool(next_token) or self._truncated_without_token} + if keys: + resp["Contents"] = [{"Key": k} for k in keys] + if next_token: + resp["NextContinuationToken"] = next_token + return resp + + def delete_objects(self, **kwargs): + self.delete_calls.append(kwargs) + sent = [o["Key"] for o in kwargs["Delete"]["Objects"]] + errors = [{"Key": k, "Code": "AccessDenied"} for k in sent if k in self._errors_for] + return {"Errors": errors} if errors else {} + + def _delete(self, s3, uri="s3://b/p/", **kwargs): + from sagemaker.train.training_session._utils import ( + delete_s3_objects_under, + ) + + return delete_s3_objects_under(s3, uri, **kwargs) + + def test_truncated_without_a_token_stops(self): + """A truncated page with no continuation token would otherwise re-request + page one forever.""" + s3 = self._S3([(["p/1"], None)], truncated_without_token=True) + + deleted, failed = self._delete(s3) + + assert (deleted, failed) == (1, []) + assert len(s3.list_calls) == 1 + + def test_failures_accumulate_across_pages(self): + """Each page is deleted as it is listed, so a later page's failures must + join the earlier ones rather than replace them.""" + s3 = self._S3([(["p/1", "p/2"], "tok"), (["p/3"], None)], errors_for=["p/1", "p/3"]) + + deleted, failed = self._delete(s3) + + assert deleted == 1 + assert sorted(failed) == ["p/1", "p/3"] + + def test_empty_page_deletes_nothing(self): + s3 = self._S3([([], None)]) + + assert self._delete(s3) == (0, []) + assert s3.delete_calls == [] + + def test_no_data_type_is_treated_as_a_prefix(self): + """A package that carries a location but no S3DataType is swept, not read + as a single key: the location is a directory in every case this SDK + writes.""" + s3 = self._S3([(["p/1"], None)]) + + self._delete(s3, s3_data_type=None) + + assert s3.list_calls == [{"Bucket": "b", "Prefix": "p/"}] + + def test_single_object_failure_comes_back(self): + s3 = self._S3([], errors_for=["p/model.tar.gz"]) + + deleted, failed = self._delete(s3, uri="s3://b/p/model.tar.gz", s3_data_type="S3Object") + + assert (deleted, failed) == (0, ["p/model.tar.gz"]) + assert s3.list_calls == [] + + @pytest.mark.parametrize("uri", ["s3://b", "s3://b/"]) + def test_bucket_root_refused_before_any_call(self, uri): + s3 = self._S3([]) + + with pytest.raises(ValueError, match="bucket root"): + self._delete(s3, uri=uri) + + assert s3.list_calls == [] + assert s3.delete_calls == [] + + def test_a_token_on_an_untruncated_page_is_not_followed(self): + """The loop stops when either signal says "done", so a final page that + still carries a continuation token ends the sweep instead of re-listing + from that token.""" + from sagemaker.train.training_session._utils import ( + delete_s3_objects_under, + ) + + class _S3: + def __init__(self): + self.list_calls = 0 + + def list_objects_v2(self, **kwargs): + self.list_calls += 1 + return { + "Contents": [{"Key": "p/1"}], + "IsTruncated": False, + "NextContinuationToken": "tok", + } + + def delete_objects(self, **kwargs): + return {} + + s3 = _S3() + deleted, failed = delete_s3_objects_under(s3, "s3://b/p/") + + assert (deleted, failed) == (1, []) + assert s3.list_calls == 1 + + def test_a_failure_without_a_key_reads_as_deleted(self): + """``DeleteObjects`` keys every error by object; an entry with no ``Key`` + names nothing to retry or report, so it is dropped rather than becoming a + ``None`` in ``failed_keys`` that the caller would have to filter. Its key + therefore counts as deleted, documented here because it means + ``deleted`` is "not reported failed", not "confirmed gone".""" + from sagemaker.train.training_session._utils import ( + delete_s3_objects_under, + ) + + class _S3: + def list_objects_v2(self, **kwargs): + return {"Contents": [{"Key": "p/1"}], "IsTruncated": False} + + def delete_objects(self, **kwargs): + return {"Errors": [{"Code": "InternalError"}]} + + deleted, failed = delete_s3_objects_under(_S3(), "s3://b/p/") + + assert failed == [] + assert deleted == 1 + + +class TestMixedBatchStartingWithATokenId: + """A batch is only a bare token list if every element is a token. + + One that merely starts with an int skipped the mixed-kind check and was read + as a single token list, putting the other elements inside a token array and + sending that. + """ + + def test_a_string_after_a_token_id_raises(self): + with pytest.raises(TypeError, match="only token IDs"): + batch_to_datum_list([10, "hello"], tokenizer=None, loss_fn_inputs=None) + + def test_the_message_names_what_it_found(self): + with pytest.raises(TypeError, match="str"): + batch_to_datum_list([10, "hello"], tokenizer=None, loss_fn_inputs=None) + + def test_a_list_after_a_token_id_raises(self): + with pytest.raises(TypeError, match="only token IDs"): + batch_to_datum_list([10, [11, 12]], tokenizer=None, loss_fn_inputs=None) + + def test_an_all_token_batch_is_still_one_datum(self): + datums = batch_to_datum_list([1, 2, 3], tokenizer=None, loss_fn_inputs=None) + + assert len(datums) == 1 + assert datums[0]["modelInput"]["chunks"][0]["tokens"] == [1, 2, 3] + + +class TestChunkedMetricsUseRealDatumCounts: + """The combined loss is a datum-count-weighted mean over the chunks. + + The counts used to be inferred from per-datum loss outputs, which are absent + unless the caller asked for them, so every chunk weighed the same and an + uneven split reported a plain mean instead. + """ + + @staticmethod + def _result(loss): + from sagemaker.train.training_session.types import ForwardBackwardOutput + + return ForwardBackwardOutput._from_api({"forwardBackward": {"metrics": {"loss": loss}}}) + + def test_uneven_chunks_are_weighted_by_size(self): + results = [self._result(1.0), self._result(1.0), self._result(11.0)] + + combined = combine_forward_backward_results(results, datum_counts=[1024, 1024, 100]) + + expected = (1.0 * 1024 + 1.0 * 1024 + 11.0 * 100) / 2148 + assert combined.metrics["loss"] == pytest.approx(expected) + + def test_without_counts_it_falls_back_to_equal_weight(self): + """Documented fallback: no counts and no per-datum outputs means a mean.""" + results = [self._result(1.0), self._result(1.0), self._result(11.0)] + + combined = combine_forward_backward_results(results) + + assert combined.metrics["loss"] == pytest.approx(13.0 / 3) + + def test_grad_norm_still_comes_from_the_last_chunk(self): + from sagemaker.train.training_session.types import ForwardBackwardOutput + + results = [ + ForwardBackwardOutput._from_api({"forwardBackward": {"metrics": {"grad_norm": 1.0}}}), + ForwardBackwardOutput._from_api({"forwardBackward": {"metrics": {"grad_norm": 7.0}}}), + ] + + combined = combine_forward_backward_results(results, datum_counts=[1000, 1]) + + assert combined.metrics["grad_norm"] == 7.0 + + +class TestHubContentArnPartition: + """Every ARN this package builds resolves its partition rather than assuming.""" + + def test_the_partition_comes_from_the_model_package_arn(self): + assert _arn_partition("arn:aws-cn:sagemaker:cn-north-1:1:model-package/g/1") == "aws-cn" + assert _arn_partition("arn:aws-us-gov:sagemaker:us-gov-west-1:1:x/y") == "aws-us-gov" + + def test_an_unreadable_arn_falls_back_to_the_commercial_partition(self): + assert _arn_partition(None) == "aws" + assert _arn_partition("") == "aws" + assert _arn_partition("not-an-arn") == "aws" + + +class TestMissingLossFnInputsAreRefused: + """A datum must carry the loss inputs its loss function requires. + + The service refuses an operation whose datums lack them, and that refusal + names neither the datum nor the key. Catching it here can name both, before + the operation is accepted and queued. + """ + + def test_a_subset_names_what_is_missing(self): + with pytest.raises(ValueError, match="logprobs"): + batch_to_datum_list( + [[5, 6]], + loss_fn_inputs={"advantages": [1.0, 1.0]}, + loss_fn_name="ppo", + ) + + def test_the_message_names_the_datum_and_what_is_present(self): + with pytest.raises(ValueError, match=r"datum\[0\].*advantages"): + batch_to_datum_list( + [[5, 6]], + loss_fn_inputs={"advantages": [1.0, 1.0]}, + loss_fn_name="ppo", + ) + + def test_no_loss_inputs_at_all_is_refused(self): + """Absent is the case the quick starts hit, so it has to raise here. + + The shape leaves ``lossFnInputs`` optional, but a loss function's own + requirement is not optional: the operation is refused for a datum + without them. + """ + with pytest.raises(ValueError, match="requires"): + batch_to_datum_list([[5, 6]], loss_fn_name="cross_entropy") + + def test_a_batch_without_a_named_loss_is_still_left_alone(self): + """No loss named, nothing to require: the caller has not chosen one yet.""" + datums = batch_to_datum_list([[5, 6]]) + + assert "lossFnInputs" not in datums[0] + + def test_a_complete_set_passes(self): + datums = batch_to_datum_list( + [[5, 6]], + loss_fn_inputs={ + "target_tokens": [6, 0], + "weights": [0.0, 1.0], + "logprobs": [-0.1, -0.2], + "advantages": [1.0, 1.0], + }, + loss_fn_name="ppo", + ) + + assert sorted(datums[0]["lossFnInputs"]) == [ + "advantages", + "logprobs", + "target_tokens", + "weights", + ] + + def test_cross_entropy_needs_only_its_own_two(self): + datums = batch_to_datum_list( + [[5, 6]], + loss_fn_inputs={"target_tokens": [6, 0], "weights": [0.0, 1.0]}, + loss_fn_name="cross_entropy", + ) + + assert datums[0]["lossFnInputs"] + + def test_an_unnamed_loss_function_is_not_second_guessed(self): + """No name given, nothing is assumed about which entries are wanted.""" + datums = batch_to_datum_list([[5, 6]], loss_fn_inputs={"advantages": [1.0, 1.0]}) + + assert datums[0]["lossFnInputs"] + + +class TestHomogeneousDatumDictBatches: + """A batch of datum dicts is read in every spelling ``from_any`` accepts. + + Only the wire spelling (``modelInput``) had a homogeneous branch, so a list + of snake_case dicts was refused outright while the very same dict inside a + mixed batch was read correctly. + """ + + _EXPECTED = [ + {"modelInput": {"chunks": [{"tokens": [1, 2, 3], "type": "encoded_text"}]}}, + {"modelInput": {"chunks": [{"tokens": [4, 5, 6], "type": "encoded_text"}]}}, + ] + + def test_snake_case_keys(self): + batch = [{"model_input": [1, 2, 3]}, {"model_input": [4, 5, 6]}] + + assert batch_to_datum_list(batch) == self._EXPECTED + + def test_the_prompt_shorthand(self): + batch = [{"prompt": [1, 2, 3]}, {"prompt": [4, 5, 6]}] + + assert batch_to_datum_list(batch) == self._EXPECTED + + def test_the_wire_spelling_is_still_passed_through_unchanged(self): + batch = [dict(d) for d in self._EXPECTED] + + assert batch_to_datum_list(batch) == self._EXPECTED + + def test_loss_inputs_still_reach_a_snake_case_batch(self): + batch = [{"model_input": [1, 2]}, {"model_input": [3, 4]}] + + datums = batch_to_datum_list(batch, loss_fn_inputs={"weights": [1.0, 1.0]}) + + assert all("lossFnInputs" in d for d in datums) + + def test_a_batch_mixing_the_two_dict_spellings_reads_each_element(self): + """Which spelling comes first must not decide how the rest are read. + + Keying the batch off its first element sent the later ones through the + wrong reading of their own keys, so one order was coerced and the other + refused. + """ + wire = {"modelInput": {"chunks": [{"tokens": [1], "type": "encoded_text"}]}} + snake = {"model_input": [2, 3]} + expected = [ + wire, + {"modelInput": {"chunks": [{"tokens": [2, 3], "type": "encoded_text"}]}}, + ] + + assert batch_to_datum_list([wire, snake]) == expected + assert batch_to_datum_list([snake, wire]) == list(reversed(expected)) + + def test_a_dict_naming_no_input_is_still_refused(self): + with pytest.raises(TypeError, match="Unsupported batch type"): + batch_to_datum_list([{"nothing_useful": 1}, {"nothing_useful": 2}]) + + def test_loss_inputs_alone_do_not_make_a_datum(self): + """``input_ids``/``messages`` are not spellings this accepts. + + Reading them as datums built an input-less example that only the + service would refuse, so the batch is rejected here instead, where the + error can name the shapes a batch may hold. + """ + with pytest.raises(TypeError, match="Unsupported batch type"): + batch_to_datum_list( + [{"input_ids": [1, 2], "loss_fn_inputs": {"weights": [1.0, 1.0]}}] * 2 + ) + with pytest.raises(TypeError, match="Unsupported batch type"): + batch_to_datum_list([{"messages": [{"role": "user"}]}] * 2) + + +class TestChunkedMetricsKeepEveryReportedKey: + """Combining read only the last chunk's metric keys. + + A metric the earlier chunks reported and the last one omitted vanished from + the combined result, so a chunked batch silently reported fewer metrics than + the same batch sent whole. + """ + + @staticmethod + def _result(metrics): + from sagemaker.train.training_session.types import ForwardBackwardOutput + + return ForwardBackwardOutput._from_api({"forwardBackward": {"metrics": metrics}}) + + def test_a_metric_only_an_earlier_chunk_reports_survives(self): + results = [self._result({"loss": 1.0, "aux": 7.0}), self._result({"loss": 3.0})] + + combined = combine_forward_backward_results(results, datum_counts=[1, 1]) + + assert combined.metrics["aux"] == pytest.approx(7.0) + assert combined.metrics["loss"] == pytest.approx(2.0) + + def test_it_is_weighted_over_only_the_chunks_that_reported_it(self): + """Weighting over every chunk would dilute it toward zero.""" + results = [self._result({"aux": 2.0}), self._result({"aux": 4.0}), self._result({})] + + combined = combine_forward_backward_results(results, datum_counts=[1, 3, 96]) + + assert combined.metrics["aux"] == pytest.approx((2.0 * 1 + 4.0 * 3) / 4) + + def test_a_namespaced_grad_norm_comes_from_the_final_chunk(self): + """The reported key is namespaced, so the suffix is what identifies it. + + Matching the bare name left the real one to be averaged, and the mean of + two gradient norms is not a gradient norm. + """ + results = [ + self._result({"actor/grad_norm": 1.0}), + self._result({"actor/grad_norm": 7.0}), + ] + + combined = combine_forward_backward_results(results, datum_counts=[1000, 1]) + + assert combined.metrics["actor/grad_norm"] == 7.0 + + def test_a_token_count_is_totalled_rather_than_averaged(self): + """Each chunk counts part of one batch, so the batch value is the sum. + + Averaging a count reports a number no chunk measured: for the split + below it gives 11502, which is lower than the first chunk's own 12288 + even though the batch held more tokens than that chunk did. + """ + results = [ + self._result({"global_token_num": 12288.0}), + self._result({"global_token_num": 912.0}), + ] + + combined = combine_forward_backward_results(results, datum_counts=[1024, 76]) + + assert combined.metrics["global_token_num"] == 13200.0 + + def test_a_token_count_from_a_single_chunk_is_unchanged(self): + """One chunk is not combined at all, so the count passes through.""" + results = [self._result({"global_token_num": 192.0})] + + combined = combine_forward_backward_results(results, datum_counts=[16]) + + assert combined.metrics["global_token_num"] == 192.0 + + def test_a_per_datum_metric_is_still_averaged(self): + """The default stays a mean; only the named counts are totalled.""" + results = [ + self._result({"actor/loss": 2.0, "global_token_num": 100.0}), + self._result({"actor/loss": 4.0, "global_token_num": 10.0}), + ] + + combined = combine_forward_backward_results(results, datum_counts=[1, 1]) + + assert combined.metrics["actor/loss"] == 3.0 + assert combined.metrics["global_token_num"] == 110.0 + + def test_a_grad_norm_only_an_earlier_chunk_reports_is_omitted(self): + """A partial norm is worse to report than none. + + The key union keeps every other metric an earlier chunk reported, but a + gradient norm describes the gradient accumulated across all chunks, so + one measured partway through is not the whole-batch value. + """ + results = [ + self._result({"actor/grad_norm": 9.0, "actor/loss": 1.0}), + self._result({"actor/loss": 3.0}), + ] + + combined = combine_forward_backward_results(results, datum_counts=[1, 1]) + + assert "actor/grad_norm" not in combined.metrics + assert combined.metrics["actor/loss"] == 2.0 diff --git a/v3-examples/training-examples/training_session_example.ipynb b/v3-examples/training-examples/training_session_example.ipynb new file mode 100644 index 0000000000..f57a823197 --- /dev/null +++ b/v3-examples/training-examples/training_session_example.ipynb @@ -0,0 +1,536 @@ +{ + "cells": [ + { + "cell_type": "markdown", + "id": "cell-0", + "metadata": {}, + "source": [ + "# Interactive Training Session\n", + "\n", + "SageMaker **interactive training sessions** give you a live, client-driven\n", + "training loop over a hosted model: you submit low-level primitives\n", + "(`forward_backward`, `optim_step`, `forward`), checkpoint with\n", + "`save_weights_for_sampler` / `save_state` / `load_state`, and sample from a\n", + "checkpoint with\n", + "`sample` — all from your own Python process.\n", + "\n", + "The SDK assigns operation sequence numbers for you (so calls can be pipelined)\n", + "and signs every request with your own AWS credentials (SigV4). You only need\n", + "standard `boto3`-resolvable credentials configured (env vars, profile, or an\n", + "attached role) — there is no separate API key or token to manage.\n", + "\n", + "This notebook has two parts:\n", + "1. **Quick start** — the minimal end-to-end loop.\n", + "2. **Going further** — train on real text with a tokenizer, inspect metrics,\n", + " checkpoint/resume, and sample." + ] + }, + { + "cell_type": "markdown", + "id": "cell-1", + "metadata": {}, + "source": [ + "## Prerequisites\n", + "\n", + "Before running, you need:\n", + "- A SageMaker **model package group** to register checkpoints under.\n", + "- An **IAM execution role** SageMaker assumes to access your resources.\n", + "- Permissions on **your own identity** for the training-session actions, plus\n", + " `iam:PassRole` for the execution role above, since these calls hand that role\n", + " to the service.\n", + "- An **S3 output path** for checkpoints and results.\n", + "- AWS credentials + region configured for `boto3`.\n", + "- `pip install 'sagemaker-train[tokenizer]'` (the tokenizer is used for text prompts).\n", + "\n", + "If you don't yet have a model package group, execution role, or bucket, the\n", + "**Optional: create prerequisites** section below shows how to create them." + ] + }, + { + "cell_type": "code", + "id": "b7f081bf", + "source": [ + "# Install the SDK (provides the `sagemaker.train.*` namespace) and the\n", + "# tokenizer used for text prompts. Restart the kernel after installing.\n", + "%pip install 'sagemaker-train[tokenizer]'" + ], + "metadata": {}, + "execution_count": null, + "outputs": [] + }, + { + "cell_type": "markdown", + "id": "cell-2", + "metadata": {}, + "source": [ + "## Configuration\n", + "\n", + "Update these values for your account before running." + ] + }, + { + "cell_type": "code", + "execution_count": null, + "id": "cell-3", + "metadata": {}, + "outputs": [], + "source": [ + "# Region to run in. The base models on offer differ by region; see\n", + "# \"Find a base model\" below.\n", + "REGION = \"us-west-2\"\n", + "\n", + "# A name you choose for this session (any string):\n", + "TRAINING_SESSION_NAME = \"\"\n", + "\n", + "# Base model to fine-tune from — an id from list_base_models() (see \"Find a\n", + "# base model\" below). The SDK resolves it to the latest Hub model ARN in REGION.\n", + "# To pin a specific version, set a full Hub model ARN here instead.\n", + "BASE_MODEL = \"\"\n", + "\n", + "# Resources in your account (create them in the next section if needed):\n", + "MODEL_PACKAGE_GROUP_ARN = (\n", + " f\"arn:aws:sagemaker:{REGION}:\"\n", + " \":model-package-group/\"\n", + ")\n", + "ROLE_ARN = \"arn:aws:iam:::role/\"\n", + "OUTPUT_S3_PATH = \"s3:///output/\"\n", + "\n", + "LORA_RANK = 8\n", + "LORA_ALPHA = 16" + ] + }, + { + "cell_type": "markdown", + "id": "cell-4", + "metadata": {}, + "source": [ + "## Optional: create prerequisites\n", + "\n", + "Skip this if you already have a model package group, execution role, and S3\n", + "bucket. These helper cells create them with `boto3`. Run only the ones you need,\n", + "then copy the resulting ARNs / paths into the **Configuration** cell above." + ] + }, + { + "cell_type": "code", + "execution_count": null, + "id": "cell-5", + "metadata": {}, + "source": [ + "import boto3\n", + "\n", + "sm = boto3.client(\"sagemaker\", region_name=REGION)\n", + "s3 = boto3.client(\"s3\", region_name=REGION)\n", + "sts = boto3.client(\"sts\")\n", + "ACCOUNT_ID = sts.get_caller_identity()[\"Account\"]\n", + "\n", + "# 1) Model package group\n", + "MPG_NAME = \"my-training-session-models\"\n", + "try:\n", + " resp = sm.create_model_package_group(ModelPackageGroupName=MPG_NAME)\n", + " print(\"created:\", resp[\"ModelPackageGroupArn\"])\n", + "except sm.exceptions.ClientError as e:\n", + " print(\"(may already exist)\", e)\n", + "\n", + "# 2) S3 bucket for outputs\n", + "# Named sagemaker-* so the managed policy's S3 permissions cover it.\n", + "BUCKET = f\"sagemaker-training-session-{ACCOUNT_ID}-{REGION}\"\n", + "try:\n", + " kwargs = {\"Bucket\": BUCKET}\n", + " if REGION != \"us-east-1\":\n", + " kwargs[\"CreateBucketConfiguration\"] = {\"LocationConstraint\": REGION}\n", + " s3.create_bucket(**kwargs)\n", + " print(\"created bucket:\", BUCKET)\n", + "except Exception as e:\n", + " print(\"(may already exist)\", e)\n", + "\n", + "print(\"\\nSet in Configuration above:\")\n", + "print(f' MODEL_PACKAGE_GROUP_ARN = \"arn:aws:sagemaker:{REGION}:{ACCOUNT_ID}:model-package-group/{MPG_NAME}\"')\n", + "print(f' OUTPUT_S3_PATH = \"s3://{BUCKET}/output/\"')" + ], + "outputs": [] + }, + { + "cell_type": "markdown", + "id": "cell-6", + "metadata": {}, + "source": [ + "The execution role needs a permissions policy allowing `sagemaker:CreateModelPackage`,\n", + "`sagemaker:DescribeHubContent`, `sagemaker:DescribeModelPackageGroup`, and S3\n", + "read/write on your bucket, plus a trust policy letting\n", + "`training-session.sagemaker.amazonaws.com` assume it (`sts:AssumeRole` +\n", + "`sts:TagSession`). See the package README for the exact JSON. Set `ROLE_ARN`\n", + "above to the role you create." + ] + }, + { + "cell_type": "code", + "id": "41578937", + "source": "# Create an execution role SageMaker assumes for the training session.\n# Skip if you already have one — just set ROLE_ARN above.\nimport json\n\niam = boto3.client(\"iam\")\nROLE_NAME = \"SageMakerTrainingSessionRole\"\n\ntrust = {\n \"Version\": \"2012-10-17\",\n \"Statement\": [{\n \"Effect\": \"Allow\",\n \"Principal\": {\"Service\": \"training-session.sagemaker.amazonaws.com\"},\n \"Action\": [\"sts:AssumeRole\", \"sts:TagSession\"],\n }],\n}\nperms = {\n \"Version\": \"2012-10-17\",\n \"Statement\": [\n {\"Effect\": \"Allow\",\n \"Action\": [\"sagemaker:CreateModelPackage\", \"sagemaker:DescribeHubContent\",\n \"sagemaker:DescribeModelPackageGroup\"],\n \"Resource\": \"*\"},\n {\"Effect\": \"Allow\",\n \"Action\": [\"s3:ListBucket\", \"s3:GetObject\", \"s3:PutObject\"],\n \"Resource\": [f\"arn:aws:s3:::{BUCKET}\", f\"arn:aws:s3:::{BUCKET}/*\"]},\n ],\n}\ntry:\n r = iam.create_role(RoleName=ROLE_NAME, AssumeRolePolicyDocument=json.dumps(trust))\n iam.put_role_policy(RoleName=ROLE_NAME, PolicyName=\"training-session-access\",\n PolicyDocument=json.dumps(perms))\n print(\"created role:\", r[\"Role\"][\"Arn\"])\nexcept iam.exceptions.EntityAlreadyExistsException:\n print(\"role already exists:\", f\"arn:aws:iam::{ACCOUNT_ID}:role/{ROLE_NAME}\")\n\nprint(f'\\nSet: ROLE_ARN = \"arn:aws:iam::{ACCOUNT_ID}:role/{ROLE_NAME}\"')", + "metadata": {}, + "execution_count": null, + "outputs": [] + }, + { + "cell_type": "markdown", + "id": "8756129a", + "source": "### Find a base model\n\n`BASE_MODEL` is the model you fine-tune **from**. Pass an id from\n`svc.list_base_models()` (listed below) and the SDK resolves it to the latest\nHub model ARN in your region. To pin a specific model version, pass that full\nHub model ARN as `BASE_MODEL` instead.\n\nCheckpoints this session saves are registered as model package *versions* in\nyour model package group and are addressed by ARN, so there is no name to\nchoose for them.", + "metadata": {} + }, + { + "cell_type": "code", + "id": "1e95ab81", + "source": "# Base models available to fine-tune in your region (pass one as BASE_MODEL above).\n# Queries the SageMaker Public Hub in REGION and returns each model's own\n# model id. Models are region-specific, so this can differ by region.\nfrom sagemaker.train.training_session import ServiceClient\n\nfor model_id in ServiceClient(region=REGION, role_arn=ROLE_ARN).list_base_models():\n print(model_id)", + "metadata": {}, + "execution_count": null, + "outputs": [] + }, + { + "cell_type": "markdown", + "id": "cell-7", + "metadata": {}, + "source": [ + "---\n", + "# Part 1 — Quick start\n", + "\n", + "## Create a Training Session\n", + "\n", + "`create_lora_training_client` creates the session and blocks until it reaches\n", + "`Running`, then returns a `TrainingClient` bound to it." + ] + }, + { + "cell_type": "code", + "execution_count": null, + "id": "cell-8", + "metadata": {}, + "outputs": [], + "source": [ + "from sagemaker.train.training_session import (\n", + " ServiceClient,\n", + " AdamParams,\n", + " Datum,\n", + " LossFunction,\n", + " ModelInput,\n", + " SamplingParams,\n", + " TensorData,\n", + " TensorDtype,\n", + ")\n", + "\n", + "svc = ServiceClient(\n", + " region=REGION,\n", + " role_arn=ROLE_ARN,\n", + " model_package_group_arn=MODEL_PACKAGE_GROUP_ARN,\n", + " s3_output_path=OUTPUT_S3_PATH,\n", + ")\n", + "tc = svc.create_lora_training_client(\n", + " base_model=BASE_MODEL,\n", + " training_session_name=TRAINING_SESSION_NAME,\n", + " rank=LORA_RANK,\n", + " alpha=LORA_ALPHA,\n", + ")\n", + "print(f\"Session running: {tc.training_session_name}\")" + ] + }, + { + "cell_type": "markdown", + "id": "cell-9", + "metadata": {}, + "source": [ + "## A minimal training step\n", + "\n", + "Pass plain text and let the SDK tokenize it (it loads the base model's tokenizer\n", + "on first use). Submit a `forward_backward` then an `optim_step`; each call\n", + "returns a future you await with `.result()`." + ] + }, + { + "cell_type": "code", + "execution_count": null, + "id": "cell-10", + "metadata": {}, + "outputs": [], + "source": [ + "# cross_entropy needs per-token loss inputs: which token to predict at each\n", + "# position (target_tokens) and how much each position counts (weights).\n", + "# Both run the full length of the tokenized input.\n", + "tokens = tc.get_tokenizer().encode(\n", + " \"Reinforcement learning teaches an agent to maximize reward.\"\n", + ")\n", + "\n", + "fb = tc.forward_backward(\n", + " [\n", + " Datum(\n", + " model_input=ModelInput.from_ints(tokens),\n", + " loss_fn_inputs={\n", + " \"target_tokens\": TensorData(\n", + " data=tokens[1:] + [tokens[-1]],\n", + " dtype=TensorDtype.INT64,\n", + " shape=[len(tokens)],\n", + " ),\n", + " \"weights\": TensorData(\n", + " data=[1.0] * len(tokens),\n", + " dtype=TensorDtype.FLOAT32,\n", + " shape=[len(tokens)],\n", + " ),\n", + " },\n", + " )\n", + " ],\n", + " loss_fn=\"cross_entropy\",\n", + ")\n", + "print(\"forward_backward metrics:\", fb.result().metrics)\n", + "\n", + "optim = tc.optim_step(AdamParams(learning_rate=1e-4))\n", + "print(\"optim_step metrics:\", optim.result().metrics)" + ] + }, + { + "cell_type": "markdown", + "id": "cell-11", + "metadata": {}, + "source": [ + "---\n", + "# Part 2 — Going further\n", + "\n", + "## Pipeline a real training loop over text\n", + "\n", + "Operations are pipelined: submit `forward_backward` and `optim_step` for a step\n", + "without blocking, then await both. The SDK orders them via auto-assigned\n", + "sequence numbers. Swap `loss_fn` for any supported objective\n", + "(`cross_entropy`, `importance_sampling`, `ppo`, `cispo`)." + ] + }, + { + "cell_type": "code", + "execution_count": null, + "id": "cell-12", + "metadata": {}, + "outputs": [], + "source": [ + "corpus = [\n", + " \"The capital of France is Paris.\",\n", + " \"Gradient descent minimizes a loss function step by step.\",\n", + " \"A transformer attends over all positions in a sequence.\",\n", + " \"LoRA adapts a base model with a small number of trainable parameters.\",\n", + "]\n", + "\n", + "def to_datum(text):\n", + " \"\"\"Build one cross_entropy training datum from a line of text.\"\"\"\n", + " ids = tc.get_tokenizer().encode(text)\n", + " return Datum(\n", + " model_input=ModelInput.from_ints(ids),\n", + " loss_fn_inputs={\n", + " \"target_tokens\": TensorData(\n", + " data=ids[1:] + [ids[-1]], dtype=TensorDtype.INT64, shape=[len(ids)]\n", + " ),\n", + " \"weights\": TensorData(\n", + " data=[1.0] * len(ids), dtype=TensorDtype.FLOAT32, shape=[len(ids)]\n", + " ),\n", + " },\n", + " )\n", + "\n", + "\n", + "for step, text in enumerate(corpus):\n", + " fb = tc.forward_backward([to_datum(text)], loss_fn=LossFunction.CROSS_ENTROPY)\n", + " opt = tc.optim_step(AdamParams(learning_rate=1e-4))\n", + " fb_res, opt_res = fb.result(), opt.result()\n", + " print(f\"step {step}: {fb_res.metrics} {text[:40]!r}\")" + ] + }, + { + "cell_type": "markdown", + "id": "cell-13", + "metadata": {}, + "source": [ + "## Inspect a forward-only pass\n", + "\n", + "`forward` runs the model without a gradient update — useful for evaluating loss\n", + "or reading `loss_fn_outputs` (e.g. per-token logprobs) on held-out text." + ] + }, + { + "cell_type": "code", + "execution_count": null, + "id": "cell-14", + "metadata": {}, + "outputs": [], + "source": [ + "# A forward-only pass still needs the loss inputs its loss function requires,\n", + "# built the same way as the training pass above.\n", + "eval_tokens = tc.get_tokenizer().encode(\"Paris is the capital of France.\")\n", + "\n", + "fwd = tc.forward(\n", + " [\n", + " Datum(\n", + " model_input=ModelInput.from_ints(eval_tokens),\n", + " loss_fn_inputs={\n", + " \"target_tokens\": TensorData(\n", + " data=eval_tokens[1:] + [eval_tokens[-1]],\n", + " dtype=TensorDtype.INT64,\n", + " ),\n", + " \"weights\": TensorData(data=[1.0] * len(eval_tokens)),\n", + " },\n", + " )\n", + " ],\n", + " loss_fn=LossFunction.CROSS_ENTROPY,\n", + ")\n", + "fwd_res = fwd.result()\n", + "print(\"eval metrics:\", fwd_res.metrics)" + ] + }, + { + "cell_type": "markdown", + "id": "cell-15", + "metadata": {}, + "source": [ + "## Checkpoint and resume\n", + "\n", + "`save_weights_for_sampler` registers inference weights as a model package.\n", + "`save_state` additionally saves optimizer state so training can resume;\n", + "`load_state` restores it." + ] + }, + { + "cell_type": "code", + "execution_count": null, + "id": "cell-16", + "metadata": {}, + "outputs": [], + "source": [ + "save_state = tc.save_state()\n", + "model_package_arn = save_state.result().model_package_arn\n", + "print(\"save_state model_package_arn:\", model_package_arn)\n", + "\n", + "tc.load_state_with_optimizer(model_package_arn).result()\n", + "print(\"load_state_with_optimizer: success\")" + ] + }, + { + "cell_type": "markdown", + "id": "cell-17", + "metadata": {}, + "source": [ + "## Sample from the trained checkpoint\n", + "\n", + "`save_weights_and_get_sampling_client` saves weights and returns a\n", + "`SamplingClient` bound to that checkpoint. `sample` returns a future. Use\n", + "`get_tokenizer()` to turn text into token IDs and to decode the result back to\n", + "text; you can also pass token IDs directly." + ] + }, + { + "cell_type": "code", + "execution_count": null, + "id": "cell-18", + "metadata": {}, + "outputs": [], + "source": "sc = tc.save_weights_and_get_sampling_client()\ntokenizer = sc.get_tokenizer()\nprompt_tokens = tokenizer.encode(\"What is reinforcement learning?\")\n\nresult = sc.sample(\n prompt=prompt_tokens,\n num_samples=1,\n sampling_params=SamplingParams(max_tokens=512, temperature=0.7),\n).result()\n\nfor seq in result.sequences:\n print(\"text:\", tokenizer.decode(seq.tokens))\n print(\"stop_reason:\", seq.stop_reason)" + }, + { + "cell_type": "markdown", + "id": "cell-19", + "metadata": {}, + "source": [ + "## Resume from a checkpoint\n", + "\n", + "After a notebook restart, resume from a `save_state` checkpoint in a new\n", + "session. `create_training_client_from_state_with_optimizer` restores weights\n", + "and optimizer state; `create_training_client_from_state` restores weights only." + ] + }, + { + "cell_type": "code", + "execution_count": null, + "id": "cell-20", + "metadata": {}, + "outputs": [], + "source": [ + "from sagemaker.train.training_session import ModelPackageKind\n", + "\n", + "ckpt = svc.get_latest_model_package(\n", + " ModelPackageKind.CHECKPOINT,\n", + " model_package_group_arn=MODEL_PACKAGE_GROUP_ARN,\n", + ")\n", + "resumed_tc = svc.create_training_client_from_state_with_optimizer(\n", + " model_package_arn=ckpt.model_package_arn,\n", + ")\n", + "print(f\"Resumed in a new session: {resumed_tc.training_session_name}\")" + ] + }, + { + "cell_type": "markdown", + "id": "cell-21", + "metadata": {}, + "source": [ + "## Cleanup\n", + "\n", + "Stop and delete every session this notebook created, so none keeps holding GPU capacity. `stop_if_active=True` stops a\n", + "running session and waits for a terminal state before deleting." + ] + }, + { + "cell_type": "code", + "execution_count": null, + "id": "cell-22", + "metadata": {}, + "outputs": [], + "source": [ + "# Both sessions this notebook created: the one from Part 1 and the one the\n", + "# resume section started. Leaving either running keeps billing GPU capacity.\n", + "# The resume section is optional, so its client may not exist. Collect the\n", + "# names first: naming a missing one here would raise before anything is\n", + "# deleted, leaving every session running.\n", + "names = [TRAINING_SESSION_NAME]\n", + "if \"resumed_tc\" in dir():\n", + " names.append(resumed_tc.training_session_name)\n", + "\n", + "for name in names:\n", + " svc.delete_session(name, stop_if_active=True)\n", + " print(f\"Deleted session: {name}\")" + ] + }, + { + "cell_type": "markdown", + "id": "f568435a", + "source": "## Sweep leaked sessions (optional)\n\nA safety net if a run failed before its cleanup cell and left sessions active.\nThis lists active sessions via `list_sessions()` and deletes them with\n`delete_session(stop_if_active=True)`. It is **dry-run by default** — set\n`DRY_RUN = False` to actually delete.", + "metadata": {} + }, + { + "cell_type": "code", + "id": "d9548bd5", + "source": [ + "DRY_RUN = True # set to False to actually stop + delete\n", + "\n", + "from sagemaker.train.training_session import SessionStatus\n", + "\n", + "ACTIVE = {SessionStatus.CREATED, SessionStatus.SCHEDULED, SessionStatus.RUNNING}\n", + "active_sessions = [s for s in svc.list_sessions() if s.status in ACTIVE]\n", + "\n", + "if not active_sessions:\n", + " print(\"No active sessions to clean up.\")\n", + "for s in active_sessions:\n", + " name = s.training_session_name\n", + " if DRY_RUN:\n", + " print(f\"[dry-run] would delete {name!r} (status={s.status})\")\n", + " else:\n", + " svc.delete_session(name, stop_if_active=True)\n", + " print(f\"Deleted {name!r}\")" + ], + "metadata": {}, + "execution_count": null, + "outputs": [] + } + ], + "metadata": { + "kernelspec": { + "display_name": "Python 3 (ipykernel)", + "language": "python", + "name": "python3" + }, + "language_info": { + "name": "python" + } + }, + "nbformat": 4, + "nbformat_minor": 5 +} \ No newline at end of file