Repository navigation
Conversation
…m ModelTrainer.train() `ModelTrainer.train()` returned `None` and kept the created job on the private `_latest_training_job`. A caller that needed the job name or ARN after `train(wait=False)` had to read the private attribute or page `ListTrainingJobs` by name prefix. The SDK's own docs, seven example notebooks, and `ModelBuilder` read the private attribute too, and `DPOTrainer` stored its job on a different attribute than the other serverless trainers, so `show_metrics()` and `stream_logs()` could not find a DPO job. Add a public `latest_training_job` property with a typed setter on `ModelTrainer` and `BaseTrainer`, return the created `TrainingJob` from `ModelTrainer.train()` for a SageMaker training job (`None` for a dry run, a pipeline capture, and local mode), store the DPO job on the shared attribute, and point the docs and error messages at the public name. --- X-AI-Prompt: How do I get the training job that ModelTrainer.train() created without reading a private attribute? X-AI-Tool: claude-code
…s every deploy path
`deploy()` took `tags` only through `**kwargs`, so `inspect.signature` could
not validate it, and two deploy paths dropped it: the fine-tuned model package
path read only `self._tags`, and the Nova path did not read the kwarg at all.
The accumulated tags also had two shapes in one SDK. `add_tags` compared
`tag["Key"]` and raised `KeyError: 'Key'` for the `{"key": ..., "value": ...}`
form that the `sagemaker.core.resources` classes take, and raised `TypeError`
for a `sagemaker.core.shapes.Tag`, while the fine-tuned path converted the
stored tags back to lowercase by hand.
Add `tags` as a named `deploy()` parameter. Normalize every accepted form (a
`{key: value}` mapping, boto-style dicts, core-style dicts, `Tag` objects) once
in `format_tags`, and make `add_tags`, `tag_exists`, `_validate_new_tags`, and
`remove_tag_with_key` go through it. Merge the caller's tags into the builder
before dispatch so the built-model, inference-component, recommendation,
fine-tuned, restricted-package, and Nova paths all tag the endpoint, and
replace the two hand-written case conversions with one helper.
---
X-AI-Prompt: Why does ModelBuilder.deploy drop my tags on a fine-tuned model, and why does add_tags raise KeyError: 'Key'?
X-AI-Tool: claude-code
…isted hub content `HubContent` had no `get_all`. The resource generator skips a List API when the list summary lacks a member the Describe input requires, and `HubContentInfo` omits `HubName`, the name the list call was made with. So every caller paged `list_hub_contents` on the boto3 client by hand: six places in the SDK and three scripts in the SageMaker agent skills repeat the same loop. `HubContentInfo` also carries `OriginalCreationTime`, the model-level catalog date, which `DescribeHubContent` does not return, and the `HubContent` class had no field for it, so `ResourceIterator` dropped it. The same gap made `HubContent.get_all_versions()` build each version without `hub_name`, so the first `refresh()` failed with `Missing required parameter in input: "HubName"`. Teach the generator two per-resource tables: list-input members to copy onto every listed resource (`HubName` -> `hub_name`), and summary-only members to keep as optional attributes (`OriginalCreationTime`). `ResourceIterator` takes the copied members as `extra_init_data`. `resources.py` is the regenerated output: `HubContent.get_all(hub_name, hub_content_type, name_contains, max_schema_version, creation_time_before, creation_time_after, sort_by, sort_order, session, region)`, the `original_creation_time` attribute, and the `hub_name` binding on `get_all_versions()`. --- X-AI-Prompt: Why does sagemaker-core have no HubContent.get_all, and why does get_all_versions() fail on refresh with a missing HubName? X-AI-Tool: claude-code
… and document
A hub model's task, license, provider, size, context window, languages,
modalities, model type, capabilities, and recipes live in
`HubContentSearchKeywords` as prefixed strings (`@task:`, `@license:`,
`@recipe:finetuning_sft_lora`, ...) and in the `HubContentDocument` JSON. The
`HubContent` resource exposed the two raw attributes and nothing else, so
`jumpstart/search.py`, `recipe_utils.py`, and `training_session/_utils.py` each
split the keywords for the one field they use, and the SageMaker agent skills
repeat the full parse in four scripts. The hub writes at most 50 keywords per
model, so a model with many languages loses its `@model-size:` and
`@context-window:` keywords while the document still carries the values. No
consumer handled that.
Add `HubContentMixin` in `sagemaker.core.hub_content` and teach the resource
generator a `RESOURCE_MIXINS` table so `resources.py` emits
`class HubContent(HubContentMixin, Base)`. The mixin adds no model field and
no state. Its properties read `hub_content_search_keywords` and
`hub_content_document` at call time: a keyword is authoritative and the
document fills what the cap dropped, so the same resource answers before and
after `refresh()`. `HubContentRecipe` parses `@recipe:` keywords and
`RecipeCollection` entries into type, technique, and strategy, and the resource
answers `supports_technique("sft")` and `fine_tuning_techniques()` directly.
---
X-AI-Prompt: How do I read a hub model's license, size, and supported fine-tuning techniques from HubContent without parsing the search keywords myself?
X-AI-Tool: claude-code
A caller that wants the job name or ARN of a long run has two choices today. `train(wait=True)` returns after the job ends, so the handle arrives hours late. `train(wait=False)` returns the job at once, but the caller then has to wait on it alone, and the shared fine-tuning wait lives in the private `sagemaker.train.common_utils.trainer_wait.wait`. Five SageMaker agent skill templates import that private function for this one reason. Add `on_submit: Optional[Callable[[TrainingJob], None]]` to `BaseTrainer.train`, `ModelTrainer.train`, and every trainer `train()` override. The trainer calls it once with the created job as soon as the service accepts the job and before any wait. A `dry_run` and a `PipelineSession` capture submit nothing and call nothing. The serverful path submits through `ModelTrainer` first and waits on the returned job afterwards, so the callback sees the job before the wait on that path too. `_record_submitted_job` stores `latest_training_job` and fires the callback in one place, so the serverless, serverful, and HyperPod paths cannot drift. --- X-AI-Prompt: How do I get the training job name or ARN right after submit when I call train(wait=True)? X-AI-Tool: claude-code
![github-actions[bot]](https://avatars.githubusercontent.com/in/15368?s=60&v=4){kind=small}
There was a problem hiding this comment.
🤖 AI code review
This PR adds four public handles the SDK previously lacked: HubContent.get_all() plus a HubContentMixin with typed catalog accessors, a code-generator mechanism (RESOURCE_MIXINS, LIST_SUMMARY_EXTRA_ATTRIBUTES, GET_ALL_INPUT_ATTRIBUTES) that produces them, a public latest_training_job property/setter and a return value from ModelTrainer.train(), and a ModelBuilder.deploy(tags=...) parameter backed by unified tag normalization (_normalize_tag/format_tags). The implementation is careful and well tested, with new regression tests for every change, and I found no correctness or compatibility regressions — the format_tags rewrite still satisfies the existing test_common_utils.py cases. Only a few minor doc/annotation nits below.
Reviewed commit d352fe5. Automated review; verify before acting.
|
|
||
| Raises: | ||
| botocore.exceptions.ClientError: This exception is raised for AWS service related errors. | ||
| The error message and error code can be parsed from the exception as follows: |
There was a problem hiding this comment.
🟡 Low: The get_all docstring documents max_results and next_token parameters, but neither exists in the method signature (ResourceIterator handles pagination internally). These stray entries will confuse readers; drop them from the docstring.
There was a problem hiding this comment.
Dropped. The generator now excludes the paging members from every get_all docstring as it excludes them from the signature (ResourceIterator pages for the caller), which removes the two stray lines here and 161 more across resources.py. Commit a20a1e6.
| } | ||
|
|
||
|
|
||
| def _normalize_tag(tag: Any) -> TagsDict: |
There was a problem hiding this comment.
🟡 Low: format_tags is annotated -> List[TagsDict] but now returns None unchanged when tags is None. Consider Optional[List[TagsDict]] to match the actual contract (several callers like _deploy_recommendation rely on the None passthrough).
There was a problem hiding this comment.
Changed to Optional[List[TagsDict]], with Optional[Tags] on the parameter. Commit a20a1e6.
| tags = kwargs.get("tags", getattr(self, "_tags", None)) | ||
| if "tags" in kwargs: | ||
| self._tags = tags | ||
| if "tags" in kwargs and tags is not None: |
There was a problem hiding this comment.
🟡 Low: _deploy_core_endpoint now calls self.add_tags(tags) here, and then the pre-existing block at lines 4001-4005 calls self.add_tags(tags) again on the same (already-merged) tags before re-reading self._tags. It's harmless because _validate_new_tags de-dups by key, but the redundant merge is dead work; the earlier block could just set tags without re-adding.
There was a problem hiding this comment.
Removed the second merge. _deploy_core_endpoint now merges the kwarg tags once and reads format_tags(self._tags) for every request. Commit a20a1e6.
… an unknown kwarg `deploy()` took `data_capture_config`, `endpoint_type`, and `inference_component_name` through `**kwargs`, and only the core endpoint path read them, so a fine-tuned or Nova deploy dropped `data_capture_config` without a word. `fallback_instance_types` and `tracing_config` were read nowhere, and a misspelled option deployed an endpoint without it. The `update_endpoint` path passed `create_endpoint_config` eleven keyword arguments the method does not take. Make the five options named parameters, validated once against `inference_config` in `_EndpointOptions`, and applied on every path: - `data_capture_config` reaches the endpoint configuration of a single-model, serverless, async, fine-tuned, restricted, and Nova deploy. An inference-component path raises `ValueError`, because the service rejects data capture there. - `fallback_instance_types` builds `InstancePools` with `instance_type` at priority 1 and up to 4 fallbacks after it. `session_helper.production_variant` takes `instance_pools` and emits `InstancePools` in place of `InstanceType`, as the service contract requires. - `tracing_config` reaches `CreateEndpointConfig` through the session, which raises `ValueError` when the installed botocore service model has no `TracingConfig` member instead of letting botocore drop it. The generated `EndpointConfig.create` gains the member once sagemaker-core is regenerated from a service model that has it. - `endpoint_type` and `inference_component_name` keep their meaning and now raise on a path that cannot apply them. - A kwarg outside `DEPLOY_KWARGS` raises `TypeError` before any service call. - `update_endpoint` creates the new endpoint configuration from the built production variant with the arguments `create_endpoint_config` takes. --- X-AI-Prompt: Why did ModelBuilder.deploy(data_capture_config=...) deploy my fine-tuned model without data capture, and how do I set fallback instance types? X-AI-Tool: claude-code
The V3 structured-configuration example passes `use_spot_instances=True` and `max_wait_time_in_seconds=7200` to `Compute`. `Compute` forbids extra fields, so the example fails on 3.27.0 with a pydantic `extra_forbidden` error. The spot flag is `Compute(enable_managed_spot_training=True)` and the wait limit lives on `StoppingCondition(max_wait_time_in_seconds=7200)`, next to `max_runtime_in_seconds`. --- X-AI-Prompt: Why does the migration guide's Compute(use_spot_instances=True) example fail with extra_forbidden? X-AI-Tool: claude-code
… API
The service indexes hub content for `Search(Resource="HubContent")` and
botocore 1.43.100 added the typed `SearchRecord.HubContent` record, but the
SDK never read it. A caller who wanted the public hub narrowed by license,
provider, task, or model size either walked every model with
`ListHubContents` or wrote the `Search` request, the `Metadata.<Field>` names,
the 20-filter and 1024-character limits, the paging, and the result parsing by
hand. The SageMaker agent skills carry a 955-line adapter for exactly this.
Add `HubContent.search(hub_name, filters, hub_content_type, session, region)`
to `HubContentMixin`. It builds the `And` expression of `HubName`,
`HubContentType`, and the caller's `HubContentFilter` entries, sorts by
`HubContentName` so the pages are stable, pages until the service returns no
`NextToken` or repeats one, and yields each record as a `HubContent` bound to
the session and region, with its keywords and document. `HubContentFilter`
holds one service filter. `HubContentFilter.metadata("license", "apache-2.0")`
maps a `HubContent` attribute name to its `Metadata.<Field>` name, and a
sequence of values selects the `In` operator. The filter checks the operator,
the value length, and the filter count against the service limits before the
request.
---
X-AI-Prompt: How do I search the SageMaker public hub for Apache-2.0 models from Meta without listing every model?
X-AI-Tool: claude-code
Commit 6e9df41 made `format_tags` normalize every tag through `_normalize_tag`, which raised `TypeError` for a `(key, value)` tuple. `Processor(tags=[("Key", "Value")])`, `ModelBiasMonitor(tags=...)`, and the feature processor `to_pipeline(tags=[(key, value)])` callers pass pairs, and three unit tests in `test_processing.py` and `test_clarify_model_monitoring.py` failed with it. Accept a two-item tuple or list as one tag. --- X-AI-Prompt: Why does Processor(tags=[("Key", "Value")]) raise TypeError after the deploy(tags=) change? X-AI-Tool: claude-code
Callers that want training-loss values, not a PNG, set the MLflow tracking URI from `training_job.mlflow_config.mlflow_resource_arn` and call `MlflowClient().get_metric_history` with `training_job.mlflow_details.mlflow_run_id` themselves. The SDK repeated those lines in `plot_training_metrics` and `get_available_metrics`, and each caller copied the recipe-specific metric names (`total_loss` for LLMFT recipes, `train/loss` for Verl recipes). `get_metric_history` returns `Dict[str, List[MetricPoint]]` sorted by step. It accepts a job object or a job name, defaults to every logged metric, and resolves caller names through `DEFAULT_METRIC_ALIASES` (`loss`, `val_loss`), first logged candidate wins. An unknown metric raises `ValueError` that lists the metrics the run has. The tracking-URI and run-id resolution now lives in one private helper that the three public functions share. The lazy `__getattr__` in `sagemaker.train` exports `get_metric_history` and `MetricPoint`, and loses its duplicate `plot_training_metrics` through `get_mlflow_url` branches. --- X-AI-Prompt: How do I read the training loss history of a SageMaker training job as data instead of a plot? X-AI-Tool: claude-code
…time `verify_reward_function` accepted any dict as a sample, sent every sample in one invocation, and selected the response format only through `is_nova`. The evaluation custom-scorer contract differs: one invocation per sample, with `id`, `model_response`, `query`, `response`, and `reference_answer` fields, and the `statusCode` and `body` envelope for Nova too. The function and `RewardFunctionOutput` were importable only from `sagemaker.train.common_utils.rlvr_reward_verifier`. `Evaluator.create` deployed a reward Lambda from a file with `python3.9` (`EVALUATOR_DEFAULT_RUNTIME`) while `CustomAgentLambda.create` carried its own `python3.12` literal. `RewardSample` and `EvalRewardSample` are the typed inputs. The `envelope` option (`oss`, `nova`, `eval_oss`, `eval_nova`) selects the request shape, the response shape, and the sample model, and `None` keeps the `is_nova` behavior. The options live in `RewardVerificationOptions`, which the function takes as an object or as keyword arguments, so `verify_reward_function(reward_function, sample_data, is_nova=False)` still works. `sagemaker.train` exports `verify_reward_function`, `RewardVerificationOptions`, `RewardFunctionOutput`, `RewardSample`, and `EvalRewardSample`. `EVALUATOR_DEFAULT_RUNTIME` is `python3.12`, `Evaluator.create(runtime=...)` reaches the Lambda create call, and `CustomAgentLambda.create` defaults to the same constant. --- X-AI-Prompt: How do I test my evaluation custom-scorer Lambda locally with the exact payload the evaluation sends, and why does Evaluator.create deploy python3.9? X-AI-Tool: claude-code
`LLMAsJudgeEvaluator._validate_custom_metrics_json` only parsed the JSON. A metric whose
`instructions` had no `{{prompt}}` or `{{prediction}}`, an array of 11 metrics, a `ratingScale`
`definition` of 101 characters, or a `value` with neither or the two of `floatValue` and
`stringValue` reached the Bedrock evaluation job and failed there, after the pipeline started.
`LLMAsJudgeEvaluator.validate_custom_metrics(custom_metrics)` is a static method that returns
every violation as a message, an empty list means valid. It accepts the list or its JSON
string. The rules: 1 to 10 metrics, each with a `customMetricDefinition` that has a string
`name`, `instructions` of at most 5000 characters that contain `{{prompt}}` or `{{prediction}}`,
and an optional `ratingScale` whose entries have a `definition` of at most 100 characters and
a `value` with exactly one of `floatValue` or `stringValue`, one value type across the scale.
`_validate_custom_metrics_json` calls it and raises `ValueError` with the joined messages, so
`evaluate()` stops before any upload. Three existing tests passed `[{"name": "test"}]` as a
valid payload and now use a valid metric.
---
X-AI-Prompt: Can the SDK check my LLM-as-Judge custom metrics JSON before I start the evaluation job?
X-AI-Tool: claude-code
…ad metadata `S3Downloader.read_bytes` downloaded the whole object through `s3_resource.Bucket(bucket).download_fileobj`, so a caller that wanted the first megabyte of a dataset, one line at a time, or the object size had to drop to boto3. Fifteen SageMaker agent skill scripts split `s3://` URIs by hand, eight call `get_object(Range=...)`, and four call `head_object`, each without the `ExpectedBucketOwner` spot check that the SDK applies on the default bucket. Add `byte_range=(start, end)` to `read_bytes`, inclusive like an HTTP `Range` header, served through `get_object`. Add `iter_lines`, which decodes `Body.iter_lines()` and never holds the whole object, and `head`, which returns a frozen `S3ObjectHead(content_length, last_modified, etag, content_type)` from `head_object`. The three methods share `parse_s3_url`, the session's `s3_client` with the same fallback as `Session.read_s3_file`, and the `_get_account_id_if_default_bucket` spot check. A negative start or an end before the start raises `ValueError` before any request. --- X-AI-Prompt: How do I read the first megabyte of an S3 object or stream its lines with the SageMaker SDK instead of boto3? X-AI-Tool: claude-code
`DatasetFormatDetector.validate_dataset` answered only `True` or `False` for a local path. It stopped at the first line, named no format, read no `s3://` source, and knew none of the Nova RLVR, PPO, and CPT shapes. Two SageMaker agent skill scripts filled the gap with their own detectors (11 and 9 format names, their own `s3://` parsing, their own `get_object(Range=)` call), and the two disagree on names for the same shape. Add `detect_dataset_format(source, options)` in `dataset_format_detector.py`. It reads at most `sample_size_bytes` (1 MiB default) through `S3Downloader.read_bytes(byte_range=)` for `s3://` or from the local file, drops the last partial line, and matches every complete line against `FORMAT_MATCHERS`: the 7 SDK JSON schemas through `jsonschema`, plus `nova_rlvr`, `nova_ppo`, and `nova_cpt` as small pure functions. The frozen `FormatDetectionResult` carries the format, the matched-line ratio, the sampled line count, `"<line>: <reason>"` entries for the lines that failed, and every candidate with its ratio. `strict=True` demands a full match. `strict=False` picks the best candidate above a 0.5 ratio. `.parquet`, `.csv`, and `.json` keep the detector's whole-file checks. `SKILL_FORMAT_ALIASES` maps the 19 skill names onto the matcher names. Export the function, the options, and the result from `sagemaker.ai_registry`. --- X-AI-Prompt: How do I find out which fine-tuning format a JSONL dataset in S3 uses before I create a DataSet from it? X-AI-Tool: claude-code
…ndle Problem: the skill templates call bedrock and bedrock-runtime by hand around BedrockModelBuilder. get_bedrock_candidates.py pages list_foundation_models and list_inference_profiles and joins the two on models[].modelArn. deploy-oss-bedrock.py loops invoke_model for 12 minutes on ModelNotReadyException after the import job completes. deploy-nova-bedrock.py calls converse with the deployment ARN. The SDK opens three Bedrock clients and lists nothing and invokes nothing. deploy() returns a response dict, so the caller picks the ARNs out of it by hand. Solution: add sagemaker.serve.bedrock_deployment. list_invocable_foundation_models() pages ListFoundationModels filtered by output modality and ListInferenceProfiles filtered to SYSTEM_DEFINED, then joins profile ids to model ids. It keeps a model when inferenceTypesSupported holds ON_DEMAND or an active cross-region profile routes to it. BedrockDeployment.wait_until_ready() probes Converse for a custom model deployment and InvokeModel for an imported model. It retries only ModelNotReadyException and raises a TimeoutError that names the ARN. invoke() does a JSON round trip, and converse() passes inferenceConfig and system through. BedrockModelBuilder.deploy() keeps its return value and sets deployment on every successful path (importedModelArn on the OSS path, modelArn plus customModelDeploymentArn on the Nova path). latest_deployment() raises ValueError before the first deploy(). Field names come from the botocore 1.43.118 bedrock and bedrock-runtime service models. Tests: sagemaker-serve/tests/unit/test_bedrock_deployment.py fails on master with ModuleNotFoundError: No module named 'sagemaker.serve.bedrock_deployment'. On this branch test_bedrock_model_builder.py and test_bedrock_deployment.py report 124 passed. --- X-AI-Prompt: Can BedrockModelBuilder list the Bedrock models I can invoke, wait until an import is ready, and invoke it without raw boto3 calls? X-AI-Tool: claude-code
ModelBuilder held two private copies of the same EC2 lookup. `_get_instance_resources` and `_infer_accelerator_count_from_instance_type` each created an EC2 client, stripped the `ml.` prefix with `str.replace`, called `describe_instance_types` for one name, and read a different slice of the response. No public SDK function listed the SageMaker instance type names or described their hardware, so every caller outside ModelBuilder wrote its own EC2 call. `sagemaker.core.instance_types` gains `InstanceSpec`, `list_instance_types`, `describe_instance_type`, and `describe_instance_types`. `list_instance_types` reads the `ProductionVariantInstanceType` enum of the loaded SageMaker service model and makes no request. `describe_instance_types` strips the `ml.` prefix, sends at most 100 names per EC2 request (the `RequestInstanceTypeList` limit), reads every page, and maps vCPUs, memory, GPU count, GPU memory, and GPU name. EC2 rejects a whole request with `InvalidInstanceType` when one name does not exist, so the function then asks for each name alone. An unknown type gets `gpu_count=0`, `None` elsewhere, and a warning. Every other `ClientError` propagates. The two ModelBuilder helpers call `describe_instance_type` with the session's `boto_session` and keep their signatures and fallbacks: `(None, None)` and `None` on an unknown type or a failed call. The existing `test_get_instance_resources` stub now carries `InstanceType`, which the real EC2 response always holds and the mapper keys on. The module joins the `sagemaker-core` API reference. --- X-AI-Prompt: Can sagemaker-core list and describe SageMaker instance types, so ModelBuilder drops its two private EC2 lookups? X-AI-Tool: claude-code
`DescribeModelPackageGroup` reports a missing group as a `ValidationException`
whose message ends in `does not exist`, not as a `ResourceNotFound` error, and
`ModelPackageGroup` offered only the generated `get` and `create`. So every
"use the group, or create it" flow parsed `ClientError` codes by hand. Ten
fine-tuning templates in the SageMaker agent skills wrap `create` in a `try`
block and match `e.response["Error"]["Code"]` against
`("ResourceInUse", "ValidationException")` to fall back to `get`, which also
treats an unrelated validation error as "the group exists".
Add `sagemaker.core.model_package_group.ModelPackageGroupMixin` and attach it
to the generated class through the `RESOURCE_MIXINS` table, so `resources.py`
emits `class ModelPackageGroup(ModelPackageGroupMixin, Base)`.
`get_or_create(name, spec=None, session=None, region=None)` calls the
generated `get`, and on a `ValidationException` whose message says the group
is absent, or on a code that starts with `ResourceNotFound`, calls the
generated `create` with the set fields of a `ModelPackageGroupSpec`
(description, tags, managed configuration). Every other error propagates. The
mixin adds no model field and no state, and the returned group keeps the
session and region it was loaded with.
---
X-AI-Prompt: Can ModelPackageGroup get or create a group in one call instead of parsing ClientError codes?
X-AI-Tool: claude-code
…cessor._start_new `Processor._start_new` built `ProcessingJob(**transformed)` from the create request. The job carried no `processing_job_arn`, no `processing_job_status`, and no client context, so `processor.latest_job.refresh()` built its client from the process default session and polled the default region (aws#5796). Four processing templates in the SageMaker agent skills call `describe_processing_job` on a boto3 client after `run()` to read the ARN. After the submit through the configured session, bind that session and its region to the job with `_set_client_context` and describe the job once with `refresh()`, so the returned job has its ARN and status and every later `refresh()`, `wait()`, and `stop()` talks to the account and region the job was created in. The `PipelineSession` path still returns the captured step arguments with no describe call. A `LocalSession` job exists only in the session's `LocalSagemakerClient`, so that path returns the request-built job as before and `run()` keeps polling it through the local client. --- X-AI-Prompt: Why does processor.latest_job have no ARN and refresh() poll the wrong region after run()? X-AI-Tool: claude-code
A generated `wait()` raised `FailedStatusError` only for a status that contains `failed`. `Stopped` is a terminal state too, so `TrainingJob.wait()` returned on a stopped job as if the job completed, and the caller had to read `training_job_status` after the wait to tell the two apart. Three training templates in the SageMaker agent skills do that and raise on `Stopped` themselves. Add `raise_on_stopped: bool = False` after `timeout` to `WAIT_METHOD_TEMPLATE` and a `STOPPED_STATUS_ERROR_TEMPLATE` block that the generator formats with the same failure-reason reference as the failed block. With `wait(raise_on_stopped=True)` a status that contains `stopped` raises `FailedStatusError` with the resource type, the status, and the failure reason. The default keeps the old return. Regenerated `resources.py`: all 15 `wait()` methods take the parameter. --- X-AI-Prompt: Can TrainingJob.wait() raise when the job is stopped instead of returning like it completed? X-AI-Tool: claude-code
When an endpoint misbehaves, an operator reads the endpoint status and failure reason, the per-variant instance type and counts, the 11 CloudWatch metrics of the endpoint dashboard, and the latest lines of the `/aws/sagemaker/Endpoints/<name>` log group. `Endpoint` exposed none of that beyond `refresh()`, so the SageMaker agent skill `collect_diagnostics.py` rebuilt the collection with raw boto3 clients from a region string. That script reads `InstanceType` from `ProductionVariantSummary`, where `DescribeEndpoint` never returns it, so the instance type is always `None`. It fetches one unpaged page of 100 log events, and it loses the session a caller bound the resource with. Add `EndpointMixin` in `sagemaker.core.endpoint` and register it in the generator's `RESOURCE_MIXINS` table, so `resources.py` emits `class Endpoint(EndpointMixin, Base)`. `diagnostics()` refreshes the resource, reads the instance type per variant from `EndpointConfig.get(endpoint_config_name)` matched by variant name, sends the 11 dashboard queries in one `GetMetricData` call keyed `"<MetricName>/<Stat>"`, and collects the log lines. A CloudWatch, Logs, or endpoint config failure lands in `errors` instead of raising. Only `DescribeEndpoint` raises. `recent_logs()` pages `FilterLogEvents` until the next token is absent, the page is empty, or the limit is reached, and formats `"<ISO timestamp> <log stream> <message>"` lines. The CloudWatch and Logs clients come from the session, region, and config the resource is bound to, so `Endpoint.get(name, session=..., region=...)` reads the metrics and logs of that account and region. The mixin adds no model field and no state. --- X-AI-Prompt: How do I get the status, metrics, and recent logs of a SageMaker endpoint from the Endpoint resource without writing my own boto3 calls? X-AI-Tool: claude-code
HyperparameterTuner builds a HyperParameterTrainingJobDefinition, which has no serverless_job_config, so AMT cannot run SFTTrainer, DPOTrainer, RLVRTrainer, or RLAIFTrainer trials. Users wrote their own sweep loops with a JSON manifest, a thread pool, and MLflow polling. FineTuningSweep runs a trainer whose train(wait=False) returns a TrainingJob or sets latest_training_job. SweepConfig selects a grid, random, or ASHA schedule with a bounded number of parallel jobs. ASHA reads the rung metric from the job's MLflow run and stops the bottom 1 - 1/reduction_factor share of each batch. With manifest_s3_uri the sweep writes its state to S3 after every trial change, and resume() skips terminal trials and attaches to running jobs by name. Trials poll training_job_status instead of TrainingJob.wait(), because rich 13 permits one live display per console and parallel waits raise LiveError. --- X-AI-Prompt: How can I run a hyperparameter sweep for SFTTrainer when HyperparameterTuner only supports ModelTrainer? X-AI-Tool: claude-code
There was a problem hiding this comment.
🤖 AI code review
This large PR adds sixteen public SDK surfaces across sagemaker-core and sagemaker-train: HubContent metadata/search/caching helpers, EndpointMixin diagnostics, ModelPackageGroupMixin.get_or_create, instance-type lookups, a reworked format_tags, trainer latest_training_job/on_submit handles, metric history, reward-function envelopes, FineTuningSweep, and ModelBuilder.deploy options. The code is well-structured (hand-written behavior isolated into mixin modules) and ships extensive unit tests. Overall it looks solid; my comments are a correctness subtlety in the telemetry-decorated generator and a couple of minor documentation/robustness notes.
Reviewed commit 8f9ea7f. Automated review; verify before acting.
| hub_name: str, | ||
| filters: Optional[Sequence[HubContentFilter]] = None, | ||
| hub_content_type: str = "Model", | ||
| session: Optional[Any] = None, |
There was a problem hiding this comment.
🟡 Low: search is a generator function (it yields), but it is wrapped by @_telemetry_emitter. The telemetry wrapper calls response = func(*args, **kwargs) and measures latency/success around that call — for a generator this returns immediately without running the body, so every search is recorded as a near-zero-latency success and any exception raised during iteration (including the ValueError from _search_expression/to_request) is never seen by the telemetry path. It also means the documented Raises: ValueError for bad filters is deferred until the caller starts iterating, not raised from the search() call itself, unlike the eager validation in get_all_cached. Consider validating the expression eagerly (e.g. build _search_expression before entering the generator, or split into an eager wrapper + inner generator) so the telemetry and the documented Raises behave as described.
There was a problem hiding this comment.
Changed to a materialized list, see the thread on the search signature. Commit a081545.
| """Typed accessors over ``hub_content_search_keywords`` and ``hub_content_document``. | ||
|
|
||
| Mixed into the generated ``sagemaker.core.resources.HubContent``. Every accessor reads | ||
| the two attributes at call time, so a ``refresh()`` that fills the document changes the |
There was a problem hiding this comment.
🟡 Low: The docstring example HubContentFilter.metadata("license", "apache-2.0") uses a lowercase value, but the team precedent and migration.md in this same PR note that Metadata.<Field> values are case-sensitive and must match the document spelling exactly (Apache-2.0). The docs/training/index.rst example here also uses apache-2.0 while migration.md uses Apache-2.0. Aligning these examples would avoid users building filters that silently return nothing.
Source: #6533
There was a problem hiding this comment.
Aligned on the document spelling, see the thread on docs/training/index.rst. Commit a081545.
| Returns: | ||
| dict[str, InstanceSpec]: One ``InstanceSpec`` per requested name, keyed by that name. | ||
|
|
||
| Raises: |
There was a problem hiding this comment.
🟡 Low: describe_instance_types builds the EC2 client with _boto_session(session).client("ec2") and no region. When neither the passed session nor the environment supplies a default region, boto3 raises NoRegionError here, but unlike list_instance_types the docstring's Raises: section does not mention it. Either pass a region/document the behavior for consistency with list_instance_types.
There was a problem hiding this comment.
Added region to describe_instance_types and describe_instance_type, passed to the EC2 client as list_instance_types does, and NoRegionError to the Raises sections. Test test_describe_instance_types_without_a_region_raises_no_region_error. Commit a081545.
…t cache shape `HubContent.get_all_cached` kept its own cache design: a file per hub read on every call, a `max_age` argument, and no in-memory layer. The JumpStart models cache (`sagemaker.core.jumpstart.cache.JumpStartModelsCache`) is the SDK precedent for cached catalog metadata. It is an `LRUCache` from `sagemaker.core.utilities.cache` with a bounded item count (20) and an expiration horizon (6 hours). The `LRUCache` calls a retrieval function on a miss, or on an expired entry with the old value. A static `JumpStartModelsAccessor` holds the process cache and takes new kwargs through `set_cache_kwargs`. Rebuild `HubCatalogCache` on that shape. The `LRUCache` holds one entry per hub and content type with the JumpStart defaults. The retrieval function lists the hub and reuses every item whose `HubContentVersion` the list still reports. The rest are described 4 at a time. The JumpStart retrieval function reads the manifest again only when its ETag changed, in the same way. The described items persist to one JSON file per hub, content type, region, and account, so a new process starts from the previous catalog. `HubCatalogAccessor` holds the process cache, and `HubContent.set_cache_kwargs(cache_expiration_horizon=..., max_cache_items=..., max_workers=..., cache_dir=...)` rebuilds it. The `LRUCache` is not thread safe, so its reads and writes run under one lock, and a retrieval runs outside it under a per-key lock: a slow hub blocks no hit on another hub, and two threads that miss the same hub share one list-and-describe pass. `LRUCache.put` on a present key read `curr_value.element` from an `Element` whose attribute is `value` and raised `AttributeError`. The JumpStart cache never reached that line, because it refreshes through `get`. The fix reads `curr_value.value`, with a regression test. Measured on the public hub (766 models): the first call takes 59.6 seconds (8 list calls, 766 describes). A call inside the horizon takes 0.3 seconds with no request. `refresh_cache` on the unchanged hub takes 1.6 seconds (8 list calls, 0 describes). A new process reads the 32.9 MB file in 2.5 seconds. --- X-AI-Prompt: Is the hub content cache in line with how the SageMaker SDK caches JumpStart model metadata? X-AI-Tool: claude-code
There was a problem hiding this comment.
🤖 AI code review
This large PR adds many public SDK surfaces: HubContent metadata accessors / search / cached listing (hub_content.py, hub_catalog.py), Endpoint.diagnostics/recent_logs (endpoint.py), ModelPackageGroup.get_or_create, instance-type lookups, BedrockDeployment, typed ModelBuilder.deploy() options, a rewritten format_tags, BaseTrainer.latest_training_job as a validating property with on_submit, and session-helper tracing/instance-pool support. The implementation is consistent with existing patterns (JumpStart cache shape, resource mixins) and is backed by extensive unit tests. Overall the code is high quality; I found one telemetry-correctness issue on the generator-based search, plus two lower-severity notes worth the author's time.
Reviewed commit f529db2. Automated review; verify before acting.
| """``True`` when the hub tags the model with the ``customization`` capability.""" | ||
| return CUSTOMIZATION_CAPABILITY in self.capabilities | ||
|
|
||
| @property |
There was a problem hiding this comment.
🟡 Low: HubContent.search is a generator function (it uses yield), but it is wrapped with @_telemetry_emitter. The telemetry wrapper calls func(*args, **kwargs) and treats the result as the finished return value. For a generator function that call returns a generator object immediately without executing any of the body, so:
- the telemetry event is recorded as SUCCESS with ~0 latency before any
SearchAPI call is actually made; - any
ClientError/ValueErrorraised while the caller iterates (e.g. a bad filter fromto_request(), or_search_expressionexceedingMAX_SEARCH_FILTERS) happens after the wrapper has already returned, so it is never captured by the FAILURE path.
The search itself still works, so this is telemetry correctness only. Consider making search return a materialized list (as get_all_cached does) or wrapping the eager validation (_search_expression) in a non-generator entry point so the decorated call executes the request-building work. get_all_cached/refresh_cache return lists and are unaffected.
There was a problem hiding this comment.
Changed to a materialized list, see the thread on the search signature. Commit a081545.
| environment: Optional[Dict[str, str]] = None | ||
| training_image: Optional[str] = None | ||
| latest_training_job: Optional[TrainingJob] = None | ||
|
|
There was a problem hiding this comment.
🟡 Low: latest_training_job changes from a plain class attribute (Optional[TrainingJob] = None) to a property with a setter that raises TypeError unless the value is a TrainingJob or None. This is a behavioral/breaking change for any consumer that previously assigned a job name string (the old docs and the removed model_customization.rst snippet did exactly trainer._latest_training_job = "my-hp-job"). The docs in this PR are updated to use TrainingJob.get(...), so the intent is clear, but it's worth calling out in the migration notes since the failure mode (TypeError at assignment) is new.
There was a problem hiding this comment.
Documented: migration.md has a row for attaching a trainer to an existing job (trainer.latest_training_job = TrainingJob.get(job_name)), with the setter contract (a TrainingJob or None, TypeError for a job name string). Commit a081545.
| """Drop the in-memory entries, and the cache file of ``key`` when one is given. | ||
|
|
||
| Returns ``True`` when a file was deleted. | ||
| """ |
There was a problem hiding this comment.
🟡 Low: _previous_value reaches into self._content_cache._lru_cache.get(key). Note that when HubCatalogCache.get() first calls self._content_cache.get(key, data_source_fallback=False) on an expired entry, LRUCache._get_item pops the key before raising KeyError and does not put it back, so by the time _previous_value runs the in-memory entry is already gone and this returns None. The retrieval then falls back to cache_file.load(), so correctness is preserved, but the in-memory "describe only changed items" shortcut is effectively never taken for expired entries — every expired-in-memory refresh re-reads the JSON file. If that disk read was meant to be avoidable, capture previous before the failing get() pops it.
There was a problem hiding this comment.
Changed: get() reads the old value before the LRUCache.get call, so an expired entry refreshes from memory and describes only the changed items without the file read. Test test_an_expired_entry_reuses_the_in_memory_items_without_the_cache_file deletes the cache file before the refresh. Commit a081545.
…e default thread pool `HubContent.get_all` described every listed item on every call, 4 at a time through `ResourceIterator(max_workers=4)`: 766 `DescribeHubContent` calls and 56 to 72 seconds for the public hub. The cache lived behind a second method, `get_all_cached`, with its own signature and no list filters, so a caller had to pick between the filters and the cache. Make `get_all` the cached method. The generator emits no `get_all` for `HubContent` (`HAND_WRITTEN_GET_ALL_RESOURCES`), and `HubContentMixin` provides it with the `ListHubContents` parameters of the generated method. The `HubCatalogCache` serves the catalog from memory inside the 6 hour horizon, from the cache file in a new process, and lists the hub and describes only the changed items after the horizon. The list filters (`name_contains`, `max_schema_version`, the creation time bounds, `sort_by`, `sort_order`) apply to the cached catalog through `HubCatalogFilters`, so a filtered call costs no request. The describes run on a `ThreadPoolExecutor` of the host default size (`min(32, cpu_count + 4)`) instead of a fixed `max_workers`. The describe client retries in the botocore `adaptive` mode with a 32 connection pool (`catalog_client_config`), passed through a new `config` argument of `Base.get_sagemaker_client`. On the public hub the default pool failed 462 of 766 describes in the standard mode and finished with 0 failures in the adaptive mode: 60.3 s cold with 44 throttle retries, 0.4 s inside the horizon, 1.9 s in a new process, 0.2 s for a filtered call. `ResourceIterator` returns to its `master` shape, with no `max_workers`, and `get_all_versions` loses the parameter. --- X-AI-Prompt: Can HubContent.get_all cache the described hub contents by default and use the default thread pool size? X-AI-Tool: claude-code
There was a problem hiding this comment.
🤖 AI code review
This large PR adds many public surfaces to sagemaker-core: HubContent.get_all/search/caching (hub_catalog.py, hub_content.py), Endpoint.diagnostics()/recent_logs() (endpoint.py), ModelPackageGroup.get_or_create(), list_instance_types/describe_instance_type, a rewritten format_tags/tag-merge path, instance_pools/tracing_config plumbing in session_helper.py, a ProcessingJob session-binding fix, and raise_on_stopped on generated wait(). I reviewed the hand-written modules in full; the generated resources.py and the test additions are mostly in the truncated portion of the 650KB diff, so I could not confirm the generator wiring (e.g. the original_creation_time field and mixin inheritance) end-to-end. The hand-written code is internally consistent. A few concrete issues below, mostly low severity.
Reviewed commit 625a4f0. Automated review; verify before acting.
| ) -> Iterator["HubContentMixin"]: | ||
| """Search a hub server side and yield each matching ``HubContent``. | ||
|
|
||
| Runs ``Search(Resource="HubContent")`` with an ``And`` expression of ``HubName``, |
There was a problem hiding this comment.
🟡 Low: search is a generator (it uses yield), but it is wrapped with @_telemetry_emitter. The decorator does response = func(*args, **kwargs) and then records SUCCESS/latency and returns response. For a generator function, calling it only constructs the generator object — none of the search body (the client.search(...) paging, and even the _search_expression(...) validation) runs until the caller iterates, which happens outside the decorator's try/except. So telemetry will always record immediate SUCCESS with ~0 latency and will never capture failures or the real duration of the paged search. get_all (which returns a list) is unaffected. Consider making search a non-generator that returns a list, or emitting telemetry from an inner helper that is fully consumed inside the decorated call.
There was a problem hiding this comment.
Changed: search() returns the matching HubContent list after the last page, so the filter validation, the paging, and the telemetry run inside the decorated call. A ValueError from a bad filter is raised from search() and recorded as a failure. Test test_search_raises_from_the_call_and_emits_a_failure. Commit a081545.
|
|
||
| @property | ||
| def output_modalities(self) -> List[str]: | ||
| """Every ``@output-modality:`` keyword, extended with document ``OutputModalities``.""" |
There was a problem hiding this comment.
🟡 Low: get_all/search/refresh_cache take session (boto3) and region keyword arguments, not sagemaker_session. The _telemetry_emitter decorator only looks for sagemaker_session on args[0]/kwargs and otherwise falls back to _get_default_sagemaker_session(). So even when a caller passes an explicit session, telemetry is emitted against a default session, and every decorated call constructs a default session (and resolves the opt-out config) regardless of the caller's intent. Worth confirming this is acceptable for these classmethods.
There was a problem hiding this comment.
Confirmed, and kept for this PR. The decorator reads sagemaker_session only, and every other session-taking classmethod in the SDK (AIRHub.list_hub_content, AIRHub.import_hub_content) has the same behavior. The default session costs about 0.2 s per call on this host (the Session constructor builds its clients), which a cache hit of get_all pays. Teaching the decorator to read a boto session is a telemetry-wide change, outside this PR.
| "SageMakerPublicHub", | ||
| filters=[ | ||
| HubContentFilter.metadata("license", "apache-2.0"), | ||
| HubContentFilter.metadata("provider", "Meta"), |
There was a problem hiding this comment.
🟡 Low: The HubContentFilter.metadata("license", "apache-2.0") example (and the same lowercase value in hub_content.py's HubContentFilter.metadata docstring) uses apache-2.0, while migration.md uses Apache-2.0. Per the team precedent, HubContentFilter.metadata() field names and values are case-sensitive and must match the document spelling exactly (e.g. Apache-2.0, not apache-2.0) (Source: #6533). The lowercase examples will silently return no results. Align the examples on the document casing.
There was a problem hiding this comment.
Aligned on the document spelling: Apache-2.0 and meta in docs/training/index.rst, the HubContentFilter.metadata docstring, the module docstring, and the tests. Against the public hub, Apache-2.0 matches 112 models and apache-2.0 none, meta matches 74 and Meta none. The metadata docstring names the rule and points a keyword value to HubContentFilter.keyword. Commit a081545.
| The answer is kept per set of credentials for the life of the process, so the STS | ||
| call happens once per principal, not once per cache read. | ||
| """ | ||
| # common_utils imports resources, which imports this module, so the import runs here. |
There was a problem hiding this comment.
🟡 Low: _catalog_resources sets resource.original_creation_time = summary["OriginalCreationTime"] directly on a pydantic model whose Base config has validate_assignment=True and extra="forbid". This relies on the generator having added an original_creation_time field to HubContent with a type that accepts the raw boto summary value (a datetime, or an ISO string when reloaded from the JSON cache file). That field addition is in the truncated generated resources.py; please confirm the field exists and its type accepts both representations, otherwise this assignment raises a validation error at get_all time.
There was a problem hiding this comment.
Confirmed. The generator adds original_creation_time: Optional[datetime.datetime] = Unassigned() to HubContent (LIST_SUMMARY_EXTRA_ATTRIBUTES in resources_codegen.py, line 13857 of resources.py). The value comes from the ListHubContents summary of the current retrieval, a datetime from boto3, never from the JSON file: the file stores only the described resource. The cached resource itself is rebuilt with model_dump(mode="json") strings, which pydantic parses back into datetime. test_a_new_process_starts_from_the_cache_file asserts original_creation_time and creation_time after a reload. Commit a081545.
| """``latest_training_job`` is the public handle for the job a trainer created. | ||
|
|
||
| Before this property the only handle was the private ``_latest_training_job``, so callers | ||
| that needed the job name or ARN after ``train(wait=False)`` had to read the private | ||
| attribute or page ``ListTrainingJobs`` by name prefix. ``ModelTrainer.train()`` also | ||
| returned ``None``, so the job could not be chained from the call that created it. | ||
|
|
||
| ``on_submit`` hands the created job to the caller as soon as the service accepts it and | ||
| before any wait, so a long ``train(wait=True)`` can record the job name or ARN without the | ||
| private ``trainer_wait.wait`` import that ``wait=False`` plus a hand-rolled wait needed. | ||
| """ |
There was a problem hiding this comment.
The docstring reads more like a changelog entry than a description of what the module tests. That history is useful, but it belongs in the PR description or migration.md. Consider trimming this to what the tests cover: the latest_training_job property and setter, the return value of train(), and the on_submit ordering contract across ModelTrainer, the serverless trainers, and HyperPod.
There was a problem hiding this comment.
Rewritten. The module docstring states what the tests cover: the latest_training_job property and setter, the TrainingJob that train() returns, and the on_submit ordering contract across ModelTrainer, the serverless trainers, and the HyperPod path. The history moved to the PR description. Commit 4042f61.
|
|
||
|
|
||
| def test_model_trainer_latest_training_job_setter_rejects_a_job_name(model_trainer): | ||
| with pytest.raises(TypeError, match="TrainingJob.get"): |
There was a problem hiding this comment.
The match argument to pytest.raises is a regex, so the dot in TrainingJob.get matches any character. It still passes, but it is looser than it looks. Wrapping with re.escape, or matching a longer distinctive phrase from the message, would make the intent clearer. Same pattern at lines 157 and 166/168.
There was a problem hiding this comment.
Changed: the matches are re.escape of a longer phrase from each message (Use TrainingJob.get(<name>) to attach an existing job. and trainer.latest_training_job = TrainingJob.get('<job-name>')), as module constants SETTER_ERROR and ATTACH_HINT. Commit 4042f61.
| for active in patches: | ||
| active.start() |
There was a problem hiding this comment.
The patch start loop runs before the try block, so if any single patch.start() raises (for example because a target attribute is renamed), the patches that already started never get stopped and will leak into every later test in the session. Moving the starts inside the try, or replacing the manual start/stop with contextlib.ExitStack, would make this safe and would also remove the repeated try/finally in both tests.
There was a problem hiding this comment.
Changed: the serverless and HyperPod cases run under one ExitStack fixture (serverless, hyperpod_trainer), so a patch.start() that raises releases every patch already started and the manual start and stop loops are gone. Commit 4042f61.
| @pytest.mark.parametrize( | ||
| "trainer_cls, module", | ||
| [ | ||
| (SFTTrainer, "sagemaker.train.sft_trainer"), | ||
| (DPOTrainer, "sagemaker.train.dpo_trainer"), | ||
| ], | ||
| ) |
There was a problem hiding this comment.
This parametrize covers SFTTrainer and DPOTrainer only. On this branch rlvr_trainer.py, rlaif_trainer.py and cpt_trainer.py also take on_submit and set latest_training_job, so a regression in any of those three would not be caught here. Could they join the list (here and at line 255)? If one of them is deliberately left out, a short comment saying why would stop a reader from assuming the omission is accidental.
There was a problem hiding this comment.
Changed: the serverless cases cover SFTTrainer, DPOTrainer, RLVRTrainer (with an evaluator ARN and skip_reward_validation), and RLAIFTrainer, in the four tests (created job, on_submit before the wait, on_submit with wait=False, dry_run). CPTTrainer is HyperPod only, so it has its own case through _train_hyperpod. Commit 4042f61.
| model_trainer.train(wait=True, logs=False, on_submit=lambda job: order.append(job)) | ||
|
|
||
| assert order == [created, "wait"] | ||
|
|
There was a problem hiding this comment.
This nicely pins the ordering. One case I do not see covered anywhere: what happens when the on_submit callback itself raises? Should train() still set latest_training_job and still wait, or should the exception propagate and skip the wait? Either answer is reasonable, but it should be a deliberate choice with a test, since a user callback that fails is the most likely real-world misuse of this hook. Related: on_submit is tested with wait=True and dry_run for ModelTrainer, but not with wait=False, which the docstring names as the original motivation.
There was a problem hiding this comment.
Made deliberate and tested: a callback that raises propagates from train(), the job stays on latest_training_job (it is already running), and the wait is skipped (test_a_raising_on_submit_keeps_the_job_and_skips_the_wait). on_submit with wait=False is covered for ModelTrainer and the four serverless trainers. Commit 4042f61.
| def _hyperpod_sft_trainer(): | ||
| trainer = SFTTrainer.__new__(SFTTrainer) | ||
| trainer.sagemaker_session = MagicMock() | ||
| trainer.compute = SimpleNamespace( | ||
| cluster_name="my-cluster", | ||
| namespace="kubeflow", | ||
| instance_type="ml.p5.48xlarge", | ||
| node_count=2, | ||
| ) | ||
| trainer._model_name = "amazon.nova-lite-v2" | ||
| trainer._customization_technique = "SFT" | ||
| trainer.training_type = "LORA" | ||
| trainer.training_image = "123.dkr.ecr.us-west-2.amazonaws.com/img:latest" | ||
| trainer.base_job_name = None | ||
| trainer.training_dataset = None | ||
| trainer.validation_dataset = None | ||
| trainer.s3_output_path = "s3://bucket/out/" | ||
| trainer.hyperparameters = None | ||
| trainer._recipe_path = "recipe-name" | ||
| trainer._overrides = None | ||
| return trainer |
There was a problem hiding this comment.
Building the trainer with SFTTrainer.new and then setting private attributes directly couples this test to the names _model_name, _customization_technique, _recipe_path and _overrides. A rename of any of those breaks the test without saying anything about the on_submit behavior under test. If the constructor is too heavy to call here, a small fixture that goes through the real constructor with the HyperPod validators patched would keep the test honest about the state the trainer ends up in. Also, training_type is the string "LORA" rather than TrainingType.LORA, which may not match what the constructor produces.
There was a problem hiding this comment.
Changed: the HyperPod fixture builds the trainer through the real SFTTrainer constructor with HyperPodCompute(cluster_name, instance_type, node_count), training_type=TrainingType.LORA by default, and patches the Hub, IAM, recipe, and CLI calls at the names the modules use. No private attribute is set by hand. Commit 4042f61.
| ): | ||
| job_name = trainer._train_hyperpod(wait=False, on_submit=submitted.append) | ||
|
|
||
| assert job_name == "my-job-123" |
There was a problem hiding this comment.
This asserts _train_hyperpod returns the string "my-job-123", while every other path in this file asserts train() returns a TrainingJob. base_trainer.py line 765 confirms the -> str return. That means the public return type of train() depends on the compute target, which is awkward for callers who want to chain the job. Is the string intentional for HyperPod, or should it return the same TrainingJob it hands to on_submit?
There was a problem hiding this comment.
Fixed: _train_hyperpod returns the TrainingJob it hands to on_submit and stores in latest_training_job, so train() returns the same type on every compute target. CPTTrainer.train documents it, FineTuningSweep drops its job-name branch, the two HyperPod notebooks print training_job_name, and the three tests that asserted the string now read training_job_name. Commit 4042f61.
| assert [job.training_job_name for job in submitted] == ["my-job-123"] | ||
| assert trainer.latest_training_job is submitted[0] | ||
| start_cmd = mock_subprocess.run.call_args_list[-1].args[0] | ||
| assert json.loads(start_cmd[start_cmd.index("--override-parameters") + 1]) |
There was a problem hiding this comment.
This only checks that the parsed override-parameters JSON is truthy. With _overrides set to None and flatten_resolved_recipe mocked to return an empty dict, it is not obvious what the payload is expected to hold. Asserting a specific key or shape would make clear what regression this line protects against; if it only guards that the flag is present, a comment saying so would help.
There was a problem hiding this comment.
Changed: the test asserts the exact payload, {"instance_type": "ml.p5.48xlarge", "container": <image>}, which are the two override parameters the trainer sends when model_source is unset. Commit 4042f61.
…and the cache refresh The automated review of PR aws#6533 left eight open findings on the branch. `HubContent.search` was a generator under `@_telemetry_emitter`, so the decorator timed the construction of the generator, recorded SUCCESS before the first `Search` call, and never saw a `ValueError` from the filter validation. `search` now returns the matching `HubContent` list, so the validation, the paging, and the telemetry run inside the call. The `HubContentFilter.metadata` examples used `apache-2.0` and `Meta`, which match no model: the service compares a `Metadata.<Field>` value with the document spelling (`Apache-2.0` matches 112 models and `apache-2.0` none, `meta` matches 74 and `Meta` none). The docstrings, `docs/training`, and the tests use the document spelling, and the `metadata` docstring names the rule and points to `HubContentFilter.keyword` for a keyword value. `HubCatalogCache.get` read the old value after `LRUCache.get` had popped the expired element, so an expired entry fell back to the cache file on every refresh. The old value is read first, and a test refreshes an expired entry with the cache file deleted. `describe_instance_types` and `describe_instance_type` take a `region`, as `list_instance_types` does, and the docstrings name `NoRegionError`. `_tag_text` turns a missing tag value into the empty string instead of `None`, because the service model marks `Tag.Value` as a necessary member. `verify_reward_function` keeps `validate_format`, `compute`, and `is_nova` as its third to fifth positional parameters, so a first-release call such as `verify_reward_function(fn, data, False, None, False)` still works, and `options` moves behind them. `migration.md` names the `latest_training_job` setter contract: a `TrainingJob` or `None`, with `TypeError` for a job name string. --- X-AI-Prompt: Can you review the latest comments on PR aws#6533 and address them? X-AI-Tool: claude-code
There was a problem hiding this comment.
🤖 AI code review
This large PR adds ~eighteen public surfaces to sagemaker-core: HubContent.get_all/search/refresh_cache backed by a new HubCatalogCache (modeled on the JumpStart cache), typed HubContent metadata accessors, EndpointMixin.diagnostics()/recent_logs(), ModelPackageGroup.get_or_create(), instance-type lookups, S3Downloader.read_bytes(byte_range=)/iter_lines()/head(), a raise_on_stopped param on generated wait() methods, a rewritten format_tags, and a public latest_training_job. It is wired through the code generator via RESOURCE_MIXINS/LIST_SUMMARY_EXTRA_ATTRIBUTES/GET_ALL_INPUT_ATTRIBUTES. The implementation is careful and well-documented; the generator template and the generated resources.py were kept in sync (e.g. the new config= param, _set_client_context). It also fixes a latent curr_value.element → .value bug in LRUCache.put. I found no blocking correctness issues in the areas I reviewed; the notes below are minor.
[{"path": "sagemaker-core/src/sagemaker/core/hub_content.py", "line": 300, "severity": "low", "body": "HubContent.get_all/search/refresh_cache accept session= and region=, but the _telemetry_emitter decorator only resolves a session from an instance (args[0]) or a sagemaker_session kwarg. For these classmethods args[0] is cls and the arg is named session, so telemetry always falls back to _get_default_sagemaker_session() — emitting telemetry against the default region/credentials instead of the caller's. This can trigger an extra default-session/credential resolution (and surprising failures in environments with no default region) even though the caller passed an explicit session/region. Consider forwarding the caller's session to telemetry (or exempting these methods)."}, {"path": "sagemaker-core/src/sagemaker/core/hub_catalog.py", "line": 388, "severity": "low", "body": "_describe_all yields zip(summaries, executor.map(describe, summaries)), and executor.map re-raises the first exception in submission order while consuming results in order. If an early describe raises, later describes that already completed are never consumed/stored, so the finally: cache_file.save(entries) persists only the items up to the failure. The docstring ("A describe that raises still leaves every described item on disk") overstates this. Not a correctness bug, but the comment is misleading; consider as_completed if you truly want to persist every successful describe before propagating."}]
Reviewed commit a081545. Automated review; verify before acting.
| :members: | ||
| :undoc-members: | ||
|
|
||
| .. automodule:: sagemaker.core.instance_types |
There was a problem hiding this comment.
Verse uses a different repo aws:master-data-preparation right?
There was a problem hiding this comment.
Yes. Verse installs the preview SDK. The SageMakerAIAgentSkills build declares SageMakerPythonSDKPreview = 3.x_master_data_preparation as a build tool, and install_preview_sdk.sh installs the five wheels the preview distribution pipeline writes from that staging branch. The staging branch and public master have diverged. master-data-preparation is 66 commits ahead of and 44 behind staging master, which is identical to public master (4042a24f). The data-preparation sync PR #2345 names the plan to port the branch to a release branch off public master. This PR targets public master, so Verse gets it when the next sync carries public master into master-data-preparation (the last one, #2345 on 2026-10-26, brought staging master up to 3.27.0). The 7 files of this PR that the staging branch also touches (common_utils.py, resources.py, telemetry_logging.py, base_trainer.py, model_builder.py, and two tests) are in that sync path. I can open the port to master-data-preparation after this merges, if that is the route you want.
…iner test review The human review of PR aws#6533 left eight threads on `test_latest_training_job.py`, and the automated review two on the hub content modules. `_train_hyperpod` returned the job name string while every other path returned the `TrainingJob`, so the return type of `train()` depended on the compute target. It returns the `TrainingJob` it hands to `on_submit` and stores in `latest_training_job`. `CPTTrainer.train` documents the handle, `FineTuningSweep` drops the job-name branch, and the two HyperPod notebooks print `training_job_name`. The test module states what it covers instead of the change history. The `pytest.raises` matches use `re.escape`. The serverless cases cover the SFT, DPO, RLVR, and RLAIF trainers through one `ExitStack` fixture, so a patch that fails to start stops nothing from being released. The HyperPod cases build the trainer through the real `SFTTrainer` constructor with `HyperPodCompute`, assert the exact `--override-parameters` payload, and cover `dry_run` and `CPTTrainer`. New cases pin `on_submit` with `wait=False` and a callback that raises: the job stays on `latest_training_job`, the error propagates, and the wait is skipped. `HubCatalogCache._describe_all` consumed `executor.map` in order, so an early failure dropped the results of later workers before the file was written. It collects with `as_completed`, stores every result, writes the file, then raises the first failure. `_telemetry_emitter` read only `sagemaker_session`, so a classmethod call with a boto `session` keyword sent its telemetry on the default credentials and built a default `Session` on every call. The decorator wraps a `boto3.Session` passed as `session` into a `Session`, and falls back to the default only without one. --- X-AI-Prompt: Can you make sure the latest review comments on PR aws#6533 are addressed? X-AI-Tool: claude-code
Brings in the three commits on master since the last merge: the vLLM v0.30 and v0.31 image metadata (aws#6526), the DLC SCL RBA config (aws#6507), and the MTRL SOP-Bench AgentCore example (aws#6532). No file is changed on the two sides. --- X-AI-Prompt: Can you refresh the PR so the latest changes from master are synced in? X-AI-Tool: claude-code
There was a problem hiding this comment.
🤖 AI code review
This large PR adds eighteen public surfaces to SageMaker SDK v3: HubContent.get_all/search/cache (hub_catalog.py, hub_content.py), Endpoint.diagnostics/recent_logs (endpoint.py), ModelPackageGroup.get_or_create, list_instance_types/describe_instance_type, trainer job handles (train()/_train_hyperpod() now return TrainingJob, latest_training_job, on_submit), typed deploy options, format_tags normalization, a TracingConfig endpoint-config member, and raise_on_stopped on generated wait() methods. The code is carefully written and well-documented, with regression tests for the training-side changes. My review of the hub catalog cache, tag handling, endpoint diagnostics, instance-type helpers, and the session-helper changes found the implementation sound; a few lower-severity items are noted below. Note the diff was truncated (~675k chars); I could not fully review the ModelBuilder.deploy, Bedrock, S3 and hub-content unit-test additions, so those are not covered here.
Reviewed commit 5ea7b09. Automated review; verify before acting.
| @@ -2124,38 +2186,43 @@ def tag_exists(tag: TagsDict, curr_tags: Optional[Tags]) -> bool: | |||
| if curr_tags is None: | |||
There was a problem hiding this comment.
🟡 Low: _validate_new_tags now calls format_tags(curr_tags) once and then tag_exists(new_tag, merged) inside the loop, and tag_exists itself calls format_tags(curr_tags) on the already-boto-form merged for every new tag. For N current and M new tags this re-normalizes the full list M times (O(N*M)). Harmless for typical tag counts, but tag_exists could take the already-normalized key list to avoid the repeated format_tags pass.
| return tags | ||
| def format_tags(tags: Optional[Tags]) -> Optional[List[TagsDict]]: | ||
| """Process tags to turn them into the expected format for Sagemaker. | ||
|
|
There was a problem hiding this comment.
🟡 Low: format_tags is now stricter: inputs the old implementation returned unchanged (e.g. a list containing a non-dict/non-pair element, or a scalar) now raise TypeError via _normalize_tag. This is the intended behavior, but because format_tags/tag_exists/remove_tag_with_key are shared helpers invoked on stored curr_tags, any existing caller that previously stashed tags in an unexpected shape will now raise instead of silently passing through. Worth confirming all internal callers (model_builder, model_builder_utils) only ever hold the accepted representations.
| query = queries_by_id.get(result.get("Id")) | ||
| if not query: | ||
| continue | ||
| previous = series.get(query.key) |
There was a problem hiding this comment.
🟡 Low: _boto_client constructs a fresh SageMakerClient(session=..., region_name=...) on every call, which eagerly builds the four SageMaker-family clients (sagemaker, runtime, featurestore-runtime, metrics) just to borrow the resolved session/region/config for a cloudwatch or logs client. diagnostics() calls this twice per invocation. Consider resolving the session/region/config once (or caching) to avoid repeatedly instantiating those unused clients.
akappagan
left a comment
akappagan
left a comment
There was a problem hiding this comment.
What's the testing strategy for these changes? How do we ensure that this changes won't break Verse or production?
|
Closing. The change lands on the branch the Verse preview SDK is built from, |
Eighteen changes to the surfaces an agent uses to pick a hub model, train it, evaluate it, and deploy it. Each closes a gap where the SageMaker agent skills (
aws-ai-ml) drop to rawboto3or a private attribute because the SDK has no public handle. One commit per change, each with a regression test that fails onmasterand passes here. Nothing regeneratessagemaker-corefrom a newer botocore. In each code block, a# master:comment states whatorigin/master(4042a24f) offers, and a# This PR:comment states what the lines below it add. Every "Live run" line is from this branch inus-west-2on 2026-10-25 and 2026-10-26. The hub runs targetSageMakerPublicHub.Hub Service: list, read, and search models
Before this PR, the SDK had no
HubContent.get_all, no typed reader forHubContentSearchKeywords, and no binding forSearch(Resource="HubContent"). The agent skills pagedlist_hub_contentsby hand in three scripts, parsed@license:and@model-size:prefixes by hand, and built theSearchexpression by hand.List the catalog (
c2db676d,cb1a6ea6,f529db22,625a4f05)Live run: 766 bound
HubContentresources in 60.3 s cold (8 list calls, 766 describes on the default pool of 32 workers, 44 throttle retries, 0 failures), against 95.7 s for one describe at a time. The 32.9 MB cache file lives inSAGEMAKER_HUB_CACHE_DIRor~/.cache/sagemaker/hub-content. Inside the horizon the call takes 0.4 s with no request, a new process 1.9 s from the file, andrefresh_cache1.4 s withlisted=766 reused=766 described=0 removed=0.name_contains="qwen3"returns 36 in 0.2 s with no request.get_all_versions()returns 4 versions ofhuggingface-vlm-qwen3-8-27b, wheremasterfails withMissing required parameter in input: "HubName".Read the metadata of a model (
d352fe5d)Live run:
license,provider,model_size, andcontext_windowareapache-2.0,qwen,1b-10b, and8k-32k.capabilitiesis['bedrock_console', 'customization'],fine_tuning_techniques()is['dpo', 'rlaif', 'rlvr', 'sft'], andsmall_sftholds 12 names. The document fills what the 50-keyword cap dropped. 79 models are at the cap, and the document alone supplies 30 licenses, 13 model sizes, and extra languages for 20. The mixin adds no model field:model_dump()is unchanged.Search server side (
93c0a67c,a0815453)Live run:
fine_tunable_smallreturns 9 in 0.8 s,apache_on_bedrock45 in 1.0 s,llama34 in 0.3 s, and an unfiltered search all 766 in 8.5 s. A coldget_alltakes 60.3 s.HubContentFilter.metadatamaps the attribute to the indexed field (licensetoMetadata.License) and rejects an unknown attribute. The 20-filter and 1024-character limits raiseValueErrorfrom thesearch()call, before a request.Metadata.<Field>matches the document spelling:Apache-2.0matches 112 models,apache-2.00,meta74,Meta0. The docstrings and docs use the document spelling.search()returns a list, not a generator, so the@_telemetry_emitterdecorator times the paged calls and records aValueErroras a failure (a0815453). Search records carry noOriginalCreationTime, soget_allstays the call for catalog dates.Training jobs
ModelTrainer.train()returns the createdTrainingJobandlatest_training_jobis public on every trainer (450c866d). The setter takes aTrainingJoborNoneand raisesTypeErrorfor a job name string, whichmigration.mdnames.on_submitruns as soon as the service accepts the job, before any wait, on the serverless, serverful, and HyperPod paths (8c8dd9ad).The HyperPod path returns the same
TrainingJobit hands toon_submit, not the job name string. The return type oftrain()no longer depends on the compute target (4042f612). A callback that raises leaves the job onlatest_training_job, propagates, and skips the wait.get_metric_historyreturns MLflow points as data, withlossresolving tototal_lossortrain/lossby recipe (fb7c6167).FineTuningSweepruns grid, random, or ASHA sweeps over the serverless trainers with a resumable S3 manifest, which AMT cannot do (HyperParameterTrainingJobDefinitionhas noserverless_job_config) (4b89520e).Evaluation
verify_reward_functionis exported with typedRewardSampleandEvalRewardSampleand anenvelopeofoss,nova,eval_oss, oreval_nova.validate_format,compute, andis_novastay the third to fifth positional parameters, so a first-release call still works (a0815453).Evaluator.createandCustomAgentLambdashare one Lambda runtime,python3.12(a4c27861).validate_custom_metricsreports every violation before the job starts, for example['Metric [0]: instructions must contain {{prompt}} or {{prediction}}.'](5a134b93).detect_dataset_formatreads a local file or one S3 byte range. It names the format from the 7 SDK schemas plus the Nova names of the skill. Four committed fixtures returnhf_prompt_completion,dpo,verl, andgenqaat confidence 1.0 (cd594911).Deployment
ModelBuilder.deploynamestags,data_capture_config,tracing_config,fallback_instance_types,endpoint_type, andinference_component_name, and applies each on every deploy path or raises. An unknown kwarg raisesTypeErrorinstead of deploying without it.fallback_instance_typesbecomesInstancePoolswith the primary type at priority 1 (6e9df413,b04f4c93,097192a5).format_tagsturns a missing tag value into"", the empty string the service model accepts, notNone(a0815453).Endpoint.diagnostics()returns the status, failure reason, per-variant instance type from the endpoint config, 11 CloudWatch metrics, and recent logs. Live run on aFailedendpoint:failure_reasonstartsUnable to provision requested ML compute capacity,metricshas the 11 keys, and the missing log group lands inerrorsasFilterLogEvents failed: An error occurred (ResourceNotFoundException)instead of raising (4c9bedb0).list_invocable_foundation_models()joinsListFoundationModelswith inference profiles: 89 of 92TEXTmodels inus-west-2, 43 through a cross-region profile.BedrockDeploymentwaits throughModelNotReadyException, theninvokeandconverse.BedrockModelBuilder.deploy()records it aslatest_deployment()(41aae122).Core resources
S3Downloader.read_bytes(byte_range=),iter_lines, andhead(2c09cf5f).ModelPackageGroup.get_or_create(6e820f61).Processor._start_newreturns a session-boundProcessingJobwith its ARN, solatest_job.refresh()no longer hits the global default region (#5796) (0727c476). Every generatedwait()takesraise_on_stopped(c6294e14).list_instance_types()returns theProductionVariantInstanceTypeenum of the loaded service model (277 names in botocore 1.43.118) with no request.describe_instance_type("ml.g5.xlarge"), with an optionalregionaslist_instance_types, returnsInstanceSpec(name='ml.g5.xlarge', vcpus=4, memory_mib=16384, gpu_count=1, gpu_memory_mib=22888, gpu_name='A10G'). The two replace two private EC2 lookups inmodel_builder.py(d9069112).migration.mdnamesCompute(enable_managed_spot_training=True)andStoppingCondition(max_wait_time_in_seconds=)in place of theCompute(use_spot_instances=...)example that fails withextra_forbidden(b6959d83).HubCatalogCache(hub_catalog.py) followsJumpStartModelsCache. It is the samesagemaker.core.utilities.cache.LRUCachewith the same defaults (20 items, 6 hour horizon). TheLRUCachecalls a retrieval function on a miss or an expired entry. A staticHubCatalogAccessorholds the process cache and takesset_cache_kwargs. The retrieval function describes only the items whose version changed, as the JumpStart one reads the manifest only when its ETag changed. The items persist to a JSON file, so a new process starts from the previous catalog. TheLRUCachereads and writes run under one lock. A retrieval runs under a per-key lock, so a slow hub blocks no hit on another hub.getreads the old value beforeLRUCache.getpops an expired element, so an expired entry refreshes from memory without a file read (a0815453). The describes complete throughas_completed, so one failure drops no result of another worker: every described item is written to the file, then the first failure is raised (4042f612).LRUCache.puton a present key raisedAttributeError(curr_value.element). The fix readscurr_value.value(f529db22).The describes run on a
ThreadPoolExecutorof the host default size (min(32, cpu_count + 4)). The client retries in the botocoreadaptivemode with a 32-connection pool (catalog_client_config), passed through a newconfigargument ofBase.get_sagemaker_client. On the public hub the default pool failed 462 of 766 describes in thestandardmode and 0 in theadaptivemode (625a4f05).The generator emits no
get_allforHubContent(HAND_WRITTEN_GET_ALL_RESOURCES).HubContentMixinprovides it with theListHubContentsparameters of a generatedget_all, andHubCatalogFiltersapplies them to the cached catalog. Hand-written behavior for a generated class lives in its own module (hub_content.py,hub_catalog.py,endpoint.py,model_package_group.py). ARESOURCE_MIXINStable in the generator mixes it in, soresources.pyemitsclass HubContent(HubContentMixin, Base)and the two others.HubContent.search,get_all, andrefresh_cachecarry the@_telemetry_emitterdecorator thatCONTRIBUTING.mdasks of public V3 entry points (8f9ea7fb).The decorator wraps a
boto3.Sessionpassed assessioninto the telemetrySession. A classmethod call reports on the account and region of the caller and builds no default session (4042f612).python src/sagemaker/core/tools/codegen.pyregeneratesresources.pyto the committed content.Tests
Twenty-two new test files, each failing on
origin/masterat collection or with the old behavior:sagemaker-core/tests/unit/test_hub_content_get_all.py,test_hub_content_mixin.py,test_hub_content_search.py,test_hub_catalog_cache.py,test_lru_cache.py,test_endpoint_diagnostics.py,test_model_package_group_get_or_create.py,test_processing_start_new_arn.py,test_wait_raise_on_stopped.py,test_instance_types.py,s3/test_client_read_primitives.py;sagemaker-train/tests/unit/train/test_latest_training_job.py,test_fine_tuning_sweep.py,common_utils/test_get_metric_history.py,common_utils/test_reward_function_contract.py,evaluate/test_validate_custom_metrics.py,ai_registry/test_evaluator_runtime.py,ai_registry/test_detect_dataset_format.py;sagemaker-serve/tests/unit/test_model_builder_deploy_tags.py,test_model_builder_deploy_options.py,test_bedrock_deployment.py,test_model_builder_instance_types.py.Full unit suites on this branch after the merge of
master(4042a24f),pytest -n 8:sagemaker-core4109 passed, 24 skipped, 5 failedandsagemaker-train3648 passed, 19 skipped, 2 failed(at4042f612),sagemaker-serve2227 passed, 25 skipped, 59 failed. Every failure fails identically onorigin/masteron the same host. The 5test_common_utils.pypath tests fail withDID NOT RAISE ValueError.sm_recipes/test_utils.py::test_load_base_recipe_types[test_case2]fails withURLError: no host given.training_session/test_submitter.py::test_close_returns_the_loops_descriptors_rather_than_waiting_for_gccounts file descriptors (assert 15 == 16) under-n 8and passes alone. The 59sagemaker-servetests importtorch,tensorflow, ortritonclient, or read JumpStart specs from S3 with expired host credentials (ExpiredToken). Thecodestyle-doc-tests (sagemaker-core)hook,tox -e flake8,pylint,docstyle,black-check, is green on this branch:black --check ./leaves 522 files unchanged andpydocstyle src/sagemakeris clean. The hook failed onmasterand on every open PR: the generator emitted two blank lines beforeloggerand no module docstring inconfig_schema.pyandshape_dag.py. The fix is in the generator, so the next botocore sync keeps the files clean. Theunit-tests (sagemaker-serve)hook failed on pandas 3.1.0 (ImportError: C extension: None not built). The merge ofmasterbrings thepandas<3.1.0test pin from #6524.Merge Checklist
By submitting this pull request, I confirm that my contribution is made under the terms of the Apache 2.0 license.