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