diff --git a/docs/api/sagemaker_core.rst b/docs/api/sagemaker_core.rst index e78b6f8ca4..a3336a4797 100644 --- a/docs/api/sagemaker_core.rst +++ b/docs/api/sagemaker_core.rst @@ -52,6 +52,18 @@ Utilities :members: :undoc-members: +.. automodule:: sagemaker.core.instance_types + :members: + :undoc-members: + +.. automodule:: sagemaker.core.hub_content + :members: + :undoc-members: + +.. automodule:: sagemaker.core.hub_catalog + :members: + :undoc-members: + Exceptions ---------- diff --git a/docs/model_customization/model_customization.rst b/docs/model_customization/model_customization.rst index 6e0095184d..5d36c2a6eb 100644 --- a/docs/model_customization/model_customization.rst +++ b/docs/model_customization/model_customization.rst @@ -154,6 +154,7 @@ Key Features # Re-attach (HyperPod — needs cluster name for log group resolution) from sagemaker.train import SFTTrainer + from sagemaker.core.resources import TrainingJob from sagemaker.core.training.configs import HyperPodCompute trainer = SFTTrainer( @@ -161,7 +162,7 @@ Key Features training_dataset="s3://unused", compute=HyperPodCompute(cluster_name="my-cluster", instance_type="ml.p5.48xlarge") ) - trainer._latest_training_job = "my-hp-job" + trainer.latest_training_job = TrainingJob.get("my-hp-job") df = trainer.show_metrics() **Monitoring: stream_logs()** diff --git a/docs/training/index.rst b/docs/training/index.rst index e2eabfae54..e23092e933 100644 --- a/docs/training/index.rst +++ b/docs/training/index.rst @@ -299,6 +299,48 @@ Train pre-configured models from the SageMaker JumpStart hub using ``ModelTraine # Complex queries with filters text_gen = search_public_hub_models(query="@task:text-generation") +**Reading a Hub Model's Metadata:** + +.. code-block:: python + + from datetime import timedelta + + from sagemaker.core.resources import HubContent + + # Every model of the public hub, with the fields the hub tags it with. + # get_all describes 4 models at a time (max_workers=4). + small_sft_models = [ + content.hub_content_name + for content in HubContent.get_all(hub_name="SageMakerPublicHub", hub_content_type="Model") + if content.supports_technique("sft") and content.model_size in ("<1b", "1b-10b") + ] + + # The same list from the hub catalog cache, built like the JumpStart models cache: an + # LRU cache of hubs in memory (6 hour horizon) backed by one JSON file per hub. The first + # call describes every model, a call inside the horizon makes no request, and a call + # after it lists the hub and describes only the models whose version changed. + models = HubContent.get_all_cached("SageMakerPublicHub") + HubContent.set_cache_kwargs(cache_expiration_horizon=timedelta(days=1), max_workers=8) + + # One model; the document fills in fields the 50-keyword cap dropped + content = HubContent.get( + hub_name="SageMakerPublicHub", + hub_content_type="Model", + hub_content_name="meta-textgeneration-llama-3-2-1b", + ) + print(content.license, content.provider, content.tasks, content.fine_tuning_techniques()) + + # Server-side search on the document fields, paged for you + from sagemaker.core.hub_content import HubContentFilter + + apache_meta_models = HubContent.search( + "SageMakerPublicHub", + filters=[ + HubContentFilter.metadata("license", "apache-2.0"), + HubContentFilter.metadata("provider", "Meta"), + ], + ) + Key points: - ``from_jumpstart_config()`` auto-configures training image, instance type, and default hyperparameters @@ -306,6 +348,9 @@ Key points: - JumpStart provides built-in datasets so you can start training immediately - Supports HuggingFace, XGBoost, CatBoost, LightGBM, and many more frameworks - Use ``list_jumpstart_models()`` and ``search_public_hub_models()`` to discover available models +- ``HubContent`` exposes ``tasks``, ``license``, ``provider``, ``model_size``, ``context_window``, ``languages``, ``capabilities``, and ``recipes`` read from its search keywords and document, so no caller parses ``HubContentSearchKeywords`` by hand +- ``HubContent.search()`` filters a hub server side on ``Metadata.`` through ``HubContentFilter``, inside the 20-filter and 1024-character service limits +- ``HubContent.get_all_cached()`` serves a hub from ``HubCatalogCache``, the same ``LRUCache`` shape as ``JumpStartModelsCache`` (20 hubs, 6 hour horizon), backed by one JSON file per hub under ``SAGEMAKER_HUB_CACHE_DIR`` or the user cache directory. It describes only the contents whose ``HubContentVersion`` changed. ``refresh_cache()`` lists again and reports the counts, ``clear_cache()`` drops the memory and the file, and ``set_cache_kwargs()`` tunes the horizon, the item cap, the worker count, and the directory :doc:`Full example notebook <../v3-examples/training-examples/jumpstart-training-example>` @@ -434,7 +479,7 @@ Provide an ordered list of candidate instance types and the platform launches th .. code-block:: python - training_job = model_trainer._latest_training_job + training_job = model_trainer.latest_training_job training_job.refresh() resource_config = training_job.resource_config diff --git a/migration.md b/migration.md index 0ea3e9bef1..43db6bb15c 100644 --- a/migration.md +++ b/migration.md @@ -531,6 +531,9 @@ pipeline = Pipeline( | Fine-tune foundation models | `JumpStartEstimator(model_id=...).fit()` | V3: jumpstart-training-example.ipynb | 🔄 REPLACED | | E2E training + inference | Separate steps | jumpstart-e2e-training-example.ipynb | 🔄 REPLACED | | Marketplace algorithms | `sagemaker.algorithm.AlgorithmEstimator` | `sagemaker.core.resources.Algorithm` (13 classes) | 🔄 REPLACED | +| List hub models | `list_jumpstart_models()` (manifest) | `HubContent.get_all(hub_name, hub_content_type)` pages `ListHubContents` and describes 4 models at a time. `HubContent.get_all_cached(hub_name)` serves the hub from an `LRUCache` with the JumpStart cache defaults (20 hubs, 6 hour horizon), backed by a JSON file per hub, and describes only the models whose version changed | 🆕 NEW IN V3 | +| Read model metadata | `JumpStartModelSpecs` fields | `HubContent.license`, `.provider`, `.model_size`, `.tasks`, `.capabilities`, `.recipes`, `.supports_technique("sft")` read from the hub keywords and document | 🆕 NEW IN V3 | +| Search hub models server side | None (`search_public_hub_models` filters client side) | `HubContent.search(hub_name, filters=[HubContentFilter.metadata("license", "Apache-2.0")])` runs `Search(Resource="HubContent")` | 🆕 NEW IN V3 | ## Configuration Management @@ -559,8 +562,7 @@ compute = Compute( instance_type="ml.m5.xlarge", instance_count=1, volume_size_in_gb=30, - use_spot_instances=True, - max_wait_time_in_seconds=7200 + enable_managed_spot_training=True ) networking = Networking( subnets=["subnet-12345"], @@ -568,7 +570,8 @@ networking = Networking( enable_inter_container_traffic_encryption=True ) stopping_condition = StoppingCondition( - max_runtime_in_seconds=3600 + max_runtime_in_seconds=3600, + max_wait_time_in_seconds=7200 ) model_trainer = ModelTrainer( training_image="...", diff --git a/sagemaker-core/src/sagemaker/core/common_utils.py b/sagemaker-core/src/sagemaker/core/common_utils.py index f76faba505..b94aad9848 100644 --- a/sagemaker-core/src/sagemaker/core/common_utils.py +++ b/sagemaker-core/src/sagemaker/core/common_utils.py @@ -1688,12 +1688,68 @@ def create_paginator_config(max_items: int = None, page_size: int = None) -> Dic } -def format_tags(tags: Tags) -> List[TagsDict]: - """Process tags to turn them into the expected format for Sagemaker.""" - if isinstance(tags, dict): - return [{"Key": str(k), "Value": str(v)} for k, v in tags.items()] +def _tag_text(value: Any) -> Union[str, PipelineVariable]: + """Return a tag key or value as the service string, keeping a pipeline variable intact.""" + if value is None or is_pipeline_variable(value): + return value + return str(value) + + +def _tag_dict(key: Any, value: Any) -> TagsDict: + return {"Key": _tag_text(key), "Value": _tag_text(value)} + + +def _normalize_tag(tag: Any) -> TagsDict: + """Return one tag as a boto-style ``{"Key": ..., "Value": ...}`` dict of strings. + + Accepts the boto form, the ``{"key": ..., "value": ...}`` form that the + ``sagemaker.core.resources`` classes take, a ``(key, value)`` pair, and a ``Tag`` + shape object (anything with ``key`` and ``value`` attributes). The key and the value + come out as ``str`` whatever the input representation, except a ``PipelineVariable``, + which the request serializer resolves later. Any other input raises ``TypeError`` so a + malformed tag fails at the call site instead of as a ``KeyError`` deep in a deploy + path. + """ + if isinstance(tag, dict): + if "Key" in tag: + return _tag_dict(tag["Key"], tag.get("Value")) + if "key" in tag: + return _tag_dict(tag["key"], tag.get("value")) + raise TypeError( + "A tag dict must have a 'Key' (or 'key') entry, " + f"got {sorted(tag)}. Pass a {{key: value}} mapping to set many tags at once." + ) + if isinstance(tag, (tuple, list)) and len(tag) == 2: + return _tag_dict(tag[0], tag[1]) + if hasattr(tag, "key") and hasattr(tag, "value"): + return _tag_dict(tag.key, tag.value) + raise TypeError( + "A tag must be a {'Key': ..., 'Value': ...} dict, a {'key': ..., 'value': ...} dict, " + f"a (key, value) pair, or a Tag object, got {type(tag).__name__}." + ) + - return tags +def format_tags(tags: Optional[Tags]) -> Optional[List[TagsDict]]: + """Process tags to turn them into the expected format for Sagemaker. + + Every accepted representation comes out as a list of boto-style + ``{"Key": ..., "Value": ...}`` dicts with string keys and values: + + * a ``{key: value}`` mapping, one tag per entry; + * a single tag dict or ``Tag`` object; + * a list of tag dicts in either casing, ``(key, value)`` pairs, ``Tag`` objects, or a mix. + + ``None`` is returned unchanged. + """ + if tags is None: + return tags + if isinstance(tags, dict): + if "Key" in tags or "key" in tags: + return [_normalize_tag(tags)] + return [_tag_dict(key, value) for key, value in tags.items()] + if isinstance(tags, (list, tuple)): + return [_normalize_tag(tag) for tag in tags] + return [_normalize_tag(tags)] def _get_resolved_path(path): @@ -2124,38 +2180,43 @@ def tag_exists(tag: TagsDict, curr_tags: Optional[Tags]) -> bool: if curr_tags is None: return False - for curr_tag in curr_tags: - if tag["Key"] == curr_tag["Key"]: + key = _normalize_tag(tag)["Key"] + for curr_tag in format_tags(curr_tags): + if key == curr_tag["Key"]: return True return False -def _validate_new_tags(new_tags: Optional[Tags], curr_tags: Optional[Tags]) -> Optional[Tags]: - """Validates new tags against existing tags. +def _validate_new_tags( + new_tags: Optional[Tags], curr_tags: Optional[Tags] +) -> Optional[List[TagsDict]]: + """Merge new tags into the current tags, without a duplicate key. + + The two inputs take any representation ``format_tags`` accepts. The result is the + merged list in boto form, so a caller that stores it (``ModelBuilder.add_tags``) + holds one shape whatever shape it received. ``None`` comes back only when the two + inputs are ``None``. Args: new_tags (Optional[Tags]): The new tags. curr_tags (Optional[Tags]): The current tags. Returns: - Optional[Tags]: The updated tags. + Optional[List[TagsDict]]: The merged tags in boto form. """ - if curr_tags is None: - return new_tags + if new_tags is None: + return format_tags(curr_tags) if curr_tags is not None else None - if curr_tags and isinstance(curr_tags, dict): - curr_tags = [curr_tags] + if curr_tags is None: + return format_tags(new_tags) - if isinstance(new_tags, dict): - if not tag_exists(new_tags, curr_tags): - curr_tags.append(new_tags) - elif isinstance(new_tags, list): - for new_tag in new_tags: - if not tag_exists(new_tag, curr_tags): - curr_tags.append(new_tag) + merged = format_tags(curr_tags) + for new_tag in format_tags(new_tags): + if not tag_exists(new_tag, merged): + merged.append(new_tag) - return curr_tags + return merged def remove_tag_with_key(key: str, tags: Optional[Tags]) -> Optional[Tags]: @@ -2170,11 +2231,9 @@ def remove_tag_with_key(key: str, tags: Optional[Tags]) -> Optional[Tags]: """ if tags is None: return tags - if isinstance(tags, dict): - tags = [tags] updated_tags = [] - for tag in tags: + for tag in format_tags(tags): if tag["Key"] != key: updated_tags.append(tag) diff --git a/sagemaker-core/src/sagemaker/core/config_schema.py b/sagemaker-core/src/sagemaker/core/config_schema.py index c064177f7b..d9b87a6e77 100644 --- a/sagemaker-core/src/sagemaker/core/config_schema.py +++ b/sagemaker-core/src/sagemaker/core/config_schema.py @@ -1,3 +1,5 @@ +"""JSON schema of the SageMaker Python SDK configuration file.""" + SAGEMAKER_PYTHON_SDK_CONFIG_SCHEMA = { "$schema": "https://json-schema.org/draft/2020-12/schema", "type": "object", diff --git a/sagemaker-core/src/sagemaker/core/endpoint.py b/sagemaker-core/src/sagemaker/core/endpoint.py new file mode 100644 index 0000000000..4308559d0f --- /dev/null +++ b/sagemaker-core/src/sagemaker/core/endpoint.py @@ -0,0 +1,382 @@ +# 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. +"""Hand-written behavior of the generated ``sagemaker.core.resources.Endpoint`` class. + +``EndpointMixin`` gives ``Endpoint`` one call that collects what an operator reads first when +an endpoint misbehaves: ``diagnostics()`` returns the endpoint status and failure reason, one +``VariantDiagnostics`` per production variant with the instance type from the endpoint +config, the 11 CloudWatch metric series of the SageMaker endpoint dashboard, and the latest +lines of the ``/aws/sagemaker/Endpoints/`` log group. ``recent_logs()`` returns the log +lines alone. + +``DescribeEndpoint`` does not return the instance type. ``ProductionVariantSummary`` carries +counts, weights, and statuses, and the instance type lives on the endpoint config, so the +mixin reads ``EndpointConfig.get(endpoint_config_name)`` and matches the variants by name. +The mixin keeps no state and adds no model field. Every client call uses the session and +region the resource was bound with, so ``Endpoint.get(name, session=..., region=...)`` +reads the metrics and logs of that account and region. +""" + +from __future__ import absolute_import + +from dataclasses import dataclass +from datetime import datetime, timedelta, timezone +from typing import Any, Dict, List, Mapping, NamedTuple, Optional, Tuple + +from botocore.exceptions import BotoCoreError, ClientError + +from sagemaker.core.utils.utils import SageMakerClient, Unassigned + +SAGEMAKER_NAMESPACE = "AWS/SageMaker" +ENDPOINT_INSTANCE_NAMESPACE = "/aws/sagemaker/Endpoints" +LOG_GROUP_PREFIX = "/aws/sagemaker/Endpoints/" +METRIC_PERIOD_SECONDS = 60 +DEFAULT_VARIANT_NAME = "AllTraffic" +DEFAULT_METRIC_WINDOW_MINUTES = 5 +DEFAULT_LOG_WINDOW_MINUTES = 15 +DEFAULT_DIAGNOSTICS_LOG_LIMIT = 50 +DEFAULT_LOG_LIMIT = 100 +MAX_LOG_EVENTS_PER_PAGE = 10000 + + +class MetricQuery(NamedTuple): + """One ``GetMetricData`` query of the endpoint dashboard.""" + + query_id: str + namespace: str + metric_name: str + stat: str + + @property + def key(self) -> str: + """The ``EndpointDiagnostics.metrics`` key, ``"/"``.""" + return f"{self.metric_name}/{self.stat}" + + +METRIC_QUERIES = ( + MetricQuery("invocations", SAGEMAKER_NAMESPACE, "Invocations", "Sum"), + MetricQuery("errors_4xx", SAGEMAKER_NAMESPACE, "Invocation4XXErrors", "Sum"), + MetricQuery("errors_5xx", SAGEMAKER_NAMESPACE, "Invocation5XXErrors", "Sum"), + MetricQuery("invocation_model_errors", SAGEMAKER_NAMESPACE, "InvocationModelErrors", "Sum"), + MetricQuery("model_latency_avg", SAGEMAKER_NAMESPACE, "ModelLatency", "Average"), + MetricQuery("model_latency_p99", SAGEMAKER_NAMESPACE, "ModelLatency", "p99"), + MetricQuery("overhead_latency", SAGEMAKER_NAMESPACE, "OverheadLatency", "Average"), + MetricQuery("cpu_util", ENDPOINT_INSTANCE_NAMESPACE, "CPUUtilization", "Average"), + MetricQuery("memory_util", ENDPOINT_INSTANCE_NAMESPACE, "MemoryUtilization", "Average"), + MetricQuery("gpu_util", ENDPOINT_INSTANCE_NAMESPACE, "GPUUtilization", "Average"), + MetricQuery("gpu_memory_util", ENDPOINT_INSTANCE_NAMESPACE, "GPUMemoryUtilization", "Average"), +) + + +class TimeWindow(NamedTuple): + """A UTC window that ends at the time it was built.""" + + start: datetime + end: datetime + + @classmethod + def last(cls, minutes: int) -> "TimeWindow": + """The window of the last ``minutes`` minutes.""" + end = datetime.now(timezone.utc) + return cls(start=end - timedelta(minutes=minutes), end=end) + + @property + def start_millis(self) -> int: + """The start as epoch milliseconds, the unit of ``FilterLogEvents``.""" + return int(self.start.timestamp() * 1000) + + @property + def end_millis(self) -> int: + """The end as epoch milliseconds, the unit of ``FilterLogEvents``.""" + return int(self.end.timestamp() * 1000) + + +@dataclass(frozen=True) +class VariantDiagnostics: + """The state of one production variant, with the instance type from the endpoint config.""" + + variant_name: str + instance_type: Optional[str] + current_instance_count: Optional[int] + desired_instance_count: Optional[int] + current_weight: Optional[float] + variant_status: List[str] + + +@dataclass(frozen=True) +class MetricSeries: + """One CloudWatch metric statistic over the diagnostics window, oldest value first.""" + + namespace: str + metric_name: str + stat: str + values: List[float] + timestamps: List[datetime] + + +@dataclass(frozen=True) +class EndpointDiagnostics: + """What ``Endpoint.diagnostics()`` collected. + + ``metrics`` is keyed ``"/"``, for example ``"ModelLatency/p99"``. + ``errors`` holds one line per data source that failed. The endpoint config, CloudWatch, + and CloudWatch Logs never raise through ``diagnostics()``. + """ + + endpoint_name: str + endpoint_status: str + failure_reason: Optional[str] + endpoint_config_name: Optional[str] + variants: List[VariantDiagnostics] + metrics: Dict[str, MetricSeries] + recent_logs: List[str] + errors: List[str] + + +class EndpointMixin: + """``diagnostics()`` and ``recent_logs()`` for the generated ``Endpoint`` resource. + + Mixed into ``sagemaker.core.resources.Endpoint``. The methods read the generated + attributes at call time and build every boto3 client from the session and region the + resource is bound to, so the mixin adds no model field and no state. + """ + + endpoint_name: Any + endpoint_config_name: Any + endpoint_status: Any + failure_reason: Any + production_variants: Any + + def diagnostics( + self, + window_minutes: int = DEFAULT_METRIC_WINDOW_MINUTES, + variant_name: Optional[str] = None, + log_limit: int = DEFAULT_DIAGNOSTICS_LOG_LIMIT, + ) -> EndpointDiagnostics: + """Collect the endpoint state, the dashboard metrics, and the latest log lines. + + Args: + window_minutes: The metric window, counted back from now. The log window is + this value or 15 minutes, whichever is longer. + variant_name: The variant whose metrics to read. Defaults to the first + production variant, or ``AllTraffic`` when the endpoint has none. + log_limit: The maximum number of log lines. + + Returns: + The collected ``EndpointDiagnostics``. A CloudWatch, CloudWatch Logs, or + endpoint config failure lands in ``errors``. + + Raises: + botocore.exceptions.ClientError: ``DescribeEndpoint`` failed. + """ + self.refresh() + variants = _assigned(self.production_variants) or [] + instance_types, config_error = self._instance_types() + metric_variant = variant_name or _first_variant_name(variants) + metrics, metric_error = self._metrics(metric_variant, TimeWindow.last(window_minutes)) + log_window_minutes = max(window_minutes, DEFAULT_LOG_WINDOW_MINUTES) + recent_logs, log_error = self._recent_logs_or_error(log_window_minutes, log_limit) + return EndpointDiagnostics( + endpoint_name=self.endpoint_name, + endpoint_status=_assigned(self.endpoint_status) or "", + failure_reason=_assigned(self.failure_reason), + endpoint_config_name=_assigned(self.endpoint_config_name), + variants=[_variant_diagnostics(summary, instance_types) for summary in variants], + metrics=metrics, + recent_logs=recent_logs, + errors=[error for error in (config_error, metric_error, log_error) if error], + ) + + def recent_logs( + self, + window_minutes: int = DEFAULT_LOG_WINDOW_MINUTES, + limit: int = DEFAULT_LOG_LIMIT, + ) -> List[str]: + """The latest lines of the ``/aws/sagemaker/Endpoints/`` log group. + + Args: + window_minutes: The log window, counted back from now. + limit: The maximum number of lines. Pages until the log group has no next + token, returns an empty page, or the limit is reached. + + Returns: + One ``" "`` line per event. + + Raises: + botocore.exceptions.ClientError: ``FilterLogEvents`` failed, for example + because the log group does not exist yet. + """ + window = TimeWindow.last(window_minutes) + client = self._boto_client("logs") + request = { + "logGroupName": LOG_GROUP_PREFIX + self.endpoint_name, + "startTime": window.start_millis, + "endTime": window.end_millis, + } + lines: List[str] = [] + next_token: Optional[str] = None + while len(lines) < limit: + page_request = dict(request, limit=min(limit - len(lines), MAX_LOG_EVENTS_PER_PAGE)) + if next_token: + page_request["nextToken"] = next_token + page = client.filter_log_events(**page_request) + events = page.get("events", []) + lines.extend(_format_log_event(event) for event in events) + next_token = page.get("nextToken") + if not events or not next_token: + break + return lines[:limit] + + def _instance_types(self) -> Tuple[Dict[str, Optional[str]], Optional[str]]: + """The instance type per variant name from the endpoint config, or the failure line.""" + # resources.py imports this module, so the resource class is imported at call time. + from sagemaker.core.resources import EndpointConfig + + config_name = _assigned(self.endpoint_config_name) + if not config_name: + return {}, None + try: + config = EndpointConfig.get( + endpoint_config_name=config_name, session=self._session, region=self._region + ) + except (BotoCoreError, ClientError) as error: + return {}, f"DescribeEndpointConfig failed: {error}" + config_variants = _assigned(config.production_variants) or [] + return { + variant.variant_name: _assigned(variant.instance_type) for variant in config_variants + }, None + + def _metrics( + self, variant_name: str, window: TimeWindow + ) -> Tuple[Dict[str, MetricSeries], Optional[str]]: + """The 11 dashboard metric series of ``variant_name``, or the failure line.""" + try: + results = self._metric_data_results(variant_name, window) + except (BotoCoreError, ClientError) as error: + return {}, f"GetMetricData failed: {error}" + return _metric_series_by_key(results), None + + def _metric_data_results(self, variant_name: str, window: TimeWindow) -> List[Dict[str, Any]]: + """Every ``MetricDataResults`` entry of the window, across all ``GetMetricData`` pages.""" + client = self._boto_client("cloudwatch") + request = { + "MetricDataQueries": metric_data_queries(self.endpoint_name, variant_name), + "StartTime": window.start, + "EndTime": window.end, + "ScanBy": "TimestampAscending", + } + results: List[Dict[str, Any]] = [] + next_token: Optional[str] = None + while True: + page_request = dict(request, NextToken=next_token) if next_token else request + page = client.get_metric_data(**page_request) + results.extend(page.get("MetricDataResults", [])) + next_token = page.get("NextToken") + if not next_token: + return results + + def _recent_logs_or_error( + self, window_minutes: int, limit: int + ) -> Tuple[List[str], Optional[str]]: + try: + return self.recent_logs(window_minutes=window_minutes, limit=limit), None + except (BotoCoreError, ClientError) as error: + return [], f"FilterLogEvents failed: {error}" + + def _boto_client(self, service_name: str) -> Any: + """A boto3 client on the session, region, and config this resource is bound to. + + ``SageMakerClient`` caches only the SageMaker family of clients, so this method builds + the CloudWatch and CloudWatch Logs clients from its resolved session, region, and config. + """ + sagemaker_client = SageMakerClient(session=self._session, region_name=self._region) + return sagemaker_client.session.client( + service_name, sagemaker_client.region_name, config=sagemaker_client.config + ) + + +def metric_data_queries(endpoint_name: str, variant_name: str) -> List[Dict[str, Any]]: + """The ``MetricDataQueries`` of the endpoint dashboard for one variant. + + SageMaker publishes every endpoint metric with the ``EndpointName`` and ``VariantName`` + dimensions, and CloudWatch matches a dimension set exactly, so each query names the two. + """ + dimensions = [ + {"Name": "EndpointName", "Value": endpoint_name}, + {"Name": "VariantName", "Value": variant_name}, + ] + return [ + { + "Id": query.query_id, + "MetricStat": { + "Metric": { + "Namespace": query.namespace, + "MetricName": query.metric_name, + "Dimensions": dimensions, + }, + "Period": METRIC_PERIOD_SECONDS, + "Stat": query.stat, + }, + } + for query in METRIC_QUERIES + ] + + +def _assigned(value: Any) -> Any: + """Return ``value`` unless it is the generated ``Unassigned()`` sentinel or ``None``.""" + if value is None or isinstance(value, Unassigned): + return None + return value + + +def _first_variant_name(variants: List[Any]) -> str: + return variants[0].variant_name if variants else DEFAULT_VARIANT_NAME + + +def _variant_diagnostics( + summary: Any, instance_types: Mapping[str, Optional[str]] +) -> VariantDiagnostics: + statuses = _assigned(summary.variant_status) or [] + return VariantDiagnostics( + variant_name=summary.variant_name, + instance_type=instance_types.get(summary.variant_name), + current_instance_count=_assigned(summary.current_instance_count), + desired_instance_count=_assigned(summary.desired_instance_count), + current_weight=_assigned(summary.current_weight), + variant_status=[status.status for status in statuses], + ) + + +def _metric_series_by_key(results: List[Mapping[str, Any]]) -> Dict[str, MetricSeries]: + """Group the results by ``"/"``. A paged query id appends to its series.""" + queries_by_id = {query.query_id: query for query in METRIC_QUERIES} + series: Dict[str, MetricSeries] = {} + for result in results: + query = queries_by_id.get(result.get("Id")) + if not query: + continue + previous = series.get(query.key) + series[query.key] = MetricSeries( + namespace=query.namespace, + metric_name=query.metric_name, + stat=query.stat, + values=(previous.values if previous else []) + list(result.get("Values", [])), + timestamps=(previous.timestamps if previous else []) + + list(result.get("Timestamps", [])), + ) + return series + + +def _format_log_event(event: Mapping[str, Any]) -> str: + timestamp = datetime.fromtimestamp(event.get("timestamp", 0) / 1000, tz=timezone.utc) + return f"{timestamp.isoformat()} {event.get('logStreamName', '')} {event.get('message', '')}" diff --git a/sagemaker-core/src/sagemaker/core/helper/session_helper.py b/sagemaker-core/src/sagemaker/core/helper/session_helper.py index e8eab15519..a847982d60 100644 --- a/sagemaker-core/src/sagemaker/core/helper/session_helper.py +++ b/sagemaker-core/src/sagemaker/core/helper/session_helper.py @@ -962,6 +962,7 @@ def endpoint_from_production_variants( vpc_config=None, enable_network_isolation=None, role=None, + tracing_config_dict=None, ): """Create an SageMaker ``Endpoint`` from a list of production variants. @@ -996,6 +997,9 @@ def endpoint_from_production_variants( artifacts. After the endpoint is created, the inference code might use the IAM role if it needs to access some AWS resources. (default: None). + tracing_config_dict (dict): The ``TracingConfig`` member of the endpoint + configuration, in service form. Raises ``ValueError`` when the installed + botocore service model has no such member (default: None). Returns: str: The name of the created ``Endpoint``. """ @@ -1011,6 +1015,7 @@ def endpoint_from_production_variants( vpc_config=vpc_config, enable_network_isolation=enable_network_isolation, role=role, + tracing_config_dict=tracing_config_dict, ) endpoint_tags = _append_project_tags(format_tags(tags)) @@ -1037,6 +1042,7 @@ def _build_endpoint_config_request( vpc_config=None, enable_network_isolation=None, role=None, + tracing_config_dict=None, ): """Build the ``CreateEndpointConfig`` request, applying config defaults. @@ -1047,11 +1053,9 @@ def _build_endpoint_config_request( dict: The ``CreateEndpointConfig`` request. """ supports_kms = any( - [ - instance_supports_kms(production_variant["InstanceType"]) - for production_variant in production_variants - if "InstanceType" in production_variant - ] + instance_supports_kms(instance_type) + for production_variant in production_variants + for instance_type in _production_variant_instance_types(production_variant) ) update_list_of_dicts_with_values_from_config( @@ -1128,9 +1132,29 @@ def _build_endpoint_config_request( config_options["EnableNetworkIsolation"] = enable_network_isolation if role is not None: config_options["ExecutionRoleArn"] = role + if tracing_config_dict is not None: + self._require_create_endpoint_config_member("TracingConfig") + config_options["TracingConfig"] = tracing_config_dict return config_options + def _require_create_endpoint_config_member(self, member_name): + """Raise ``ValueError`` when the installed service model lacks a request member. + + botocore drops a request member its service model does not know, so a config such + as ``TracingConfig`` would vanish without this check. + """ + operation_model = self.sagemaker_client.meta.service_model.operation_model( + "CreateEndpointConfig" + ) + if member_name in operation_model.input_shape.members: + return + raise ValueError( + f"{member_name} is not a CreateEndpointConfig member in the installed botocore " + f"{botocore.__version__} service model, so the endpoint configuration cannot carry " + "it. Upgrade botocore to a release whose SageMaker model has it." + ) + def create_endpoint_config( self, name, @@ -1143,6 +1167,7 @@ def create_endpoint_config( vpc_config=None, enable_network_isolation=None, role=None, + tracing_config_dict=None, ): """Create an Amazon SageMaker endpoint configuration and nothing else. @@ -1168,6 +1193,9 @@ def create_endpoint_config( containers deployed to the endpoint. role (str): An execution role ARN, required when the production variants carry no model name, as with inference components. + tracing_config_dict (dict): The ``TracingConfig`` member in service + form. Raises ``ValueError`` when the installed botocore service + model has no such member. Returns: str: Name of the endpoint configuration created. Under a @@ -1185,6 +1213,7 @@ def create_endpoint_config( vpc_config=vpc_config, enable_network_isolation=enable_network_isolation, role=role, + tracing_config_dict=tracing_config_dict, ) def submit(request): @@ -2859,6 +2888,13 @@ def update_args(args: Dict[str, Any], **kwargs): args.update({key: value}) +def _production_variant_instance_types(production_variant): + """Return the instance types a production variant provisions, one or the pool list.""" + if "InstanceType" in production_variant: + return [production_variant["InstanceType"]] + return [pool["InstanceType"] for pool in production_variant.get("InstancePools", [])] + + def production_variant( model_name=None, instance_type=None, @@ -2873,6 +2909,7 @@ def production_variant( managed_instance_scaling=None, routing_config=None, inference_ami_version=None, + instance_pools=None, ): """Create a production variant description suitable for use in a ``ProductionVariant`` list. @@ -2906,6 +2943,10 @@ def production_variant( https://docs.aws.amazon.com/sagemaker/latest/dg/your-algorithms-inference-code.html #your-algorithms -inference-algo-ping-requests + instance_pools (list[dict]): Ordered ``InstancePools`` entries + (``{"InstanceType": ..., "Priority": ...}``), the primary type first. The service + takes ``InstancePools`` in place of ``InstanceType``, so the variant carries the pools + and no ``InstanceType`` when this is set. Ignored for a serverless variant. Returns: dict[str, str]: An SageMaker ``ProductionVariant`` description @@ -2928,7 +2969,10 @@ def production_variant( else: initial_instance_count = initial_instance_count or 1 production_variant_configuration["InitialInstanceCount"] = initial_instance_count - production_variant_configuration["InstanceType"] = instance_type + if instance_pools: + production_variant_configuration["InstancePools"] = instance_pools + else: + production_variant_configuration["InstanceType"] = instance_type update_args( production_variant_configuration, VolumeSizeInGB=volume_size, diff --git a/sagemaker-core/src/sagemaker/core/hub_catalog.py b/sagemaker-core/src/sagemaker/core/hub_catalog.py new file mode 100644 index 0000000000..ac9f4a2956 --- /dev/null +++ b/sagemaker-core/src/sagemaker/core/hub_catalog.py @@ -0,0 +1,408 @@ +# 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. +"""A cache of the described contents of a hub, in the shape of the JumpStart models cache. + +``HubContent.get_all`` costs one ``ListHubContents`` call per page plus one +``DescribeHubContent`` call per item, about 770 describes for the public hub. The list +is cheap (8 pages, about 1 second) and the describes are the cost (60 to 95 seconds). + +``HubCatalogCache`` follows ``sagemaker.core.jumpstart.cache.JumpStartModelsCache``. Its +in-memory layer is the same ``sagemaker.core.utilities.cache.LRUCache``, with one entry per +hub and content type, a bounded item count, and an expiration horizon. The ``LRUCache`` +calls the retrieval function on a miss, or on an expired entry with the old value. The +retrieval function lists the hub and keeps every described item whose ``HubContentVersion`` +still matches the list summary. Only the new or changed items are described. The JumpStart +retrieval function reads the manifest again only when its ETag changed, in the same way. + +On top of the JumpStart shape, the described items persist to one JSON file per hub, +content type, region, and account. A new process starts from the previous catalog instead +of describing every item again. The file lives under the ``platformdirs`` user cache +directory (``~/.cache/sagemaker/hub-content`` on Linux) or ``SAGEMAKER_HUB_CACHE_DIR``. It +holds the service response for catalog metadata and no credentials. A file is written to +a temporary name and renamed into place, so a reader never sees a partial file. + +``HubCatalogAccessor`` holds the process-wide cache instance, as +``JumpStartModelsAccessor`` holds the JumpStart one. +""" + +from __future__ import absolute_import + +import datetime +import json +import logging +import os +import re +import tempfile +import threading +from concurrent.futures import ThreadPoolExecutor +from dataclasses import dataclass +from typing import Any, Callable, Dict, List, Optional, Tuple + +from platformdirs import user_cache_dir + +from sagemaker.core.utilities.cache import LRUCache + +logger = logging.getLogger(__name__) + +CACHE_DIR_ENV_VAR = "SAGEMAKER_HUB_CACHE_DIR" +CACHE_APP_NAME = "sagemaker" +CACHE_SUBDIRECTORY = "hub-content" +CACHE_FORMAT_VERSION = 2 +HUB_CATALOG_DEFAULT_MAX_CACHE_ITEMS = 20 +HUB_CATALOG_DEFAULT_CACHE_EXPIRATION_HORIZON = datetime.timedelta(hours=6) +HUB_CATALOG_DEFAULT_DESCRIBE_WORKERS = 4 +LIST_PAGE_SIZE = 100 +_UNSAFE_FILE_NAME_CHARACTERS = re.compile(r"[^A-Za-z0-9._-]+") + +Summary = Dict[str, Any] +Entry = Dict[str, Any] + + +def default_cache_dir() -> str: + """Return the directory of the hub content cache files. + + ``SAGEMAKER_HUB_CACHE_DIR`` wins when set. The default is the ``hub-content`` + directory under the ``platformdirs`` user cache directory of the ``sagemaker`` + application, ``~/.cache/sagemaker/hub-content`` on Linux. + """ + override = os.environ.get(CACHE_DIR_ENV_VAR) + if override: + return os.path.expanduser(override) + return os.path.join(user_cache_dir(CACHE_APP_NAME), CACHE_SUBDIRECTORY) + + +@dataclass(frozen=True) +class HubCatalogKey: + """The key of one cache entry: one hub and content type in one region and account.""" + + hub_name: str + hub_content_type: str + region: str + account_id: str + + def file_name(self) -> str: + """Return the cache file name, with every part reduced to file-system-safe text.""" + parts = (self.account_id, self.region, self.hub_name, self.hub_content_type) + safe_parts = [_UNSAFE_FILE_NAME_CHARACTERS.sub("_", part) for part in parts] + return "-".join(safe_parts) + ".json" + + +DescribeFunction = Callable[[HubCatalogKey, Summary], Any] +ListFunction = Callable[[HubCatalogKey], List[Summary]] + + +@dataclass(frozen=True) +class HubCatalogSource: + """The service calls of one retrieval: list a hub, and describe one listed summary. + + The calls carry the session of the caller. They are passed per call and not stored on + the cache, so one process cache serves every session. + """ + + list_function: ListFunction + describe_function: DescribeFunction + + +@dataclass(frozen=True) +class HubCatalogRefresh: + """What one retrieval did: how many items it listed, reused, described, and dropped.""" + + listed: int + reused: int + described: int + removed: int + cache_path: str + + +@dataclass(frozen=True) +class HubCatalogValue: + """The value of one cache entry: the described items and the refresh that produced them. + + ``entries`` maps a content name to ``{"listed_version": ..., "resource": ...}``, where + ``resource`` is the ``model_dump`` of the described ``HubContent``. ``summaries`` is the + ``ListHubContents`` output of the retrieval, in service order. + """ + + entries: Dict[str, Entry] + summaries: List[Summary] + refresh: HubCatalogRefresh + + +class HubCatalogFile: + """Read and write the on-disk copy of one :class:`HubCatalogKey`.""" + + def __init__(self, key: HubCatalogKey, cache_dir: Optional[str] = None): + self.key = key + self.cache_dir = cache_dir or default_cache_dir() + self.path = os.path.join(self.cache_dir, key.file_name()) + + def load(self) -> Dict[str, Entry]: + """Return the entries by name, or an empty dict for a missing or unreadable file.""" + try: + with open(self.path, "r", encoding="utf-8") as cache_file: + payload = json.load(cache_file) + except FileNotFoundError: + return {} + except (OSError, ValueError) as error: + logger.warning("Ignoring the unreadable hub content cache %s: %s", self.path, error) + return {} + if payload.get("format_version") != CACHE_FORMAT_VERSION: + return {} + entries = payload.get("entries") + return entries if isinstance(entries, dict) else {} + + def save(self, entries: Dict[str, Entry]) -> None: + """Write the entries, through a temporary file and a rename.""" + os.makedirs(self.cache_dir, exist_ok=True) + payload = { + "format_version": CACHE_FORMAT_VERSION, + "written_at": datetime.datetime.now(datetime.timezone.utc).isoformat(), + "key": { + "hub_name": self.key.hub_name, + "hub_content_type": self.key.hub_content_type, + "region": self.key.region, + "account_id": self.key.account_id, + }, + "entries": entries, + } + descriptor, temporary_path = tempfile.mkstemp( + dir=self.cache_dir, prefix=self.key.file_name(), suffix=".tmp" + ) + try: + with os.fdopen(descriptor, "w", encoding="utf-8") as cache_file: + json.dump(payload, cache_file) + os.replace(temporary_path, self.path) + except BaseException: + _remove_quietly(temporary_path) + raise + + def clear(self) -> bool: + """Delete the file. Return ``True`` when a file was deleted.""" + try: + os.remove(self.path) + except FileNotFoundError: + return False + return True + + +class HubCatalogCache: + """Cache of the described contents of hubs, one ``LRUCache`` entry per hub and type. + + Built like ``JumpStartModelsCache``: an ``LRUCache`` with a bounded item count and an + expiration horizon. The ``LRUCache`` calls the retrieval function on a miss, or on an + expired entry with the old value. The retrieval function lists the hub and reuses the + old items whose ``HubContentVersion`` the list still reports. The rest are described + ``max_workers`` at a time, and the result persists to the cache file of the key. + + Args: + max_cache_items: Maximum number of hubs to hold in memory. Default: 20. + cache_expiration_horizon: Maximum age of an in-memory entry before the hub is + listed again. Default: 6 hours. + max_workers: How many describe calls to keep in flight. Default: 4. + cache_dir: The directory of the cache files. Default: :func:`default_cache_dir`. + """ + + def __init__( + self, + max_cache_items: int = HUB_CATALOG_DEFAULT_MAX_CACHE_ITEMS, + cache_expiration_horizon: datetime.timedelta = HUB_CATALOG_DEFAULT_CACHE_EXPIRATION_HORIZON, + max_workers: int = HUB_CATALOG_DEFAULT_DESCRIBE_WORKERS, + cache_dir: Optional[str] = None, + ) -> None: + if max_workers < 1: + raise ValueError(f"max_workers must be at least 1, got {max_workers}.") + self._max_workers = max_workers + self._cache_dir = cache_dir + # The LRUCache is not thread safe, so its reads and writes run under this lock. A + # retrieval runs outside it, with one lock per key so two threads that miss the + # same hub share one list-and-describe pass while other hubs stay unblocked. + self._cache_lock = threading.Lock() + self._key_locks: Dict[HubCatalogKey, threading.Lock] = {} + self._content_cache = LRUCache[HubCatalogKey, HubCatalogValue]( + max_cache_items=max_cache_items, + expiration_horizon=cache_expiration_horizon, + retrieval_function=self._retrieval_function, + ) + + def get(self, key: HubCatalogKey, source: HubCatalogSource) -> HubCatalogValue: + """Return the catalog of a hub from memory, or retrieve it when absent or expired.""" + with self._cache_lock: + try: + value, _ = self._content_cache.get(key, data_source_fallback=False) + return value + except KeyError: + previous = self._previous_value(key) + with self._key_lock(key): + with self._cache_lock: + try: + value, _ = self._content_cache.get(key, data_source_fallback=False) + return value + except KeyError: + pass + return self._retrieve_and_store(key, previous, source) + + def refresh(self, key: HubCatalogKey, source: HubCatalogSource) -> HubCatalogValue: + """List the hub again, describe the changed items, and return the new catalog.""" + with self._cache_lock: + previous = self._previous_value(key) + with self._key_lock(key): + return self._retrieve_and_store(key, previous, source) + + def clear(self, key: Optional[HubCatalogKey] = None) -> bool: + """Drop the in-memory entries, and the cache file of ``key`` when one is given. + + Returns ``True`` when a file was deleted. + """ + with self._cache_lock: + self._content_cache.clear() + if key is None: + return False + return HubCatalogFile(key, self._cache_dir).clear() + + def _previous_value(self, key: HubCatalogKey) -> Optional[HubCatalogValue]: + """Return the stored value of ``key``, expired or not, without a retrieval.""" + element = self._content_cache._lru_cache.get(key) # pylint: disable=protected-access + return element.value if element is not None else None + + def _key_lock(self, key: HubCatalogKey) -> threading.Lock: + with self._cache_lock: + return self._key_locks.setdefault(key, threading.Lock()) + + def _retrieve_and_store( + self, key: HubCatalogKey, previous: Optional[HubCatalogValue], source: HubCatalogSource + ) -> HubCatalogValue: + """Run the retrieval outside the cache lock, then store the result under it.""" + value = self._retrieval_function(key, previous, source) + with self._cache_lock: + self._content_cache.put(key, value) + return value + + def _retrieval_function( + self, + key: HubCatalogKey, + value: Optional[HubCatalogValue], + source: Optional[HubCatalogSource] = None, + ) -> HubCatalogValue: + """Bring the catalog of ``key`` up to date, from the old value and the cache file. + + The ``LRUCache`` holds this function as its ``retrieval_function`` and calls it with + ``key`` and ``value`` only. :meth:`get` and :meth:`refresh` call it themselves with + the ``source`` of the caller, so the ``LRUCache`` path raises for a missing source. + """ + if source is None: + raise RuntimeError("HubCatalogCache retrieves only through get() or refresh().") + cache_file = HubCatalogFile(key, self._cache_dir) + previous = value.entries if value is not None else cache_file.load() + summaries = source.list_function(key) + reused = reusable_entries(previous, summaries) + missing = [s for s in summaries if s.get("HubContentName") not in reused] + entries = dict(reused) + try: + for summary, resource in self._describe_all(source, key, missing): + entries[summary["HubContentName"]] = resource_entry( + resource, summary.get("HubContentVersion") + ) + finally: + # A describe that raises still leaves every described item on disk, so the + # next retrieval describes only the rest. + cache_file.save(entries) + refresh = HubCatalogRefresh( + listed=len(summaries), + reused=len(reused), + described=len(missing), + removed=len(set(previous) - set(entries)), + cache_path=cache_file.path, + ) + return HubCatalogValue(entries=entries, summaries=summaries, refresh=refresh) + + def _describe_all(self, source: HubCatalogSource, key: HubCatalogKey, summaries): + """Yield ``(summary, resource)`` for each summary, described on a thread pool.""" + + def describe(summary: Summary) -> Any: + return source.describe_function(key, summary) + + with ThreadPoolExecutor( + max_workers=self._max_workers, thread_name_prefix="sagemaker-hub-describe" + ) as executor: + yield from zip(summaries, executor.map(describe, summaries)) + + +class HubCatalogAccessor: + """Static holder of the process-wide :class:`HubCatalogCache`. + + Mirrors ``JumpStartModelsAccessor``: one cache per process, built from the stored + cache kwargs, rebuilt when :meth:`set_cache_kwargs` changes them, and dropped through + :meth:`reset_cache`. + """ + + _cache: Optional[HubCatalogCache] = None + _cache_kwargs: Dict[str, Any] = {} + _lock = threading.Lock() + + @staticmethod + def get_cache() -> HubCatalogCache: + """Return the process cache, built on the first call with the stored kwargs.""" + with HubCatalogAccessor._lock: + if HubCatalogAccessor._cache is None: + HubCatalogAccessor._cache = HubCatalogCache(**HubCatalogAccessor._cache_kwargs) + return HubCatalogAccessor._cache + + @staticmethod + def set_cache_kwargs(**cache_kwargs: Any) -> None: + """Store new ``HubCatalogCache`` kwargs and drop the cache, so the next call rebuilds it.""" + with HubCatalogAccessor._lock: + HubCatalogAccessor._cache_kwargs = dict(cache_kwargs) + HubCatalogAccessor._cache = None + + @staticmethod + def reset_cache() -> None: + """Drop the process cache and its kwargs, so the next call builds a default one.""" + with HubCatalogAccessor._lock: + HubCatalogAccessor._cache = None + HubCatalogAccessor._cache_kwargs = {} + + +def resource_entry(resource: Any, listed_version: str) -> Entry: + """Return the cache entry of one described resource.""" + return { + "listed_version": listed_version, + "resource": resource.model_dump(exclude_unset=True, mode="json"), + } + + +def reusable_entries(entries: Dict[str, Entry], summaries: List[Summary]) -> Dict[str, Entry]: + """Return the entries whose listed version matches the current list summary.""" + reusable = {} + for summary in summaries: + name = summary.get("HubContentName") + entry = entries.get(name) + if entry and entry.get("listed_version") == summary.get("HubContentVersion"): + reusable[name] = entry + return reusable + + +def ordered_entries(value: HubCatalogValue) -> List[Tuple[Summary, Entry]]: + """Return ``(summary, entry)`` pairs in the order the service listed the hub.""" + pairs = [] + for summary in value.summaries: + entry = value.entries.get(summary.get("HubContentName")) + if entry is not None: + pairs.append((summary, entry)) + return pairs + + +def _remove_quietly(path: str) -> None: + try: + os.remove(path) + except OSError: + pass diff --git a/sagemaker-core/src/sagemaker/core/hub_content.py b/sagemaker-core/src/sagemaker/core/hub_content.py new file mode 100644 index 0000000000..1d44b93c3c --- /dev/null +++ b/sagemaker-core/src/sagemaker/core/hub_content.py @@ -0,0 +1,669 @@ +# 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. +"""Hand-written behavior of the generated ``sagemaker.core.resources.HubContent`` class. + +A hub model describes itself in two places. ``HubContentSearchKeywords`` holds up to 50 +prefixed strings (``@task:``, ``@license:``, ``@provider:``, ``@model-size:``, +``@context-window:``, ``@framework:``, ``@language:``, ``@data-type:``, +``@input-modality:``, ``@output-modality:``, ``@model-type:``, ``@capability:``, +``@recipe:``), and ``HubContentDocument`` holds the same values as JSON fields with their +original casing and no cap. ``ListHubContents`` and ``Search`` return the keywords, +``DescribeHubContent`` and ``Search`` return the document. + +``HubContentMixin`` gives ``HubContent`` typed accessors over the two: ``tasks``, +``license``, ``provider``, ``model_size``, ``context_window``, ``languages``, +``capabilities``, ``recipes``, and so on. A keyword, when present, is authoritative. The +document fills in only what the keyword cap dropped. The mixin keeps no state, so it +works on a resource from ``HubContent.get`` (document, no ``OriginalCreationTime``), from +``HubContent.get_all`` (keywords and, after the per-item refresh, the document), and on a +``HubContent`` built by hand from a boto response. + +``HubContent.search`` runs the service-side ``Search`` API with ``Resource="HubContent"``. +The service indexes the document fields as ``Metadata.`` (``Metadata.Task``, +``Metadata.License``, ``Metadata.Provider``, ``Metadata.ModelSize``, ...), so a filter such +as ``HubContentFilter.metadata("license", "apache-2.0")`` narrows the public hub without +listing it. Each result record is a full ``HubContent`` with its keywords and document. +""" + +from __future__ import absolute_import + +import json +import threading +from dataclasses import dataclass +from typing import Any, Dict, Iterable, Iterator, List, Mapping, Optional, Sequence, Union + +from sagemaker.core.hub_catalog import ( + HubCatalogAccessor, + HubCatalogKey, + HubCatalogRefresh, + HubCatalogSource, + HubCatalogValue, + ordered_entries, +) +from sagemaker.core.telemetry.constants import Feature +from sagemaker.core.telemetry.telemetry_logging import _telemetry_emitter + +KEYWORD_PREFIXES = { + "tasks": "@task:", + "license": "@license:", + "provider": "@provider:", + "model_size": "@model-size:", + "context_window": "@context-window:", + "framework": "@framework:", + "languages": "@language:", + "data_types": "@data-type:", + "input_modalities": "@input-modality:", + "output_modalities": "@output-modality:", + "model_type": "@model-type:", +} +CAPABILITY_PREFIX = "@capability:" +RECIPE_PREFIX = "@recipe:" +CUSTOMIZATION_CAPABILITY = "customization" +BEDROCK_CAPABILITY = "bedrock_console" +FINE_TUNING_RECIPE_TYPE = "finetuning" + +DOCUMENT_FIELDS = { + "tasks": "Task", + "license": "License", + "provider": "Provider", + "model_size": "ModelSize", + "context_window": "ContextWindow", + "framework": "Framework", + "languages": "Languages", + "data_types": "DataType", + "input_modalities": "InputModalities", + "output_modalities": "OutputModalities", + "model_type": "ModelTypes", +} +LIST_FIELDS = frozenset( + ("tasks", "languages", "data_types", "input_modalities", "output_modalities") +) + +# Search indexes the hub content document as ``Metadata.``. The two extra names are +# document fields that have no search keyword. +METADATA_FIELDS = dict(DOCUMENT_FIELDS, capabilities="Capabilities", recipes="RecipeCollection") +SEARCH_OPERATORS = frozenset( + ( + "Equals", + "NotEquals", + "GreaterThan", + "GreaterThanOrEqualTo", + "LessThan", + "LessThanOrEqualTo", + "Contains", + "Exists", + "NotExists", + "In", + ) +) +# Service limits of a SearchExpression: 20 filters per expression, 1024 characters per value. +MAX_SEARCH_FILTERS = 20 +MAX_SEARCH_FILTER_VALUE_LENGTH = 1024 +SEARCH_PAGE_SIZE = 100 +LIST_PAGE_SIZE = 100 +# The account of a set of credentials, so a cache hit on an unchanged hub makes no STS call. +# Keyed by the frozen credentials, which change when the caller assumes another role. +_ACCOUNT_IDS: Dict[Any, str] = {} +_ACCOUNT_IDS_LOCK = threading.Lock() + + +@dataclass(frozen=True) +class HubContentRecipe: + """One ``@recipe:`` keyword or ``RecipeCollection`` entry of a hub model. + + The hub tags a recipe as ``@recipe:{type}_{technique}`` or + ``@recipe:{type}_{technique}_{strategy}``, all lowercase. ``@recipe:finetuning_sft_lora`` + is a FineTuning recipe for the SFT technique with the LoRA strategy, and + ``@recipe:finetuning_cpt`` has no strategy. + """ + + recipe_type: str + technique: str + strategy: Optional[str] = None + + @classmethod + def from_keyword(cls, keyword: str) -> Optional["HubContentRecipe"]: + """Parse a ``@recipe:`` keyword, or return ``None`` for any other keyword.""" + if not keyword.lower().startswith(RECIPE_PREFIX): + return None + parts = keyword[len(RECIPE_PREFIX) :].lower().split("_", 2) + if len(parts) < 2 or not parts[0] or not parts[1]: + return None + strategy = parts[2] if len(parts) == 3 and parts[2] else None + return cls(recipe_type=parts[0], technique=parts[1], strategy=strategy) + + @classmethod + def from_document(cls, recipe: Mapping[str, Any]) -> Optional["HubContentRecipe"]: + """Parse one ``RecipeCollection`` entry of a hub content document.""" + recipe_type = recipe.get("Type") or recipe.get("RecipeType") + technique = recipe.get("CustomizationTechnique") or recipe.get("EvaluationType") + if not recipe_type or not technique: + return None + strategy = recipe.get("Strategy") or recipe.get("TrainingType") + return cls( + recipe_type=str(recipe_type).lower(), + technique=str(technique).lower(), + strategy=str(strategy).lower() if strategy else None, + ) + + +@dataclass(frozen=True) +class HubContentFilter: + """One ``Filter`` of a ``Search(Resource="HubContent")`` expression. + + ``name`` is a search property such as ``HubContentName``, ``HubContentSearchKeywords``, + or ``Metadata.License``. ``operator`` is one of ``SEARCH_OPERATORS``. ``value`` is the + filter value, or a sequence of values for ``In`` (joined with commas, as the service + expects). ``Exists`` and ``NotExists`` take no value. + """ + + name: str + operator: str = "Equals" + value: Optional[Union[str, Sequence[str]]] = None + + @classmethod + def metadata( + cls, attribute: str, value: Union[str, Sequence[str]], operator: str = "Equals" + ) -> "HubContentFilter": + """Build a ``Metadata.`` filter from a ``HubContent`` attribute name. + + ``HubContentFilter.metadata("license", "apache-2.0")`` is + ``HubContentFilter("Metadata.License", "Equals", "apache-2.0")``. A sequence of + values selects the ``In`` operator. + """ + if attribute not in METADATA_FIELDS: + raise ValueError( + f"{attribute!r} is not a hub content metadata attribute. " + f"Choose one of: {', '.join(sorted(METADATA_FIELDS))}." + ) + if not isinstance(value, str) and operator == "Equals": + operator = "In" + return cls(name=f"Metadata.{METADATA_FIELDS[attribute]}", operator=operator, value=value) + + @classmethod + def keyword(cls, keyword: str) -> "HubContentFilter": + """Match a hub content whose search keywords hold ``keyword``, such as ``@task:summarization``.""" + return cls(name="HubContentSearchKeywords", operator="Equals", value=keyword) + + def to_request(self) -> Dict[str, str]: + """Return the filter in the service form, after a check of the operator and value.""" + if self.operator not in SEARCH_OPERATORS: + raise ValueError( + f"{self.operator!r} is not a Search operator. " + f"Choose one of: {', '.join(sorted(SEARCH_OPERATORS))}." + ) + request = {"Name": self.name, "Operator": self.operator} + if self.operator in ("Exists", "NotExists"): + return request + if self.value is None: + raise ValueError(f"A {self.operator} filter on {self.name} takes a value.") + value = self.value if isinstance(self.value, str) else ",".join(self.value) + if not value or len(value) > MAX_SEARCH_FILTER_VALUE_LENGTH: + raise ValueError( + f"A Search filter value holds 1 to {MAX_SEARCH_FILTER_VALUE_LENGTH} characters, " + f"the value for {self.name} has {len(value)}." + ) + request["Value"] = value + return request + + +class HubContentMixin: # pylint: disable=too-many-public-methods + """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 + answer, and the mixin adds no model field and no state. + """ + + hub_content_search_keywords: Any + hub_content_document: Any + + @property + def tasks(self) -> List[str]: + """Every ``@task:`` keyword, extended with document ``Task`` values the cap dropped.""" + return self._list_field("tasks") + + @property + def license(self) -> Optional[str]: + """The ``@license:`` keyword, or the document ``License``.""" + return self._scalar_field("license") + + @property + def provider(self) -> Optional[str]: + """The ``@provider:`` keyword, or the document ``Provider``.""" + return self._scalar_field("provider") + + @property + def model_size(self) -> Optional[str]: + """The raw ``@model-size:`` value (``"7b"``, ``"1b-10b"``, ``"<1b"``), or the document ``ModelSize``.""" + return self._scalar_field("model_size") + + @property + def context_window(self) -> Optional[str]: + """The raw ``@context-window:`` value, or the document ``ContextWindow``.""" + return self._scalar_field("context_window") + + @property + def framework(self) -> Optional[str]: + """The ``@framework:`` keyword, or the document ``Framework``.""" + return self._scalar_field("framework") + + @property + def languages(self) -> List[str]: + """Every ``@language:`` keyword, extended with document ``Languages``.""" + return self._list_field("languages") + + @property + def data_types(self) -> List[str]: + """Every ``@data-type:`` keyword, extended with document ``DataType``.""" + return self._list_field("data_types") + + @property + def input_modalities(self) -> List[str]: + """Every ``@input-modality:`` keyword, extended with document ``InputModalities``.""" + return self._list_field("input_modalities") + + @property + def output_modalities(self) -> List[str]: + """Every ``@output-modality:`` keyword, extended with document ``OutputModalities``.""" + return self._list_field("output_modalities") + + @property + def model_type(self) -> Optional[str]: + """The ``@model-type:`` keyword, or the first document ``ModelTypes`` value.""" + return self._scalar_field("model_type") + + @property + def capabilities(self) -> List[str]: + """Every ``@capability:`` value plus the document ``Capabilities``, lowercase.""" + values = [v.lower() for v in self._keyword_values().get("capabilities", [])] + for value in _document_values(self._document(), "Capabilities"): + if value.lower() not in values: + values.append(value.lower()) + return values + + @property + def customization_supported(self) -> bool: + """``True`` when the hub tags the model with the ``customization`` capability.""" + return CUSTOMIZATION_CAPABILITY in self.capabilities + + @property + def bedrock_supported(self) -> bool: + """``True`` when the hub tags the model with the ``bedrock_console`` capability.""" + return BEDROCK_CAPABILITY in self.capabilities + + @property + def recipes(self) -> List[HubContentRecipe]: + """Every ``@recipe:`` keyword plus the document ``RecipeCollection``, de-duplicated.""" + recipes = [ + recipe + for recipe in (HubContentRecipe.from_keyword(k) for k in self._keywords()) + if recipe + ] + for raw in self._document().get("RecipeCollection") or []: + if isinstance(raw, Mapping): + recipe = HubContentRecipe.from_document(raw) + if recipe and recipe not in recipes: + recipes.append(recipe) + return recipes + + @classmethod + @_telemetry_emitter(feature=Feature.MODEL_CUSTOMIZATION, func_name="HubContent.search") + def search( + cls, + hub_name: str, + filters: Optional[Sequence[HubContentFilter]] = None, + hub_content_type: str = "Model", + session: Optional[Any] = None, + region: Optional[str] = None, + ) -> Iterator["HubContentMixin"]: + """Search a hub server side and yield each matching ``HubContent``. + + Runs ``Search(Resource="HubContent")`` with an ``And`` expression of ``HubName``, + ``HubContentType``, and ``filters``, sorted by ``HubContentName`` so the pages are + stable, and pages until the service returns no ``NextToken``. Each result is a full + ``HubContent`` bound to ``session`` and ``region``, with its search keywords and + document, so ``content.license`` and ``content.recipes`` work without a second call. + + Args: + hub_name (str): The hub to search, for example ``SageMakerPublicHub``. + filters (Optional[Sequence[HubContentFilter]]): Up to 18 extra filters, for + example ``[HubContentFilter.metadata("license", "apache-2.0")]``. + hub_content_type (str): ``Model``, ``Notebook``, ``ModelReference``, or another + hub content type (default: ``Model``). + session (Optional[boto3.session.Session]): Boto3 session. + region (Optional[str]): Region name. + + Returns: + Iterator[HubContent]: The matching hub content, one page at a time. + + Raises: + ValueError: If more than ``MAX_SEARCH_FILTERS`` filters result, or a filter has an + unknown operator or an out-of-range value. + """ + expression = _search_expression(hub_name, hub_content_type, filters or ()) + client = cls.get_sagemaker_client(session=session, region_name=region) + next_token = None + while True: + request = { + "Resource": "HubContent", + "SearchExpression": expression, + "SortBy": "HubContentName", + "SortOrder": "Ascending", + "MaxResults": SEARCH_PAGE_SIZE, + } + if next_token: + request["NextToken"] = next_token + response = client.search(**request) + for record in response.get("Results", []): + content = cls._from_search_record(record.get("HubContent", {}), hub_name) + if content is not None: + yield content._set_client_context(session=session, region=region) + page_token = response.get("NextToken") + if not page_token or page_token == next_token: + return + next_token = page_token + + @classmethod + def _from_search_record( + cls, record: Mapping[str, Any], hub_name: str + ) -> Optional["HubContentMixin"]: + """Build a ``HubContent`` from one ``SearchRecord.HubContent`` dict.""" + if not record: + return None + from sagemaker.core.utils.code_injection.codec import transform + + transformed = transform(dict(record), "HubContent") + transformed.setdefault("hub_name", hub_name) + known = {name: value for name, value in transformed.items() if name in cls.model_fields} + return cls(**known) + + def fine_tuning_techniques(self) -> List[str]: + """The techniques of the model's FineTuning recipes (``sft``, ``dpo``, ...), sorted.""" + return sorted( + {r.technique for r in self.recipes if r.recipe_type == FINE_TUNING_RECIPE_TYPE} + ) + + def supports_technique(self, technique: str) -> bool: + """``True`` when a FineTuning recipe for ``technique`` exists (case-insensitive).""" + return technique.lower() in self.fine_tuning_techniques() + + @classmethod + @_telemetry_emitter(feature=Feature.MODEL_CUSTOMIZATION, func_name="HubContent.get_all_cached") + def get_all_cached( + cls, + hub_name: str, + hub_content_type: str = "Model", + session: Optional[Any] = None, + region: Optional[str] = None, + ) -> List["HubContentMixin"]: + """Return the contents of a hub from the hub catalog cache. + + The cache is a ``sagemaker.core.hub_catalog.HubCatalogCache``, built like the + JumpStart models cache. In memory it is an ``LRUCache`` of up to 20 hubs with a 6 + hour expiration horizon. On disk it is one JSON file per hub, content type, region, + and account. The first call for a hub lists it and describes every item, 4 at a + time. A call inside the horizon returns the in-memory catalog with no request. A + call after the horizon, or in a new process, lists the hub and describes only the + items whose ``HubContentVersion`` changed. Each returned ``HubContent`` is bound to + ``session`` and ``region`` and carries ``original_creation_time`` from the list + summary. + + Tune the cache with :meth:`set_cache_kwargs`, for example + ``HubContent.set_cache_kwargs(cache_expiration_horizon=timedelta(days=1), max_workers=8)``. + + Args: + hub_name (str): The hub to list, for example ``SageMakerPublicHub``. + hub_content_type (str): ``Model``, ``Notebook``, ``ModelReference``, ``DataSet``, + or ``JsonDoc`` (default: ``Model``). + session (Optional[boto3.session.Session]): Boto3 session. + region (Optional[str]): Region name. + + Returns: + List[HubContent]: Every content of the hub, in the order the service listed it. + """ + key = cls._catalog_key(hub_name, hub_content_type, session, region) + value = HubCatalogAccessor.get_cache().get(key, cls._catalog_source(session, region)) + return cls._catalog_resources(value, session, region) + + @classmethod + @_telemetry_emitter(feature=Feature.MODEL_CUSTOMIZATION, func_name="HubContent.refresh_cache") + def refresh_cache( + cls, + hub_name: str, + hub_content_type: str = "Model", + session: Optional[Any] = None, + region: Optional[str] = None, + ) -> HubCatalogRefresh: + """List the hub again, describe the changed items, and report what the call did. + + Same cache as :meth:`get_all_cached`, with the expiration horizon ignored, for a + caller that wants the counts (listed, reused, described, removed) and the cache + file path. + """ + key = cls._catalog_key(hub_name, hub_content_type, session, region) + source = cls._catalog_source(session, region) + return HubCatalogAccessor.get_cache().refresh(key, source).refresh + + @classmethod + def clear_cache( + cls, + hub_name: Optional[str] = None, + hub_content_type: str = "Model", + session: Optional[Any] = None, + region: Optional[str] = None, + ) -> bool: + """Drop the in-memory hub catalogs, and the cache file of ``hub_name`` when given. + + Returns ``True`` when a cache file was deleted. + """ + cache = HubCatalogAccessor.get_cache() + if hub_name is None: + return cache.clear() + return cache.clear(cls._catalog_key(hub_name, hub_content_type, session, region)) + + @classmethod + def set_cache_kwargs(cls, **cache_kwargs: Any) -> None: + """Rebuild the hub catalog cache with new settings and drop its in-memory entries. + + Accepts the keyword arguments of ``HubCatalogCache``: ``max_cache_items``, + ``cache_expiration_horizon``, ``max_workers``, and ``cache_dir``. The cache files + on disk stay in place. + """ + HubCatalogAccessor.set_cache_kwargs(**cache_kwargs) + + @classmethod + def _catalog_source(cls, session, region) -> HubCatalogSource: + """Return the list and describe calls of one retrieval, bound to the caller session.""" + + def list_hub(key: HubCatalogKey) -> List[Dict[str, Any]]: + client = cls.get_sagemaker_client(session=session, region_name=key.region) + return _list_hub_contents(client, key.hub_name, key.hub_content_type) + + def describe(key: HubCatalogKey, summary: Dict[str, Any]): + return cls._describe_summary(key.hub_name, summary, session, region) + + return HubCatalogSource(list_function=list_hub, describe_function=describe) + + @classmethod + def _catalog_key(cls, hub_name, hub_content_type, session, region) -> HubCatalogKey: + """Build the cache key of a hub from the call arguments and the resolved account.""" + from boto3.session import Session + + resolved_region = region or (session.region_name if session is not None else None) + if not resolved_region: + resolved_region = Session().region_name or "default" + return HubCatalogKey( + hub_name=hub_name, + hub_content_type=hub_content_type, + region=str(resolved_region), + account_id=cls._account_id(session, region), + ) + + @classmethod + def _catalog_resources(cls, value: HubCatalogValue, session, region) -> List["HubContentMixin"]: + """Build bound ``HubContent`` resources from a catalog value, in list order.""" + resources = [] + for summary, entry in ordered_entries(value): + resource = cls(**entry["resource"]) + if summary.get("OriginalCreationTime") is not None: + resource.original_creation_time = summary["OriginalCreationTime"] + resources.append(resource._set_client_context(session=session, region=region)) + return resources + + @classmethod + def _describe_summary(cls, hub_name, summary, session, region): + """Build a ``HubContent`` from one list summary, bind it, and describe it.""" + from sagemaker.core.utils.code_injection.codec import transform + + init_data = transform(dict(summary), "HubContentInfo") + init_data.setdefault("hub_name", hub_name) + known = {name: value for name, value in init_data.items() if name in cls.model_fields} + resource = cls(**known) + resource._set_client_context(session=session, region=region) + resource.refresh() + return resource + + @classmethod + def _account_id(cls, session, region) -> str: + """Return the account of the caller, so two accounts never share a cache file. + + 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. + """ + from boto3.session import Session + + from sagemaker.core.common_utils import sts_regional_endpoint + + boto_session = session or Session(region_name=region) + resolved_region = region or boto_session.region_name + credentials = boto_session.get_credentials() + cache_key = credentials.get_frozen_credentials() if credentials is not None else None + with _ACCOUNT_IDS_LOCK: + account_id = _ACCOUNT_IDS.get(cache_key) if cache_key is not None else None + if account_id is None: + sts_client = boto_session.client( + "sts", + region_name=resolved_region, + endpoint_url=sts_regional_endpoint(resolved_region), + ) + account_id = sts_client.get_caller_identity()["Account"] + if cache_key is not None: + with _ACCOUNT_IDS_LOCK: + _ACCOUNT_IDS[cache_key] = account_id + return account_id + + def _keywords(self) -> List[str]: + keywords = _assigned(self.hub_content_search_keywords) or [] + return [k for k in keywords if isinstance(k, str)] + + def _keyword_values(self) -> Dict[str, List[str]]: + return _keyword_values(self._keywords()) + + def _document(self) -> Dict[str, Any]: + return _parse_document(_assigned(self.hub_content_document)) + + def _scalar_field(self, attribute: str) -> Optional[str]: + keyword_hits = self._keyword_values().get(attribute, []) + if keyword_hits: + return keyword_hits[0] + document_hits = _document_values(self._document(), DOCUMENT_FIELDS[attribute]) + return document_hits[0] if document_hits else None + + def _list_field(self, attribute: str) -> List[str]: + merged = list(self._keyword_values().get(attribute, [])) + seen = {v.lower() for v in merged} + for value in _document_values(self._document(), DOCUMENT_FIELDS[attribute]): + if value.lower() not in seen: + merged.append(value) + seen.add(value.lower()) + return merged + + +def _assigned(value: Any) -> Any: + """Return ``value`` unless it is the generated ``Unassigned()`` sentinel or ``None``.""" + if value is None or type(value).__name__ == "Unassigned": + return None + return value + + +def _list_hub_contents(client, hub_name: str, hub_content_type: str) -> List[Dict[str, Any]]: + """Return every ``ListHubContents`` summary of the hub, in service order.""" + summaries: List[Dict[str, Any]] = [] + next_token = None + while True: + request = { + "HubName": hub_name, + "HubContentType": hub_content_type, + "MaxResults": LIST_PAGE_SIZE, + } + if next_token: + request["NextToken"] = next_token + response = client.list_hub_contents(**request) + summaries.extend(response.get("HubContentSummaries", [])) + page_token = response.get("NextToken") + if not page_token or page_token == next_token: + return summaries + next_token = page_token + + +def _search_expression( + hub_name: str, hub_content_type: str, filters: Sequence[HubContentFilter] +) -> Dict[str, Any]: + """Build the ``And`` expression of a hub content search, inside the service limits.""" + request_filters = [ + HubContentFilter("HubName", "Equals", hub_name).to_request(), + HubContentFilter("HubContentType", "Equals", hub_content_type).to_request(), + ] + request_filters.extend(search_filter.to_request() for search_filter in filters) + if len(request_filters) > MAX_SEARCH_FILTERS: + raise ValueError( + f"A SearchExpression holds at most {MAX_SEARCH_FILTERS} filters, HubName and " + f"HubContentType included. Got {len(request_filters)}." + ) + return {"Operator": "And", "Filters": request_filters} + + +def _parse_document(document: Any) -> Dict[str, Any]: + if isinstance(document, Mapping): + return dict(document) + if isinstance(document, str) and document: + try: + parsed = json.loads(document) + except ValueError: + return {} + return parsed if isinstance(parsed, dict) else {} + return {} + + +def _keyword_values(keywords: Iterable[str]) -> Dict[str, List[str]]: + """Group keyword values by attribute, in hub order, keeping the hub's spelling.""" + values: Dict[str, List[str]] = {} + for keyword in keywords: + lowered = keyword.lower() + if lowered.startswith(CAPABILITY_PREFIX): + values.setdefault("capabilities", []).append(keyword[len(CAPABILITY_PREFIX) :]) + continue + for attribute, prefix in KEYWORD_PREFIXES.items(): + if lowered.startswith(prefix): + value = keyword[len(prefix) :] + if value: + values.setdefault(attribute, []).append(value) + break + return values + + +def _document_values(document: Mapping[str, Any], field_name: str) -> List[str]: + value = document.get(field_name) + items = value if isinstance(value, list) else [value] + return [item for item in items if isinstance(item, str) and item] diff --git a/sagemaker-core/src/sagemaker/core/instance_types.py b/sagemaker-core/src/sagemaker/core/instance_types.py index 2a76af28a8..f0b43636fb 100644 --- a/sagemaker-core/src/sagemaker/core/instance_types.py +++ b/sagemaker-core/src/sagemaker/core/instance_types.py @@ -10,12 +10,22 @@ # 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. -"""Accessors to retrieve instance types.""" +"""Accessors to retrieve instance types. + +``retrieve`` and ``retrieve_default`` read the instance types a JumpStart model supports. +``list_instance_types``, ``describe_instance_type``, and ``describe_instance_types`` read the +SageMaker hosting instance-type names from the installed botocore model and their hardware +from EC2 ``DescribeInstanceTypes``. +""" from __future__ import absolute_import import logging -from typing import List, Optional +from dataclasses import dataclass +from typing import Any, Dict, List, Optional, Sequence + +import boto3 +from botocore.exceptions import ClientError from sagemaker.core.jumpstart import utils as jumpstart_utils from sagemaker.core.jumpstart import artifacts @@ -25,6 +35,33 @@ logger = logging.getLogger(__name__) +SAGEMAKER_INSTANCE_TYPE_PREFIX = "ml." +PRODUCTION_VARIANT_INSTANCE_TYPE_SHAPE = "ProductionVariantInstanceType" +# ``RequestInstanceTypeList`` in the EC2 service model allows 100 names per request. +DESCRIBE_INSTANCE_TYPES_BATCH_SIZE = 100 +INVALID_INSTANCE_TYPE_ERROR_CODE = "InvalidInstanceType" + + +@dataclass(frozen=True) +class InstanceSpec: + """Hardware of one SageMaker instance type as EC2 ``DescribeInstanceTypes`` reports it. + + Attributes: + name: The SageMaker instance type, for example ``ml.g5.xlarge``. + vcpus: The default vCPU count, or ``None`` when EC2 does not know the type. + memory_mib: The instance memory in MiB, or ``None`` when EC2 does not know the type. + gpu_count: The number of GPUs. ``0`` for a CPU instance and for an unknown type. + gpu_memory_mib: The GPU memory in MiB summed over all GPUs, or ``None`` without GPUs. + gpu_name: The GPU model name, for example ``A10G``, or ``None`` without GPUs. + """ + + name: str + vcpus: Optional[int] + memory_mib: Optional[int] + gpu_count: int + gpu_memory_mib: Optional[int] + gpu_name: Optional[str] + def retrieve_default( region: Optional[str] = None, @@ -162,3 +199,178 @@ def retrieve( sagemaker_session=sagemaker_session, training_instance_type=training_instance_type, ) + + +def _boto_session(session: Optional[boto3.Session]) -> boto3.Session: + """Returns the given session, or a new default session.""" + if session is not None: + return session + return boto3.Session() + + +def _ec2_instance_type(instance_type: str) -> str: + """Returns the EC2 name of a SageMaker instance type (``ml.g5.xlarge`` to ``g5.xlarge``).""" + if instance_type.startswith(SAGEMAKER_INSTANCE_TYPE_PREFIX): + return instance_type[len(SAGEMAKER_INSTANCE_TYPE_PREFIX) :] + return instance_type + + +def _unknown_spec(instance_type: str) -> InstanceSpec: + """Returns the spec of an instance type that EC2 does not know, and logs a warning.""" + logger.warning( + "EC2 DescribeInstanceTypes does not know instance type %s. " + "Its InstanceSpec has gpu_count=0 and no other hardware.", + instance_type, + ) + return InstanceSpec( + name=instance_type, + vcpus=None, + memory_mib=None, + gpu_count=0, + gpu_memory_mib=None, + gpu_name=None, + ) + + +def _spec_from_ec2(instance_type: str, instance_info: Dict[str, Any]) -> InstanceSpec: + """Builds an ``InstanceSpec`` from one EC2 ``InstanceTypeInfo`` element.""" + gpus = instance_info.get("GpuInfo", {}).get("Gpus", []) + gpu_count = sum(gpu.get("Count", 0) for gpu in gpus) + first_gpu = gpus[0] if gpus else {} + memory_per_gpu_mib = first_gpu.get("MemoryInfo", {}).get("SizeInMiB") + gpu_memory_mib = None + if gpu_count > 0 and memory_per_gpu_mib is not None: + gpu_memory_mib = memory_per_gpu_mib * gpu_count + return InstanceSpec( + name=instance_type, + vcpus=instance_info.get("VCpuInfo", {}).get("DefaultVCpus"), + memory_mib=instance_info.get("MemoryInfo", {}).get("SizeInMiB"), + gpu_count=gpu_count, + gpu_memory_mib=gpu_memory_mib, + gpu_name=first_gpu.get("Name"), + ) + + +def _describe_ec2_instance_types( + ec2_client: Any, instance_types: List[str] +) -> Dict[str, Dict[str, Any]]: + """Calls EC2 ``DescribeInstanceTypes`` for the names and reads every page of the response.""" + request: Dict[str, Any] = { + "InstanceTypes": [_ec2_instance_type(instance_type) for instance_type in instance_types] + } + described: Dict[str, Dict[str, Any]] = {} + while True: + response = ec2_client.describe_instance_types(**request) + for instance_info in response.get("InstanceTypes", []): + described[instance_info["InstanceType"]] = instance_info + next_token = response.get("NextToken") + # A real token is a non-empty string. The type check also stops a client stub + # whose ``get`` returns a truthy Mock from looping forever. + if not isinstance(next_token, str) or not next_token: + return described + request["NextToken"] = next_token + + +def _describe_known(ec2_client: Any, instance_types: List[str]) -> Dict[str, Dict[str, Any]]: + """Returns the EC2 ``InstanceTypeInfo`` of each name that EC2 knows, keyed by EC2 name. + + EC2 answers ``InvalidInstanceType`` for the whole request when one name does not exist. + The function then asks for each name alone, so one unknown name hides no other. + """ + try: + return _describe_ec2_instance_types(ec2_client, instance_types) + except ClientError as error: + if error.response.get("Error", {}).get("Code") != INVALID_INSTANCE_TYPE_ERROR_CODE: + raise + if len(instance_types) == 1: + return {} + described: Dict[str, Dict[str, Any]] = {} + for instance_type in instance_types: + described.update(_describe_known(ec2_client, [instance_type])) + return described + + +def _describe_batch(ec2_client: Any, instance_types: List[str]) -> Dict[str, InstanceSpec]: + """Describes one request of at most 100 instance types and maps each name to a spec.""" + described = _describe_known(ec2_client, instance_types) + specs: Dict[str, InstanceSpec] = {} + for instance_type in instance_types: + instance_info = described.get(_ec2_instance_type(instance_type)) + if instance_info is None: + specs[instance_type] = _unknown_spec(instance_type) + else: + specs[instance_type] = _spec_from_ec2(instance_type, instance_info) + return specs + + +def list_instance_types( + region: Optional[str] = None, session: Optional[boto3.Session] = None +) -> List[str]: + """Lists the SageMaker hosting instance types that the loaded service model knows. + + The names come from the ``ProductionVariantInstanceType`` enum of the SageMaker service + model that the session loads, so the call makes no request. The list holds the + service-wide set and keeps types the region does not offer. + + Args: + region (str): The region of the SageMaker client. Defaults to the session region. + session (boto3.Session): The session that creates the client. Defaults to a new session. + + Returns: + list[str]: The instance type names in model order, for example ``ml.t2.medium`` first. + + Raises: + botocore.exceptions.NoRegionError: If neither ``region`` nor the session has a region. + """ + sagemaker_client = _boto_session(session).client("sagemaker", region_name=region) + shape = sagemaker_client.meta.service_model.shape_for(PRODUCTION_VARIANT_INSTANCE_TYPE_SHAPE) + return list(shape.enum) + + +def describe_instance_types( + instance_types: Sequence[str], session: Optional[boto3.Session] = None +) -> Dict[str, InstanceSpec]: + """Describes the hardware of SageMaker instance types through EC2 ``DescribeInstanceTypes``. + + The function strips the ``ml.`` prefix before the EC2 call and sends 100 names per request, + the EC2 limit. EC2 rejects a whole request when one name does not exist, so the function + then describes each name of that request alone. An instance type that EC2 does not know + (a retired type, or a type the region does not offer) gets an ``InstanceSpec`` with + ``gpu_count=0`` and ``None`` in every other field. The function logs a warning for it. + + Args: + instance_types (Sequence[str]): The SageMaker instance types, for example + ``["ml.g5.xlarge", "ml.m5.large"]``. + session (boto3.Session): The session that creates the EC2 client. Defaults to a new + session. + + Returns: + dict[str, InstanceSpec]: One ``InstanceSpec`` per requested name, keyed by that name. + + Raises: + botocore.exceptions.ClientError: If EC2 rejects a request for a reason other than an + unknown instance type, for example missing ``ec2:DescribeInstanceTypes`` permission. + """ + ec2_client = _boto_session(session).client("ec2") + specs: Dict[str, InstanceSpec] = {} + for start in range(0, len(instance_types), DESCRIBE_INSTANCE_TYPES_BATCH_SIZE): + batch = list(instance_types[start : start + DESCRIBE_INSTANCE_TYPES_BATCH_SIZE]) + specs.update(_describe_batch(ec2_client, batch)) + return specs + + +def describe_instance_type( + instance_type: str, session: Optional[boto3.Session] = None +) -> InstanceSpec: + """Describes the hardware of one SageMaker instance type through EC2. + + Args: + instance_type (str): The SageMaker instance type, for example ``ml.g5.xlarge``. + session (boto3.Session): The session that creates the EC2 client. Defaults to a new + session. + + Returns: + InstanceSpec: The hardware. An instance type that EC2 does not know gets + ``gpu_count=0`` and ``None`` in every other field. See :func:`describe_instance_types`. + """ + return describe_instance_types([instance_type], session=session)[instance_type] diff --git a/sagemaker-core/src/sagemaker/core/model_package_group.py b/sagemaker-core/src/sagemaker/core/model_package_group.py new file mode 100644 index 0000000000..ca588a10ed --- /dev/null +++ b/sagemaker-core/src/sagemaker/core/model_package_group.py @@ -0,0 +1,129 @@ +# 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. +"""Hand-written behavior of the generated ``sagemaker.core.resources.ModelPackageGroup`` class. + +``DescribeModelPackageGroup`` reports a missing group as a ``ValidationException`` whose +message ends in ``does not exist``, not as a ``ResourceNotFound`` error. So every caller +that wants "use the group, or create it" parsed ``ClientError`` codes and messages by hand. +``ModelPackageGroupMixin.get_or_create`` does that once. It returns the existing group, or +creates the group from a ``ModelPackageGroupSpec`` when the describe call says the group is +absent. Every other error propagates. +""" + +from __future__ import absolute_import + +import logging +from dataclasses import dataclass +from typing import TYPE_CHECKING, Any, Dict, List, Optional + +from boto3.session import Session +from botocore.exceptions import ClientError + +from sagemaker.core.shapes import ManagedConfiguration, Tag + +if TYPE_CHECKING: + from sagemaker.core.resources import ModelPackageGroup + +logger = logging.getLogger(__name__) + +NOT_FOUND_ERROR_CODE_PREFIX = "ResourceNotFound" +VALIDATION_ERROR_CODE = "ValidationException" +NOT_FOUND_MESSAGE_MARKERS = ("does not exist", "not found", "could not find") + + +@dataclass(frozen=True) +class ModelPackageGroupSpec: + """The create-only fields of a model package group. + + ``get_or_create`` passes the set fields to ``ModelPackageGroup.create`` when the group is + absent, and ignores the spec when the group exists. + """ + + model_package_group_description: Optional[str] = None + tags: Optional[List[Tag]] = None + managed_configuration: Optional[ManagedConfiguration] = None + + def create_kwargs(self) -> Dict[str, Any]: + """Return the set fields as keyword arguments of ``ModelPackageGroup.create``.""" + fields = { + "model_package_group_description": self.model_package_group_description, + "tags": self.tags, + "managed_configuration": self.managed_configuration, + } + return {name: value for name, value in fields.items() if value is not None} + + +def is_not_found_error(error: ClientError) -> bool: + """Return ``True`` when a describe ``ClientError`` means the resource does not exist. + + ``DescribeModelPackageGroup`` returns ``ValidationException`` with a message that ends in + ``does not exist`` for a missing group. A code that starts with ``ResourceNotFound`` means + the same on other APIs. + """ + error_details = error.response.get("Error", {}) + code = error_details.get("Code", "") + if code.startswith(NOT_FOUND_ERROR_CODE_PREFIX): + return True + message = error_details.get("Message", "").lower() + return code == VALIDATION_ERROR_CODE and any( + marker in message for marker in NOT_FOUND_MESSAGE_MARKERS + ) + + +class ModelPackageGroupMixin: + """``get_or_create`` for the generated ``sagemaker.core.resources.ModelPackageGroup``. + + Mixed into the generated class. The mixin adds no model field and no state. It calls the + generated ``get`` and ``create`` classmethods, so the returned group keeps the session and + region it was loaded with. + """ + + @classmethod + def get_or_create( + cls, + model_package_group_name: str, + spec: Optional[ModelPackageGroupSpec] = None, + session: Optional[Session] = None, + region: Optional[str] = None, + ) -> "ModelPackageGroup": + """Return the model package group, or create it when the describe call says it is absent. + + Args: + model_package_group_name: The name of the model group. + spec: The description, tags, and managed configuration of a group this call + creates. Ignored when the group exists. + session: Boto3 session. + region: Region name. + + Returns: + The existing or the created ModelPackageGroup resource. + + Raises: + botocore.exceptions.ClientError: Every describe or create error other than a + missing group. + """ + try: + return cls.get( + model_package_group_name=model_package_group_name, session=session, region=region + ) + except ClientError as error: + if not is_not_found_error(error): + raise + logger.info("Model package group %s does not exist. Creating it.", model_package_group_name) + create_kwargs = (spec or ModelPackageGroupSpec()).create_kwargs() + return cls.create( + model_package_group_name=model_package_group_name, + session=session, + region=region, + **create_kwargs, + ) diff --git a/sagemaker-core/src/sagemaker/core/processing.py b/sagemaker-core/src/sagemaker/core/processing.py index 912c0d1075..44fe55a62f 100644 --- a/sagemaker-core/src/sagemaker/core/processing.py +++ b/sagemaker-core/src/sagemaker/core/processing.py @@ -718,7 +718,20 @@ def submit(request): transformed = transform(serialized_request, "CreateProcessingJobRequest") # Remove tags from transformed dict as ProcessingJob resource doesn't accept it transformed.pop("tags", None) - return ProcessingJob(**transformed) + job = ProcessingJob(**transformed) + if isinstance(self.sagemaker_session, LocalSession): + # A local job exists only in the session's LocalSagemakerClient, so a + # describe over AWS would fail. run() polls it through that client. + return job + # The create request carries no ARN and no status, and a session-less + # ProcessingJob builds its client from the process default session, so + # latest_job.refresh() polled the default region (#5796). Bind the session + # the job was submitted with, then describe once to fill the ARN and status. + job._set_client_context( + session=self.sagemaker_session.boto_session, + region=self.sagemaker_session.boto_region_name, + ) + return job.refresh() def _get_process_args(self, inputs, outputs, experiment_config): """Gets a dict of arguments for a new Amazon SageMaker processing job.""" diff --git a/sagemaker-core/src/sagemaker/core/resources.py b/sagemaker-core/src/sagemaker/core/resources.py index c0780dc553..7627503798 100644 --- a/sagemaker-core/src/sagemaker/core/resources.py +++ b/sagemaker-core/src/sagemaker/core/resources.py @@ -45,7 +45,9 @@ from typing import ClassVar from sagemaker.core.serializers.base import BaseSerializer from sagemaker.core.deserializers.base import BaseDeserializer - +from sagemaker.core.endpoint import EndpointMixin +from sagemaker.core.hub_content import HubContentMixin +from sagemaker.core.model_package_group import ModelPackageGroupMixin logger = get_textual_rich_logger(__name__) @@ -478,6 +480,7 @@ def wait( self, poll: int = 5, timeout: Optional[int] = None, + raise_on_stopped: bool = False, ) -> None: """ Wait for a AIBenchmarkJob resource. @@ -485,10 +488,11 @@ def wait( Parameters: poll: The number of seconds to wait between each poll. timeout: The maximum number of seconds to wait before timing out. + raise_on_stopped: Whether to raise FailedStatusError when the resource stops instead of completing. Raises: TimeoutExceededError: If the resource does not reach a terminal state before the timeout. - FailedStatusError: If the resource reaches a failed state. + FailedStatusError: If the resource reaches a failed state, or a stopped state when raise_on_stopped is True. WaiterError: Raised when an error occurs while waiting. """ @@ -526,6 +530,13 @@ def wait( reason=self.failure_reason, ) + if raise_on_stopped and "stopped" in current_status.lower(): + raise FailedStatusError( + resource_type="AIBenchmarkJob", + status=current_status, + reason=self.failure_reason, + ) + return if timeout is not None and time.time() - start_time >= timeout: @@ -553,8 +564,6 @@ def get_all( Get all AIBenchmarkJob resources Parameters: - max_results: The maximum number of benchmark jobs to return in the response. - next_token: If the previous call to ListAIBenchmarkJobs didn't return the full set of jobs, the call returns a token for getting the next set. name_contains: A string in the job name. This filter returns only jobs whose name contains the specified string. status_equals: A filter that returns only benchmark jobs with the specified status. creation_time_after: A filter that returns only jobs created after the specified time. @@ -938,6 +947,7 @@ def wait( self, poll: int = 5, timeout: Optional[int] = None, + raise_on_stopped: bool = False, ) -> None: """ Wait for a AIRecommendationJob resource. @@ -945,10 +955,11 @@ def wait( Parameters: poll: The number of seconds to wait between each poll. timeout: The maximum number of seconds to wait before timing out. + raise_on_stopped: Whether to raise FailedStatusError when the resource stops instead of completing. Raises: TimeoutExceededError: If the resource does not reach a terminal state before the timeout. - FailedStatusError: If the resource reaches a failed state. + FailedStatusError: If the resource reaches a failed state, or a stopped state when raise_on_stopped is True. WaiterError: Raised when an error occurs while waiting. """ @@ -986,6 +997,13 @@ def wait( reason=self.failure_reason, ) + if raise_on_stopped and "stopped" in current_status.lower(): + raise FailedStatusError( + resource_type="AIRecommendationJob", + status=current_status, + reason=self.failure_reason, + ) + return if timeout is not None and time.time() - start_time >= timeout: @@ -1013,8 +1031,6 @@ def get_all( Get all AIRecommendationJob resources Parameters: - max_results: The maximum number of recommendation jobs to return in the response. - next_token: If the previous call to ListAIRecommendationJobs didn't return the full set of jobs, the call returns a token for getting the next set. name_contains: A string in the job name. This filter returns only jobs whose name contains the specified string. status_equals: A filter that returns only recommendation jobs with the specified status. creation_time_after: A filter that returns only jobs created after the specified time. @@ -1315,8 +1331,6 @@ def get_all( Get all AIWorkloadConfig resources Parameters: - max_results: The maximum number of AI workload configurations to return in the response. - next_token: If the previous call to ListAIWorkloadConfigs didn't return the full set of configurations, the call returns a token for getting the next set of configurations. name_contains: A string in the configuration name. This filter returns only configurations whose name contains the specified string. creation_time_after: A filter that returns only configurations created after the specified time. creation_time_before: A filter that returns only configurations created before the specified time. @@ -1697,8 +1711,6 @@ def get_all( created_before: A filter that returns only actions created on or before the specified time. sort_by: The property used to sort results. The default value is CreationTime. sort_order: The sort order. The default value is Descending. - next_token: If the previous call to ListActions didn't return the full set of actions, the call returns a token for getting the next set of actions. - max_results: The maximum number of actions to return in the response. The default value is 10. session: Boto3 session. region: Region name. @@ -2147,9 +2159,7 @@ def get_all( Parameters: creation_time_after: A filter that returns only algorithms created after the specified time (timestamp). creation_time_before: A filter that returns only algorithms created before the specified time (timestamp). - max_results: The maximum number of algorithms to return in the response. name_contains: A string in the algorithm name. This filter returns only algorithms whose name contains the specified string. - next_token: If the response to a previous ListAlgorithms request was truncated, the response includes a NextToken. To retrieve the next set of algorithms, use the token in the next request. sort_by: The parameter by which to sort the results. The default is CreationTime. sort_order: The sort order for the results. The default is Ascending. session: Boto3 session. @@ -2620,8 +2630,6 @@ def get_all( Get all App resources Parameters: - next_token: If the previous response was truncated, you will receive this token. Use it in your next request to receive the next set of results. - max_results: This parameter defines the maximum number of results that can be return in a single response. The MaxResults parameter is an upper bound, not a target. If there are more results available than the value specified, a NextToken is provided in the response. The NextToken indicates that the user should get the next set of results by providing this token as a part of a subsequent call. The default value for MaxResults is 10. sort_order: The sort order for the results. The default is Ascending. sort_by: The parameter by which to sort the results. The default is CreationTime. domain_id_equals: A parameter to search for the domain ID. @@ -2970,8 +2978,6 @@ def get_all( Get all AppImageConfig resources Parameters: - max_results: The total number of items to return in the response. If the total number of items available is more than the value specified, a NextToken is provided in the response. To resume pagination, provide the NextToken value in the as part of a subsequent call. The default value is 10. - next_token: If the previous call to ListImages didn't return the full set of AppImageConfigs, the call returns a token for getting the next set of AppImageConfigs. name_contains: A filter that returns only AppImageConfigs whose name contains the specified string. creation_time_before: A filter that returns only AppImageConfigs created on or before the specified time. creation_time_after: A filter that returns only AppImageConfigs created on or after the specified time. @@ -3345,8 +3351,6 @@ def get_all( created_before: A filter that returns only artifacts created on or before the specified time. sort_by: The property used to sort results. The default value is CreationTime. sort_order: The sort order. The default value is Descending. - next_token: If the previous call to ListArtifacts didn't return the full set of artifacts, the call returns a token for getting the next set of artifacts. - max_results: The maximum number of artifacts to return in the response. The default value is 10. session: Boto3 session. region: Region name. @@ -3501,8 +3505,6 @@ def get_all( created_before: A filter that returns only associations created on or before the specified time. sort_by: The property used to sort results. The default value is CreationTime. sort_order: The sort order. The default value is Descending. - next_token: If the previous call to ListAssociations didn't return the full set of associations, the call returns a token for getting the next set of associations. - max_results: The maximum number of associations to return in the response. The default value is 10. session: Boto3 session. region: Region name. @@ -3916,6 +3918,7 @@ def wait( self, poll: int = 5, timeout: Optional[int] = None, + raise_on_stopped: bool = False, ) -> None: """ Wait for a AutoMLJob resource. @@ -3923,10 +3926,11 @@ def wait( Parameters: poll: The number of seconds to wait between each poll. timeout: The maximum number of seconds to wait before timing out. + raise_on_stopped: Whether to raise FailedStatusError when the resource stops instead of completing. Raises: TimeoutExceededError: If the resource does not reach a terminal state before the timeout. - FailedStatusError: If the resource reaches a failed state. + FailedStatusError: If the resource reaches a failed state, or a stopped state when raise_on_stopped is True. WaiterError: Raised when an error occurs while waiting. """ @@ -3964,6 +3968,13 @@ def wait( reason=self.failure_reason, ) + if raise_on_stopped and "stopped" in current_status.lower(): + raise FailedStatusError( + resource_type="AutoMLJob", + status=current_status, + reason=self.failure_reason, + ) + return if timeout is not None and time.time() - start_time >= timeout: @@ -4001,8 +4012,6 @@ def get_all( status_equals: Request a list of jobs, using a filter for status. sort_order: The sort order for the results. The default is Descending. sort_by: The parameter by which to sort the results. The default is Name. - max_results: Request a list of jobs up to a specified limit. - next_token: If the previous response was truncated, you receive this token. Use it in your next request to receive the next set of results. session: Boto3 session. region: Region name. @@ -4069,8 +4078,6 @@ def get_all_candidates( candidate_name_equals: List the candidates for the job and filter by candidate name. sort_order: The sort order for the results. The default is Ascending. sort_by: The parameter by which to sort the results. The default is Descending. - max_results: List the job's candidates up to a specified limit. - next_token: If the previous response was truncated, you receive this token. Use it in your next request to receive the next set of results. session: Boto3 session. region: Region name. @@ -4406,6 +4413,7 @@ def wait( self, poll: int = 5, timeout: Optional[int] = None, + raise_on_stopped: bool = False, ) -> None: """ Wait for a AutoMLJobV2 resource. @@ -4413,10 +4421,11 @@ def wait( Parameters: poll: The number of seconds to wait between each poll. timeout: The maximum number of seconds to wait before timing out. + raise_on_stopped: Whether to raise FailedStatusError when the resource stops instead of completing. Raises: TimeoutExceededError: If the resource does not reach a terminal state before the timeout. - FailedStatusError: If the resource reaches a failed state. + FailedStatusError: If the resource reaches a failed state, or a stopped state when raise_on_stopped is True. WaiterError: Raised when an error occurs while waiting. """ @@ -4454,6 +4463,13 @@ def wait( reason=self.failure_reason, ) + if raise_on_stopped and "stopped" in current_status.lower(): + raise FailedStatusError( + resource_type="AutoMLJobV2", + status=current_status, + reason=self.failure_reason, + ) + return if timeout is not None and time.time() - start_time >= timeout: @@ -4977,9 +4993,7 @@ def get_all( Parameters: creation_time_after: Set a start time for the time range during which you want to list SageMaker HyperPod clusters. Timestamps are formatted according to the ISO 8601 standard. Acceptable formats include: YYYY-MM-DDThh:mm:ss.sssTZD (UTC), for example, 2014-10-01T20:30:00.000Z YYYY-MM-DDThh:mm:ss.sssTZD (with offset), for example, 2014-10-01T12:30:00.000-08:00 YYYY-MM-DD, for example, 2014-10-01 Unix time in seconds, for example, 1412195400. This is also referred to as Unix Epoch time and represents the number of seconds since midnight, January 1, 1970 UTC. For more information about the timestamp format, see Timestamp in the Amazon Web Services Command Line Interface User Guide. creation_time_before: Set an end time for the time range during which you want to list SageMaker HyperPod clusters. A filter that returns nodes in a SageMaker HyperPod cluster created before the specified time. The acceptable formats are the same as the timestamp formats for CreationTimeAfter. For more information about the timestamp format, see Timestamp in the Amazon Web Services Command Line Interface User Guide. - max_results: Specifies the maximum number of clusters to evaluate for the operation (not necessarily the number of matching items). After SageMaker processes the number of clusters up to MaxResults, it stops the operation and returns the matching clusters up to that point. If all the matching clusters are desired, SageMaker will go through all the clusters until NextToken is empty. name_contains: Set the maximum number of instances to print in the list. - next_token: Set the next token to retrieve the list of SageMaker HyperPod clusters. sort_by: The field by which to sort results. The default value is CREATION_TIME. sort_order: The sort order for results. The default value is Ascending. training_plan_arn: The Amazon Resource Name (ARN); of the training plan to filter clusters by. For more information about reserving GPU capacity for your SageMaker HyperPod clusters using Amazon SageMaker Training Plan, see CreateTrainingPlan . @@ -5099,8 +5113,6 @@ def get_all_nodes( creation_time_after: A filter that returns nodes in a SageMaker HyperPod cluster created after the specified time. Timestamps are formatted according to the ISO 8601 standard. Acceptable formats include: YYYY-MM-DDThh:mm:ss.sssTZD (UTC), for example, 2014-10-01T20:30:00.000Z YYYY-MM-DDThh:mm:ss.sssTZD (with offset), for example, 2014-10-01T12:30:00.000-08:00 YYYY-MM-DD, for example, 2014-10-01 Unix time in seconds, for example, 1412195400. This is also referred to as Unix Epoch time and represents the number of seconds since midnight, January 1, 1970 UTC. For more information about the timestamp format, see Timestamp in the Amazon Web Services Command Line Interface User Guide. creation_time_before: A filter that returns nodes in a SageMaker HyperPod cluster created before the specified time. The acceptable formats are the same as the timestamp formats for CreationTimeAfter. For more information about the timestamp format, see Timestamp in the Amazon Web Services Command Line Interface User Guide. instance_group_name_contains: A filter that returns the instance groups whose name contain a specified string. - max_results: The maximum number of nodes to return in the response. - next_token: If the result of the previous ListClusterNodes request was truncated, the response includes a NextToken. To retrieve the next set of cluster nodes, use the token in the next request. sort_by: The field by which to sort results. The default value is CREATION_TIME. sort_order: The sort order for results. The default value is Ascending. include_node_logical_ids: Specifies whether to include nodes that are still being provisioned in the response. When set to true, the response includes all nodes regardless of their provisioning status. When set to False (default), only nodes with assigned InstanceIds are returned. @@ -5729,8 +5741,6 @@ def get_all( status: Filter for status. sort_by: Filter for sorting the list by a given value. For example, sort by name, creation time, or status. sort_order: The order of the list. By default, listed in Descending order according to by SortBy. To change the list order, you can specify SortOrder to be Ascending. - next_token: If the previous response was truncated, you will receive this token. Use it in your next request to receive the next set of results. - max_results: The maximum number of cluster policies to list. session: Boto3 session. region: Region name. @@ -6062,9 +6072,7 @@ def get_all( creation_time_before: A filter that returns only Git repositories that were created before the specified time. last_modified_time_after: A filter that returns only Git repositories that were last modified after the specified time. last_modified_time_before: A filter that returns only Git repositories that were last modified before the specified time. - max_results: The maximum number of Git repositories to return in the response. name_contains: A string in the Git repositories name. This filter returns only repositories whose name contains the specified string. - next_token: If the result of a ListCodeRepositoriesOutput request was truncated, the response includes a NextToken. To get the next set of Git repositories, use the token in the next request. sort_by: The field to sort results by. The default is Name. sort_order: The sort order for results. The default is Ascending. session: Boto3 session. @@ -6440,6 +6448,7 @@ def wait( self, poll: int = 5, timeout: Optional[int] = None, + raise_on_stopped: bool = False, ) -> None: """ Wait for a CompilationJob resource. @@ -6447,10 +6456,11 @@ def wait( Parameters: poll: The number of seconds to wait between each poll. timeout: The maximum number of seconds to wait before timing out. + raise_on_stopped: Whether to raise FailedStatusError when the resource stops instead of completing. Raises: TimeoutExceededError: If the resource does not reach a terminal state before the timeout. - FailedStatusError: If the resource reaches a failed state. + FailedStatusError: If the resource reaches a failed state, or a stopped state when raise_on_stopped is True. WaiterError: Raised when an error occurs while waiting. """ @@ -6488,6 +6498,13 @@ def wait( reason=self.failure_reason, ) + if raise_on_stopped and "stopped" in current_status.lower(): + raise FailedStatusError( + resource_type="CompilationJob", + status=current_status, + reason=self.failure_reason, + ) + return if timeout is not None and time.time() - start_time >= timeout: @@ -6517,8 +6534,6 @@ def get_all( Get all CompilationJob resources Parameters: - next_token: If the result of the previous ListCompilationJobs request was truncated, the response includes a NextToken. To retrieve the next set of model compilation jobs, use the token in the next request. - max_results: The maximum number of model compilation jobs to return in the response. creation_time_after: A filter that returns the model compilation jobs that were created after a specified time. creation_time_before: A filter that returns the model compilation jobs that were created before a specified time. last_modified_time_after: A filter that returns the model compilation jobs that were modified after a specified time. @@ -7054,8 +7069,6 @@ def get_all( cluster_arn: Filter for ARN of the cluster. sort_by: Filter for sorting the list by a given value. For example, sort by name, creation time, or status. sort_order: The order of the list. By default, listed in Descending order according to by SortBy. To change the list order, you can specify SortOrder to be Ascending. - next_token: If the previous response was truncated, you will receive this token. Use it in your next request to receive the next set of results. - max_results: The maximum number of compute allocation definitions to list. session: Boto3 session. region: Region name. @@ -7421,8 +7434,6 @@ def get_all( created_before: A filter that returns only contexts created on or before the specified time. sort_by: The property used to sort results. The default value is CreationTime. sort_order: The sort order. The default value is Descending. - next_token: If the previous call to ListContexts didn't return the full set of contexts, the call returns a token for getting the next set of contexts. - max_results: The maximum number of contexts to return in the response. The default value is 10. session: Boto3 session. region: Region name. @@ -7786,8 +7797,6 @@ def get_all( endpoint_name: A filter that lists the data quality job definitions associated with the specified endpoint. sort_by: The field to sort results by. The default is CreationTime. sort_order: Whether to sort the results in Ascending or Descending order. The default is Descending. - next_token: If the result of the previous ListDataQualityJobDefinitions request was truncated, the response includes a NextToken. To retrieve the next set of transform jobs, use the token in the next request.> - max_results: The maximum number of data quality monitoring job definitions to return in the response. name_contains: A string in the data quality monitoring job definition name. This filter returns only data quality monitoring job definitions whose name contains the specified string. creation_time_before: A filter that returns only data quality monitoring job definitions created before the specified time. creation_time_after: A filter that returns only data quality monitoring job definitions created after the specified time. @@ -8000,8 +8009,6 @@ def get_all( Get all Device resources Parameters: - next_token: The response from the last list when returning a list large enough to need tokening. - max_results: Maximum number of results to select. latest_heartbeat_after: Select fleets where the job was updated after X model_name: A filter that searches devices that contains this name in any of their models. device_fleet_name: Filter for fleets containing this name in their device fleet name. @@ -8379,8 +8386,6 @@ def get_all( Get all DeviceFleet resources Parameters: - next_token: The response from the last list when returning a list large enough to need tokening. - max_results: The maximum number of results to select. creation_time_after: Filter fleets where packaging job was created after specified time. creation_time_before: Filter fleets where the edge packaging job was created before specified time. last_modified_time_after: Select fleets where the job was updated after X @@ -9474,8 +9479,6 @@ def get_all( Get all EdgeDeploymentPlan resources Parameters: - next_token: The response from the last list when returning a list large enough to need tokening. - max_results: The maximum number of results to select (50 by default). creation_time_after: Selects edge deployment plans created after this time. creation_time_before: Selects edge deployment plans created before this time. last_modified_time_after: Selects edge deployment plans that were last updated after this time. @@ -9709,7 +9712,6 @@ def get_all_stage_devices( Parameters: stage_name: The name of the stage in the deployment. - max_results: The maximum number of requests to select. exclude_devices_deployed_in_other_stage: Toggle for excluding devices deployed in other stages. session: Boto3 session. region: Region name. @@ -10033,6 +10035,7 @@ def wait( self, poll: int = 5, timeout: Optional[int] = None, + raise_on_stopped: bool = False, ) -> None: """ Wait for a EdgePackagingJob resource. @@ -10040,10 +10043,11 @@ def wait( Parameters: poll: The number of seconds to wait between each poll. timeout: The maximum number of seconds to wait before timing out. + raise_on_stopped: Whether to raise FailedStatusError when the resource stops instead of completing. Raises: TimeoutExceededError: If the resource does not reach a terminal state before the timeout. - FailedStatusError: If the resource reaches a failed state. + FailedStatusError: If the resource reaches a failed state, or a stopped state when raise_on_stopped is True. WaiterError: Raised when an error occurs while waiting. """ @@ -10081,6 +10085,13 @@ def wait( reason=self.edge_packaging_job_status_message, ) + if raise_on_stopped and "stopped" in current_status.lower(): + raise FailedStatusError( + resource_type="EdgePackagingJob", + status=current_status, + reason=self.edge_packaging_job_status_message, + ) + return if timeout is not None and time.time() - start_time >= timeout: @@ -10111,8 +10122,6 @@ def get_all( Get all EdgePackagingJob resources Parameters: - next_token: The response from the last list when returning a list large enough to need tokening. - max_results: Maximum number of results to select. creation_time_after: Select jobs where the job was created after specified time. creation_time_before: Select jobs where the job was created before specified time. last_modified_time_after: Select jobs where the job was updated after specified time. @@ -10172,7 +10181,7 @@ def get_all( ) -class Endpoint(Base): +class Endpoint(EndpointMixin, Base): """ Class representing resource Endpoint @@ -10647,8 +10656,6 @@ def get_all( Parameters: sort_by: Sorts the list of results. The default is CreationTime. sort_order: The sort order for results. The default is Descending. - next_token: If the result of a ListEndpoints request was truncated, the response includes a NextToken. To retrieve the next set of endpoints, use the token in the next request. - max_results: The maximum number of endpoints to return in the response. This value defaults to 10. name_contains: A string in endpoint names. This filter returns only endpoints whose name contains the specified string. creation_time_before: A filter that returns only endpoints that were created before the specified time (timestamp). creation_time_after: A filter that returns only endpoints with a creation time greater than or equal to the specified time (timestamp). @@ -11319,8 +11326,6 @@ def get_all( Parameters: sort_by: The field to sort results by. The default is CreationTime. sort_order: The sort order for results. The default is Descending. - next_token: If the result of the previous ListEndpointConfig request was truncated, the response includes a NextToken. To retrieve the next set of endpoint configurations, use the token in the next request. - max_results: The maximum number of training jobs to return in the response. name_contains: A string in the endpoint configuration name. This filter returns only endpoint configurations whose name contains the specified string. creation_time_before: A filter that returns only endpoint configurations created before the specified time (timestamp). creation_time_after: A filter that returns only endpoint configurations with a creation time greater than or equal to the specified time (timestamp). @@ -11668,8 +11673,6 @@ def get_all( created_before: A filter that returns only experiments created before the specified time. sort_by: The property used to sort results. The default value is CreationTime. sort_order: The sort order. The default value is Descending. - next_token: If the previous call to ListExperiments didn't return the full set of experiments, the call returns a token for getting the next set of experiments. - max_results: The maximum number of experiments to return in the response. The default value is 10. session: Boto3 session. region: Region name. @@ -12211,8 +12214,6 @@ def get_all( creation_time_before: Use this parameter to search for FeatureGroupss created before a specific date and time. sort_order: The order in which feature groups are listed. sort_by: The value on which the feature group list is sorted. - max_results: The maximum number of results returned by ListFeatureGroups. - next_token: A token to resume pagination of ListFeatureGroups results. session: Boto3 session. region: Region name. @@ -13252,8 +13253,6 @@ def get_all( creation_time_after: A filter that returns only flow definitions with a creation time greater than or equal to the specified timestamp. creation_time_before: A filter that returns only flow definitions that were created before the specified timestamp. sort_order: An optional value that specifies whether you want the results sorted in Ascending or Descending order. - next_token: A token to resume pagination. - max_results: The total number of items to return. If the total number of available items is more than the value specified in MaxResults, then a NextToken will be provided in the output that you can use to resume pagination. session: Boto3 session. region: Region name. @@ -13755,8 +13754,6 @@ def get_all( last_modified_time_after: Only list hubs that were last modified after the time specified. sort_by: Sort hubs by either name or creation time. sort_order: Sort hubs by ascending or descending order. - max_results: The maximum number of hubs to list. - next_token: If the response to a previous ListHubs request was truncated, the response includes a NextToken. To retrieve the next set of hubs, use the token in the next request. session: Boto3 session. region: Region name. @@ -13805,7 +13802,7 @@ def get_all( ) -class HubContent(Base): +class HubContent(HubContentMixin, Base): """ Class representing resource HubContent @@ -13830,6 +13827,7 @@ class HubContent(Base): hub_content_dependencies: The location of any dependencies that the hub content has, such as scripts, model artifacts, datasets, or notebooks. failure_reason: The failure reason if importing hub content failed. last_modified_time: The last modified time of the hub content. + original_creation_time: The date and time when the hub content was originally created, before any updates or revisions. """ @@ -13854,6 +13852,7 @@ class HubContent(Base): creation_time: Optional[datetime.datetime] = Unassigned() last_modified_time: Optional[datetime.datetime] = Unassigned() hub_name: Optional[str] = Unassigned() + original_creation_time: Optional[datetime.datetime] = Unassigned() def get_name(self) -> str: attributes = vars(self) @@ -14210,6 +14209,86 @@ def load( region=region, ) + @classmethod + @Base.add_validate_call + def get_all( + cls, + hub_name: StrPipeVar, + hub_content_type: StrPipeVar, + name_contains: Optional[StrPipeVar] = Unassigned(), + max_schema_version: Optional[StrPipeVar] = Unassigned(), + creation_time_before: Optional[datetime.datetime] = Unassigned(), + creation_time_after: Optional[datetime.datetime] = Unassigned(), + sort_by: Optional[StrPipeVar] = Unassigned(), + sort_order: Optional[StrPipeVar] = Unassigned(), + session: Optional[Session] = None, + region: Optional[StrPipeVar] = None, + max_workers: int = 4, + ) -> ResourceIterator["HubContent"]: + """ + Get all HubContent resources + + Parameters: + hub_name: The name of the hub to list the contents of. + hub_content_type: The type of hub content to list. + name_contains: Only list hub content if the name contains the specified string. + max_schema_version: The upper bound of the hub content schema verion. + creation_time_before: Only list hub content that was created before the time specified. + creation_time_after: Only list hub content that was created after the time specified. + sort_by: Sort hub content versions by either name or creation time. + sort_order: Sort hubs by ascending or descending order. + session: Boto3 session. + region: Region name. + max_workers: How many listed resources to describe at a time, on a thread pool. Defaults to 4. Keep it at or below the max_pool_connections of the client (10 by default). + + Returns: + Iterator for listed HubContent resources. + + 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: + ``` + try: + # AWS service call here + except botocore.exceptions.ClientError as e: + error_message = e.response['Error']['Message'] + error_code = e.response['Error']['Code'] + ``` + ResourceNotFound: Resource being access is not found. + """ + + client = Base.get_sagemaker_client( + session=session, region_name=region, service_name="sagemaker" + ) + + operation_input_args = { + "HubName": hub_name, + "HubContentType": hub_content_type, + "NameContains": name_contains, + "MaxSchemaVersion": max_schema_version, + "CreationTimeBefore": creation_time_before, + "CreationTimeAfter": creation_time_after, + "SortBy": sort_by, + "SortOrder": sort_order, + } + + # serialize the input request + operation_input_args = serialize(operation_input_args) + logger.debug(f"Serialized input request: {operation_input_args}") + + return ResourceIterator( + client=client, + session=session, + region=region, + list_method="list_hub_contents", + summaries_key="HubContentSummaries", + summary_name="HubContentInfo", + resource_cls=HubContent, + extra_init_data={"hub_name": hub_name}, + max_workers=max_workers, + list_method_kwargs=operation_input_args, + ) + @Base.add_validate_call def get_all_versions( self, @@ -14221,6 +14300,7 @@ def get_all_versions( sort_order: Optional[StrPipeVar] = Unassigned(), session: Optional[Session] = None, region: Optional[str] = None, + max_workers: int = 4, ) -> ResourceIterator["HubContent"]: """ List hub content versions. @@ -14232,10 +14312,9 @@ def get_all_versions( creation_time_after: Only list hub content versions that were created after the time specified. sort_by: Sort hub content versions by either name or creation time. sort_order: Sort hub content versions by ascending or descending order. - max_results: The maximum number of hub content versions to list. - next_token: If the response to a previous ListHubContentVersions request was truncated, the response includes a NextToken. To retrieve the next set of hub content versions, use the token in the next request. session: Boto3 session. region: Region name. + max_workers: How many listed resources to describe at a time, on a thread pool. Defaults to 4. Keep it at or below the max_pool_connections of the client (10 by default). Returns: Iterator for listed HubContent. @@ -14279,6 +14358,8 @@ def get_all_versions( summary_name="HubContentInfo", resource_cls=HubContent, list_method_kwargs=operation_input_args, + extra_init_data={"hub_name": self.hub_name}, + max_workers=max_workers, ) @@ -14941,8 +15022,6 @@ def get_all( creation_time_after: A filter that returns only human task user interfaces with a creation time greater than or equal to the specified timestamp. creation_time_before: A filter that returns only human task user interfaces that were created before the specified timestamp. sort_order: An optional value that specifies whether you want the results sorted in Ascending or Descending order. - next_token: A token to resume pagination. - max_results: The total number of items to return. If the total number of available items is more than the value specified in MaxResults, then a NextToken will be provided in the output that you can use to resume pagination. session: Boto3 session. region: Region name. @@ -15320,6 +15399,7 @@ def wait( self, poll: int = 5, timeout: Optional[int] = None, + raise_on_stopped: bool = False, ) -> None: """ Wait for a HyperParameterTuningJob resource. @@ -15327,10 +15407,11 @@ def wait( Parameters: poll: The number of seconds to wait between each poll. timeout: The maximum number of seconds to wait before timing out. + raise_on_stopped: Whether to raise FailedStatusError when the resource stops instead of completing. Raises: TimeoutExceededError: If the resource does not reach a terminal state before the timeout. - FailedStatusError: If the resource reaches a failed state. + FailedStatusError: If the resource reaches a failed state, or a stopped state when raise_on_stopped is True. WaiterError: Raised when an error occurs while waiting. """ @@ -15368,6 +15449,13 @@ def wait( reason=self.failure_reason, ) + if raise_on_stopped and "stopped" in current_status.lower(): + raise FailedStatusError( + resource_type="HyperParameterTuningJob", + status=current_status, + reason=self.failure_reason, + ) + return if timeout is not None and time.time() - start_time >= timeout: @@ -15460,8 +15548,6 @@ def get_all( Get all HyperParameterTuningJob resources Parameters: - next_token: If the result of the previous ListHyperParameterTuningJobs request was truncated, the response includes a NextToken. To retrieve the next set of tuning jobs, use the token in the next request. - max_results: The maximum number of tuning jobs to return. The default value is 10. sort_by: The field to sort results by. The default is Name. sort_order: The sort order for results. The default is Ascending. name_contains: A string in the tuning job name. This filter returns only tuning jobs whose name contains the specified string. @@ -15531,8 +15617,6 @@ def get_all_training_jobs( Gets a list of TrainingJobSummary objects that describe the training jobs that a hyperparameter tuning job launched. Parameters: - next_token: If the result of the previous ListTrainingJobsForHyperParameterTuningJob request was truncated, the response includes a NextToken. To retrieve the next set of training jobs, use the token in the next request. - max_results: The maximum number of training jobs to return. The default value is 10. status_equals: A filter that returns only training jobs with the specified status. sort_by: The field to sort results by. The default is Name. If the value of this field is FinalObjectiveMetricValue, any training jobs that did not return an objective metric are not listed. sort_order: The sort order for results. The default is Ascending. @@ -16040,9 +16124,7 @@ def get_all( creation_time_before: A filter that returns only images created on or before the specified time. last_modified_time_after: A filter that returns only images modified on or after the specified time. last_modified_time_before: A filter that returns only images modified on or before the specified time. - max_results: The maximum number of images to return in the response. The default value is 10. name_contains: A filter that returns only images whose name contains the specified string. - next_token: If the previous call to ListImages didn't return the full set of images, the call returns a token for getting the next set of images. sort_by: The property used to sort results. The default value is CREATION_TIME. sort_order: The sort order. The default value is DESCENDING. session: Boto3 session. @@ -16106,8 +16188,6 @@ def get_all_aliases( Parameters: alias: The alias of the image version. version: The version of the image. If image version is not specified, the aliases of all versions of the image are listed. - max_results: The maximum number of aliases to return. - next_token: If the previous call to ListAliases didn't return the full set of aliases, the call returns a token for retrieving the next set of aliases. session: Boto3 session. region: Region name. @@ -17082,8 +17162,6 @@ def get_all( Parameters: sort_by: The field by which to sort the inference components in the response. The default is CreationTime. sort_order: The sort order for results. The default is Descending. - next_token: A token that you use to get the next set of results following a truncated response. If the response to the previous request was truncated, that response provides the value for this token. - max_results: The maximum number of inference components to return in the response. This value defaults to 10. name_contains: Filters the results to only those inference components with a name that contains the specified string. creation_time_before: Filters the results to only those inference components that were created before the specified time. creation_time_after: Filters the results to only those inference components that were created after the specified time. @@ -17695,8 +17773,6 @@ def get_all( last_modified_time_before: Selects inference experiments which were last modified before this timestamp. sort_by: The column by which to sort the listed inference experiments. sort_order: The direction of sorting (ascending or descending). - next_token: The response from the last list when returning a list large enough to need tokening. - max_results: The maximum number of results to select. session: Boto3 session. region: Region name. @@ -18029,6 +18105,7 @@ def wait( self, poll: int = 5, timeout: Optional[int] = None, + raise_on_stopped: bool = False, ) -> None: """ Wait for a InferenceRecommendationsJob resource. @@ -18036,10 +18113,11 @@ def wait( Parameters: poll: The number of seconds to wait between each poll. timeout: The maximum number of seconds to wait before timing out. + raise_on_stopped: Whether to raise FailedStatusError when the resource stops instead of completing. Raises: TimeoutExceededError: If the resource does not reach a terminal state before the timeout. - FailedStatusError: If the resource reaches a failed state. + FailedStatusError: If the resource reaches a failed state, or a stopped state when raise_on_stopped is True. WaiterError: Raised when an error occurs while waiting. """ @@ -18077,6 +18155,13 @@ def wait( reason=self.failure_reason, ) + if raise_on_stopped and "stopped" in current_status.lower(): + raise FailedStatusError( + resource_type="InferenceRecommendationsJob", + status=current_status, + reason=self.failure_reason, + ) + return if timeout is not None and time.time() - start_time >= timeout: @@ -18183,8 +18268,6 @@ def get_all( status_equals: A filter that retrieves only inference recommendations jobs with a specific status. sort_by: The parameter by which to sort the results. sort_order: The sort order for the results. - next_token: If the response to a previous ListInferenceRecommendationsJobsRequest request was truncated, the response includes a NextToken. To retrieve the next set of recommendations, use the token in the next request. - max_results: The maximum number of recommendations to return in the response. model_name_equals: A filter that returns only jobs that were created for this model. model_package_version_arn_equals: A filter that returns only jobs that were created for this versioned model package. session: Boto3 session. @@ -18249,8 +18332,6 @@ def get_all_steps( Parameters: step_type: A filter to return details about the specified type of subtask. BENCHMARK: Evaluate the performance of your model on different instance types. - max_results: The maximum number of results to return. - next_token: A token that you can specify to return more results from the list. Specify this field if you have a token that was returned from a previous request. session: Boto3 session. region: Region name. @@ -18601,6 +18682,7 @@ def wait( self, poll: int = 5, timeout: Optional[int] = None, + raise_on_stopped: bool = False, ) -> None: """ Wait for a Job resource. @@ -18608,10 +18690,11 @@ def wait( Parameters: poll: The number of seconds to wait between each poll. timeout: The maximum number of seconds to wait before timing out. + raise_on_stopped: Whether to raise FailedStatusError when the resource stops instead of completing. Raises: TimeoutExceededError: If the resource does not reach a terminal state before the timeout. - FailedStatusError: If the resource reaches a failed state. + FailedStatusError: If the resource reaches a failed state, or a stopped state when raise_on_stopped is True. WaiterError: Raised when an error occurs while waiting. """ @@ -18647,6 +18730,11 @@ def wait( resource_type="Job", status=current_status, reason=self.failure_reason ) + if raise_on_stopped and "stopped" in current_status.lower(): + raise FailedStatusError( + resource_type="Job", status=current_status, reason=self.failure_reason + ) + return if timeout is not None and time.time() - start_time >= timeout: @@ -18739,8 +18827,6 @@ def get_all( Parameters: job_category: The category of jobs to list. - next_token: If the previous response was truncated, this token retrieves the next set of results. - max_results: The maximum number of jobs to return in the response. The default value is 50. creation_time_after: A filter that returns only jobs created after the specified time. creation_time_before: A filter that returns only jobs created before the specified time. last_modified_time_after: A filter that returns only jobs modified after the specified time. @@ -19107,6 +19193,7 @@ def wait( self, poll: int = 5, timeout: Optional[int] = None, + raise_on_stopped: bool = False, ) -> None: """ Wait for a LabelingJob resource. @@ -19114,10 +19201,11 @@ def wait( Parameters: poll: The number of seconds to wait between each poll. timeout: The maximum number of seconds to wait before timing out. + raise_on_stopped: Whether to raise FailedStatusError when the resource stops instead of completing. Raises: TimeoutExceededError: If the resource does not reach a terminal state before the timeout. - FailedStatusError: If the resource reaches a failed state. + FailedStatusError: If the resource reaches a failed state, or a stopped state when raise_on_stopped is True. WaiterError: Raised when an error occurs while waiting. """ @@ -19155,6 +19243,13 @@ def wait( reason=self.failure_reason, ) + if raise_on_stopped and "stopped" in current_status.lower(): + raise FailedStatusError( + resource_type="LabelingJob", + status=current_status, + reason=self.failure_reason, + ) + return if timeout is not None and time.time() - start_time >= timeout: @@ -19188,8 +19283,6 @@ def get_all( creation_time_before: A filter that returns only labeling jobs created before the specified time (timestamp). last_modified_time_after: A filter that returns only labeling jobs modified after the specified time (timestamp). last_modified_time_before: A filter that returns only labeling jobs modified before the specified time (timestamp). - max_results: The maximum number of labeling jobs to return in each page of the response. - next_token: If the result of the previous ListLabelingJobs request was truncated, the response includes a NextToken. To retrieve the next set of labeling jobs, use the token in the next request. name_contains: A string in the labeling job name. This filter returns only labeling jobs whose name contains the specified string. sort_by: The field to sort results by. The default is CreationTime. sort_order: The sort order for results. The default is Ascending. @@ -19392,8 +19485,6 @@ def get_all( created_before: A timestamp to filter against lineage groups created before a certain point in time. sort_by: The parameter by which to sort the results. The default is CreationTime. sort_order: The sort order for the results. The default is Ascending. - next_token: If the response is truncated, SageMaker returns this token. To retrieve the next set of algorithms, use it in the subsequent request. - max_results: The maximum number of endpoints to return in the response. This value defaults to 10. session: Boto3 session. region: Region name. @@ -19976,8 +20067,6 @@ def get_all( account_default_status: Filter for MLflow Apps with the specified AccountDefaultStatus. sort_by: Filter for MLflow Apps sorting by name, creation time, or creation status. sort_order: Change the order of the listed MLflow Apps. By default, MLflow Apps are listed in Descending order by creation time. To change the list order, specify SortOrder to be Ascending. - next_token: If the previous response was truncated, use this token in your next request to receive the next set of results. - max_results: The maximum number of MLflow Apps to list. session: Boto3 session. region: Region name. @@ -20608,8 +20697,6 @@ def get_all( mlflow_version: Filter for tracking servers using the specified MLflow version. sort_by: Filter for trackings servers sorting by name, creation time, or creation status. sort_order: Change the order of the listed tracking servers. By default, tracking servers are listed in Descending order by creation time. To change the list order, you can specify SortOrder to be Ascending. - next_token: If the previous response was truncated, you will receive this token. Use it in your next request to receive the next set of results. - max_results: The maximum number of tracking servers to list. session: Boto3 session. region: Region name. @@ -20948,8 +21035,6 @@ def get_all( Parameters: sort_by: Sorts the list of results. The default is CreationTime. sort_order: The sort order for results. The default is Descending. - next_token: If the response to a previous ListModels request was truncated, the response includes a NextToken. To retrieve the next set of models, use the token in the next request. - max_results: The maximum number of models to return in the response. name_contains: A string in the model name. This filter returns only models whose name contains the specified string. creation_time_before: A filter that returns only models created before the specified time (timestamp). creation_time_after: A filter that returns only models with a creation time greater than or equal to the specified time (timestamp). @@ -21010,8 +21095,6 @@ def get_all_metadata( Parameters: search_expression: One or more filters that searches for the specified resource or resources in a search. All resource objects that satisfy the expression's condition are included in the search results. Specify the Framework, FrameworkVersion, Domain or Task to filter supported. Filter names and values are case-sensitive. - next_token: If the response to a previous ListModelMetadataResponse request was truncated, the response includes a NextToken. To retrieve the next set of model metadata, use the token in the next request. - max_results: The maximum number of models to return in the response. session: Boto3 session. region: Region name. @@ -21366,8 +21449,6 @@ def get_all( endpoint_name: Name of the endpoint to monitor for model bias. sort_by: Whether to sort results by the Name or CreationTime field. The default is CreationTime. sort_order: Whether to sort the results in Ascending or Descending order. The default is Descending. - next_token: The token returned if the response is truncated. To retrieve the next set of job executions, use it in the next request. - max_results: The maximum number of model bias jobs to return in the response. The default value is 10. name_contains: Filter for model bias jobs whose name contains a specified string. creation_time_before: A filter that returns only model bias jobs created before a specified time. creation_time_after: A filter that returns only model bias jobs created after a specified time. @@ -21805,10 +21886,8 @@ def get_all( Parameters: creation_time_after: Only list model cards that were created after the time specified. creation_time_before: Only list model cards that were created before the time specified. - max_results: The maximum number of model cards to list. name_contains: Only list model cards with names that contain the specified string. model_card_status: Only list model cards with the specified approval status. - next_token: If the response to a previous ListModelCards request was truncated, the response includes a NextToken. To retrieve the next set of model cards, use the token in the next request. sort_by: Sort model cards by either name or creation time. Sorts by creation time by default. sort_order: Sort model cards by ascending or descending order. session: Boto3 session. @@ -21873,8 +21952,6 @@ def get_all_versions( Parameters: creation_time_after: Only list model card versions that were created after the time specified. creation_time_before: Only list model card versions that were created before the time specified. - max_results: The maximum number of model card versions to list. - next_token: If the response to a previous ListModelCardVersions request was truncated, the response includes a NextToken. To retrieve the next set of model card versions, use the token in the next request. sort_by: Sort listed model card versions by version. Sorts by version by default. sort_order: Sort model card versions by ascending or descending order. session: Boto3 session. @@ -22152,6 +22229,7 @@ def wait( self, poll: int = 5, timeout: Optional[int] = None, + raise_on_stopped: bool = False, ) -> None: """ Wait for a ModelCardExportJob resource. @@ -22159,10 +22237,11 @@ def wait( Parameters: poll: The number of seconds to wait between each poll. timeout: The maximum number of seconds to wait before timing out. + raise_on_stopped: Whether to raise FailedStatusError when the resource stops instead of completing. Raises: TimeoutExceededError: If the resource does not reach a terminal state before the timeout. - FailedStatusError: If the resource reaches a failed state. + FailedStatusError: If the resource reaches a failed state, or a stopped state when raise_on_stopped is True. WaiterError: Raised when an error occurs while waiting. """ @@ -22200,6 +22279,13 @@ def wait( reason=self.failure_reason, ) + if raise_on_stopped and "stopped" in current_status.lower(): + raise FailedStatusError( + resource_type="ModelCardExportJob", + status=current_status, + reason=self.failure_reason, + ) + return if timeout is not None and time.time() - start_time >= timeout: @@ -22237,8 +22323,6 @@ def get_all( status_equals: Only list model card export jobs with the specified status. sort_by: Sort model card export jobs by either name or creation time. Sorts by creation time by default. sort_order: Sort model card export jobs by ascending or descending order. - next_token: If the response to a previous ListModelCardExportJobs request was truncated, the response includes a NextToken. To retrieve the next set of model card export jobs, use the token in the next request. - max_results: The maximum number of model card export jobs to list. session: Boto3 session. region: Region name. @@ -22609,8 +22693,6 @@ def get_all( endpoint_name: Name of the endpoint to monitor for model explainability. sort_by: Whether to sort results by the Name or CreationTime field. The default is CreationTime. sort_order: Whether to sort the results in Ascending or Descending order. The default is Descending. - next_token: The token returned if the response is truncated. To retrieve the next set of job executions, use it in the next request. - max_results: The maximum number of jobs to return in the response. The default value is 10. name_contains: Filter for model explainability jobs whose name contains a specified string. creation_time_before: A filter that returns only model explainability jobs created before a specified time. creation_time_after: A filter that returns only model explainability jobs created after a specified time. @@ -23284,12 +23366,10 @@ def get_all( Parameters: creation_time_after: A filter that returns only model packages created after the specified time (timestamp). creation_time_before: A filter that returns only model packages created before the specified time (timestamp). - max_results: The maximum number of model packages to return in the response. name_contains: A string in the model package name. This filter returns only model packages whose name contains the specified string. model_approval_status: A filter that returns only the model packages with the specified approval status. model_package_group_name: A filter that returns only model versions that belong to the specified model group. model_package_type: A filter that returns only the model packages of the specified type. This can be one of the following values. UNVERSIONED - List only unversioined models. This is the default value if no ModelPackageType is specified. VERSIONED - List only versioned models. BOTH - List both versioned and unversioned models. - next_token: If the response to a previous ListModelPackages request was truncated, the response includes a NextToken. To retrieve the next set of model packages, use the token in the next request. sort_by: The parameter by which to sort the results. The default is CreationTime. sort_order: The sort order for the results. The default is Ascending. session: Boto3 session. @@ -23387,7 +23467,7 @@ def batch_get( return BatchDescribeModelPackageOutput(**transformed_response) -class ModelPackageGroup(Base): +class ModelPackageGroup(ModelPackageGroupMixin, Base): """ Class representing resource ModelPackageGroup @@ -23758,9 +23838,7 @@ def get_all( Parameters: creation_time_after: A filter that returns only model groups created after the specified time. creation_time_before: A filter that returns only model groups created before the specified time. - max_results: The maximum number of results to return in the response. name_contains: A string in the model group name. This filter returns only model groups whose name contains the specified string. - next_token: If the result of the previous ListModelPackageGroups request was truncated, the response includes a NextToken. To retrieve the next set of model groups, use the token in the next request. sort_by: The field to sort results by. The default is CreationTime. sort_order: The sort order for results. The default is Ascending. cross_account_filter_option: A filter that returns either model groups shared with you or model groups in your own account. When the value is CrossAccount, the results show the resources made discoverable to you from other accounts. When the value is SameAccount or null, the results show resources from your account. The default is SameAccount. @@ -24249,8 +24327,6 @@ def get_all( endpoint_name: A filter that returns only model quality monitoring job definitions that are associated with the specified endpoint. sort_by: The field to sort results by. The default is CreationTime. sort_order: Whether to sort the results in Ascending or Descending order. The default is Descending. - next_token: If the result of the previous ListModelQualityJobDefinitions request was truncated, the response includes a NextToken. To retrieve the next set of model quality monitoring job definitions, use the token in the next request. - max_results: The maximum number of results to return in a call to ListModelQualityJobDefinitions. name_contains: A string in the transform job name. This filter returns only model quality monitoring job definitions whose name contains the specified string. creation_time_before: A filter that returns only model quality monitoring job definitions created before the specified time. creation_time_after: A filter that returns only model quality monitoring job definitions created after the specified time. @@ -24408,8 +24484,6 @@ def get_all( Parameters: monitoring_schedule_name: The name of a monitoring schedule. - next_token: If the result of the previous ListMonitoringAlerts request was truncated, the response includes a NextToken. To retrieve the next set of alerts in the history, use the token in the next request. - max_results: The maximum number of results to display. The default is 100. session: Boto3 session. region: Region name. @@ -24595,8 +24669,6 @@ def get_all( endpoint_name: Name of a specific endpoint to fetch jobs for. sort_by: Whether to sort the results by the Status, CreationTime, or ScheduledTime field. The default is CreationTime. sort_order: Whether to sort the results in Ascending or Descending order. The default is Descending. - next_token: The token returned if the response is truncated. To retrieve the next set of job executions, use it in the next request. - max_results: The maximum number of jobs to return in the response. The default value is 10. scheduled_time_before: Filter for jobs scheduled before a specified time. scheduled_time_after: Filter for jobs scheduled after a specified time. creation_time_before: A filter that returns only jobs created before a specified time. @@ -25134,8 +25206,6 @@ def get_all( endpoint_name: Name of a specific endpoint to fetch schedules for. sort_by: Whether to sort the results by the Status, CreationTime, or ScheduledTime field. The default is CreationTime. sort_order: Whether to sort the results in Ascending or Descending order. The default is Descending. - next_token: The token returned if the response is truncated. To retrieve the next set of job executions, use it in the next request. - max_results: The maximum number of jobs to return in the response. The default value is 10. name_contains: Filter for monitoring schedules whose name contains a specified string. creation_time_before: A filter that returns only monitoring schedules created before a specified time. creation_time_after: A filter that returns only monitoring schedules created after a specified time. @@ -25823,8 +25893,6 @@ def get_all( Get all NotebookInstance resources Parameters: - next_token: If the previous call to the ListNotebookInstances is truncated, the response includes a NextToken. You can use this token in your subsequent ListNotebookInstances request to fetch the next set of notebook instances. You might specify a filter or a sort order in your request. When response is truncated, you must use the same values for the filer and sort order in the next request. - max_results: The maximum number of notebook instances to return. sort_by: The field to sort results by. The default is Name. sort_order: The sort order for results. name_contains: A string in the notebook instances' name. This filter returns only notebook instances whose name contains the specified string. @@ -26180,8 +26248,6 @@ def get_all( Get all NotebookInstanceLifecycleConfig resources Parameters: - next_token: If the result of a ListNotebookInstanceLifecycleConfigs request was truncated, the response includes a NextToken. To get the next set of lifecycle configurations, use the token in the next request. - max_results: The maximum number of lifecycle configurations to return in the response. sort_by: Sorts the list of results. The default is CreationTime. sort_order: The sort order for results. name_contains: A string in the lifecycle configuration name. This filter returns only lifecycle configurations whose name contains the specified string. @@ -26576,6 +26642,7 @@ def wait( self, poll: int = 5, timeout: Optional[int] = None, + raise_on_stopped: bool = False, ) -> None: """ Wait for a OptimizationJob resource. @@ -26583,10 +26650,11 @@ def wait( Parameters: poll: The number of seconds to wait between each poll. timeout: The maximum number of seconds to wait before timing out. + raise_on_stopped: Whether to raise FailedStatusError when the resource stops instead of completing. Raises: TimeoutExceededError: If the resource does not reach a terminal state before the timeout. - FailedStatusError: If the resource reaches a failed state. + FailedStatusError: If the resource reaches a failed state, or a stopped state when raise_on_stopped is True. WaiterError: Raised when an error occurs while waiting. """ @@ -26624,6 +26692,13 @@ def wait( reason=self.failure_reason, ) + if raise_on_stopped and "stopped" in current_status.lower(): + raise FailedStatusError( + resource_type="OptimizationJob", + status=current_status, + reason=self.failure_reason, + ) + return if timeout is not None and time.time() - start_time >= timeout: @@ -26654,8 +26729,6 @@ def get_all( Get all OptimizationJob resources Parameters: - next_token: A token that you use to get the next set of results following a truncated response. If the response to the previous request was truncated, that response provides the value for this token. - max_results: The maximum number of optimization jobs to return in the response. The default is 50. creation_time_after: Filters the results to only those optimization jobs that were created after the specified time. creation_time_before: Filters the results to only those optimization jobs that were created before the specified time. last_modified_time_after: Filters the results to only those optimization jobs that were updated after the specified time. @@ -27817,8 +27890,6 @@ def get_all( created_before: A filter that returns the pipelines that were created before a specified time. sort_by: The field by which to sort results. The default is CreatedTime. sort_order: The sort order for results. - next_token: If the result of the previous ListPipelines request was truncated, the response includes a NextToken. To retrieve the next set of pipelines, use the token in the next request. - max_results: The maximum number of pipelines to return in the response. session: Boto3 session. region: Region name. @@ -28226,8 +28297,6 @@ def get_all( created_before: A filter that returns the pipeline executions that were created before a specified time. sort_by: The field by which to sort results. The default is CreatedTime. sort_order: The sort order for results. - next_token: If the result of the previous ListPipelineExecutions request was truncated, the response includes a NextToken. To retrieve the next set of pipeline executions, use the token in the next request. - max_results: The maximum number of pipeline executions to return in the response. session: Boto3 session. region: Region name. @@ -28330,8 +28399,6 @@ def get_all_steps( Gets a list of PipeLineExecutionStep objects. Parameters: - next_token: If the result of the previous ListPipelineExecutionSteps request was truncated, the response includes a NextToken. To retrieve the next set of pipeline execution steps, use the token in the next request. - max_results: The maximum number of pipeline execution steps to return in the response. sort_order: The field by which to sort results. The default is CreatedTime. session: Boto3 session. region: Region name. @@ -28383,8 +28450,6 @@ def get_all_parameters( Gets a list of parameters for a pipeline execution. Parameters: - next_token: If the result of the previous ListPipelineParametersForExecution request was truncated, the response includes a NextToken. To retrieve the next set of parameters, use the token in the next request. - max_results: The maximum number of parameters to return in the response. session: Boto3 session. region: Region name. @@ -29290,6 +29355,7 @@ def wait( self, poll: int = 5, timeout: Optional[int] = None, + raise_on_stopped: bool = False, logs: Optional[bool] = False, ) -> None: """ @@ -29298,11 +29364,12 @@ def wait( Parameters: poll: The number of seconds to wait between each poll. timeout: The maximum number of seconds to wait before timing out. + raise_on_stopped: Whether to raise FailedStatusError when the resource stops instead of completing. logs: Whether to print logs while waiting. Raises: TimeoutExceededError: If the resource does not reach a terminal state before the timeout. - FailedStatusError: If the resource reaches a failed state. + FailedStatusError: If the resource reaches a failed state, or a stopped state when raise_on_stopped is True. WaiterError: Raised when an error occurs while waiting. """ @@ -29360,6 +29427,13 @@ def wait( reason=self.failure_reason, ) + if raise_on_stopped and "stopped" in current_status.lower(): + raise FailedStatusError( + resource_type="ProcessingJob", + status=current_status, + reason=self.failure_reason, + ) + return if timeout is not None and time.time() - start_time >= timeout: @@ -29397,8 +29471,6 @@ def get_all( status_equals: A filter that retrieves only processing jobs with a specific status. sort_by: The field to sort results by. The default is CreationTime. sort_order: The sort order for results. The default is Ascending. - next_token: If the result of the previous ListProcessingJobs request was truncated, the response includes a NextToken. To retrieve the next set of processing jobs, use the token in the next request. - max_results: The maximum number of processing jobs to return in the response. session: Boto3 session. region: Region name. @@ -29833,9 +29905,7 @@ def get_all( Parameters: creation_time_after: A filter that returns the projects that were created after a specified time. creation_time_before: A filter that returns the projects that were created before a specified time. - max_results: The maximum number of projects to return in the response. name_contains: A filter that returns the projects whose name contains a specified string. - next_token: If the result of the previous ListProjects request was truncated, the response includes a NextToken. To retrieve the next set of projects, use the token in the next request. sort_by: The field by which to sort results. The default is CreationTime. sort_order: The sort order for results. The default is Ascending. session: Boto3 session. @@ -29938,8 +30008,6 @@ def get_all( creation_time_before: Use this parameter to search for ResourceCatalogs created before a specific date and time. sort_order: The order in which the resource catalogs are listed. sort_by: The value on which the resource catalog list is sorted. - max_results: The maximum number of results returned by ListResourceCatalogs. - next_token: A token to resume pagination of ListResourceCatalogs results. session: Boto3 session. region: Region name. @@ -30550,8 +30618,6 @@ def get_all( Get all Space resources Parameters: - next_token: If the previous response was truncated, you will receive this token. Use it in your next request to receive the next set of results. - max_results: This parameter defines the maximum number of results that can be return in a single response. The MaxResults parameter is an upper bound, not a target. If there are more results available than the value specified, a NextToken is provided in the response. The NextToken indicates that the user should get the next set of results by providing this token as a part of a subsequent call. The default value for MaxResults is 10. sort_order: The sort order for the results. The default is Ascending. sort_by: The parameter by which to sort the results. The default is CreationTime. domain_id_equals: A parameter to search for the domain ID. @@ -30852,8 +30918,6 @@ def get_all( Get all StudioLifecycleConfig resources Parameters: - max_results: The total number of items to return in the response. If the total number of items available is more than the value specified, a NextToken is provided in the response. To resume pagination, provide the NextToken value in the as part of a subsequent call. The default value is 10. - next_token: If the previous call to ListStudioLifecycleConfigs didn't return the full set of Lifecycle Configurations, the call returns a token for getting the next set of Lifecycle Configurations. name_contains: A string in the Lifecycle Configuration name. This filter returns only Lifecycle Configurations whose name contains the specified string. app_type_equals: A parameter to search for the App Type to which the Lifecycle Configuration is attached. creation_time_before: A filter that returns only Lifecycle Configurations created on or before the specified time. @@ -31040,8 +31104,6 @@ def get_all( Parameters: name_contains: A string in the work team name. This filter returns only work teams whose name contains the specified string. - next_token: If the result of the previous ListSubscribedWorkteams request was truncated, the response includes a NextToken. To retrieve the next set of labeling jobs, use the token in the next request. - max_results: The maximum number of work teams to return in each page of the response. session: Boto3 session. region: Region name. @@ -31126,8 +31188,6 @@ def get_all( Parameters: resource_arn: The Amazon Resource Name (ARN) of the resource whose tags you want to retrieve. - next_token: If the response to the previous ListTags request is truncated, SageMaker returns this token. To retrieve the next set of tags, use it in the subsequent request. - max_results: Maximum number of tags to return. session: Boto3 session. region: Region name. @@ -31788,6 +31848,7 @@ def wait( self, poll: int = 5, timeout: Optional[int] = None, + raise_on_stopped: bool = False, logs: Optional[bool] = False, ) -> None: """ @@ -31796,11 +31857,12 @@ def wait( Parameters: poll: The number of seconds to wait between each poll. timeout: The maximum number of seconds to wait before timing out. + raise_on_stopped: Whether to raise FailedStatusError when the resource stops instead of completing. logs: Whether to print logs while waiting. Raises: TimeoutExceededError: If the resource does not reach a terminal state before the timeout. - FailedStatusError: If the resource reaches a failed state. + FailedStatusError: If the resource reaches a failed state, or a stopped state when raise_on_stopped is True. WaiterError: Raised when an error occurs while waiting. """ @@ -31870,6 +31932,13 @@ def wait( reason=self.failure_reason, ) + if raise_on_stopped and "stopped" in current_status.lower(): + raise FailedStatusError( + resource_type="TrainingJob", + status=current_status, + reason=self.failure_reason, + ) + return if timeout is not None and time.time() - start_time >= timeout: @@ -31964,8 +32033,6 @@ def get_all( Get all TrainingJob resources Parameters: - next_token: If the result of the previous ListTrainingJobs request was truncated, the response includes a NextToken. To retrieve the next set of training jobs, use the token in the next request. - max_results: The maximum number of training jobs to return in the response. creation_time_after: A filter that returns only training jobs created after the specified time (timestamp). creation_time_before: A filter that returns only training jobs created before the specified time (timestamp). last_modified_time_after: A filter that returns only training jobs modified after the specified time (timestamp). @@ -32321,8 +32388,6 @@ def get_all( Get all TrainingPlan resources Parameters: - next_token: A token to continue pagination if more results are available. - max_results: The maximum number of results to return in the response. start_time_after: Filter to list only training plans with an actual start time after this date. start_time_before: Filter to list only training plans with an actual start time before this date. sort_by: The training plan field to sort the results by (e.g., StartTime, Status). @@ -33117,6 +33182,7 @@ def wait( self, poll: int = 5, timeout: Optional[int] = None, + raise_on_stopped: bool = False, logs: Optional[bool] = False, ) -> None: """ @@ -33125,11 +33191,12 @@ def wait( Parameters: poll: The number of seconds to wait between each poll. timeout: The maximum number of seconds to wait before timing out. + raise_on_stopped: Whether to raise FailedStatusError when the resource stops instead of completing. logs: Whether to print logs while waiting. Raises: TimeoutExceededError: If the resource does not reach a terminal state before the timeout. - FailedStatusError: If the resource reaches a failed state. + FailedStatusError: If the resource reaches a failed state, or a stopped state when raise_on_stopped is True. WaiterError: Raised when an error occurs while waiting. """ @@ -33187,6 +33254,13 @@ def wait( reason=self.failure_reason, ) + if raise_on_stopped and "stopped" in current_status.lower(): + raise FailedStatusError( + resource_type="TransformJob", + status=current_status, + reason=self.failure_reason, + ) + return if timeout is not None and time.time() - start_time >= timeout: @@ -33224,8 +33298,6 @@ def get_all( status_equals: A filter that retrieves only transform jobs with a specific status. sort_by: The field to sort results by. The default is CreationTime. sort_order: The sort order for results. The default is Descending. - next_token: If the result of the previous ListTransformJobs request was truncated, the response includes a NextToken. To retrieve the next set of transform jobs, use the token in the next request. - max_results: The maximum number of transform jobs to return in the response. The default value is 10. session: Boto3 session. region: Region name. @@ -33581,8 +33653,6 @@ def get_all( created_before: A filter that returns only trials created before the specified time. sort_by: The property used to sort results. The default value is CreationTime. sort_order: The sort order. The default value is Descending. - max_results: The maximum number of trials to return in the response. The default value is 10. - next_token: If the previous call to ListTrials didn't return the full set of trials, the call returns a token for getting the next set of trials. session: Boto3 session. region: Region name. @@ -34050,8 +34120,6 @@ def get_all( created_before: A filter that returns only components created before the specified time. sort_by: The property used to sort results. The default value is CreationTime. sort_order: The sort order. The default value is Descending. - max_results: The maximum number of components to return in the response. The default value is 10. - next_token: If the previous call to ListTrialComponents didn't return the full set of components, the call returns a token for getting the next set of components. session: Boto3 session. region: Region name. @@ -34775,8 +34843,6 @@ def get_all( Get all UserProfile resources Parameters: - next_token: If the previous response was truncated, you will receive this token. Use it in your next request to receive the next set of results. - max_results: This parameter defines the maximum number of results that can be return in a single response. The MaxResults parameter is an upper bound, not a target. If there are more results available than the value specified, a NextToken is provided in the response. The NextToken indicates that the user should get the next set of results by providing this token as a part of a subsequent call. The default value for MaxResults is 10. sort_order: The sort order for the results. The default is Ascending. sort_by: The parameter by which to sort the results. The default is CreationTime. domain_id_equals: A parameter by which to filter the results. @@ -35260,8 +35326,6 @@ def get_all( sort_by: Sort workforces using the workforce name or creation date. sort_order: Sort workforces in ascending or descending order. name_contains: A filter you can use to search for workforces using part of the workforce name. - next_token: A token to resume pagination. - max_results: The maximum number of workforces returned in the response. session: Boto3 session. region: Region name. @@ -35604,8 +35668,6 @@ def get_all( sort_by: The field to sort results by. The default is CreationTime. sort_order: The sort order for results. The default is Ascending. name_contains: A string in the work team's name. This filter returns only work teams whose name contains the specified string. - next_token: If the result of the previous ListWorkteams request was truncated, the response includes a NextToken. To retrieve the next set of labeling jobs, use the token in the next request. - max_results: The maximum number of work teams to return in each page of the response. session: Boto3 session. region: Region name. @@ -35666,8 +35728,6 @@ def get_all_labeling_jobs( Parameters: workteam_arn: The Amazon Resource Name (ARN) of the work team for which you want to see labeling jobs for. - max_results: The maximum number of labeling jobs to return in each page of the response. - next_token: If the result of the previous ListLabelingJobsForWorkteam request was truncated, the response includes a NextToken. To retrieve the next set of labeling jobs, use the token in the next request. creation_time_after: A filter that returns only labeling jobs created after the specified time (timestamp). creation_time_before: A filter that returns only labeling jobs created before the specified time (timestamp). job_reference_code_contains: A filter the limits jobs to only the ones whose job reference code contains the specified string. diff --git a/sagemaker-core/src/sagemaker/core/s3/__init__.py b/sagemaker-core/src/sagemaker/core/s3/__init__.py index 6343b79cf6..700674529c 100644 --- a/sagemaker-core/src/sagemaker/core/s3/__init__.py +++ b/sagemaker-core/src/sagemaker/core/s3/__init__.py @@ -18,6 +18,7 @@ from sagemaker.core.s3.client import ( # noqa: F401 S3Uploader, S3Downloader, + S3ObjectHead, parse_s3_url, is_s3_url, s3_path_join, @@ -35,6 +36,7 @@ __all__ = [ "S3Uploader", "S3Downloader", + "S3ObjectHead", "parse_s3_url", "is_s3_url", "s3_path_join", diff --git a/sagemaker-core/src/sagemaker/core/s3/client.py b/sagemaker-core/src/sagemaker/core/s3/client.py index 15a9a8bb5f..303c143eee 100644 --- a/sagemaker-core/src/sagemaker/core/s3/client.py +++ b/sagemaker-core/src/sagemaker/core/s3/client.py @@ -17,9 +17,10 @@ import logging import io -from typing import Union +from dataclasses import dataclass +from datetime import datetime from functools import reduce -from typing import Optional +from typing import Iterator, Optional, Tuple, Union from six.moves.urllib.parse import urlparse from sagemaker.core.helper.session_helper import Session @@ -27,6 +28,64 @@ logger = logging.getLogger("sagemaker") +@dataclass(frozen=True) +class S3ObjectHead: + """Object metadata from ``HeadObject``. + + Attributes: + content_length (int): Object size in bytes. + last_modified (datetime): Last modification time. + etag (str): The entity tag as S3 returns it, quotes included. + content_type (Optional[str]): The MIME type, or ``None`` when S3 returns none. + """ + + content_length: int + last_modified: datetime + etag: str + content_type: Optional[str] + + +def _s3_client(sagemaker_session): + """Return the session's S3 client, or create one from its boto session. + + ``Session.s3_client`` is ``None`` for a session built without ``_initialize``. This is + the same fallback ``Session.read_s3_file`` uses. + """ + if sagemaker_session.s3_client is not None: + return sagemaker_session.s3_client + return sagemaker_session.boto_session.client( + "s3", region_name=sagemaker_session.boto_region_name + ) + + +def _expected_owner_args(sagemaker_session, bucket) -> dict: + """Return ``{"ExpectedBucketOwner": }`` for the SDK-generated default bucket. + + Spot check: enforce ownership only on the session's default bucket. Cross-account reads + (JumpStart, shared buckets) stay untouched. + """ + expected_owner = sagemaker_session._get_account_id_if_default_bucket(bucket) + if expected_owner: + return {"ExpectedBucketOwner": expected_owner} + return {} + + +def _range_header(byte_range: Tuple[int, int]) -> str: + """Return the HTTP ``Range`` header for an inclusive ``(start, end)`` byte range.""" + start, end = byte_range + if start < 0 or end < start: + raise ValueError( + "byte_range must satisfy 0 <= start <= end, got ({}, {}).".format(start, end) + ) + return "bytes={}-{}".format(start, end) + + +def _decode_lines(body, encoding: str) -> Iterator[str]: + """Yield the decoded lines of a streaming ``GetObject`` body without their line ends.""" + for raw_line in body.iter_lines(): + yield raw_line.decode(encoding) + + class S3Uploader(object): """Contains static methods for uploading directories or files to S3.""" @@ -184,7 +243,9 @@ def read_file(s3_uri, sagemaker_session=None) -> str: return sagemaker_session.read_s3_file(bucket=bucket, key_prefix=object_key) @staticmethod - def read_bytes(s3_uri, sagemaker_session=None) -> bytes: + def read_bytes( + s3_uri, sagemaker_session=None, byte_range: Optional[Tuple[int, int]] = None + ) -> bytes: """Static method that returns the contents of a s3 object as bytes. Args: @@ -193,28 +254,82 @@ def read_bytes(s3_uri, sagemaker_session=None) -> bytes: manages interactions with Amazon SageMaker APIs and any other AWS services needed. If not specified, one is created using the default AWS configuration chain. + byte_range (Optional[Tuple[int, int]]): Inclusive ``(start, end)`` offsets, as in + an HTTP ``Range`` header. ``(0, 1023)`` reads the first 1024 bytes through + ``GetObject``. ``None`` downloads the whole object. Returns: - bytes: The body of the file. + bytes: The body of the file, or the requested part of it. + + Raises: + ValueError: If ``start`` is negative or ``end`` is less than ``start``. """ sagemaker_session = sagemaker_session or Session() - bucket, object_key = parse_s3_url(s3_uri) + extra_args = _expected_owner_args(sagemaker_session, bucket) - # Spot check: enforce ownership only when reading from the session's default - # bucket. Cross-account reads (e.g. JumpStart) are left untouched. - extra_args = None - expected_owner = sagemaker_session._get_account_id_if_default_bucket(bucket) - if expected_owner: - extra_args = {"ExpectedBucketOwner": expected_owner} + if byte_range is not None: + range_header = _range_header(byte_range) + response = _s3_client(sagemaker_session).get_object( + Bucket=bucket, Key=object_key, Range=range_header, **extra_args + ) + return response["Body"].read() bytes_io = io.BytesIO() sagemaker_session.s3_resource.Bucket(bucket).download_fileobj( - object_key, bytes_io, ExtraArgs=extra_args + object_key, bytes_io, ExtraArgs=extra_args or None ) bytes_io.seek(0) return bytes_io.read() + @staticmethod + def iter_lines(s3_uri, sagemaker_session=None, encoding: str = "utf-8") -> Iterator[str]: + """Static method that streams the lines of a text object one at a time. + + The body streams through ``GetObject``. The whole object is never held in memory. + + Args: + s3_uri (str): An S3 uri that refers to a s3 object. + sagemaker_session (sagemaker.core.helper.session_helper.Session): Session object which + manages interactions with Amazon SageMaker APIs and other AWS services. If not + specified, one is created from the default AWS configuration chain. + encoding (str): The text encoding of the object. Defaults to ``utf-8``. + + Returns: + Iterator[str]: The lines of the object without their line ends. + """ + sagemaker_session = sagemaker_session or Session() + bucket, object_key = parse_s3_url(s3_uri) + response = _s3_client(sagemaker_session).get_object( + Bucket=bucket, Key=object_key, **_expected_owner_args(sagemaker_session, bucket) + ) + return _decode_lines(response["Body"], encoding) + + @staticmethod + def head(s3_uri, sagemaker_session=None) -> S3ObjectHead: + """Static method that returns the metadata of a s3 object through ``HeadObject``. + + Args: + s3_uri (str): An S3 uri that refers to a s3 object. + sagemaker_session (sagemaker.core.helper.session_helper.Session): Session object which + manages interactions with Amazon SageMaker APIs and other AWS services. If not + specified, one is created from the default AWS configuration chain. + + Returns: + S3ObjectHead: Size, last modification time, entity tag, and content type. + """ + sagemaker_session = sagemaker_session or Session() + bucket, object_key = parse_s3_url(s3_uri) + response = _s3_client(sagemaker_session).head_object( + Bucket=bucket, Key=object_key, **_expected_owner_args(sagemaker_session, bucket) + ) + return S3ObjectHead( + content_length=response["ContentLength"], + last_modified=response["LastModified"], + etag=response["ETag"], + content_type=response.get("ContentType"), + ) + @staticmethod def list(s3_uri, sagemaker_session=None): """Static method that lists the contents of an S3 uri. diff --git a/sagemaker-core/src/sagemaker/core/tools/constants.py b/sagemaker-core/src/sagemaker/core/tools/constants.py index 6d5580820c..fe07f0f0a4 100644 --- a/sagemaker-core/src/sagemaker/core/tools/constants.py +++ b/sagemaker-core/src/sagemaker/core/tools/constants.py @@ -78,8 +78,7 @@ # language governing permissions and limitations under the License. """ -LOGGER_STRING = """ -logger = get_textual_rich_logger(__name__) +LOGGER_STRING = """logger = get_textual_rich_logger(__name__) """ diff --git a/sagemaker-core/src/sagemaker/core/tools/resources_codegen.py b/sagemaker-core/src/sagemaker/core/tools/resources_codegen.py index 22428f3006..a59cd09ba4 100644 --- a/sagemaker-core/src/sagemaker/core/tools/resources_codegen.py +++ b/sagemaker-core/src/sagemaker/core/tools/resources_codegen.py @@ -71,6 +71,7 @@ CREATE_METHOD_TEMPLATE_WITHOUT_DEFAULTS, IMPORT_METHOD_TEMPLATE, FAILED_STATUS_ERROR_TEMPLATE, + STOPPED_STATUS_ERROR_TEMPLATE, GET_NAME_METHOD_TEMPLATE, GET_ALL_METHOD_NO_ARGS_TEMPLATE, GET_ALL_METHOD_WITH_ARGS_TEMPLATE, @@ -106,6 +107,64 @@ "ModelPackage": {"ModelPackageName": "model_package_arn"}, } +# Some List APIs return a summary member that the Describe API does not. A resource class +# is generated from the Describe output, so ResourceIterator drops such members when it +# builds the resource from the summary (`extra="forbid"` leaves no other option). This +# maps a resource to the summary members it keeps as optional attributes. ListHubContents +# returns OriginalCreationTime, the model-level catalog date, while DescribeHubContent +# returns only the version-level CreationTime. +LIST_SUMMARY_EXTRA_ATTRIBUTES = { + "HubContent": {"OriginalCreationTime": "Optional[datetime.datetime] = Unassigned()"}, +} + +# Hand-written behavior for a generated resource class. The generator emits +# ``class (, Base)`` so the mixin's properties and methods sit beside the +# generated API methods on the one public class, and the hand-written code stays in its +# own module where it can be tested without regenerating resources.py. +RESOURCE_MIXINS = { + "Endpoint": ("sagemaker.core.endpoint", "EndpointMixin"), + "HubContent": ("sagemaker.core.hub_content", "HubContentMixin"), + "ModelPackageGroup": ("sagemaker.core.model_package_group", "ModelPackageGroupMixin"), +} + +# Some List APIs take an input member that the resource's Describe input requires but the +# list summary does not return (ListHubContents takes HubName, HubContentInfo omits it). +# get_all() cannot build a resource from such a summary unless the list argument is +# copied onto every resource. This maps a resource to the list-input members to copy, as +# {summary-less member: resource attribute}. +GET_ALL_INPUT_ATTRIBUTES = { + "HubContent": {"HubName": "hub_name"}, +} + +# ResourceIterator describes every listed resource before it returns it, one call per +# resource. A list API whose catalog runs to hundreds of items (ListHubContents returns +# the public hub, about 770 models) takes minutes one describe at a time, so the +# generated list methods of these resources take ``max_workers`` and hand it to the +# iterator, which keeps that many describes in flight. The value is the default worker +# count: 4 describes of the public hub run 1.7 to 2.8 times faster than 1, and 8 run +# into DescribeHubContent throttling and gain nothing more. +CONCURRENT_DESCRIBE_RESOURCES = {"HubContent": 4} +MAX_WORKERS_ARG_TEMPLATE = "max_workers: int = {default}," +MAX_WORKERS_DOCSTRING_TEMPLATE = ( + "max_workers: How many listed resources to describe at a time, on a thread pool. " + "Defaults to {default}. Keep it at or below the max_pool_connections of the client " + "(10 by default).\n" +) + + +def _max_workers_code(resource_name: str): + """Return the signature line and the docstring line of ``max_workers`` for a resource. + + The two are empty for a resource that ``CONCURRENT_DESCRIBE_RESOURCES`` does not name. + """ + default = CONCURRENT_DESCRIBE_RESOURCES.get(resource_name) + if default is None: + return "", "" + return ( + add_indent(MAX_WORKERS_ARG_TEMPLATE.format(default=default), 4), + add_indent(MAX_WORKERS_DOCSTRING_TEMPLATE.format(default=default)), + ) + class ResourcesCodeGen: """A class for generating resources based on a service JSON file. @@ -208,6 +267,9 @@ def generate_imports(self) -> str: "from sagemaker.core.serializers.base import BaseSerializer", "from sagemaker.core.deserializers.base import BaseDeserializer", ] + imports.extend( + f"from {module} import {mixin}" for module, mixin in sorted(RESOURCE_MIXINS.values()) + ) formated_imports = "\n".join(imports) formated_imports += "\n\n" @@ -361,7 +423,7 @@ def generate_resource_class( ) # Start defining the class - resource_class = f"class {resource_name}(Base):\n" + resource_class = f"class {resource_name}({self._resource_bases(resource_name)}):\n" class_documentation_string = f"Class representing resource {resource_name}\n\n" class_documentation_string += "Attributes:\n" @@ -461,7 +523,7 @@ def generate_resource_class( else: # If there's no 'get' or 'list' or 'create' method, generate a class with no attributes resource_attributes = [] - resource_class = f"class {resource_name}(Base):\n" + resource_class = f"class {resource_name}({self._resource_bases(resource_name)}):\n" class_documentation_string = f"Class representing resource {resource_name}\n" resource_class += add_indent(f'"""\n{class_documentation_string}\n"""\n', 4) @@ -535,8 +597,15 @@ def _get_class_attributes(self, resource_name: str, class_methods: list) -> tupl class_attributes["hub_name"] = "Optional[str] = Unassigned()" class_attributes_string = class_attributes_string.replace("hub_name: str", "") class_attributes_string = ( - class_attributes_string + "hub_name: Optional[str] = Unassigned()" + class_attributes_string + "hub_name: Optional[str] = Unassigned()\n" ) + for member, annotation in LIST_SUMMARY_EXTRA_ATTRIBUTES.get(resource_name, {}).items(): + attribute = convert_to_snake_case(member) + class_attributes[attribute] = annotation + class_attributes_string += f"{attribute}: {annotation}\n" + attributes_and_documentation[member] = self.shapes_extractor.combined_shapes[ + self._list_summary_shape(resource_name) + ]["members"][member].get("documentation") if resource_name == "Endpoint": class_attributes["serializer"] = "Optional[BaseSerializer] = None" class_attributes_string = class_attributes_string.replace( @@ -1026,6 +1095,7 @@ def _generate_docstring( return_string: str = None, include_intelligent_defaults_errors: bool = False, exclude_resource_attrs: list = None, + extra_parameters_docstring: str = "", ) -> str: """Generate the docstring for a method of a resource. @@ -1039,6 +1109,8 @@ def _generate_docstring( return_string (str): The return string. include_intelligent_defaults_errors (bool): Whether to include intelligent defaults errors. exclude_resource_attrs (list): A list of attributes to exclude from the docstring. + extra_parameters_docstring (str): Parameter lines to add after ``session`` and + ``region``, for a method argument that is not an operation input member. Returns: str: The generated docstring for the IMPORT method. @@ -1060,6 +1132,7 @@ def _generate_docstring( if not _shape_attr_documentation_string: docstring += "\nParameters:\n" docstring += add_indent("session: Boto3 session.\nregion: Region name.\n") + docstring += extra_parameters_docstring if include_return_resource_docstring: docstring += f"\nReturns:\n" f" The {resource_name} resource.\n" @@ -1714,6 +1787,26 @@ def generate_additional_get_all_method(self, method: Method, resource_attributes "list_method_kwargs=operation_input_args", ] + # An object method lists resources of its own type under its own identifiers + # (ListHubContentVersions lists the versions of one HubContent). The summaries + # omit the identifiers the list call was made with, so copy them onto every + # resource or the first refresh() fails with a missing required parameter. + if ( + method.method_type == MethodType.OBJECT.value + and method.return_type == method.resource_name + ): + input_attributes = GET_ALL_INPUT_ATTRIBUTES.get(method.resource_name, {}) + if input_attributes: + extra_init = ", ".join( + f"'{attribute}': self.{attribute}" for attribute in input_attributes.values() + ) + resource_iterator_args_list.append(f"extra_init_data={{{extra_init}}}") + + extra_args, extra_parameters_docstring = _max_workers_code(method.return_type) + if extra_args: + method_args += "\n" + extra_args + resource_iterator_args_list.append("max_workers=max_workers") + resource_iterator_args = ",\n".join(resource_iterator_args_list) resource_iterator_args = add_indent(resource_iterator_args, 8) serialize_operation_input = SERIALIZE_INPUT_TEMPLATE.format( @@ -1734,7 +1827,8 @@ def generate_additional_get_all_method(self, method: Method, resource_attributes operation_input_shape_name=operation_input_shape_name, include_session_region=True, return_string=return_string, - exclude_resource_attrs=exclude_resource_attrs, + exclude_resource_attrs=exclude_list + (exclude_resource_attrs or []), + extra_parameters_docstring=extra_parameters_docstring, ) return GENERIC_METHOD_TEMPLATE.format( @@ -1829,6 +1923,10 @@ def generate_wait_method(self, resource_name: str) -> str: resource_name=resource_name, reason=failure_reason ) formatted_failed_block = add_indent(formatted_failed_block, 16) + formatted_stopped_block = STOPPED_STATUS_ERROR_TEMPLATE.format( + resource_name=resource_name, reason=failure_reason + ) + formatted_stopped_block = add_indent(formatted_stopped_block, 16) logs_arg = "" logs_arg_doc = "" @@ -1851,6 +1949,7 @@ def generate_wait_method(self, resource_name: str) -> str: terminal_resource_states=terminal_resource_states, status_key_path=status_key_path, failed_error_block=formatted_failed_block, + stopped_error_block=formatted_stopped_block, resource_name=resource_name, logs_arg=logs_arg, logs_arg_doc=logs_arg_doc, @@ -1934,6 +2033,20 @@ def generate_wait_for_delete_method(self, resource_name: str) -> str: ) return formatted_method + @staticmethod + def _resource_bases(resource_name: str) -> str: + """Return the base-class list of a generated resource class.""" + mixin = RESOURCE_MIXINS.get(resource_name) + return f"{mixin[1]}, Base" if mixin else "Base" + + def _list_summary_shape(self, resource_name: str) -> str: + """Return the summary shape name of the resource's List API output.""" + operation_name = "List" + resource_name + "s" + output_shape = self.operations[operation_name]["output"]["shape"] + output_members = self.shapes[output_shape]["members"] + summaries_key = next(key for key in output_members if key != "NextToken") + return self.shapes[output_members[summaries_key]["shape"]]["member"]["shape"] + def generate_get_all_method(self, resource_name: str) -> str: """Auto-Generate 'get_all' class Method [list API] for a resource. @@ -1974,7 +2087,13 @@ def generate_get_all_method(self, resource_name: str) -> str: get_operation_required_input = [] custom_key_mapping_str = "" - if any(member not in summary_members for member in get_operation_required_input): + input_attributes = GET_ALL_INPUT_ATTRIBUTES.get(resource_name, {}) + missing_members = [ + member + for member in get_operation_required_input + if member not in summary_members and member not in input_attributes + ] + if missing_members: if "MonitoringJobDefinitionSummary" == summary_name: custom_key_mapping = { "monitoring_job_definition_name": "job_definition_name", @@ -2002,6 +2121,18 @@ def generate_get_all_method(self, resource_name: str) -> str: if custom_key_mapping_str: resource_iterator_args_list.append("custom_key_mapping=custom_key_mapping") + if input_attributes: + extra_init = ", ".join( + f"'{attribute}': {convert_to_snake_case(member)}" + for member, attribute in input_attributes.items() + ) + resource_iterator_args_list.append(f"extra_init_data={{{extra_init}}}") + + extra_args, extra_parameters_docstring = _max_workers_code(resource_name) + if extra_args: + extra_args += "\n" + resource_iterator_args_list.append("max_workers=max_workers") + exclude_list = ["next_token", "max_results"] get_all_args = self._generate_method_args(operation_input_shape_name, exclude_list) @@ -2015,6 +2146,8 @@ def generate_get_all_method(self, resource_name: str) -> str: operation=operation, custom_key_mapping=custom_key_mapping_str, resource_iterator_args=resource_iterator_args, + extra_args=extra_args, + extra_parameters_docstring=extra_parameters_docstring, ) return formatted_method @@ -2026,7 +2159,8 @@ def generate_get_all_method(self, resource_name: str) -> str: resource_iterator_args = ",\n".join(resource_iterator_args_list) resource_iterator_args = add_indent(resource_iterator_args, 8) - # generate docstring + # generate docstring. The paging members are not parameters: ResourceIterator + # pages for the caller, so they stay out of the docstring too. docstring = self._generate_docstring( title=f"Get all {resource_name} resources", operation_name=operation_name, @@ -2035,6 +2169,8 @@ def generate_get_all_method(self, resource_name: str) -> str: include_session_region=True, include_return_resource_docstring=False, return_string=f"Returns:\n" f" Iterator for listed {resource_name} resources.\n", + exclude_resource_attrs=exclude_list, + extra_parameters_docstring=extra_parameters_docstring, ) formatted_method = GET_ALL_METHOD_WITH_ARGS_TEMPLATE.format( @@ -2045,6 +2181,7 @@ def generate_get_all_method(self, resource_name: str) -> str: operation_input_args=operation_input_args, custom_key_mapping=custom_key_mapping_str, resource_iterator_args=resource_iterator_args, + extra_args=extra_args, ) return formatted_method @@ -2115,6 +2252,7 @@ def generate_config_schema(self) -> str: with open(output, "w") as file: # Generate and write the license to the file file.write( + '"""JSON schema of the SageMaker Python SDK configuration file."""\n\n' f"SAGEMAKER_PYTHON_SDK_CONFIG_SCHEMA = {json.dumps(combined_config_schema, indent=4)}" ) diff --git a/sagemaker-core/src/sagemaker/core/tools/shapes_extractor.py b/sagemaker-core/src/sagemaker/core/tools/shapes_extractor.py index b92e50524c..53ccefc547 100644 --- a/sagemaker-core/src/sagemaker/core/tools/shapes_extractor.py +++ b/sagemaker-core/src/sagemaker/core/tools/shapes_extractor.py @@ -43,6 +43,9 @@ def __init__(self, combined_shapes: Optional[dict] = None): self.shape_dag = self.get_shapes_dag() with open(SHAPE_DAG_FILE_PATH, "w") as f: + f.write( + '"""Generated dependency graph of the service shapes, used for serialization."""\n\n' + ) f.write("SHAPE_DAG=") f.write(textwrap.indent(pprint.pformat(self.shape_dag, width=1), "") + "\n") reformat_file_with_black(SHAPE_DAG_FILE_PATH) diff --git a/sagemaker-core/src/sagemaker/core/tools/templates.py b/sagemaker-core/src/sagemaker/core/tools/templates.py index 981a9fefc7..a82c1bad68 100644 --- a/sagemaker-core/src/sagemaker/core/tools/templates.py +++ b/sagemaker-core/src/sagemaker/core/tools/templates.py @@ -272,6 +272,11 @@ def refresh( raise FailedStatusError(resource_type="{resource_name}", status=current_status, reason={reason}) """ +STOPPED_STATUS_ERROR_TEMPLATE = """ +if raise_on_stopped and "stopped" in current_status.lower(): + raise FailedStatusError(resource_type="{resource_name}", status=current_status, reason={reason}) +""" + INIT_WAIT_LOGS_TEMPLATE = """ instance_count = {get_instance_count} if logs: @@ -303,6 +308,7 @@ def wait( self, poll: int = 5, timeout: Optional[int] = None, + raise_on_stopped: bool = False, {logs_arg} ) -> None: """ @@ -311,10 +317,11 @@ def wait( Parameters: poll: The number of seconds to wait between each poll. timeout: The maximum number of seconds to wait before timing out. + raise_on_stopped: Whether to raise FailedStatusError when the resource stops instead of completing. {logs_arg_doc} Raises: TimeoutExceededError: If the resource does not reach a terminal state before the timeout. - FailedStatusError: If the resource reaches a failed state. + FailedStatusError: If the resource reaches a failed state, or a stopped state when raise_on_stopped is True. WaiterError: Raised when an error occurs while waiting. """ @@ -346,6 +353,7 @@ def wait( if current_status in terminal_states: logger.info(f"Final Resource Status: [bold]{{current_status}}") {failed_error_block} +{stopped_error_block} return if timeout is not None and time.time() - start_time >= timeout: @@ -524,7 +532,7 @@ def get_all( {get_all_args} session: Optional[Session] = None, region: Optional[StrPipeVar] = None, -) -> ResourceIterator["{resource}"]: +{extra_args}) -> ResourceIterator["{resource}"]: {docstring} client = Base.get_sagemaker_client(session=session, region_name=region, service_name="{service_name}") @@ -548,14 +556,14 @@ def get_all( cls, session: Optional[Session] = None, region: Optional[StrPipeVar] = None, -) -> ResourceIterator["{resource}"]: +{extra_args}) -> ResourceIterator["{resource}"]: """ Get all {resource} resources. Parameters: session: Boto3 session. region: Region name. - +{extra_parameters_docstring} Returns: Iterator for listed {resource} resources. diff --git a/sagemaker-core/src/sagemaker/core/utilities/cache.py b/sagemaker-core/src/sagemaker/core/utilities/cache.py index 9c138845d2..9eedced59a 100644 --- a/sagemaker-core/src/sagemaker/core/utilities/cache.py +++ b/sagemaker-core/src/sagemaker/core/utilities/cache.py @@ -128,7 +128,7 @@ def put(self, key: KeyType, value: Optional[ValType] = None) -> None: if value is None: value = self._retrieval_function( # type: ignore - key=key, value=curr_value.element if curr_value else None + key=key, value=curr_value.value if curr_value else None ) self._lru_cache[key] = self.Element( diff --git a/sagemaker-core/src/sagemaker/core/utils/code_injection/shape_dag.py b/sagemaker-core/src/sagemaker/core/utils/code_injection/shape_dag.py index ff5f80e6be..65db26511d 100644 --- a/sagemaker-core/src/sagemaker/core/utils/code_injection/shape_dag.py +++ b/sagemaker-core/src/sagemaker/core/utils/code_injection/shape_dag.py @@ -1,3 +1,5 @@ +"""Generated dependency graph of the service shapes, used for serialization.""" + SHAPE_DAG = { "AIAdapterModelPackageEntry": { "members": [ diff --git a/sagemaker-core/src/sagemaker/core/utils/utils.py b/sagemaker-core/src/sagemaker/core/utils/utils.py index dc6b96378d..33decfba20 100644 --- a/sagemaker-core/src/sagemaker/core/utils/utils.py +++ b/sagemaker-core/src/sagemaker/core/utils/utils.py @@ -18,7 +18,8 @@ import re import subprocess import threading -from collections import OrderedDict +from collections import OrderedDict, deque +from concurrent.futures import ThreadPoolExecutor from boto3.session import Session from botocore.config import Config @@ -533,7 +534,14 @@ def _retries_configured_outside_code(session: Session) -> bool: class ResourceIterator(Generic[T]): - """ResourceIterator class to iterate over a list of resources.""" + """ResourceIterator class to iterate over a list of resources. + + Each listed resource is described (``refresh()``) before it is returned, so a list of + N resources costs one list call per page plus N describe calls. With ``max_workers`` + above 1 the describes of the next ``max_workers`` resources run on a thread pool while + the caller consumes the current one. The resources still come back in list order, and + the list calls stay on the thread of the caller. + """ def __init__( self, @@ -546,6 +554,8 @@ def __init__( custom_key_mapping: dict = None, session=None, region: Optional[str] = None, + extra_init_data: dict = None, + max_workers: int = 1, ): """Initialize a ResourceIterator object @@ -563,7 +573,22 @@ def __init__( Each returned resource is bound to it, so its refresh and other object methods use the same account. Defaults to None (the process default client). region (str, optional): The region the list call was made in. Defaults to None. + extra_init_data (dict, optional): Attributes set on every resource in addition to + the summary, for a list argument the summary does not echo back (for example + the ``hub_name`` a ``ListHubContents`` call was made with). Defaults to None. + max_workers (int, optional): How many listed resources to describe at a time. + 1 describes each resource when the caller asks for it. A higher value keeps + up to ``max_workers`` ``refresh()`` calls in flight on a thread pool, so a + caller that stops early makes at most ``max_workers - 1`` describe calls + beyond the resources it consumed. Keep it at or below the + ``max_pool_connections`` of the client (10 by default), or urllib3 discards + connections. Defaults to 1. + + Raises: + ValueError: If ``max_workers`` is less than 1. """ + if max_workers < 1: + raise ValueError(f"max_workers must be at least 1, got {max_workers}.") self.summaries_key = summaries_key self.summary_name = summary_name self.client = client @@ -572,86 +597,128 @@ def __init__( self.custom_key_mapping = custom_key_mapping self.session = session self.region = region + self.extra_init_data = extra_init_data or {} + self.max_workers = max_workers self.resource_cls = resource_cls self.index = 0 self.summary_list = [] self.next_token = None + self._exhausted = False + self._pending = deque() + self._executor: Optional[ThreadPoolExecutor] = None def __iter__(self): """Return the iterator object.""" return self def __next__(self) -> T: + """Return the next item from the iterator.""" + if self.max_workers == 1: + return self._load(self._next_summary()) + while len(self._pending) < self.max_workers: + try: + summary = self._next_summary() + except StopIteration: + break + self._pending.append(self._get_executor().submit(self._load, summary)) + if not self._pending: + self._shutdown_executor() + raise StopIteration + return self._pending.popleft().result() + + def _get_executor(self) -> ThreadPoolExecutor: + """Return the describe pool, creating it on the first concurrent describe.""" + if self._executor is None: + self._executor = ThreadPoolExecutor( + max_workers=self.max_workers, thread_name_prefix="sagemaker-resource-refresh" + ) + return self._executor + + def _shutdown_executor(self): + """Release the describe pool once the iteration is over.""" + if self._executor is not None: + self._executor.shutdown(wait=False) + self._executor = None + + def _load(self, summary: dict) -> T: + """Build the resource for one summary, bind it to the session, and describe it.""" + # Initialize the resource object + if is_primitive_class(self.resource_cls): + # If the resource class is a primitive class, there will be only one element in the summary + return list(summary.values())[0] + + # Transform the resource summary into format to initialize object + init_data = transform(summary, self.summary_name) + + if self.custom_key_mapping: + init_data = {self.custom_key_mapping.get(k, k): v for k, v in init_data.items()} + + # Filter out the fields that are not in the resource class + fields = self.resource_cls.__annotations__ + init_data = {k: v for k, v in init_data.items() if k in fields} + for key, value in self.extra_init_data.items(): + init_data.setdefault(key, value) + + resource_object = self.resource_cls(**init_data) + + # Bind the resource to the session/region it was listed with + if hasattr(resource_object, "_set_client_context"): + resource_object._set_client_context( # pylint: disable=protected-access + session=self.session, region=self.region + ) + + # If the resource object has refresh method, refresh and return it + if hasattr(resource_object, "refresh"): + resource_object.refresh() + return resource_object + + def _next_summary(self) -> dict: + """Return the next list summary, fetching the next page when the current one is spent.""" + if self._exhausted: + raise StopIteration # If there are summaries in the summary_list, return the next summary - """Return the next item from the iterator.""" if len(self.summary_list) > 0 and self.index < len(self.summary_list): # Get the next summary from the resource summary_list summary = self.summary_list[self.index] self.index += 1 - - # Initialize the resource object - if is_primitive_class(self.resource_cls): - # If the resource class is a primitive class, there will be only one element in the summary - resource_object = list(summary.values())[0] - else: - # Transform the resource summary into format to initialize object - init_data = transform(summary, self.summary_name) - - if self.custom_key_mapping: - init_data = {self.custom_key_mapping.get(k, k): v for k, v in init_data.items()} - - # Filter out the fields that are not in the resource class - fields = self.resource_cls.__annotations__ - init_data = {k: v for k, v in init_data.items() if k in fields} - - resource_object = self.resource_cls(**init_data) - - # Bind the resource to the session/region it was listed with - if hasattr(resource_object, "_set_client_context"): - resource_object._set_client_context( # pylint: disable=protected-access - session=self.session, region=self.region - ) - - # If the resource object has refresh method, refresh and return it - if hasattr(resource_object, "refresh"): - resource_object.refresh() - return resource_object + return summary # If index reached the end of summary_list, and there is no next token, raise StopIteration - elif ( + if ( len(self.summary_list) > 0 and self.index >= len(self.summary_list) and (not self.next_token) ): + self._exhausted = True raise StopIteration # Otherwise, get the next page of summaries by calling the list method with the next token if available - else: - while True: - previous_token = self.next_token - if self.next_token: - response = getattr(self.client, self.list_method)( - NextToken=self.next_token, **self.list_method_kwargs - ) - else: - response = getattr(self.client, self.list_method)(**self.list_method_kwargs) - - self.summary_list = response.get(self.summaries_key, []) - self.next_token = response.get("NextToken", None) - self.index = 0 - - if len(self.summary_list) > 0: - return self.__next__() - - # A page can be empty and still carry a NextToken (the service - # filters results after paginating), so an empty page only ends - # the iteration when there is no further page. A NextToken that - # comes back unchanged would page forever, so it also ends the - # iteration (the same consecutive-repeat check botocore paginators do). - if not self.next_token or self.next_token == previous_token: - raise StopIteration + while True: + previous_token = self.next_token + if self.next_token: + response = getattr(self.client, self.list_method)( + NextToken=self.next_token, **self.list_method_kwargs + ) + else: + response = getattr(self.client, self.list_method)(**self.list_method_kwargs) + + self.summary_list = response.get(self.summaries_key, []) + self.next_token = response.get("NextToken", None) + self.index = 0 + + if len(self.summary_list) > 0: + return self._next_summary() + + # A page can be empty and still carry a NextToken (the service + # filters results after paginating), so an empty page only ends + # the iteration when there is no further page. A NextToken that + # comes back unchanged would page forever, so it also ends the + # iteration (the same consecutive-repeat check botocore paginators do). + if not self.next_token or self.next_token == previous_token: + self._exhausted = True + raise StopIteration def serialize(value: Any) -> Any: diff --git a/sagemaker-core/tests/unit/s3/test_client_read_primitives.py b/sagemaker-core/tests/unit/s3/test_client_read_primitives.py new file mode 100644 index 0000000000..5412a84cee --- /dev/null +++ b/sagemaker-core/tests/unit/s3/test_client_read_primitives.py @@ -0,0 +1,184 @@ +# 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. +"""``S3Downloader`` reads a byte range, streams lines, and returns ``HeadObject`` metadata. + +Before this change ``S3Downloader.read_bytes`` had no ``byte_range`` parameter and downloaded +the whole object, and the class had no ``iter_lines`` and no ``head``. Fifteen SageMaker agent +skill scripts split ``s3://`` URIs by hand, eight called ``get_object(Range=...)``, and four +called ``head_object`` directly, each without the ``ExpectedBucketOwner`` spot check. +""" + +from __future__ import absolute_import + +from datetime import datetime, timezone + +import pytest +from mock import Mock + +from sagemaker.core.s3 import client as s3 + +REGION = "us-west-2" +S3_URI = "s3://mybucket/prefix/data.jsonl" +ACCOUNT_ID = "123456789012" + + +@pytest.fixture() +def sagemaker_session(): + boto_mock = Mock(name="boto_session", region_name=REGION) + session_mock = Mock( + name="sagemaker_session", + boto_session=boto_mock, + boto_region_name=REGION, + config=None, + local_mode=False, + default_bucket_prefix=None, + ) + session_mock.s3_client = Mock(name="s3_client") + session_mock._get_account_id_if_default_bucket = Mock(return_value=None) + return session_mock + + +def _get_object_response(payload: bytes): + body = Mock(name="Body") + body.read = Mock(return_value=payload) + body.iter_lines = Mock(return_value=iter(payload.splitlines())) + return {"Body": body} + + +def test_read_bytes_with_range_sends_inclusive_range_header(sagemaker_session): + sagemaker_session.s3_client.get_object.return_value = _get_object_response(b"abc") + + result = s3.S3Downloader.read_bytes(S3_URI, sagemaker_session, byte_range=(0, 2)) + + assert result == b"abc" + sagemaker_session.s3_client.get_object.assert_called_once_with( + Bucket="mybucket", Key="prefix/data.jsonl", Range="bytes=0-2" + ) + sagemaker_session.s3_resource.Bucket.assert_not_called() + + +def test_read_bytes_with_range_passes_expected_bucket_owner(sagemaker_session): + sagemaker_session._get_account_id_if_default_bucket.return_value = ACCOUNT_ID + sagemaker_session.s3_client.get_object.return_value = _get_object_response(b"x") + + s3.S3Downloader.read_bytes(S3_URI, sagemaker_session, byte_range=(10, 1033)) + + sagemaker_session._get_account_id_if_default_bucket.assert_called_once_with("mybucket") + sagemaker_session.s3_client.get_object.assert_called_once_with( + Bucket="mybucket", + Key="prefix/data.jsonl", + Range="bytes=10-1033", + ExpectedBucketOwner=ACCOUNT_ID, + ) + + +def test_read_bytes_without_range_downloads_the_whole_object(sagemaker_session): + def _download_fileobj(key, fileobj, ExtraArgs): + fileobj.write(b"whole object") + + bucket_mock = sagemaker_session.s3_resource.Bucket.return_value + bucket_mock.download_fileobj.side_effect = _download_fileobj + + result = s3.S3Downloader.read_bytes(S3_URI, sagemaker_session) + + assert result == b"whole object" + bucket_mock.download_fileobj.assert_called_once() + assert bucket_mock.download_fileobj.call_args.kwargs == {"ExtraArgs": None} + sagemaker_session.s3_client.get_object.assert_not_called() + + +@pytest.mark.parametrize("byte_range", [(-1, 10), (5, 4), (-3, -1)]) +def test_read_bytes_rejects_bad_range(sagemaker_session, byte_range): + with pytest.raises(ValueError) as error: + s3.S3Downloader.read_bytes(S3_URI, sagemaker_session, byte_range=byte_range) + + assert "byte_range" in str(error.value) + sagemaker_session.s3_client.get_object.assert_not_called() + + +def test_iter_lines_streams_and_decodes_each_line(sagemaker_session): + response = _get_object_response("first\nzweite Zeile \u00e4\nthird".encode("utf-8")) + sagemaker_session.s3_client.get_object.return_value = response + + lines = s3.S3Downloader.iter_lines(S3_URI, sagemaker_session) + + assert list(lines) == ["first", "zweite Zeile \u00e4", "third"] + response["Body"].iter_lines.assert_called_once_with() + response["Body"].read.assert_not_called() + sagemaker_session.s3_client.get_object.assert_called_once_with( + Bucket="mybucket", Key="prefix/data.jsonl" + ) + + +def test_iter_lines_uses_the_given_encoding_and_owner_check(sagemaker_session): + sagemaker_session._get_account_id_if_default_bucket.return_value = ACCOUNT_ID + response = _get_object_response("caf\u00e9\n".encode("latin-1")) + sagemaker_session.s3_client.get_object.return_value = response + + lines = s3.S3Downloader.iter_lines(S3_URI, sagemaker_session, encoding="latin-1") + + assert list(lines) == ["caf\u00e9"] + sagemaker_session.s3_client.get_object.assert_called_once_with( + Bucket="mybucket", Key="prefix/data.jsonl", ExpectedBucketOwner=ACCOUNT_ID + ) + + +def test_head_maps_head_object_fields(sagemaker_session): + last_modified = datetime(2026, 10, 23, 12, 0, 0, tzinfo=timezone.utc) + sagemaker_session.s3_client.head_object.return_value = { + "ContentLength": 2048, + "LastModified": last_modified, + "ETag": '"d41d8cd98f00b204e9800998ecf8427e"', + "ContentType": "application/jsonlines", + "Metadata": {}, + } + + head = s3.S3Downloader.head(S3_URI, sagemaker_session) + + assert head == s3.S3ObjectHead( + content_length=2048, + last_modified=last_modified, + etag='"d41d8cd98f00b204e9800998ecf8427e"', + content_type="application/jsonlines", + ) + sagemaker_session.s3_client.head_object.assert_called_once_with( + Bucket="mybucket", Key="prefix/data.jsonl" + ) + + +def test_head_without_content_type_and_with_owner_check(sagemaker_session): + sagemaker_session._get_account_id_if_default_bucket.return_value = ACCOUNT_ID + sagemaker_session.s3_client.head_object.return_value = { + "ContentLength": 0, + "LastModified": datetime(2026, 1, 1, tzinfo=timezone.utc), + "ETag": '"etag"', + } + + head = s3.S3Downloader.head(S3_URI, sagemaker_session) + + assert head.content_type is None + assert head.content_length == 0 + sagemaker_session.s3_client.head_object.assert_called_once_with( + Bucket="mybucket", Key="prefix/data.jsonl", ExpectedBucketOwner=ACCOUNT_ID + ) + + +def test_read_primitives_fall_back_to_a_boto_client_without_s3_client(sagemaker_session): + sagemaker_session.s3_client = None + boto_client = sagemaker_session.boto_session.client.return_value + boto_client.get_object.return_value = _get_object_response(b"ab") + + result = s3.S3Downloader.read_bytes(S3_URI, sagemaker_session, byte_range=(0, 1)) + + assert result == b"ab" + sagemaker_session.boto_session.client.assert_called_once_with("s3", region_name=REGION) diff --git a/sagemaker-core/tests/unit/test_endpoint_diagnostics.py b/sagemaker-core/tests/unit/test_endpoint_diagnostics.py new file mode 100644 index 0000000000..205f0e7902 --- /dev/null +++ b/sagemaker-core/tests/unit/test_endpoint_diagnostics.py @@ -0,0 +1,444 @@ +# 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. +"""``Endpoint`` collects its own diagnostics through ``EndpointMixin``. + +Before the mixin, ``Endpoint`` had no ``diagnostics()`` and no ``recent_logs()``. The +SageMaker agent skill ``collect_diagnostics.py`` built raw boto3 clients from a region string, +read ``InstanceType`` from ``ProductionVariantSummary`` where it never exists, so the instance +type was always ``None``, fetched one unpaged page of 100 log events, and ignored the session +a caller had bound the resource to. +""" + +from __future__ import absolute_import + +from datetime import datetime, timezone +from unittest.mock import Mock, patch + +import boto3 +import pytest +from botocore.exceptions import ClientError, EndpointConnectionError + +from sagemaker.core.endpoint import ( + METRIC_QUERIES, + EndpointDiagnostics, + EndpointMixin, + MetricSeries, + VariantDiagnostics, +) +from sagemaker.core.resources import Endpoint, EndpointConfig +from sagemaker.core.shapes import ( + ProductionVariant, + ProductionVariantStatus, + ProductionVariantSummary, +) + +ENDPOINT_NAME = "ep" +CONFIG_NAME = "cfg" +VARIANT_NAME = "AllTraffic" +SHADOW_VARIANT_NAME = "Canary" +INSTANCE_TYPE = "ml.g5.xlarge" +LOG_GROUP = "/aws/sagemaker/Endpoints/ep" +TIMESTAMPS = [ + datetime(2026, 10, 23, 12, 0, tzinfo=timezone.utc), + datetime(2026, 10, 23, 12, 1, tzinfo=timezone.utc), +] +EXPECTED_QUERY_IDS = [ + "invocations", + "errors_4xx", + "errors_5xx", + "invocation_model_errors", + "model_latency_avg", + "model_latency_p99", + "overhead_latency", + "cpu_util", + "memory_util", + "gpu_util", + "gpu_memory_util", +] +EXPECTED_SAGEMAKER_METRICS = { + ("Invocations", "Sum"), + ("Invocation4XXErrors", "Sum"), + ("Invocation5XXErrors", "Sum"), + ("InvocationModelErrors", "Sum"), + ("ModelLatency", "Average"), + ("ModelLatency", "p99"), + ("OverheadLatency", "Average"), +} +EXPECTED_INSTANCE_METRICS = { + ("CPUUtilization", "Average"), + ("MemoryUtilization", "Average"), + ("GPUUtilization", "Average"), + ("GPUMemoryUtilization", "Average"), +} + + +def _endpoint(**overrides): + kwargs = { + "endpoint_name": ENDPOINT_NAME, + "endpoint_config_name": CONFIG_NAME, + "endpoint_status": "InService", + "production_variants": [ + ProductionVariantSummary( + variant_name=VARIANT_NAME, + current_instance_count=2, + desired_instance_count=3, + current_weight=1.0, + variant_status=[ + ProductionVariantStatus(status="Updating"), + ProductionVariantStatus(status="Baking"), + ], + ), + ProductionVariantSummary(variant_name=SHADOW_VARIANT_NAME, current_instance_count=1), + ], + } + kwargs.update(overrides) + return Endpoint(**kwargs) + + +def _endpoint_config(): + return EndpointConfig( + endpoint_config_name=CONFIG_NAME, + production_variants=[ + ProductionVariant(variant_name=SHADOW_VARIANT_NAME, instance_type="ml.m5.large"), + ProductionVariant(variant_name=VARIANT_NAME, instance_type=INSTANCE_TYPE), + ], + ) + + +def _metric_data_response(): + return { + "MetricDataResults": [ + {"Id": "model_latency_p99", "Values": [120.0, 480.0], "Timestamps": TIMESTAMPS}, + {"Id": "invocations", "Values": [10.0, 12.0], "Timestamps": TIMESTAMPS}, + {"Id": "not_a_query", "Values": [1.0], "Timestamps": TIMESTAMPS[:1]}, + ] + } + + +def _log_event(seconds, message): + return { + "timestamp": 1_700_000_000_000 + seconds * 1000, + "logStreamName": "AllTraffic/i-0abc", + "message": message, + } + + +def _client_error(operation): + return ClientError( + {"Error": {"Code": "AccessDeniedException", "Message": "no " + operation}}, operation + ) + + +@pytest.fixture +def clients(): + """Canned CloudWatch and Logs clients, keyed the way ``_boto_client`` is called.""" + cloudwatch = Mock(name="cloudwatch") + cloudwatch.get_metric_data.return_value = _metric_data_response() + logs = Mock(name="logs") + logs.filter_log_events.return_value = {"events": [_log_event(0, "ready")]} + return {"cloudwatch": cloudwatch, "logs": logs} + + +@pytest.fixture +def mocked_sources(clients): + with ( + patch.object(Endpoint, "refresh") as refresh, + patch.object(EndpointConfig, "get", return_value=_endpoint_config()) as get_config, + patch.object( + EndpointMixin, "_boto_client", side_effect=lambda service_name: clients[service_name] + ), + ): + yield {"refresh": refresh, "get_config": get_config, "clients": clients} + + +def test_endpoint_is_an_endpoint_mixin_without_new_fields(): + assert issubclass(Endpoint, EndpointMixin) + assert not {"diagnostics", "recent_logs"} & set(Endpoint.model_fields) + assert "diagnostics" not in _endpoint().model_dump() + + +def test_diagnostics_refreshes_and_reads_the_endpoint_state(mocked_sources): + endpoint = _endpoint(failure_reason="CannotStartContainerError") + + result = endpoint.diagnostics() + + mocked_sources["refresh"].assert_called_once_with() + assert isinstance(result, EndpointDiagnostics) + assert result.endpoint_name == ENDPOINT_NAME + assert result.endpoint_status == "InService" + assert result.failure_reason == "CannotStartContainerError" + assert result.endpoint_config_name == CONFIG_NAME + assert result.errors == [] + + +def test_diagnostics_reads_the_instance_type_from_the_endpoint_config(mocked_sources): + """``DescribeEndpoint`` never returns ``InstanceType``. The skill read ``None`` forever.""" + session = boto3.Session( + aws_access_key_id="AKIAEXAMPLE", aws_secret_access_key="secret", region_name="eu-west-1" + ) + endpoint = _endpoint()._set_client_context(session=session, region="eu-west-1") + + result = endpoint.diagnostics() + + mocked_sources["get_config"].assert_called_once_with( + endpoint_config_name=CONFIG_NAME, session=session, region="eu-west-1" + ) + assert result.variants == [ + VariantDiagnostics( + variant_name=VARIANT_NAME, + instance_type=INSTANCE_TYPE, + current_instance_count=2, + desired_instance_count=3, + current_weight=1.0, + variant_status=["Updating", "Baking"], + ), + VariantDiagnostics( + variant_name=SHADOW_VARIANT_NAME, + instance_type="ml.m5.large", + current_instance_count=1, + desired_instance_count=None, + current_weight=None, + variant_status=[], + ), + ] + + +def test_diagnostics_sends_the_eleven_dashboard_queries_for_the_first_variant(mocked_sources): + _endpoint().diagnostics(window_minutes=30) + + call = mocked_sources["clients"]["cloudwatch"].get_metric_data.call_args + queries = call.kwargs["MetricDataQueries"] + assert len(queries) == len(METRIC_QUERIES) == 11 + assert [query["Id"] for query in queries] == EXPECTED_QUERY_IDS + expected_dimensions = [ + {"Name": "EndpointName", "Value": ENDPOINT_NAME}, + {"Name": "VariantName", "Value": VARIANT_NAME}, + ] + by_namespace = {} + for query in queries: + metric_stat = query["MetricStat"] + assert metric_stat["Metric"]["Dimensions"] == expected_dimensions + assert metric_stat["Period"] == 60 + by_namespace.setdefault(metric_stat["Metric"]["Namespace"], set()).add( + (metric_stat["Metric"]["MetricName"], metric_stat["Stat"]) + ) + assert by_namespace == { + "AWS/SageMaker": EXPECTED_SAGEMAKER_METRICS, + "/aws/sagemaker/Endpoints": EXPECTED_INSTANCE_METRICS, + } + assert call.kwargs["ScanBy"] == "TimestampAscending" + window = call.kwargs["EndTime"] - call.kwargs["StartTime"] + assert window.total_seconds() == 30 * 60 + + +def test_diagnostics_targets_the_requested_variant(mocked_sources): + _endpoint().diagnostics(variant_name=SHADOW_VARIANT_NAME) + + queries = mocked_sources["clients"]["cloudwatch"].get_metric_data.call_args.kwargs[ + "MetricDataQueries" + ] + assert {query["MetricStat"]["Metric"]["Dimensions"][1]["Value"] for query in queries} == { + SHADOW_VARIANT_NAME + } + + +def test_diagnostics_defaults_to_all_traffic_without_variants(mocked_sources): + _endpoint(production_variants=[]).diagnostics() + + queries = mocked_sources["clients"]["cloudwatch"].get_metric_data.call_args.kwargs[ + "MetricDataQueries" + ] + assert queries[0]["MetricStat"]["Metric"]["Dimensions"][1]["Value"] == "AllTraffic" + + +def test_diagnostics_keys_metrics_by_name_and_stat(mocked_sources): + result = _endpoint().diagnostics() + + assert sorted(result.metrics) == ["Invocations/Sum", "ModelLatency/p99"] + assert result.metrics["ModelLatency/p99"] == MetricSeries( + namespace="AWS/SageMaker", + metric_name="ModelLatency", + stat="p99", + values=[120.0, 480.0], + timestamps=TIMESTAMPS, + ) + + +def test_diagnostics_joins_paged_metric_results(mocked_sources): + cloudwatch = mocked_sources["clients"]["cloudwatch"] + cloudwatch.get_metric_data.side_effect = [ + { + "MetricDataResults": [ + {"Id": "invocations", "Values": [1.0], "Timestamps": TIMESTAMPS[:1]} + ], + "NextToken": "page-2", + }, + { + "MetricDataResults": [ + {"Id": "invocations", "Values": [2.0], "Timestamps": TIMESTAMPS[1:]} + ] + }, + ] + + result = _endpoint().diagnostics() + + assert result.metrics["Invocations/Sum"].values == [1.0, 2.0] + assert result.metrics["Invocations/Sum"].timestamps == TIMESTAMPS + second_call = cloudwatch.get_metric_data.call_args_list[1] + assert second_call.kwargs["NextToken"] == "page-2" + + +def test_diagnostics_formats_and_limits_the_log_lines(mocked_sources): + logs = mocked_sources["clients"]["logs"] + logs.filter_log_events.return_value = { + "events": [_log_event(0, "model loaded"), _log_event(1, "worker started")], + "nextToken": "more", + } + + result = _endpoint().diagnostics(log_limit=2) + + assert result.recent_logs == [ + "2023-11-14T22:13:20+00:00 AllTraffic/i-0abc model loaded", + "2023-11-14T22:13:21+00:00 AllTraffic/i-0abc worker started", + ] + logs.filter_log_events.assert_called_once() + request = logs.filter_log_events.call_args.kwargs + assert request["logGroupName"] == LOG_GROUP + assert request["limit"] == 2 + assert request["endTime"] - request["startTime"] == 15 * 60 * 1000 + + +def test_diagnostics_widens_the_log_window_to_the_metric_window(mocked_sources): + _endpoint().diagnostics(window_minutes=60) + + request = mocked_sources["clients"]["logs"].filter_log_events.call_args.kwargs + assert request["endTime"] - request["startTime"] == 60 * 60 * 1000 + + +def test_diagnostics_collects_a_cloudwatch_error_instead_of_raising(mocked_sources): + mocked_sources["clients"]["cloudwatch"].get_metric_data.side_effect = _client_error( + "GetMetricData" + ) + + result = _endpoint().diagnostics() + + assert result.metrics == {} + assert len(result.errors) == 1 + assert result.errors[0].startswith("GetMetricData failed: ") + assert "no GetMetricData" in result.errors[0] + assert result.recent_logs == ["2023-11-14T22:13:20+00:00 AllTraffic/i-0abc ready"] + + +def test_diagnostics_collects_a_logs_error_and_a_config_error(mocked_sources): + mocked_sources["clients"]["logs"].filter_log_events.side_effect = _client_error( + "FilterLogEvents" + ) + mocked_sources["get_config"].side_effect = EndpointConnectionError( + endpoint_url="https://api.sagemaker.eu-west-1.amazonaws.com" + ) + + result = _endpoint().diagnostics() + + assert result.recent_logs == [] + assert [variant.instance_type for variant in result.variants] == [None, None] + assert [error.split(":")[0] for error in result.errors] == [ + "DescribeEndpointConfig failed", + "FilterLogEvents failed", + ] + assert len(result.metrics) == 2 + + +def test_diagnostics_skips_the_config_lookup_without_a_config_name(mocked_sources): + result = _endpoint(endpoint_config_name=None).diagnostics() + + mocked_sources["get_config"].assert_not_called() + assert result.endpoint_config_name is None + assert result.variants[0].instance_type is None + assert result.errors == [] + + +def test_diagnostics_raises_when_describe_endpoint_fails(mocked_sources): + mocked_sources["refresh"].side_effect = _client_error("DescribeEndpoint") + + with pytest.raises(ClientError, match="no DescribeEndpoint"): + _endpoint().diagnostics() + + +def test_recent_logs_pages_until_the_next_token_is_absent(mocked_sources): + logs = mocked_sources["clients"]["logs"] + logs.filter_log_events.side_effect = [ + {"events": [_log_event(0, "one")], "nextToken": "page-2"}, + {"events": [_log_event(1, "two")], "nextToken": "page-3"}, + {"events": [_log_event(2, "three")]}, + ] + + lines = _endpoint().recent_logs(window_minutes=10, limit=100) + + assert [line.split(" ", 2)[2] for line in lines] == ["one", "two", "three"] + requests = [call.kwargs for call in logs.filter_log_events.call_args_list] + assert [request.get("nextToken") for request in requests] == [None, "page-2", "page-3"] + assert [request["limit"] for request in requests] == [100, 99, 98] + assert requests[0]["endTime"] - requests[0]["startTime"] == 10 * 60 * 1000 + + +def test_recent_logs_stops_on_an_empty_page_with_a_token(mocked_sources): + """CloudWatch Logs can return a token with no events forever. The pager must not loop.""" + logs = mocked_sources["clients"]["logs"] + logs.filter_log_events.side_effect = [ + {"events": [_log_event(0, "one")], "nextToken": "page-2"}, + {"events": [], "nextToken": "page-3"}, + ] + + lines = _endpoint().recent_logs() + + assert len(lines) == 1 + assert logs.filter_log_events.call_count == 2 + + +def test_recent_logs_stops_at_the_limit(mocked_sources): + logs = mocked_sources["clients"]["logs"] + logs.filter_log_events.return_value = { + "events": [_log_event(0, "one"), _log_event(1, "two"), _log_event(2, "three")], + "nextToken": "page-2", + } + + lines = _endpoint().recent_logs(limit=3) + + assert len(lines) == 3 + logs.filter_log_events.assert_called_once() + + +def test_recent_logs_raises_a_client_error(mocked_sources): + mocked_sources["clients"]["logs"].filter_log_events.side_effect = _client_error( + "FilterLogEvents" + ) + + with pytest.raises(ClientError, match="no FilterLogEvents"): + _endpoint().recent_logs() + + +def test_boto_clients_use_the_bound_session_and_region(): + """The skill built clients from a region string and lost the caller's credentials.""" + session = boto3.Session( + aws_access_key_id="AKIABOUND0000000000", aws_secret_access_key="s", region_name="us-east-1" + ) + endpoint = _endpoint()._set_client_context(session=session, region="eu-central-1") + + cloudwatch = endpoint._boto_client("cloudwatch") + logs = endpoint._boto_client("logs") + + assert cloudwatch.meta.service_model.service_name == "cloudwatch" + assert logs.meta.service_model.service_name == "logs" + assert cloudwatch.meta.region_name == logs.meta.region_name == "eu-central-1" + assert cloudwatch._request_signer._credentials.access_key == "AKIABOUND0000000000" + assert logs._request_signer._credentials.access_key == "AKIABOUND0000000000" diff --git a/sagemaker-core/tests/unit/test_hub_catalog_cache.py b/sagemaker-core/tests/unit/test_hub_catalog_cache.py new file mode 100644 index 0000000000..89ef61489e --- /dev/null +++ b/sagemaker-core/tests/unit/test_hub_catalog_cache.py @@ -0,0 +1,446 @@ +# 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. +"""``HubContent.get_all_cached`` serves a hub from the hub catalog cache. + +Before this change every list of a hub described every item: 766 ``DescribeHubContent`` +calls and 60 to 95 seconds for the public hub, on every call. ``HubCatalogCache`` follows +``JumpStartModelsCache``: an ``LRUCache`` of hubs in memory with an expiration horizon, and +a retrieval function that lists the hub and describes only the items whose version +changed. A JSON file per hub, content type, region, and account lets a new process start +from the previous catalog. +""" + +from __future__ import absolute_import + +import datetime +import json +import threading +import time +from concurrent.futures import ThreadPoolExecutor +from unittest.mock import MagicMock, patch + +import pytest + +from sagemaker.core import hub_catalog +from sagemaker.core.hub_catalog import ( + CACHE_FORMAT_VERSION, + HUB_CATALOG_DEFAULT_CACHE_EXPIRATION_HORIZON, + HUB_CATALOG_DEFAULT_MAX_CACHE_ITEMS, + HubCatalogAccessor, + HubCatalogCache, + HubCatalogFile, + HubCatalogKey, + HubCatalogSource, + default_cache_dir, + reusable_entries, +) +from sagemaker.core.resources import Base, HubContent +from sagemaker.core.utilities.cache import LRUCache +from sagemaker.core.utils.utils import SageMakerClient + +HUB = "SageMakerPublicHub" +ACCOUNT = "123456789012" +CREATED = datetime.datetime(2026, 3, 1, tzinfo=datetime.timezone.utc) +ORIGINAL = datetime.datetime(2025, 6, 1, tzinfo=datetime.timezone.utc) + + +def _summary(name, version="1.0.0"): + return { + "HubContentName": name, + "HubContentArn": f"arn:aws:sagemaker:us-west-2:aws:hub-content/{HUB}/Model/{name}/{version}", + "HubContentVersion": version, + "HubContentType": "Model", + "DocumentSchemaVersion": "2.0.0", + "HubContentStatus": "Available", + "HubContentSearchKeywords": ["@license:apache-2.0", "@model-size:1b-10b"], + "CreationTime": CREATED, + "OriginalCreationTime": ORIGINAL, + } + + +class _Hub: + """A fake SageMaker client over a dict of name to version.""" + + def __init__(self, versions): + self.versions = dict(versions) + self.client = MagicMock() + self.client.list_hub_contents.side_effect = self._list + self.client.describe_hub_content.side_effect = self._describe + self.described = [] + self.lists = 0 + self.in_flight = 0 + self.peak_in_flight = 0 + self._lock = threading.Lock() + + def _list(self, **kwargs): + self.lists += 1 + names = sorted(self.versions) + page = kwargs.get("NextToken") + if page is None: + first, token = names[:2], ("page-2" if len(names) > 2 else None) + else: + first, token = names[2:], None + response = {"HubContentSummaries": [_summary(n, self.versions[n]) for n in first]} + if token: + response["NextToken"] = token + return response + + def _describe(self, **kwargs): + assert kwargs["HubName"] + name = kwargs["HubContentName"] + with self._lock: + self.described.append(name) + self.in_flight += 1 + self.peak_in_flight = max(self.peak_in_flight, self.in_flight) + try: + summary = _summary(name, self.versions[name]) + summary.pop("OriginalCreationTime") + return { + **summary, + "HubName": kwargs["HubName"], + "HubArn": f"arn:aws:sagemaker:us-west-2:aws:hub/{HUB}", + "HubContentDocument": json.dumps( + {"License": "Apache-2.0", "Version": self.versions[name]} + ), + } + finally: + with self._lock: + self.in_flight -= 1 + + +@pytest.fixture(autouse=True) +def _isolated_cache(tmp_path, monkeypatch): + monkeypatch.setenv(hub_catalog.CACHE_DIR_ENV_VAR, str(tmp_path / "hub-cache")) + SageMakerClient.reset() + HubCatalogAccessor.reset_cache() + with ( + patch.object(HubContent, "_account_id", return_value=ACCOUNT), + patch( + "sagemaker.core.telemetry.telemetry_logging._get_default_sagemaker_session", + return_value=MagicMock(sagemaker_config=None), + ), + patch("sagemaker.core.telemetry.telemetry_logging._send_telemetry_request"), + ): + yield tmp_path / "hub-cache" + HubCatalogAccessor.reset_cache() + SageMakerClient.reset() + + +@pytest.fixture +def telemetry_send(): + with patch("sagemaker.core.telemetry.telemetry_logging._send_telemetry_request") as send: + yield send + + +def _patched(hub): + return patch.object(Base, "get_sagemaker_client", return_value=hub.client) + + +def _cache_payload(cache_dir, region="us-west-2", hub=HUB, content_type="Model"): + return json.loads((cache_dir / f"{ACCOUNT}-{region}-{hub}-{content_type}.json").read_text()) + + +def test_first_call_describes_every_item_and_writes_the_cache_file(_isolated_cache): + hub = _Hub({"model-a": "1.0.0", "model-b": "2.0.0", "model-c": "1.0.0"}) + + with _patched(hub): + models = HubContent.get_all_cached(HUB, region="us-west-2") + + assert [m.hub_content_name for m in models] == ["model-a", "model-b", "model-c"] + assert sorted(hub.described) == ["model-a", "model-b", "model-c"] + for call in hub.client.describe_hub_content.call_args_list: + assert call.kwargs["HubName"] == HUB + assert models[0].original_creation_time == ORIGINAL + assert models[0].hub_content_document == json.dumps( + {"License": "Apache-2.0", "Version": "1.0.0"} + ) + assert models[0].license == "apache-2.0" + assert models[0]._region == "us-west-2" + payload = _cache_payload(_isolated_cache) + assert payload["format_version"] == CACHE_FORMAT_VERSION + assert sorted(payload["entries"]) == ["model-a", "model-b", "model-c"] + assert payload["entries"]["model-b"]["listed_version"] == "2.0.0" + + +def test_a_call_inside_the_horizon_makes_no_request(): + hub = _Hub({"model-a": "1.0.0", "model-b": "1.0.0"}) + with _patched(hub): + first = HubContent.get_all_cached(HUB, region="us-west-2") + hub.described.clear() + hub.lists = 0 + second = HubContent.get_all_cached(HUB, region="us-west-2") + + assert (hub.lists, hub.described) == (0, []) + assert [m.hub_content_name for m in second] == [m.hub_content_name for m in first] + assert second[0] is not first[0] + + +def test_refresh_cache_lists_again_and_describes_only_new_and_changed_items(): + hub = _Hub({"model-a": "1.0.0", "model-b": "2.0.0", "model-c": "1.0.0"}) + with _patched(hub): + HubContent.get_all_cached(HUB, region="us-west-2") + hub.described.clear() + hub.versions["model-b"] = "2.1.0" + hub.versions["model-d"] = "1.0.0" + del hub.versions["model-c"] + + with _patched(hub): + refresh = HubContent.refresh_cache(HUB, region="us-west-2") + models = HubContent.get_all_cached(HUB, region="us-west-2") + + assert sorted(hub.described) == ["model-b", "model-d"] + assert (refresh.listed, refresh.reused, refresh.described, refresh.removed) == (3, 1, 2, 1) + assert [m.hub_content_name for m in models] == ["model-a", "model-b", "model-d"] + assert models[1].hub_content_version == "2.1.0" + assert models[1].hub_content_document == json.dumps( + {"License": "Apache-2.0", "Version": "2.1.0"} + ) + assert models[0].original_creation_time == ORIGINAL + + +def test_an_expired_entry_lists_again_and_reuses_the_unchanged_items(): + hub = _Hub({"model-a": "1.0.0", "model-b": "1.0.0"}) + HubContent.set_cache_kwargs(cache_expiration_horizon=datetime.timedelta(milliseconds=50)) + with _patched(hub): + HubContent.get_all_cached(HUB, region="us-west-2") + assert hub.lists == 1 + hub.described.clear() + hub.versions["model-b"] = "1.1.0" + time.sleep(0.1) + models = HubContent.get_all_cached(HUB, region="us-west-2") + + assert hub.lists == 2 + assert hub.described == ["model-b"] + assert [m.hub_content_version for m in models] == ["1.0.0", "1.1.0"] + + +def test_a_new_process_starts_from_the_cache_file(_isolated_cache): + hub = _Hub({"model-a": "1.0.0", "model-b": "1.0.0"}) + with _patched(hub): + HubContent.get_all_cached(HUB, region="us-west-2") + HubCatalogAccessor.reset_cache() + hub.described.clear() + + with _patched(hub): + models = HubContent.get_all_cached(HUB, region="us-west-2") + + assert hub.described == [] + assert [m.hub_content_name for m in models] == ["model-a", "model-b"] + assert models[0].license == "apache-2.0" + + +def test_cache_files_are_separate_per_account_region_hub_and_type(_isolated_cache): + hub = _Hub({"model-a": "1.0.0"}) + with _patched(hub): + HubContent.get_all_cached(HUB, region="us-west-2") + HubContent.get_all_cached(HUB, region="eu-west-1") + HubContent.get_all_cached("my-hub", hub_content_type="Notebook", region="us-west-2") + + names = sorted(path.name for path in _isolated_cache.iterdir()) + assert names == [ + f"{ACCOUNT}-eu-west-1-{HUB}-Model.json", + f"{ACCOUNT}-us-west-2-{HUB}-Model.json", + f"{ACCOUNT}-us-west-2-my-hub-Notebook.json", + ] + + +def test_clear_cache_drops_memory_and_deletes_the_file_once(): + hub = _Hub({"model-a": "1.0.0"}) + with _patched(hub): + HubContent.get_all_cached(HUB, region="us-west-2") + + assert HubContent.clear_cache(HUB, region="us-west-2") is True + assert HubContent.clear_cache(HUB, region="us-west-2") is False + assert HubContent.clear_cache() is False + hub.described.clear() + HubContent.get_all_cached(HUB, region="us-west-2") + + assert hub.described == ["model-a"] + + +def test_an_unreadable_cache_file_is_ignored(_isolated_cache, caplog): + _isolated_cache.mkdir(parents=True) + (_isolated_cache / f"{ACCOUNT}-us-west-2-{HUB}-Model.json").write_text("{not json") + hub = _Hub({"model-a": "1.0.0"}) + + with _patched(hub), caplog.at_level("WARNING"): + models = HubContent.get_all_cached(HUB, region="us-west-2") + + assert [m.hub_content_name for m in models] == ["model-a"] + assert "Ignoring the unreadable hub content cache" in caplog.text + + +def test_a_failed_describe_raises_and_keeps_the_described_items_in_the_file(_isolated_cache): + hub = _Hub({"model-a": "1.0.0", "model-b": "1.0.0", "model-c": "1.0.0"}) + + def describe(**kwargs): + if kwargs["HubContentName"] == "model-c": + raise RuntimeError("describe of model-c failed") + return hub._describe(**kwargs) + + hub.client.describe_hub_content.side_effect = describe + HubContent.set_cache_kwargs(max_workers=1) + with _patched(hub): + with pytest.raises(RuntimeError, match="describe of model-c failed"): + HubContent.get_all_cached(HUB, region="us-west-2") + + assert sorted(_cache_payload(_isolated_cache)["entries"]) == ["model-a", "model-b"] + hub.client.describe_hub_content.side_effect = hub._describe + hub.described.clear() + with _patched(hub): + refresh = HubContent.refresh_cache(HUB, region="us-west-2") + + assert hub.described == ["model-c"] + assert (refresh.reused, refresh.described) == (2, 1) + + +def test_set_cache_kwargs_rebuilds_the_cache_and_rejects_bad_values(): + HubContent.set_cache_kwargs(max_workers=0) + with pytest.raises(ValueError, match="max_workers must be at least 1"): + HubCatalogAccessor.get_cache() + + HubContent.set_cache_kwargs(max_cache_items=3) + cache = HubCatalogAccessor.get_cache() + assert cache._content_cache._max_cache_items == 3 + assert HubCatalogAccessor.get_cache() is cache + + +def test_the_cache_is_an_lru_cache_with_the_jumpstart_defaults(): + cache = HubCatalogCache() + + assert isinstance(cache._content_cache, LRUCache) + assert cache._content_cache._max_cache_items == HUB_CATALOG_DEFAULT_MAX_CACHE_ITEMS == 20 + assert cache._content_cache._expiration_horizon == HUB_CATALOG_DEFAULT_CACHE_EXPIRATION_HORIZON + assert HUB_CATALOG_DEFAULT_CACHE_EXPIRATION_HORIZON == datetime.timedelta(hours=6) + + +def test_get_all_cached_and_refresh_cache_emit_telemetry(telemetry_send): + hub = _Hub({"model-a": "1.0.0"}) + + with _patched(hub): + HubContent.get_all_cached(HUB, region="us-west-2") + HubContent.refresh_cache(HUB, region="us-west-2") + + func_names = [call.args[5].split("&")[0] for call in telemetry_send.call_args_list] + assert func_names == ["HubContent.get_all_cached", "HubContent.refresh_cache"] + assert all(call.args[0] == 1 for call in telemetry_send.call_args_list) + + +def test_account_id_calls_sts_once_per_set_of_credentials(monkeypatch): + import sagemaker.core.hub_content as hub_content_module + + monkeypatch.setattr(hub_content_module, "_ACCOUNT_IDS", {}) + account_id = hub_content_module.HubContentMixin.__dict__["_account_id"].__func__ + sts = MagicMock() + sts.get_caller_identity.return_value = {"Account": ACCOUNT} + session = MagicMock() + session.region_name = "us-west-2" + session.client.return_value = sts + session.get_credentials.return_value.get_frozen_credentials.return_value = ( + "AKID", + "secret", + None, + ) + + first = account_id(HubContent, session, "us-west-2") + second = account_id(HubContent, session, "us-west-2") + + assert (first, second) == (ACCOUNT, ACCOUNT) + assert sts.get_caller_identity.call_count == 1 + assert session.client.call_args.kwargs["region_name"] == "us-west-2" + + +def test_default_cache_dir_honors_the_environment_variable(monkeypatch, tmp_path): + monkeypatch.setenv(hub_catalog.CACHE_DIR_ENV_VAR, str(tmp_path / "override")) + assert default_cache_dir() == str(tmp_path / "override") + + monkeypatch.delenv(hub_catalog.CACHE_DIR_ENV_VAR) + assert default_cache_dir().endswith("hub-content") + + +def test_cache_key_file_name_is_file_system_safe(): + key = HubCatalogKey( + hub_name="arn:aws:hub/x y", hub_content_type="Model", region="us-west-2", account_id=ACCOUNT + ) + + assert key.file_name() == f"{ACCOUNT}-us-west-2-arn_aws_hub_x_y-Model.json" + + +def test_reusable_entries_match_on_the_listed_version(): + entries = { + "same": {"listed_version": "1.0.0", "resource": {}}, + "changed": {"listed_version": "1.0.0", "resource": {}}, + "gone": {"listed_version": "1.0.0", "resource": {}}, + } + summaries = [_summary("same", "1.0.0"), _summary("changed", "1.1.0"), _summary("new", "1.0.0")] + + assert list(reusable_entries(entries, summaries)) == ["same"] + + +def test_cache_file_rejects_another_format_version(tmp_path): + key = HubCatalogKey(HUB, "Model", "us-west-2", ACCOUNT) + cache_file = HubCatalogFile(key, str(tmp_path)) + cache_file.save({"model-a": {"listed_version": "1.0.0", "resource": {"hub_content_name": "a"}}}) + payload = json.loads(open(cache_file.path).read()) + payload["format_version"] = CACHE_FORMAT_VERSION + 1 + open(cache_file.path, "w").write(json.dumps(payload)) + + assert cache_file.load() == {} + + +def test_a_slow_retrieval_of_one_hub_does_not_block_a_hit_on_another_hub(): + cache = HubCatalogCache() + fast_key = HubCatalogKey("fast-hub", "Model", "us-west-2", ACCOUNT) + slow_key = HubCatalogKey("slow-hub", "Model", "us-west-2", ACCOUNT) + started = threading.Event() + + def slow_list(key): + started.set() + time.sleep(0.5) + return [] + + fast_source = HubCatalogSource(lambda key: [], lambda key, summary: None) + slow_source = HubCatalogSource(slow_list, lambda key, summary: None) + cache.get(fast_key, fast_source) + + with ThreadPoolExecutor(max_workers=1) as executor: + slow = executor.submit(cache.get, slow_key, slow_source) + assert started.wait(timeout=5) + hit_started = time.monotonic() + cache.get(fast_key, fast_source) + hit_elapsed = time.monotonic() - hit_started + slow.result(timeout=5) + + assert hit_elapsed < 0.4 + + +def test_two_threads_that_miss_the_same_hub_share_one_retrieval(): + cache = HubCatalogCache() + key = HubCatalogKey("hub", "Model", "us-west-2", ACCOUNT) + lists = [] + + def list_hub(key): + lists.append(threading.get_ident()) + return [_summary("model-a")] + + def describe(key, summary): + resource = MagicMock() + resource.model_dump.return_value = {"hub_content_name": "model-a"} + return resource + + source = HubCatalogSource(list_hub, describe) + with ThreadPoolExecutor(max_workers=4) as executor: + values = list(executor.map(lambda _: cache.get(key, source), range(4))) + + assert len(lists) == 1 + assert all(value.refresh.listed == 1 for value in values) diff --git a/sagemaker-core/tests/unit/test_hub_content_get_all.py b/sagemaker-core/tests/unit/test_hub_content_get_all.py new file mode 100644 index 0000000000..19a11a4508 --- /dev/null +++ b/sagemaker-core/tests/unit/test_hub_content_get_all.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/ +# +# 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. +"""``HubContent.get_all`` lists a hub's contents as bound resources. + +Before this method the SDK had no resource-level ``ListHubContents``: callers paged the +boto3 client by hand (six places in the SDK did the same), ``ResourceIterator`` dropped +``OriginalCreationTime`` because the ``HubContent`` class had no such field, and +``get_all_versions()`` built resources without the ``hub_name`` the list call used, so +the first ``refresh()`` failed with ``Missing required parameter in input: "HubName"``. +""" + +from __future__ import absolute_import + +import datetime +from unittest.mock import MagicMock, patch + +import boto3 +import pytest + +from sagemaker.core.resources import Base, HubContent +from sagemaker.core.utils.utils import ResourceIterator, SageMakerClient + +HUB = "SageMakerPublicHub" +CREATED = datetime.datetime(2024, 3, 1, tzinfo=datetime.timezone.utc) +ORIGINAL = datetime.datetime(2023, 6, 1, tzinfo=datetime.timezone.utc) + + +def _summary(name, version="1.0.0"): + return { + "HubContentName": name, + "HubContentArn": f"arn:aws:sagemaker:us-west-2:aws:hub-content/{HUB}/Model/{name}/{version}", + "HubContentVersion": version, + "HubContentType": "Model", + "DocumentSchemaVersion": "2.0.0", + "HubContentStatus": "Available", + "HubContentSearchKeywords": ["@task:text-generation", "@license:apache-2.0"], + "CreationTime": CREATED, + "OriginalCreationTime": ORIGINAL, + } + + +def _describe(**kwargs): + if "HubName" not in kwargs: + raise ValueError('Missing required parameter in input: "HubName"') + summary = _summary(kwargs["HubContentName"], kwargs.get("HubContentVersion", "1.0.0")) + summary.pop("OriginalCreationTime") + return { + **summary, + "HubName": kwargs["HubName"], + "HubArn": f"arn:aws:sagemaker:us-west-2:aws:hub/{HUB}", + "HubContentDocument": "{}", + } + + +@pytest.fixture(autouse=True) +def _reset_client_cache(): + SageMakerClient.reset() + yield + SageMakerClient.reset() + + +@pytest.fixture +def client(): + fake = MagicMock() + fake.list_hub_contents.side_effect = [ + {"HubContentSummaries": [_summary("model-a")], "NextToken": "page-2"}, + {"HubContentSummaries": [_summary("model-b")]}, + ] + fake.list_hub_content_versions.return_value = { + "HubContentSummaries": [_summary("model-a", "1.0.0"), _summary("model-a", "2.0.0")] + } + fake.describe_hub_content.side_effect = _describe + with patch.object(Base, "get_sagemaker_client", return_value=fake): + yield fake + + +def test_get_all_pages_through_the_hub_and_binds_hub_name(client): + contents = list(HubContent.get_all(hub_name=HUB, hub_content_type="Model")) + + assert [content.hub_content_name for content in contents] == ["model-a", "model-b"] + assert all(content.hub_name == HUB for content in contents) + assert client.list_hub_contents.call_args_list[0].kwargs == { + "HubName": HUB, + "HubContentType": "Model", + } + assert client.list_hub_contents.call_args_list[1].kwargs["NextToken"] == "page-2" + + +def test_get_all_forwards_the_list_filters(client): + client.list_hub_contents.side_effect = [{"HubContentSummaries": []}] + + list( + HubContent.get_all( + hub_name=HUB, + hub_content_type="Model", + name_contains="llama", + sort_by="HubContentName", + sort_order="Ascending", + ) + ) + + assert client.list_hub_contents.call_args.kwargs == { + "HubName": HUB, + "HubContentType": "Model", + "NameContains": "llama", + "SortBy": "HubContentName", + "SortOrder": "Ascending", + } + + +def test_get_all_keeps_original_creation_time_from_the_summary(client): + contents = list(HubContent.get_all(hub_name=HUB, hub_content_type="Model")) + + assert contents[0].original_creation_time == ORIGINAL + assert contents[0].creation_time == CREATED + assert contents[0].hub_content_search_keywords == [ + "@task:text-generation", + "@license:apache-2.0", + ] + + +def test_get_all_refreshes_each_resource_with_the_hub_name(client): + list(HubContent.get_all(hub_name=HUB, hub_content_type="Model")) + + for call in client.describe_hub_content.call_args_list: + assert call.kwargs["HubName"] == HUB + + +def test_get_all_binds_the_session_and_region_to_each_resource(client): + session = boto3.Session( + aws_access_key_id="AKID", aws_secret_access_key="secret", region_name="eu-west-1" + ) + + content = next( + iter( + HubContent.get_all( + hub_name=HUB, hub_content_type="Model", session=session, region="eu-west-1" + ) + ) + ) + + assert content._session is session + assert content._region == "eu-west-1" + + +def test_get_all_versions_binds_the_hub_name_of_the_parent(client): + parent = HubContent(hub_name=HUB, hub_content_type="Model", hub_content_name="model-a") + + versions = list(parent.get_all_versions()) + + assert [version.hub_content_version for version in versions] == ["1.0.0", "2.0.0"] + assert all(version.hub_name == HUB for version in versions) + for call in client.describe_hub_content.call_args_list: + assert call.kwargs["HubName"] == HUB + + +def test_resource_iterator_extra_init_data_does_not_override_the_summary(): + fake = MagicMock() + fake.list_hub_contents.return_value = {"HubContentSummaries": [_summary("model-a")]} + fake.describe_hub_content.side_effect = _describe + + iterator = ResourceIterator( + client=fake, + summaries_key="HubContentSummaries", + summary_name="HubContentInfo", + resource_cls=HubContent, + list_method="list_hub_contents", + list_method_kwargs={"HubName": HUB, "HubContentType": "Model"}, + extra_init_data={"hub_name": HUB, "hub_content_name": "ignored-when-present"}, + ) + with patch.object(Base, "get_sagemaker_client", return_value=fake): + content = next(iterator) + + assert content.hub_content_name == "model-a" + assert content.hub_name == HUB diff --git a/sagemaker-core/tests/unit/test_hub_content_mixin.py b/sagemaker-core/tests/unit/test_hub_content_mixin.py new file mode 100644 index 0000000000..aa18b150cc --- /dev/null +++ b/sagemaker-core/tests/unit/test_hub_content_mixin.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/ +# +# 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. +"""``HubContent`` reads its own catalog metadata through ``HubContentMixin``. + +Before the mixin every consumer split ``HubContentSearchKeywords`` by prefix itself: +``jumpstart/search.py``, ``recipe_utils.py``, ``training_session/_utils.py``, and four +scripts in the SageMaker agent skills, each with its own partial field list and none +handling the 50-keyword cap that drops ``@model-size:`` and ``@context-window:`` on models +with many languages. +""" + +from __future__ import absolute_import + +import json + +from sagemaker.core.hub_content import HubContentMixin, HubContentRecipe +from sagemaker.core.resources import HubContent + +KEYWORDS = [ + "@task:text generation", + "@task:summarization", + "@license:apache-2.0", + "@provider:Meta", + "@model-size:1b-10b", + "@context-window:128k", + "@framework:huggingface", + "@language:english", + "@language:german", + "@data-type:text", + "@input-modality:text", + "@output-modality:text", + "@model-type:open_weights", + "@capability:customization", + "@capability:bedrock_console", + "@recipe:finetuning_sft_lora", + "@recipe:finetuning_dpo_lora", + "@recipe:finetuning_cpt", + "@recipe:evaluation_llmasjudge", + "@huggingface-id:meta-llama/Llama-3.2-1B", + "unknown keyword", +] + + +def _content(**overrides): + kwargs = { + "hub_name": "SageMakerPublicHub", + "hub_content_type": "Model", + "hub_content_name": "meta-textgeneration-llama-3-2-1b", + "hub_content_search_keywords": list(KEYWORDS), + } + kwargs.update(overrides) + return HubContent(**kwargs) + + +def test_hub_content_is_a_hub_content_mixin(): + assert issubclass(HubContent, HubContentMixin) + # The mixin adds behavior only: no model field, nothing in model_dump(). + assert not {"tasks", "license", "recipes"} & set(HubContent.model_fields) + assert "tasks" not in _content().model_dump() + + +def test_every_keyword_prefix_is_readable(): + content = _content() + + assert content.tasks == ["text generation", "summarization"] + assert content.license == "apache-2.0" + assert content.provider == "Meta" + assert content.model_size == "1b-10b" + assert content.context_window == "128k" + assert content.framework == "huggingface" + assert content.languages == ["english", "german"] + assert content.data_types == ["text"] + assert content.input_modalities == ["text"] + assert content.output_modalities == ["text"] + assert content.model_type == "open_weights" + assert content.capabilities == ["customization", "bedrock_console"] + assert content.customization_supported is True + assert content.bedrock_supported is True + + +def test_recipes_parse_type_technique_and_optional_strategy(): + content = _content() + + assert HubContentRecipe("finetuning", "sft", "lora") in content.recipes + assert HubContentRecipe("finetuning", "cpt", None) in content.recipes + assert HubContentRecipe("evaluation", "llmasjudge", None) in content.recipes + assert content.fine_tuning_techniques() == ["cpt", "dpo", "sft"] + assert content.supports_technique("SFT") is True + assert content.supports_technique("rlvr") is False + + +def test_recipe_keyword_parser_ignores_other_keywords_and_malformed_recipes(): + assert HubContentRecipe.from_keyword("@task:text generation") is None + assert HubContentRecipe.from_keyword("@recipe:finetuning") is None + assert HubContentRecipe.from_keyword("@recipe:FineTuning_SFT_LoRA") == HubContentRecipe( + "finetuning", "sft", "lora" + ) + + +def test_document_fills_fields_the_keyword_cap_dropped(): + document = { + "Task": "Speech Recognition", + "License": "MIT", + "Provider": "OpenAI", + "ModelSize": "1.5B", + "ContextWindow": "30s", + "Languages": ["English", "French"], + "Capabilities": ["CUSTOMIZATION"], + "RecipeCollection": [ + {"Type": "FineTuning", "CustomizationTechnique": "SFT", "Strategy": "LORA"}, + {"Type": "Evaluation", "EvaluationType": "ASR"}, + ], + } + content = _content( + hub_content_search_keywords=["@task:speech recognition", "@language:english"], + hub_content_document=json.dumps(document), + ) + + # A keyword is authoritative; the document only fills what is missing. + assert content.tasks == ["speech recognition"] + assert content.languages == ["english", "French"] + assert content.license == "MIT" + assert content.provider == "OpenAI" + assert content.model_size == "1.5B" + assert content.context_window == "30s" + assert content.capabilities == ["customization"] + assert content.recipes == [ + HubContentRecipe("finetuning", "sft", "lora"), + HubContentRecipe("evaluation", "asr", None), + ] + + +def test_keyword_wins_over_document_for_scalar_fields(): + content = _content( + hub_content_search_keywords=["@license:apache-2.0", "@model-size:7b"], + hub_content_document=json.dumps({"License": "Apache 2.0", "ModelSize": "8B"}), + ) + + assert content.license == "apache-2.0" + assert content.model_size == "7b" + + +def test_unassigned_keywords_and_malformed_document_read_as_empty(): + content = _content( + hub_content_search_keywords=HubContent.model_fields["hub_content_search_keywords"].default, + hub_content_document="not json", + ) + + assert content.tasks == [] + assert content.license is None + assert content.capabilities == [] + assert content.recipes == [] + assert content.customization_supported is False + + +def test_accessors_follow_a_later_refresh_of_the_document(): + content = _content(hub_content_search_keywords=["@task:text generation"]) + assert content.license is None + + content.hub_content_document = json.dumps({"License": "Llama 3.2 Community"}) + + assert content.license == "Llama 3.2 Community" diff --git a/sagemaker-core/tests/unit/test_hub_content_search.py b/sagemaker-core/tests/unit/test_hub_content_search.py new file mode 100644 index 0000000000..b447f95062 --- /dev/null +++ b/sagemaker-core/tests/unit/test_hub_content_search.py @@ -0,0 +1,208 @@ +# 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. +"""``HubContent.search`` runs ``Search(Resource="HubContent")`` with ``Metadata.*`` filters. + +Before this change the SDK had no binding for the ``HubContent`` search resource: the typed +``SearchRecord.hub_content`` shape existed since v3.25.0 and nothing read it, so a caller who +wanted the public hub narrowed by license, provider, task, or model size either listed every +model with ``ListHubContents`` or wrote the ``Search`` request, the paging, and the result +parsing by hand. +""" + +from __future__ import absolute_import + +import json +from unittest.mock import Mock, patch + +import pytest + +from sagemaker.core.hub_content import ( + MAX_SEARCH_FILTERS, + HubContentFilter, + _search_expression, +) +from sagemaker.core.resources import HubContent + +HUB = "SageMakerPublicHub" + + +@pytest.fixture(autouse=True) +def _offline_telemetry(): + """Keep the ``HubContent.search`` telemetry off the network and off the default session.""" + with ( + patch( + "sagemaker.core.telemetry.telemetry_logging._get_default_sagemaker_session", + return_value=Mock(sagemaker_config=None), + ), + patch("sagemaker.core.telemetry.telemetry_logging._send_telemetry_request") as send, + ): + yield send + + +def _record(name, license_name="apache-2.0", hub_name=HUB): + return { + "HubContent": { + "HubContentName": name, + "HubName": hub_name, + "HubContentType": "Model", + "HubContentVersion": "1.0.0", + "HubContentArn": f"arn:aws:sagemaker:us-west-2:aws:hub-content/{hub_name}/Model/{name}/1.0.0", + "HubContentSearchKeywords": [f"@license:{license_name}", "@task:text generation"], + "HubContentDocument": json.dumps({"License": license_name, "ModelSize": "1b-10b"}), + "HubContentStatus": "Available", + "Tags": [{"Key": "owner", "Value": "aws"}], + } + } + + +def _client(pages): + client = Mock() + client.search.side_effect = pages + return client + + +class TestHubContentFilter: + def test_metadata_filter_maps_an_attribute_to_its_document_field(self): + search_filter = HubContentFilter.metadata("license", "apache-2.0") + + assert search_filter.to_request() == { + "Name": "Metadata.License", + "Operator": "Equals", + "Value": "apache-2.0", + } + + def test_metadata_filter_with_a_sequence_uses_in(self): + request = HubContentFilter.metadata("model_size", ["<1b", "1b-10b"]).to_request() + + assert request == {"Name": "Metadata.ModelSize", "Operator": "In", "Value": "<1b,1b-10b"} + + def test_metadata_filter_rejects_an_unknown_attribute(self): + with pytest.raises(ValueError, match="not a hub content metadata attribute"): + HubContentFilter.metadata("flavor", "sweet") + + def test_keyword_filter_matches_a_search_keyword(self): + assert HubContentFilter.keyword("@task:summarization").to_request() == { + "Name": "HubContentSearchKeywords", + "Operator": "Equals", + "Value": "@task:summarization", + } + + def test_exists_takes_no_value_and_other_operators_require_one(self): + assert HubContentFilter("Metadata.License", "Exists").to_request() == { + "Name": "Metadata.License", + "Operator": "Exists", + } + with pytest.raises(ValueError, match="takes a value"): + HubContentFilter("Metadata.License", "Contains").to_request() + + def test_operator_and_value_limits(self): + with pytest.raises(ValueError, match="not a Search operator"): + HubContentFilter("Metadata.License", "Like", "x").to_request() + with pytest.raises(ValueError, match="1 to 1024 characters"): + HubContentFilter("Metadata.License", "Equals", "x" * 1025).to_request() + with pytest.raises(ValueError, match="1 to 1024 characters"): + HubContentFilter("Metadata.License", "Equals", "").to_request() + + +class TestSearchExpression: + def test_expression_leads_with_hub_name_and_content_type(self): + expression = _search_expression( + HUB, "Model", [HubContentFilter.metadata("provider", "Meta")] + ) + + assert expression == { + "Operator": "And", + "Filters": [ + {"Name": "HubName", "Operator": "Equals", "Value": HUB}, + {"Name": "HubContentType", "Operator": "Equals", "Value": "Model"}, + {"Name": "Metadata.Provider", "Operator": "Equals", "Value": "Meta"}, + ], + } + + def test_expression_enforces_the_filter_cap(self): + too_many = [HubContentFilter.keyword(f"@task:{i}") for i in range(MAX_SEARCH_FILTERS - 1)] + + with pytest.raises(ValueError, match=f"at most {MAX_SEARCH_FILTERS} filters"): + _search_expression(HUB, "Model", too_many) + + +class TestHubContentSearch: + def test_search_pages_and_yields_bound_hub_content(self): + client = _client( + [ + {"Results": [_record("model-a")], "NextToken": "page-2"}, + {"Results": [_record("model-b", "llama3")]}, + ] + ) + session = Mock(name="session") + + with patch.object(HubContent, "get_sagemaker_client", return_value=client): + results = list( + HubContent.search( + HUB, + filters=[HubContentFilter.metadata("license", "apache-2.0")], + session=session, + region="us-west-2", + ) + ) + + assert [content.hub_content_name for content in results] == ["model-a", "model-b"] + assert all(isinstance(content, HubContent) for content in results) + assert results[0].hub_name == HUB + assert results[0].license == "apache-2.0" + assert results[0].model_size == "1b-10b" + assert results[1].license == "llama3" + assert results[0]._session is session + assert results[0]._region == "us-west-2" + first_request = client.search.call_args_list[0].kwargs + assert first_request["Resource"] == "HubContent" + assert first_request["SortBy"] == "HubContentName" + assert first_request["SearchExpression"]["Filters"][2] == { + "Name": "Metadata.License", + "Operator": "Equals", + "Value": "apache-2.0", + } + assert "NextToken" not in first_request + assert client.search.call_args_list[1].kwargs["NextToken"] == "page-2" + + def test_search_stops_on_a_repeated_token_and_an_empty_result(self): + client = _client( + [ + {"Results": [_record("model-a")], "NextToken": "same"}, + {"Results": [], "NextToken": "same"}, + ] + ) + + with patch.object(HubContent, "get_sagemaker_client", return_value=client): + results = list(HubContent.search(HUB)) + + assert [content.hub_content_name for content in results] == ["model-a"] + assert client.search.call_count == 2 + + def test_search_drops_record_members_the_resource_lacks(self): + client = _client([{"Results": [_record("model-a")]}]) + + with patch.object(HubContent, "get_sagemaker_client", return_value=client): + (content,) = list(HubContent.search(HUB, hub_content_type="Model")) + + assert content.hub_content_status == "Available" + assert not hasattr(content, "tags") + + def test_search_emits_telemetry_once_per_call(self, _offline_telemetry): + client = _client([{"Results": [_record("model-a")]}]) + + with patch.object(HubContent, "get_sagemaker_client", return_value=client): + list(HubContent.search(HUB)) + + assert _offline_telemetry.call_count == 1 + assert _offline_telemetry.call_args.args[5].startswith("HubContent.search&") diff --git a/sagemaker-core/tests/unit/test_instance_types.py b/sagemaker-core/tests/unit/test_instance_types.py new file mode 100644 index 0000000000..cf1fdb89c0 --- /dev/null +++ b/sagemaker-core/tests/unit/test_instance_types.py @@ -0,0 +1,282 @@ +# 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.core.instance_types`` lists SageMaker instance types and describes their hardware. + +Before this change the module held only the JumpStart ``retrieve`` functions, so +``from sagemaker.core.instance_types import InstanceSpec`` raised ``ImportError``. The only +hardware lookups lived in two private ``ModelBuilder`` methods (``_get_instance_resources`` +and ``_infer_accelerator_count_from_instance_type``), each with its own EC2 +``describe_instance_types`` call, its own ``ml.`` strip, and no batching. +""" + +from __future__ import absolute_import + +import logging +from unittest.mock import Mock + +import boto3 +import pytest +from botocore.exceptions import ClientError + +from sagemaker.core.instance_types import ( + InstanceSpec, + describe_instance_type, + describe_instance_types, + list_instance_types, +) + +G5_XLARGE_INFO = { + "InstanceType": "g5.xlarge", + "VCpuInfo": {"DefaultVCpus": 4}, + "MemoryInfo": {"SizeInMiB": 16384}, + "GpuInfo": { + "Gpus": [{"Name": "A10G", "Count": 1, "MemoryInfo": {"SizeInMiB": 24576}}], + "TotalGpuMemoryInMiB": 24576, + }, +} +P4D_24XLARGE_INFO = { + "InstanceType": "p4d.24xlarge", + "VCpuInfo": {"DefaultVCpus": 96}, + "MemoryInfo": {"SizeInMiB": 1179648}, + "GpuInfo": { + "Gpus": [{"Name": "A100", "Count": 8, "MemoryInfo": {"SizeInMiB": 40960}}], + "TotalGpuMemoryInMiB": 327680, + }, +} +M5_LARGE_INFO = { + "InstanceType": "m5.large", + "VCpuInfo": {"DefaultVCpus": 2}, + "MemoryInfo": {"SizeInMiB": 8192}, +} +KNOWN_INFOS = { + info["InstanceType"]: info for info in (G5_XLARGE_INFO, P4D_24XLARGE_INFO, M5_LARGE_INFO) +} + + +def _invalid_instance_type_error(names): + return ClientError( + { + "Error": { + "Code": "InvalidInstanceType", + "Message": f"The following supplied instance types do not exist: [{names}]", + } + }, + "DescribeInstanceTypes", + ) + + +def _ec2_client(known_infos): + """An EC2 client stub that, like EC2, rejects a whole request with one unknown name.""" + + def describe_instance_types(**request): + unknown = [name for name in request["InstanceTypes"] if name not in known_infos] + if unknown: + raise _invalid_instance_type_error(", ".join(unknown)) + return {"InstanceTypes": [known_infos[name] for name in request["InstanceTypes"]]} + + client = Mock() + client.describe_instance_types.side_effect = describe_instance_types + return client + + +def _session(ec2_client): + session = Mock(spec=boto3.Session) + session.client.return_value = ec2_client + return session + + +def _requested_names(ec2_client): + return [ + call.kwargs["InstanceTypes"] for call in ec2_client.describe_instance_types.call_args_list + ] + + +def test_list_instance_types_reads_the_enum_of_the_loaded_service_model(): + sagemaker_client = Mock() + sagemaker_client.meta.service_model.shape_for.return_value.enum = [ + "ml.t2.medium", + "ml.g5.xlarge", + ] + session = _session(sagemaker_client) + + names = list_instance_types(region="eu-west-1", session=session) + + assert names == ["ml.t2.medium", "ml.g5.xlarge"] + session.client.assert_called_once_with("sagemaker", region_name="eu-west-1") + sagemaker_client.meta.service_model.shape_for.assert_called_once_with( + "ProductionVariantInstanceType" + ) + + +def test_list_instance_types_reads_the_installed_botocore_model_without_a_request(): + session = boto3.Session( + region_name="us-west-2", aws_access_key_id="unused", aws_secret_access_key="unused" + ) + + names = list_instance_types(session=session) + + assert names[0] == "ml.t2.medium" + assert "ml.g5.xlarge" in names + assert all(name.startswith("ml.") for name in names) + assert len(names) > 200 + + +def test_describe_instance_types_sends_at_most_100_names_per_ec2_request(): + names = [f"ml.c5.{index}xlarge" for index in range(150)] + known_infos = { + name[3:]: { + "InstanceType": name[3:], + "VCpuInfo": {"DefaultVCpus": 2}, + "MemoryInfo": {"SizeInMiB": 4096}, + } + for name in names + } + ec2_client = _ec2_client(known_infos) + + specs = describe_instance_types(names, session=_session(ec2_client)) + + requested = _requested_names(ec2_client) + assert [len(batch) for batch in requested] == [100, 50] + assert requested[0][0] == "c5.0xlarge" + assert requested[1][-1] == "c5.149xlarge" + assert list(specs) == names + assert specs["ml.c5.149xlarge"] == InstanceSpec( + name="ml.c5.149xlarge", + vcpus=2, + memory_mib=4096, + gpu_count=0, + gpu_memory_mib=None, + gpu_name=None, + ) + + +def test_describe_instance_types_maps_gpu_fields_and_strips_the_ml_prefix(): + ec2_client = _ec2_client(KNOWN_INFOS) + + specs = describe_instance_types( + ["ml.g5.xlarge", "ml.p4d.24xlarge", "ml.m5.large"], session=_session(ec2_client) + ) + + assert _requested_names(ec2_client) == [["g5.xlarge", "p4d.24xlarge", "m5.large"]] + assert specs["ml.g5.xlarge"] == InstanceSpec( + name="ml.g5.xlarge", + vcpus=4, + memory_mib=16384, + gpu_count=1, + gpu_memory_mib=24576, + gpu_name="A10G", + ) + assert specs["ml.p4d.24xlarge"] == InstanceSpec( + name="ml.p4d.24xlarge", + vcpus=96, + memory_mib=1179648, + gpu_count=8, + gpu_memory_mib=327680, + gpu_name="A100", + ) + assert specs["ml.m5.large"] == InstanceSpec( + name="ml.m5.large", + vcpus=2, + memory_mib=8192, + gpu_count=0, + gpu_memory_mib=None, + gpu_name=None, + ) + + +def test_describe_instance_types_reads_every_page_of_the_ec2_response(): + ec2_client = Mock() + ec2_client.describe_instance_types.side_effect = [ + {"InstanceTypes": [G5_XLARGE_INFO], "NextToken": "page-2"}, + {"InstanceTypes": [M5_LARGE_INFO]}, + ] + + specs = describe_instance_types(["ml.g5.xlarge", "ml.m5.large"], session=_session(ec2_client)) + + assert ec2_client.describe_instance_types.call_args_list[1].kwargs["NextToken"] == "page-2" + assert specs["ml.g5.xlarge"].gpu_name == "A10G" + assert specs["ml.m5.large"].vcpus == 2 + + +def test_unknown_instance_type_gets_an_empty_spec_and_a_warning(caplog): + ec2_client = _ec2_client(KNOWN_INFOS) + + with caplog.at_level(logging.WARNING, logger="sagemaker.core.instance_types"): + specs = describe_instance_types( + ["ml.g5.xlarge", "ml.retired.xlarge", "ml.m5.large"], session=_session(ec2_client) + ) + + assert specs["ml.retired.xlarge"] == InstanceSpec( + name="ml.retired.xlarge", + vcpus=None, + memory_mib=None, + gpu_count=0, + gpu_memory_mib=None, + gpu_name=None, + ) + # After the rejected request the function asks for each name alone, so the known names + # still resolve. + assert specs["ml.g5.xlarge"].gpu_count == 1 + assert specs["ml.m5.large"].vcpus == 2 + assert _requested_names(ec2_client) == [ + ["g5.xlarge", "retired.xlarge", "m5.large"], + ["g5.xlarge"], + ["retired.xlarge"], + ["m5.large"], + ] + warnings = [record for record in caplog.records if record.levelno == logging.WARNING] + assert len(warnings) == 1 + assert "ml.retired.xlarge" in warnings[0].getMessage() + + +def test_describe_instance_type_returns_the_one_spec(): + ec2_client = _ec2_client(KNOWN_INFOS) + + spec = describe_instance_type("ml.g5.xlarge", session=_session(ec2_client)) + + assert spec.name == "ml.g5.xlarge" + assert (spec.vcpus, spec.memory_mib, spec.gpu_count) == (4, 16384, 1) + assert _requested_names(ec2_client) == [["g5.xlarge"]] + + +def test_describe_instance_type_of_an_unknown_type_has_no_hardware(): + spec = describe_instance_type("ml.retired.xlarge", session=_session(_ec2_client(KNOWN_INFOS))) + + assert spec == InstanceSpec("ml.retired.xlarge", None, None, 0, None, None) + + +def test_describe_instance_types_raises_every_other_ec2_error(): + ec2_client = Mock() + ec2_client.describe_instance_types.side_effect = ClientError( + {"Error": {"Code": "UnauthorizedOperation", "Message": "not allowed"}}, + "DescribeInstanceTypes", + ) + + with pytest.raises(ClientError, match="UnauthorizedOperation"): + describe_instance_types(["ml.g5.xlarge", "ml.m5.large"], session=_session(ec2_client)) + + assert ec2_client.describe_instance_types.call_count == 1 + + +def test_describe_instance_types_with_no_names_makes_no_request(): + ec2_client = _ec2_client(KNOWN_INFOS) + + assert describe_instance_types([], session=_session(ec2_client)) == {} + ec2_client.describe_instance_types.assert_not_called() + + +def test_instance_spec_is_frozen(): + spec = InstanceSpec("ml.g5.xlarge", 4, 16384, 1, 24576, "A10G") + + with pytest.raises(AttributeError): + spec.gpu_count = 2 diff --git a/sagemaker-core/tests/unit/test_lru_cache.py b/sagemaker-core/tests/unit/test_lru_cache.py new file mode 100644 index 0000000000..6a5776028b --- /dev/null +++ b/sagemaker-core/tests/unit/test_lru_cache.py @@ -0,0 +1,76 @@ +# 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. +"""``LRUCache.put`` on a present key hands the old value to the retrieval function. + +Before this change the call raised ``AttributeError: 'Element' object has no attribute +'element'``, because ``put`` read ``curr_value.element`` from an ``Element`` whose attribute +is ``value``. ``JumpStartModelsCache`` never reached the line, because it refreshes +through ``get`` and ``_get_item``. ``HubCatalogCache.refresh`` reaches it. +""" + +from __future__ import absolute_import + +import datetime + +import pytest + +from sagemaker.core.utilities.cache import LRUCache + + +def _cache(calls, horizon=datetime.timedelta(hours=6)): + def retrieve(key, value): + calls.append((key, value)) + return f"value-{len(calls)}" + + return LRUCache(max_cache_items=2, expiration_horizon=horizon, retrieval_function=retrieve) + + +def test_put_on_a_present_key_passes_the_old_value_to_the_retrieval_function(): + calls = [] + cache = _cache(calls) + assert cache.get("key") == ("value-1", False) + + cache.put("key") + + assert cache.get("key") == ("value-2", True) + assert calls == [("key", None), ("key", "value-1")] + + +def test_put_with_a_value_stores_it_without_a_retrieval(): + calls = [] + cache = _cache(calls) + + cache.put("key", "given") + + assert cache.get("key") == ("given", True) + assert calls == [] + + +def test_put_evicts_the_least_recently_used_key_at_the_cap(): + calls = [] + cache = _cache(calls) + cache.put("first") + cache.put("second") + + cache.put("third") + + assert "first" not in cache + assert "second" in cache and "third" in cache + assert len(cache) == 2 + + +def test_get_without_fallback_raises_for_an_absent_key(): + cache = _cache([]) + + with pytest.raises(KeyError, match="not found in LRUCache"): + cache.get("absent", data_source_fallback=False) diff --git a/sagemaker-core/tests/unit/test_model_package_group_get_or_create.py b/sagemaker-core/tests/unit/test_model_package_group_get_or_create.py new file mode 100644 index 0000000000..5ca61a206e --- /dev/null +++ b/sagemaker-core/tests/unit/test_model_package_group_get_or_create.py @@ -0,0 +1,198 @@ +# 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. +"""``ModelPackageGroup.get_or_create`` returns the existing group or creates it. + +Before the change ``ModelPackageGroup`` had only the generated ``get`` and ``create``. +``DescribeModelPackageGroup`` reports a missing group as a ``ValidationException`` whose +message ends in ``does not exist``, so ten fine-tuning templates in the SageMaker agent +skills wrapped ``create`` in a ``try`` block and parsed ``e.response["Error"]["Code"]`` +against ``("ResourceInUse", "ValidationException")`` to fall back to ``get``. +""" + +from __future__ import absolute_import + +from unittest.mock import MagicMock, patch + +import boto3 +import pytest +from botocore.exceptions import ClientError + +from sagemaker.core.model_package_group import ( + ModelPackageGroupMixin, + ModelPackageGroupSpec, + is_not_found_error, +) +from sagemaker.core.resources import Base, ModelPackageGroup +from sagemaker.core.shapes import Tag + +GROUP_NAME = "my-model-group" + + +def _client_error(code, message): + return ClientError({"Error": {"Code": code, "Message": message}}, "DescribeModelPackageGroup") + + +def _not_found_error(): + return _client_error( + "ValidationException", + f"Model Package Group {GROUP_NAME} does not exist.", + ) + + +def _group(**overrides): + kwargs = { + "model_package_group_name": GROUP_NAME, + "model_package_group_arn": f"arn:aws:sagemaker:us-west-2:123456789012:model-package-group/{GROUP_NAME}", + } + kwargs.update(overrides) + return ModelPackageGroup(**kwargs) + + +def test_model_package_group_is_a_model_package_group_mixin(): + assert issubclass(ModelPackageGroup, ModelPackageGroupMixin) + # The mixin adds behavior only: no model field, nothing in model_dump(). + assert "spec" not in ModelPackageGroup.model_fields + assert "get_or_create" not in _group().model_dump() + + +def test_get_or_create_returns_the_existing_group_without_create(): + existing = _group() + with ( + patch.object(ModelPackageGroup, "get", return_value=existing) as get, + patch.object(ModelPackageGroup, "create") as create, + ): + result = ModelPackageGroup.get_or_create( + GROUP_NAME, spec=ModelPackageGroupSpec(model_package_group_description="ignored") + ) + + assert result is existing + get.assert_called_once_with(model_package_group_name=GROUP_NAME, session=None, region=None) + create.assert_not_called() + + +def test_get_or_create_creates_the_group_when_describe_says_it_does_not_exist(): + created = _group() + with ( + patch.object(ModelPackageGroup, "get", side_effect=_not_found_error()), + patch.object(ModelPackageGroup, "create", return_value=created) as create, + ): + result = ModelPackageGroup.get_or_create(GROUP_NAME) + + assert result is created + create.assert_called_once_with(model_package_group_name=GROUP_NAME, session=None, region=None) + + +def test_get_or_create_passes_the_spec_session_and_region_to_create(): + session = boto3.Session( + aws_access_key_id="AKIDTEST", aws_secret_access_key="secret", region_name="us-west-2" + ) + tags = [Tag(key="team", value="ml")] + spec = ModelPackageGroupSpec(model_package_group_description="Fine-tuned models", tags=tags) + with ( + patch.object(ModelPackageGroup, "get", side_effect=_not_found_error()) as get, + patch.object(ModelPackageGroup, "create", return_value=_group()) as create, + ): + ModelPackageGroup.get_or_create(GROUP_NAME, spec=spec, session=session, region="eu-west-1") + + get.assert_called_once_with( + model_package_group_name=GROUP_NAME, session=session, region="eu-west-1" + ) + create.assert_called_once_with( + model_package_group_name=GROUP_NAME, + session=session, + region="eu-west-1", + model_package_group_description="Fine-tuned models", + tags=tags, + ) + + +def test_get_or_create_creates_on_a_resource_not_found_code(): + error = _client_error("ResourceNotFoundException", "No such group") + with ( + patch.object(ModelPackageGroup, "get", side_effect=error), + patch.object(ModelPackageGroup, "create", return_value=_group()) as create, + ): + ModelPackageGroup.get_or_create(GROUP_NAME) + + create.assert_called_once() + + +def test_get_or_create_propagates_every_other_error(): + denied = _client_error("AccessDeniedException", "User is not authorized") + with ( + patch.object(ModelPackageGroup, "get", side_effect=denied), + patch.object(ModelPackageGroup, "create") as create, + ): + with pytest.raises(ClientError) as raised: + ModelPackageGroup.get_or_create(GROUP_NAME) + + assert raised.value is denied + create.assert_not_called() + + +def test_a_validation_exception_with_another_message_propagates(): + invalid = _client_error("ValidationException", "1 validation error detected: name too long") + with ( + patch.object(ModelPackageGroup, "get", side_effect=invalid), + patch.object(ModelPackageGroup, "create") as create, + ): + with pytest.raises(ClientError): + ModelPackageGroup.get_or_create(GROUP_NAME) + + create.assert_not_called() + + +def test_get_or_create_end_to_end_through_the_generated_client_calls(): + client = MagicMock() + client.describe_model_package_group.side_effect = [ + _not_found_error(), + { + "ModelPackageGroupName": GROUP_NAME, + "ModelPackageGroupArn": f"arn:aws:sagemaker:us-west-2:123456789012:model-package-group/{GROUP_NAME}", + "ModelPackageGroupStatus": "Completed", + }, + ] + # The generated get() and create() call Base.get_sagemaker_client, not cls. + with patch.object(Base, "get_sagemaker_client", return_value=client): + result = ModelPackageGroup.get_or_create( + GROUP_NAME, spec=ModelPackageGroupSpec(model_package_group_description="desc") + ) + + client.create_model_package_group.assert_called_once_with( + ModelPackageGroupName=GROUP_NAME, ModelPackageGroupDescription="desc" + ) + assert client.describe_model_package_group.call_count == 2 + assert result.model_package_group_status == "Completed" + + +def test_spec_create_kwargs_drops_unset_fields(): + assert ModelPackageGroupSpec().create_kwargs() == {} + assert ModelPackageGroupSpec(model_package_group_description="d").create_kwargs() == { + "model_package_group_description": "d" + } + + +@pytest.mark.parametrize( + "code, message, expected", + [ + ("ValidationException", "Model Package Group x does not exist.", True), + ("ValidationException", "Could not find model package group x", True), + ("ResourceNotFound", "Resource not found", True), + ("ResourceNotFoundException", "Resource not found", True), + ("ValidationException", "1 validation error detected", False), + ("AccessDeniedException", "does not exist", False), + ("ThrottlingException", "Rate exceeded", False), + ], +) +def test_is_not_found_error(code, message, expected): + assert is_not_found_error(_client_error(code, message)) is expected diff --git a/sagemaker-core/tests/unit/test_processing_start_new_arn.py b/sagemaker-core/tests/unit/test_processing_start_new_arn.py new file mode 100644 index 0000000000..6b3ac8b652 --- /dev/null +++ b/sagemaker-core/tests/unit/test_processing_start_new_arn.py @@ -0,0 +1,178 @@ +# 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. +"""``Processor._start_new`` returns a session-bound ``ProcessingJob`` with its ARN. + +Before the change ``_start_new`` built ``ProcessingJob(**transformed)`` from the create +request only. The job had no ``processing_job_arn``, no ``processing_job_status``, and no +client context, so ``processor.latest_job.refresh()`` polled the process default region +(issue #5796) and four processing templates in the SageMaker agent skills called +``describe_processing_job`` on a boto3 client to read the ARN. +""" + +from __future__ import absolute_import + +from unittest.mock import MagicMock, Mock, patch + +import pytest + +from sagemaker.core.local.local_session import LocalSession +from sagemaker.core.processing import Processor +from sagemaker.core.resources import Base +from sagemaker.core.utils.utils import Unassigned +from sagemaker.core.workflow.pipeline_context import PipelineSession + +JOB_NAME = "test-job" +JOB_ARN = f"arn:aws:sagemaker:eu-west-1:123456789012:processing-job/{JOB_NAME}" +ROLE = "arn:aws:iam::123456789012:role/SageMakerRole" +PROCESS_ARGS = { + "job_name": JOB_NAME, + "inputs": [], + "output_config": {"Outputs": []}, + "resources": { + "ClusterConfig": { + "InstanceCount": 1, + "InstanceType": "ml.m5.xlarge", + "VolumeSizeInGB": 30, + } + }, + "stopping_condition": None, + "app_specification": {"ImageUri": "test-image:latest"}, + "environment": None, + "network_config": None, + "role_arn": ROLE, + "tags": [], +} +DESCRIBE_RESPONSE = { + "ProcessingJobName": JOB_NAME, + "ProcessingJobArn": JOB_ARN, + "ProcessingJobStatus": "InProgress", + "AppSpecification": {"ImageUri": "test-image:latest"}, + "ProcessingResources": { + "ClusterConfig": { + "InstanceCount": 1, + "InstanceType": "ml.m5.xlarge", + "VolumeSizeInGB": 30, + } + }, + "RoleArn": ROLE, +} + + +def _session(spec=None): + session = Mock(spec=spec) if spec else Mock() + session.boto_session = Mock() + session.boto_session.region_name = "eu-west-1" + session.boto_region_name = "eu-west-1" + session.sagemaker_client = Mock() + session.sagemaker_config = {} + session._intercept_create_request = Mock( + side_effect=lambda request, create, name: create(request) + ) + return session + + +@pytest.fixture +def mock_session(): + return _session() + + +@pytest.fixture +def recorded_clients(): + """Patch Base.get_sagemaker_client; record (session, region) per call.""" + calls = [] + client = MagicMock() + client.describe_processing_job.return_value = DESCRIBE_RESPONSE + + def fake(session=None, region_name=None, service_name="sagemaker"): + calls.append((session, region_name)) + return client + + with patch.object(Base, "get_sagemaker_client", side_effect=fake): + yield calls, client + + +def _processor(session): + processor = Processor( + role=ROLE, + image_uri="test-image:latest", + instance_count=1, + instance_type="ml.m5.xlarge", + sagemaker_session=session, + ) + processor._current_job_name = JOB_NAME + return processor + + +def _start_new(processor): + with patch.object(processor, "_get_process_args", return_value=dict(PROCESS_ARGS)): + return processor._start_new([], [], None) + + +def test_start_new_returns_the_job_with_its_arn_and_status(mock_session, recorded_clients): + _, client = recorded_clients + + job = _start_new(_processor(mock_session)) + + mock_session.sagemaker_client.create_processing_job.assert_called_once() + client.describe_processing_job.assert_called_once_with(ProcessingJobName=JOB_NAME) + assert job.processing_job_name == JOB_NAME + assert job.processing_job_arn == JOB_ARN + assert job.processing_job_status == "InProgress" + + +def test_start_new_describes_through_the_session_and_its_region(mock_session, recorded_clients): + calls, _ = recorded_clients + + job = _start_new(_processor(mock_session)) + + assert calls == [(mock_session.boto_session, "eu-west-1")] + assert job._session is mock_session.boto_session + assert job._region == "eu-west-1" + + +def test_refresh_on_the_returned_job_keeps_the_session_region(mock_session, recorded_clients): + calls, client = recorded_clients + job = _start_new(_processor(mock_session)) + calls.clear() + + job.refresh() + + assert calls == [(mock_session.boto_session, "eu-west-1")] + assert client.describe_processing_job.call_count == 2 + + +def test_pipeline_session_captures_the_step_arguments_without_a_describe(recorded_clients): + calls, client = recorded_clients + pipeline_session = _session(spec=PipelineSession) + pipeline_session._intercept_create_request = Mock(return_value=None) + + result = _start_new(_processor(pipeline_session)) + + assert result is None + pipeline_session._intercept_create_request.assert_called_once() + pipeline_session.sagemaker_client.create_processing_job.assert_not_called() + client.describe_processing_job.assert_not_called() + assert calls == [] + + +def test_local_session_returns_the_request_job_without_a_describe(recorded_clients): + calls, client = recorded_clients + local_session = _session(spec=LocalSession) + + job = _start_new(_processor(local_session)) + + local_session.sagemaker_client.create_processing_job.assert_called_once() + client.describe_processing_job.assert_not_called() + assert calls == [] + assert job.processing_job_name == JOB_NAME + assert isinstance(job.processing_job_arn, Unassigned) diff --git a/sagemaker-core/tests/unit/test_resource_iterator_max_workers.py b/sagemaker-core/tests/unit/test_resource_iterator_max_workers.py new file mode 100644 index 0000000000..3fa77ae295 --- /dev/null +++ b/sagemaker-core/tests/unit/test_resource_iterator_max_workers.py @@ -0,0 +1,171 @@ +# 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. +"""``ResourceIterator(max_workers=N)`` describes N listed resources at a time. + +Before this change the iterator described one resource per ``next()`` call, so a list of +the public hub (about 770 models) spent one ``DescribeHubContent`` round trip per model +in series, about 95 seconds. ``HubContent.get_all`` and ``get_all_versions`` take +``max_workers`` and hand it to the iterator, with 4 as the default. +""" + +from __future__ import absolute_import + +import inspect +import threading +import time +from unittest.mock import MagicMock, patch + +import pytest + +from sagemaker.core.resources import Base, HubContent +from sagemaker.core.utils.utils import ResourceIterator, SageMakerClient + +HUB = "SageMakerPublicHub" +DESCRIBE_SECONDS = 0.05 + + +def _summary(name): + return { + "HubContentName": name, + "HubContentArn": f"arn:aws:sagemaker:us-west-2:aws:hub-content/{HUB}/Model/{name}/1.0.0", + "HubContentVersion": "1.0.0", + "HubContentType": "Model", + "DocumentSchemaVersion": "2.0.0", + "HubContentStatus": "Available", + } + + +class _Recorder: + """A fake client whose describe call sleeps and records its concurrency.""" + + def __init__(self, names, fail_on=None): + self.client = MagicMock() + self.client.list_hub_contents.return_value = { + "HubContentSummaries": [_summary(name) for name in names] + } + self.client.describe_hub_content.side_effect = self._describe + self.fail_on = fail_on + self.in_flight = 0 + self.peak_in_flight = 0 + self.started = [] + self._lock = threading.Lock() + + def _describe(self, **kwargs): + with self._lock: + self.in_flight += 1 + self.peak_in_flight = max(self.peak_in_flight, self.in_flight) + self.started.append(kwargs["HubContentName"]) + try: + time.sleep(DESCRIBE_SECONDS) + if kwargs["HubContentName"] == self.fail_on: + raise RuntimeError(f"describe of {self.fail_on} failed") + return { + **_summary(kwargs["HubContentName"]), + "HubName": kwargs["HubName"], + "HubContentDocument": "{}", + } + finally: + with self._lock: + self.in_flight -= 1 + + +def _iterator(recorder, max_workers): + return ResourceIterator( + client=recorder.client, + summaries_key="HubContentSummaries", + summary_name="HubContentInfo", + resource_cls=HubContent, + list_method="list_hub_contents", + list_method_kwargs={"HubName": HUB, "HubContentType": "Model"}, + extra_init_data={"hub_name": HUB}, + max_workers=max_workers, + ) + + +@pytest.fixture(autouse=True) +def _reset_client_cache(): + SageMakerClient.reset() + yield + SageMakerClient.reset() + + +def test_max_workers_describes_resources_concurrently_in_list_order(): + names = [f"model-{index:02d}" for index in range(12)] + recorder = _Recorder(names) + + with patch.object(Base, "get_sagemaker_client", return_value=recorder.client): + started = time.monotonic() + contents = list(_iterator(recorder, max_workers=4)) + elapsed = time.monotonic() - started + + assert [content.hub_content_name for content in contents] == names + assert all(content.hub_name == HUB for content in contents) + assert recorder.peak_in_flight > 1 + assert recorder.client.describe_hub_content.call_count == len(names) + assert elapsed < len(names) * DESCRIBE_SECONDS + + +def test_default_max_workers_describes_one_resource_at_a_time(): + recorder = _Recorder(["model-a", "model-b", "model-c"]) + + with patch.object(Base, "get_sagemaker_client", return_value=recorder.client): + contents = list(_iterator(recorder, max_workers=1)) + + assert [content.hub_content_name for content in contents] == ["model-a", "model-b", "model-c"] + assert recorder.peak_in_flight == 1 + + +def test_stopping_early_describes_at_most_max_workers_resources(): + names = [f"model-{index:02d}" for index in range(20)] + recorder = _Recorder(names) + + with patch.object(Base, "get_sagemaker_client", return_value=recorder.client): + iterator = _iterator(recorder, max_workers=4) + first = next(iterator) + iterator._shutdown_executor() + + assert first.hub_content_name == "model-00" + assert recorder.client.describe_hub_content.call_count <= 4 + + +def test_a_describe_error_raises_in_list_order_and_releases_the_pool(): + recorder = _Recorder(["model-a", "model-b", "model-c"], fail_on="model-b") + + with patch.object(Base, "get_sagemaker_client", return_value=recorder.client): + iterator = _iterator(recorder, max_workers=3) + assert next(iterator).hub_content_name == "model-a" + with pytest.raises(RuntimeError, match="describe of model-b failed"): + next(iterator) + iterator._shutdown_executor() + + assert iterator._executor is None + + +def test_max_workers_below_one_is_rejected(): + with pytest.raises(ValueError, match="max_workers must be at least 1"): + _iterator(_Recorder([]), max_workers=0) + + +def test_hub_content_list_methods_take_max_workers(): + for method in (HubContent.get_all, HubContent.get_all_versions): + parameter = inspect.signature(method).parameters["max_workers"] + assert parameter.default == 4, method.__qualname__ + + +def test_hub_content_get_all_hands_max_workers_to_the_iterator(): + recorder = _Recorder(["model-a"]) + + with patch.object(Base, "get_sagemaker_client", return_value=recorder.client): + iterator = HubContent.get_all(hub_name=HUB, hub_content_type="Model", max_workers=2) + + assert iterator.max_workers == 2 diff --git a/sagemaker-core/tests/unit/test_wait_raise_on_stopped.py b/sagemaker-core/tests/unit/test_wait_raise_on_stopped.py new file mode 100644 index 0000000000..39185e8454 --- /dev/null +++ b/sagemaker-core/tests/unit/test_wait_raise_on_stopped.py @@ -0,0 +1,131 @@ +# 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. +"""Generated ``wait()`` methods take ``raise_on_stopped`` and raise on a stopped resource. + +Before the change ``wait()`` raised ``FailedStatusError`` only for a status that contains +``failed``. ``Stopped`` is a terminal state too, so ``TrainingJob.wait()`` returned as if +the job completed, and three training templates in the SageMaker agent skills read +``training_job_status`` after the wait and raised on ``Stopped`` themselves. +""" + +from __future__ import absolute_import + +import inspect +from unittest.mock import MagicMock, patch + +import pytest + +from sagemaker.core import resources +from sagemaker.core.resources import Base, ProcessingJob, TrainingJob +from sagemaker.core.shapes import ProcessingClusterConfig, ProcessingResources +from sagemaker.core.utils.exceptions import FailedStatusError + +JOB_NAME = "my-job" + + +def _describe_training_job(status, failure_reason=None): + response = { + "TrainingJobName": JOB_NAME, + "TrainingJobArn": f"arn:aws:sagemaker:us-west-2:123456789012:training-job/{JOB_NAME}", + "TrainingJobStatus": status, + "SecondaryStatus": status, + } + if failure_reason: + response["FailureReason"] = failure_reason + return response + + +@pytest.fixture +def training_client(): + client = MagicMock() + with patch.object(Base, "get_sagemaker_client", return_value=client): + yield client + + +def test_every_generated_wait_takes_raise_on_stopped(): + wait_methods = [ + cls.wait + for _, cls in inspect.getmembers(resources, inspect.isclass) + if issubclass(cls, Base) and "wait" in vars(cls) + ] + + assert len(wait_methods) >= 10 + for wait in wait_methods: + parameter = inspect.signature(wait).parameters["raise_on_stopped"] + assert parameter.default is False, wait.__qualname__ + + +def test_wait_returns_on_stopped_by_default(training_client): + training_client.describe_training_job.return_value = _describe_training_job("Stopped") + + assert TrainingJob(training_job_name=JOB_NAME).wait(poll=0) is None + + +def test_wait_raises_on_stopped_when_asked(training_client): + training_client.describe_training_job.return_value = _describe_training_job( + "Stopped", failure_reason="Stopped by user" + ) + + with pytest.raises(FailedStatusError) as raised: + TrainingJob(training_job_name=JOB_NAME).wait(poll=0, raise_on_stopped=True) + + message = str(raised.value) + assert "TrainingJob" in message + assert "Stopped" in message + assert "Stopped by user" in message + + +def test_wait_with_raise_on_stopped_returns_on_completed(training_client): + training_client.describe_training_job.return_value = _describe_training_job("Completed") + + assert TrainingJob(training_job_name=JOB_NAME).wait(poll=0, raise_on_stopped=True) is None + + +def test_wait_with_raise_on_stopped_still_raises_on_failed(training_client): + training_client.describe_training_job.return_value = _describe_training_job( + "Failed", failure_reason="AlgorithmError" + ) + + with pytest.raises(FailedStatusError, match="Failed"): + TrainingJob(training_job_name=JOB_NAME).wait(poll=0, raise_on_stopped=True) + + +def test_wait_polls_through_stopping_before_it_raises(training_client): + training_client.describe_training_job.side_effect = [ + _describe_training_job("Stopping"), + _describe_training_job("Stopped"), + ] + + with pytest.raises(FailedStatusError): + TrainingJob(training_job_name=JOB_NAME).wait(poll=0, raise_on_stopped=True) + + assert training_client.describe_training_job.call_count == 2 + + +def test_processing_job_wait_raises_on_stopped_when_asked(training_client): + training_client.describe_processing_job.return_value = { + "ProcessingJobName": JOB_NAME, + "ProcessingJobStatus": "Stopped", + } + # The generated ProcessingJob.wait reads the instance count before the first refresh. + job = ProcessingJob( + processing_job_name=JOB_NAME, + processing_resources=ProcessingResources( + cluster_config=ProcessingClusterConfig( + instance_count=1, instance_type="ml.m5.large", volume_size_in_gb=30 + ) + ), + ) + + with pytest.raises(FailedStatusError, match="ProcessingJob"): + job.wait(poll=0, raise_on_stopped=True) diff --git a/sagemaker-serve/src/sagemaker/serve/bedrock_deployment.py b/sagemaker-serve/src/sagemaker/serve/bedrock_deployment.py new file mode 100644 index 0000000000..a71b57b848 --- /dev/null +++ b/sagemaker-serve/src/sagemaker/serve/bedrock_deployment.py @@ -0,0 +1,373 @@ +# 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. +"""Bedrock foundation model listing and the deployment handle of ``BedrockModelBuilder``.""" + +from __future__ import absolute_import + +import json +import logging +import time +from dataclasses import dataclass +from functools import partial +from typing import Any, Callable, Dict, List, Optional, Union + +import boto3 +from botocore.exceptions import ClientError + +logger = logging.getLogger(__name__) + +ON_DEMAND_INFERENCE_TYPE = "ON_DEMAND" +INFERENCE_PROFILE_INFERENCE_TYPE = "INFERENCE_PROFILE" +SYSTEM_DEFINED_PROFILE_TYPE = "SYSTEM_DEFINED" +ACTIVE_PROFILE_STATUS = "ACTIVE" +GLOBAL_PROFILE_PREFIX = "global." +MODEL_NOT_READY_ERROR_CODE = "ModelNotReadyException" +JSON_CONTENT_TYPE = "application/json" +PROBE_TEXT = "Hi" +# A readiness probe asks for one token so that it costs almost nothing. +PROBE_MAX_TOKENS = 1 + + +@dataclass(frozen=True) +class BedrockFoundationModel: + """One foundation model that the region can invoke. + + Attributes: + model_id: The Bedrock model id, for example ``anthropic.claude-3-haiku-20240307-v1:0``. + model_arn: The foundation model ARN. + model_name: The branded model name. + provider_name: The model provider. + input_modalities: ``TEXT``, ``IMAGE``, or ``EMBEDDING`` values the model accepts. + output_modalities: ``TEXT``, ``IMAGE``, or ``EMBEDDING`` values the model returns. + inference_types: ``ON_DEMAND`` and ``PROVISIONED`` from the model summary, plus + ``INFERENCE_PROFILE`` when a system-defined inference profile routes to the model. + inference_profile_id: The cross-region inference profile id that routes to the model, + or None. Pass it as ``modelId`` when ``ON_DEMAND`` is absent. + """ + + model_id: str + model_arn: str + model_name: str + provider_name: str + input_modalities: List[str] + output_modalities: List[str] + inference_types: List[str] + inference_profile_id: Optional[str] + + +def _collect_pages(list_call: Callable[..., Dict[str, Any]], result_key: str) -> List[Dict]: + """Follow ``nextToken`` through every page of a Bedrock list call. + + ``ListInferenceProfiles`` pages with ``nextToken``. ``ListFoundationModels`` has no + token in botocore 1.43.118 and returns one page, which this loop also handles. + + Args: + list_call: The client method with its filter arguments already bound. + result_key: The response key that holds the page items. + + Returns: + The items of every page in order. + """ + items: List[Dict] = [] + params: Dict[str, str] = {} + while True: + response = list_call(**params) + items.extend(response.get(result_key, [])) + next_token = response.get("nextToken") + if not next_token: + return items + params = {"nextToken": next_token} + + +def _profile_ids_by_model_id(profiles: List[Dict[str, Any]]) -> Dict[str, List[str]]: + """Map each model id to the sorted active system-defined profile ids that route to it. + + Args: + profiles: The ``inferenceProfileSummaries`` items. + + Returns: + Model id to sorted profile ids. The model id is the last ARN path segment of + ``models[].modelArn``, because a cross-region profile lists ARNs of other regions. + """ + profile_ids: Dict[str, set] = {} + for profile in profiles: + if profile.get("status") != ACTIVE_PROFILE_STATUS: + continue + if profile.get("type") != SYSTEM_DEFINED_PROFILE_TYPE: + continue + for routed_model in profile.get("models", []): + model_id = routed_model.get("modelArn", "").rsplit("/", 1)[-1] + profile_ids.setdefault(model_id, set()).add(profile["inferenceProfileId"]) + return {model_id: sorted(ids) for model_id, ids in profile_ids.items()} + + +def _choose_profile_id(profile_ids: List[str]) -> Optional[str]: + """Prefer the geographic profile to the global one, then the first sorted id. + + Args: + profile_ids: The sorted profile ids that route to one model. + + Returns: + The chosen profile id, or None when the list is empty. + """ + if not profile_ids: + return None + geographic_ids = [ + profile_id for profile_id in profile_ids if not profile_id.startswith(GLOBAL_PROFILE_PREFIX) + ] + return (geographic_ids or profile_ids)[0] + + +def _to_foundation_model( + summary: Dict[str, Any], profile_ids: List[str] +) -> Optional[BedrockFoundationModel]: + """Build the dataclass for one model summary, or None when nothing can invoke the model. + + Args: + summary: One ``FoundationModelSummary`` item. + profile_ids: The system-defined profile ids that route to the model. + + Returns: + The model when its inference types include ``ON_DEMAND`` or a profile routes to it, + else None. + """ + inference_types = list(summary.get("inferenceTypesSupported", [])) + profile_id = _choose_profile_id(profile_ids) + if profile_id: + inference_types.append(INFERENCE_PROFILE_INFERENCE_TYPE) + if ON_DEMAND_INFERENCE_TYPE not in inference_types and not profile_id: + return None + return BedrockFoundationModel( + model_id=summary["modelId"], + model_arn=summary.get("modelArn", ""), + model_name=summary.get("modelName", ""), + provider_name=summary.get("providerName", ""), + input_modalities=list(summary.get("inputModalities", [])), + output_modalities=list(summary.get("outputModalities", [])), + inference_types=inference_types, + inference_profile_id=profile_id, + ) + + +def list_invocable_foundation_models( + region: Optional[str] = None, + boto_session: Optional[boto3.Session] = None, + output_modality: str = "TEXT", +) -> List[BedrockFoundationModel]: + """List the foundation models that ``InvokeModel`` or ``Converse`` accepts in the region. + + Calls ``ListFoundationModels`` filtered by output modality and ``ListInferenceProfiles`` + filtered to system-defined profiles. The result keeps a model when + ``inferenceTypesSupported`` includes ``ON_DEMAND`` or an active cross-region profile + routes to it. + + Args: + region: The region to list. None keeps the session region. + boto_session: The session to use. None creates ``boto3.Session()``. + output_modality: ``TEXT``, ``IMAGE``, or ``EMBEDDING``. Defaults to ``TEXT``. + + Returns: + The invocable models in the order that ``ListFoundationModels`` returned them. + """ + session = boto_session or boto3.Session() + client = session.client("bedrock", region_name=region) + summaries = _collect_pages( + partial(client.list_foundation_models, byOutputModality=output_modality), + "modelSummaries", + ) + profiles = _collect_pages( + partial(client.list_inference_profiles, typeEquals=SYSTEM_DEFINED_PROFILE_TYPE), + "inferenceProfileSummaries", + ) + profile_ids_by_model = _profile_ids_by_model_id(profiles) + models = [] + for summary in summaries: + model = _to_foundation_model(summary, profile_ids_by_model.get(summary["modelId"], [])) + if model is not None: + models.append(model) + logger.info("%d of %d %s models are invocable", len(models), len(summaries), output_modality) + return models + + +def _error_code(error: ClientError) -> Optional[str]: + """Return the service error code of a ``ClientError``.""" + return error.response.get("Error", {}).get("Code") + + +def _is_json_media_type(media_type: str) -> bool: + """Return True when the media type is ``application/json`` with or without parameters.""" + return media_type.split(";", 1)[0].strip().lower() == JSON_CONTENT_TYPE + + +class BedrockDeployment: + """A handle on a Bedrock model that ``BedrockModelBuilder.deploy()`` made invocable. + + ``invoke`` and ``converse`` send requests to ``deployment_arn`` when it is set (a Nova + custom model deployment) and to ``model_arn`` otherwise (an imported model). The first + call creates the ``bedrock-runtime`` client. + + Args: + model_arn: The custom model ARN or the imported model ARN. + boto_session: The session that creates the ``bedrock-runtime`` client. + region: The region override. None keeps the session region. + deployment_arn: The custom model deployment ARN, or None for an imported model. + """ + + def __init__( + self, + model_arn: str, + boto_session: boto3.Session, + region: Optional[str] = None, + deployment_arn: Optional[str] = None, + ): + """Store the identifiers and the session. The constructor sends no request.""" + self._model_arn = model_arn + self._deployment_arn = deployment_arn + self._boto_session = boto_session + self._region = region + self._runtime_client = None + + @property + def model_arn(self) -> str: + """The custom model ARN or the imported model ARN.""" + return self._model_arn + + @property + def deployment_arn(self) -> Optional[str]: + """The custom model deployment ARN, or None for an imported model.""" + return self._deployment_arn + + @property + def _invocation_target(self) -> str: + """The ``modelId`` that ``InvokeModel`` and ``Converse`` take.""" + return self._deployment_arn or self._model_arn + + def _runtime(self): + """Return the ``bedrock-runtime`` client, created on first use.""" + if self._runtime_client is None: + self._runtime_client = self._boto_session.client( + "bedrock-runtime", region_name=self._region + ) + return self._runtime_client + + def wait_until_ready( + self, + timeout: int = 1800, + poll: int = 30, + probe_body: Optional[Dict[str, Any]] = None, + ) -> None: + """Block until the model answers a one-token request. + + Bedrock returns ``ModelNotReadyException`` for some minutes after an import or a + deployment. The probe calls ``Converse`` for a custom model deployment, and + ``InvokeModel`` with ``probe_body`` for an imported model. The method raises every + other service error at once, so pass a ``probe_body`` that the model accepts. + + Args: + timeout: Maximum seconds to wait. Defaults to 1800. + poll: Seconds between probes. Defaults to 30. + probe_body: The ``InvokeModel`` body for an imported model. Defaults to + ``{"prompt": "Hi", "max_tokens": 1}``, the completion schema of imported + text models. An embedding model takes ``{"inputText": "Hi"}``. + + Raises: + TimeoutError: When the model is not ready after ``timeout`` seconds. + botocore.exceptions.ClientError: When a probe fails for another reason. + """ + if probe_body is None: + probe_body = {"prompt": PROBE_TEXT, "max_tokens": PROBE_MAX_TOKENS} + elapsed = 0 + while not self._probe_ready(probe_body): + if elapsed >= timeout: + raise TimeoutError( + f"Bedrock model {self._invocation_target} was not ready after {timeout}s." + ) + logger.info( + "Bedrock model %s is not ready yet (elapsed %ds)", self._invocation_target, elapsed + ) + time.sleep(poll) + elapsed += poll + + def _probe_ready(self, probe_body: Dict[str, Any]) -> bool: + """Send one probe. Return False on ``ModelNotReadyException``, raise other errors.""" + try: + self._send_probe(probe_body) + except ClientError as error: + if _error_code(error) != MODEL_NOT_READY_ERROR_CODE: + raise + return False + return True + + def _send_probe(self, probe_body: Dict[str, Any]) -> None: + """Send the one-token probe that fits the deployment type.""" + if self._deployment_arn: + self.converse( + messages=[{"role": "user", "content": [{"text": PROBE_TEXT}]}], + inference_config={"maxTokens": PROBE_MAX_TOKENS}, + ) + return + self.invoke(probe_body) + + def invoke( + self, + body: Union[str, bytes, Dict[str, Any]], + content_type: str = JSON_CONTENT_TYPE, + accept: str = JSON_CONTENT_TYPE, + ) -> Dict[str, Any]: + """Call ``InvokeModel`` and read the response body. + + Args: + body: The request body. ``json.dumps`` serializes a dict. + content_type: The request media type. Defaults to ``application/json``. + accept: The response media type. Defaults to ``application/json``. + + Returns: + The parsed JSON body when ``accept`` is JSON. Otherwise a dict with the raw + ``body`` bytes and the response ``contentType``. + """ + payload = json.dumps(body) if isinstance(body, dict) else body + response = self._runtime().invoke_model( + modelId=self._invocation_target, + body=payload, + contentType=content_type, + accept=accept, + ) + raw_body = response["body"].read() + if _is_json_media_type(accept): + return json.loads(raw_body) + return {"body": raw_body, "contentType": response.get("contentType")} + + def converse( + self, + messages: List[Dict[str, Any]], + inference_config: Optional[Dict[str, Any]] = None, + system: Optional[List[Dict[str, Any]]] = None, + ) -> Dict[str, Any]: + """Call ``Converse`` and return its response. + + Args: + messages: The conversation, for example + ``[{"role": "user", "content": [{"text": "Hi"}]}]``. + inference_config: The ``inferenceConfig`` map (``maxTokens``, ``temperature``, + ``topP``, ``stopSequences``), or None. + system: The ``system`` content blocks, or None. + + Returns: + The ``Converse`` response with ``output``, ``stopReason``, ``usage``, and ``metrics``. + """ + params: Dict[str, Any] = {"modelId": self._invocation_target, "messages": messages} + if inference_config is not None: + params["inferenceConfig"] = inference_config + if system is not None: + params["system"] = system + return self._runtime().converse(**params) diff --git a/sagemaker-serve/src/sagemaker/serve/bedrock_model_builder.py b/sagemaker-serve/src/sagemaker/serve/bedrock_model_builder.py index 0c1693b2f6..c0ae71b1a0 100644 --- a/sagemaker-serve/src/sagemaker/serve/bedrock_model_builder.py +++ b/sagemaker-serve/src/sagemaker/serve/bedrock_model_builder.py @@ -23,6 +23,7 @@ from urllib.parse import urlparse from sagemaker.serve.utils.model_package_utils import is_restricted_model_package +from sagemaker.serve.bedrock_deployment import BedrockDeployment from sagemaker.serve.model_reuse import ( build_source_tag, find_active_bedrock_deployment_for_model, @@ -112,6 +113,14 @@ class BedrockModelBuilder: logging a warning rather than raising. To avoid redeploying the same model artifacts, opt into resource reuse with ``reuse_resources=True``. + Deployment handle: + Every successful ``deploy()`` sets ``deployment`` to a + :class:`~sagemaker.serve.bedrock_deployment.BedrockDeployment` that holds the + model ARN, plus the custom model deployment ARN on the Nova path. + ``latest_deployment()`` returns that handle and raises ``ValueError`` before the + first successful ``deploy()``. Call ``deployment.wait_until_ready()``, then + ``deployment.invoke(body)`` or ``deployment.converse(messages)``. + Args: model: The model to deploy. Can be a ModelTrainer, BaseTrainer (SFTTrainer, DPOTrainer, RLVRTrainer, RLAIFTrainer, etc.), @@ -144,6 +153,7 @@ def __init__( self._bedrock_client = None self._sagemaker_client = None self._imported_model_id = None + self.deployment: Optional[BedrockDeployment] = None self.sagemaker_session = Session() self.boto_session = self.sagemaker_session.boto_session @@ -393,10 +403,12 @@ def deploy( model_arn, existing_deployment, ) - return { + response = { "modelArn": model_arn, "customModelDeploymentArn": existing_deployment, } + self._record_deployment(response) + return response logger.info( "Reusing existing custom model %s (matched model-source tag); " "creating a new deployment on it. Pass reuse_resources=False to " @@ -408,6 +420,7 @@ def deploy( model_arn=model_arn, deployment_name=deploy_name ) response.setdefault("modelArn", model_arn) + self._record_deployment(response) return response if self._is_rmp: @@ -445,6 +458,7 @@ def deploy( deploy_name = deployment_name or f"{custom_model_name}-deployment" response = self.create_deployment(model_arn=model_arn, deployment_name=deploy_name) response.setdefault("modelArn", model_arn) + self._record_deployment(response) return response else: # Resolve and validate the Bedrock role: the provided role_arn if given, @@ -495,6 +509,7 @@ def deploy( modelIdentifier=model_arn ) self._imported_model_id = model_details.get("modelName") + self._record_deployment(model_details) return model_details # 2. Otherwise, an import job may already be running for this @@ -515,6 +530,7 @@ def deploy( jobIdentifier=job_arn ) self._imported_model_id = job_details.get("importedModelName") + self._record_deployment(job_details) return job_details # If artifacts are a tar.gz, extract to S3 first (Bedrock requires uncompressed format) @@ -575,8 +591,49 @@ def deploy( # Return the completed job details and store imported model ID job_details = self._get_bedrock_client().get_model_import_job(jobIdentifier=job_arn) self._imported_model_id = job_details.get("importedModelName") + self._record_deployment(job_details) return job_details + def latest_deployment(self) -> BedrockDeployment: + """Return the handle that the most recent successful ``deploy()`` set. + + Returns: + The :class:`~sagemaker.serve.bedrock_deployment.BedrockDeployment` for the + imported model or the Nova custom model deployment. + + Raises: + ValueError: When ``deploy()`` has not run, or its response had no model ARN. + """ + if self.deployment is None: + raise ValueError( + "No Bedrock deployment is available. Call deploy() first, then " + "latest_deployment()." + ) + return self.deployment + + def _record_deployment(self, response: Dict[str, Any]) -> None: + """Set ``self.deployment`` from the response that ``deploy()`` returns. + + Args: + response: The ``get_model_import_job`` response (``importedModelArn``), the + ``get_imported_model`` response (``modelArn``), or the Nova response + (``modelArn`` and ``customModelDeploymentArn``). A response with no model + ARN resets ``self.deployment`` to None and logs a warning. + """ + model_arn = response.get("importedModelArn") or response.get("modelArn") + if not model_arn: + logger.warning( + "The Bedrock response carries no model ARN, so latest_deployment() is " + "unavailable until the next deploy() call." + ) + self.deployment = None + return + self.deployment = BedrockDeployment( + model_arn=model_arn, + boto_session=self.boto_session, + deployment_arn=response.get("customModelDeploymentArn"), + ) + def create_deployment( self, model_arn: str, diff --git a/sagemaker-serve/src/sagemaker/serve/model_builder.py b/sagemaker-serve/src/sagemaker/serve/model_builder.py index 07e7c18268..95fe418ca7 100644 --- a/sagemaker-serve/src/sagemaker/serve/model_builder.py +++ b/sagemaker-serve/src/sagemaker/serve/model_builder.py @@ -25,7 +25,7 @@ import logging import uuid import platform -from dataclasses import dataclass, field +from dataclasses import dataclass, field, fields from pathlib import Path from typing import Any, Dict, List, Optional, TYPE_CHECKING, Union, Set from botocore.exceptions import ClientError @@ -49,15 +49,18 @@ InferenceComponentSpecification, InferenceComponentContainerSpecification, InferenceComponentRuntimeConfig, + InstancePool, ProductionVariant, ) +from sagemaker.core.shapes import DataCaptureConfig as DataCaptureConfigShape from sagemaker.core.resources import ( ModelPackage, ModelPackageGroup, ModelCard, ModelPackageModelCard, ) -from sagemaker.core.utils.utils import logger +from sagemaker.core.utils.utils import Unassigned, logger, serialize +from sagemaker.core.utils.code_injection.codec import transform from sagemaker.core.exceptions import CapacityError, UnexpectedStatusException from sagemaker.core.helper import session_helper from sagemaker.core.helper.iam_role_resolver import resolve_and_validate_role @@ -101,8 +104,10 @@ from sagemaker.core.jumpstart.configs import JumpStartConfig from sagemaker.core.jumpstart.utils import add_jumpstart_uri_tags from sagemaker.core.jumpstart.artifacts.kwargs import _retrieve_model_deploy_kwargs +from sagemaker.core.instance_types import describe_instance_type from sagemaker.core.inference_config import AsyncInferenceConfig, ServerlessInferenceConfig +from sagemaker.core.model_monitor.data_capture_config import DataCaptureConfig from sagemaker.serve.batch_inference.batch_transform_inference_config import ( BatchTransformInferenceConfig, ) @@ -173,6 +178,35 @@ from sagemaker.serve.ai_inference_recommender.workload import Workload _LOWEST_MMS_VERSION = "1.2" +# Keyword arguments that deploy() forwards to the model-server wrappers, the core endpoint +# path, and the inference-component path. A name outside this set raises TypeError, so a +# misspelled option fails at the call instead of deploying an endpoint without it. +DEPLOY_KWARGS = frozenset( + ( + "accelerator_type", + "accept_eula", + "async_inference_config", + "built_model", + "container_startup_health_check_timeout", + "deserializer", + "endpoint_logging", + "explainer_config", + "inference_ami_version", + "inference_recommendation_id", + "kms_key", + "managed_instance_scaling", + "mode", + "model_data_download_timeout", + "resources", + "role", + "routing_config", + "serializer", + "serverless_inference_config", + "volume_size", + ) +) +# InstancePools takes 1 to 5 entries, the primary instance type included. +MAX_FALLBACK_INSTANCE_TYPES = 4 SCRIPT_PARAM_NAME = "sagemaker_program" DIR_PARAM_NAME = "sagemaker_submit_directory" CONTAINER_LOG_LEVEL_PARAM_NAME = "sagemaker_container_log_level" @@ -182,6 +216,182 @@ SAGEMAKER_OUTPUT_LOCATION = "sagemaker_s3_output" +@dataclass(frozen=True) +class _EndpointOptions: + """Endpoint-configuration options of ``deploy()`` that each deploy path reads by name. + + ``deploy()`` validates the options once against ``inference_config``, then each path + applies the ones it supports and raises ``ValueError`` for the ones it cannot apply, + so no option is dropped on the floor. + """ + + data_capture_config: Optional[DataCaptureConfig] = None + tracing_config: Optional[Any] = None + fallback_instance_types: Optional[List[str]] = None + endpoint_type: Optional[EndpointType] = None + inference_component_name: Optional[str] = None + + def validate(self, inference_config: Optional[Any]) -> None: + """Reject a type or a combination with ``inference_config`` that no path can apply.""" + is_serverless = isinstance(inference_config, ServerlessInferenceConfig) + is_inference_component = isinstance(inference_config, ResourceRequirements) + if self.data_capture_config is not None and not isinstance( + self.data_capture_config, DataCaptureConfig + ): + raise TypeError( + "data_capture_config must be a " + "sagemaker.core.model_monitor.data_capture_config.DataCaptureConfig, got " + f"{type(self.data_capture_config).__name__}." + ) + if self.data_capture_config is not None and is_inference_component: + raise ValueError( + "SageMaker rejects data capture on an inference-component endpoint. Drop " + "data_capture_config or deploy without a ResourceRequirements inference_config." + ) + if self.endpoint_type is not None and not isinstance(self.endpoint_type, EndpointType): + raise TypeError( + "endpoint_type must be a sagemaker.core.enums.EndpointType, got " + f"{type(self.endpoint_type).__name__}." + ) + if self.endpoint_type == EndpointType.INFERENCE_COMPONENT_BASED and not ( + is_inference_component + ): + raise ValueError( + "endpoint_type=EndpointType.INFERENCE_COMPONENT_BASED requires " + "inference_config=ResourceRequirements(...)." + ) + if self.endpoint_type == EndpointType.MODEL_BASED and is_inference_component: + raise ValueError( + "endpoint_type=EndpointType.MODEL_BASED conflicts with a ResourceRequirements " + "inference_config, which deploys an inference component." + ) + self._validate_fallback_instance_types(is_serverless) + + def _validate_fallback_instance_types(self, is_serverless: bool) -> None: + if self.fallback_instance_types is None: + return + if is_serverless: + raise ValueError("fallback_instance_types does not apply to a serverless endpoint.") + is_name_list = isinstance(self.fallback_instance_types, list) and all( + isinstance(name, str) and name for name in self.fallback_instance_types + ) + if not is_name_list: + raise TypeError( + "fallback_instance_types must be a list of instance type names such as " + f"['ml.g6.2xlarge'], got {self.fallback_instance_types!r}." + ) + if len(self.fallback_instance_types) > MAX_FALLBACK_INSTANCE_TYPES: + raise ValueError( + f"fallback_instance_types takes at most {MAX_FALLBACK_INSTANCE_TYPES} instance " + f"types (InstancePools holds 5 with the primary type), got " + f"{len(self.fallback_instance_types)}." + ) + if len(set(self.fallback_instance_types)) != len(self.fallback_instance_types): + raise ValueError( + f"fallback_instance_types repeats an instance type: {self.fallback_instance_types}." + ) + + def set_names(self) -> List[str]: + """Return the names of the options the caller set.""" + return [option.name for option in fields(self) if getattr(self, option.name) is not None] + + def reject_for(self, path_description: str) -> None: + """Raise ``ValueError`` when any option is set, for a path that applies none of them.""" + set_names = self.set_names() + if set_names: + raise ValueError( + f"{', '.join(set_names)} cannot be applied by {path_description}. " + "Drop the option or deploy through a path that supports it." + ) + + def require_model_based(self, path_description: str) -> None: + """Raise ``ValueError`` for options that fit only an inference-component endpoint.""" + if self.endpoint_type == EndpointType.INFERENCE_COMPONENT_BASED: + raise ValueError( + f"{path_description} places the model on the production variant, so " + "endpoint_type=EndpointType.INFERENCE_COMPONENT_BASED cannot be applied." + ) + if self.inference_component_name is not None: + raise ValueError( + f"{path_description} creates no inference component, so " + "inference_component_name cannot be applied." + ) + + def require_inference_component_based(self, path_description: str) -> None: + """Raise ``ValueError`` for options that fit only a model-on-variant endpoint.""" + if self.endpoint_type == EndpointType.MODEL_BASED: + raise ValueError( + f"{path_description} hosts the model as an inference component, so " + "endpoint_type=EndpointType.MODEL_BASED cannot be applied." + ) + if self.data_capture_config is not None: + raise ValueError( + f"SageMaker rejects data capture on an inference-component endpoint, which " + f"{path_description} creates. Drop data_capture_config." + ) + + def as_deploy_kwargs(self) -> Dict[str, Any]: + """Return the set options as the keyword arguments ``_deploy_core_endpoint`` reads.""" + return {name: getattr(self, name) for name in self.set_names()} + + +def _instance_pools(instance_type: Optional[str], fallback_instance_types: List[str]) -> List[Dict]: + """Build ``InstancePools`` with ``instance_type`` at priority 1 and the fallbacks after it.""" + if not instance_type: + raise ValueError( + "fallback_instance_types requires an instance_type: the primary type is priority 1 " + "and the fallbacks follow it." + ) + ordered = [instance_type] + [name for name in fallback_instance_types if name != instance_type] + return [ + {"InstanceType": name, "Priority": priority} + for priority, name in enumerate(ordered, start=1) + ] + + +def _instance_pool_shapes(instance_type: Optional[str], fallback_instance_types: List[str]): + """Return the ``InstancePool`` shapes for a ``sagemaker.core.shapes.ProductionVariant``.""" + return [ + InstancePool(instance_type=pool["InstanceType"], priority=pool["Priority"]) + for pool in _instance_pools(instance_type, fallback_instance_types) + ] + + +def _with_instance_pools( + variant: ProductionVariant, fallback_instance_types: Optional[List[str]] +) -> ProductionVariant: + """Move a variant from ``instance_type`` to ``instance_pools`` when fallbacks are set. + + The service takes ``InstancePools`` in place of ``InstanceType``, so the primary type + becomes the priority-1 pool and ``instance_type`` is left unassigned. + """ + if not fallback_instance_types: + return variant + pools = _instance_pool_shapes(variant.instance_type, fallback_instance_types) + return variant.model_copy(update={"instance_pools": pools, "instance_type": Unassigned()}) + + +def _tracing_config_to_request_dict(tracing_config: Optional[Any]) -> Optional[Dict[str, Any]]: + """Return ``tracing_config`` in service form, from a dict, a config object, or a shape.""" + if tracing_config is None: + return None + if isinstance(tracing_config, dict): + return dict(tracing_config) + if hasattr(tracing_config, "_to_request_dict"): + return tracing_config._to_request_dict() + return serialize(tracing_config) + + +def _data_capture_config_shape( + data_capture_config: Optional[DataCaptureConfig], +) -> Optional[DataCaptureConfigShape]: + """Convert a Model Monitor ``DataCaptureConfig`` to the shape ``EndpointConfig.create`` takes.""" + if data_capture_config is None: + return None + request = data_capture_config._to_request_dict() + return DataCaptureConfigShape(**transform(request, "DataCaptureConfig")) + + @dataclass class ModelBuilder(_InferenceRecommenderMixin, _ModelBuilderServers, _ModelBuilderUtils): """Unified interface for building and deploying machine learning models. @@ -790,22 +1000,26 @@ def _fetch_hosting_configs_for_custom_model(self) -> dict: return hosting_configs def _get_instance_resources(self, instance_type: str) -> tuple: - """Get CPU and memory for an instance type by querying EC2.""" + """Get CPU and memory for an instance type by querying EC2. + + Args: + instance_type: The SageMaker instance type (for example 'ml.g5.12xlarge') + + Returns: + A ``(vcpus, memory_mib)`` tuple, or ``(None, None)`` when EC2 does not know the + instance type or the call fails + """ try: - ec2_client = self.sagemaker_session.boto_session.client("ec2") - ec2_instance_type = instance_type.replace("ml.", "") - response = ec2_client.describe_instance_types(InstanceTypes=[ec2_instance_type]) - if response["InstanceTypes"]: - instance_info = response["InstanceTypes"][0] - cpus = instance_info["VCpuInfo"]["DefaultVCpus"] - memory_mb = instance_info["MemoryInfo"]["SizeInMiB"] - return cpus, memory_mb + spec = describe_instance_type( + instance_type, session=self.sagemaker_session.boto_session + ) except Exception as e: logger.warning( f"Could not query instance type {instance_type}: {e}. " f"Unable to validate CPU requirements. Proceeding with recipe defaults." ) - return None, None + return None, None + return spec.vcpus, spec.memory_mib def _resolve_compute_requirements( self, instance_type: str, user_resource_requirements: Optional[ResourceRequirements] = None @@ -989,23 +1203,21 @@ def _infer_accelerator_count_from_instance_type(self, instance_type: str) -> Opt """Infer the number of accelerator devices by querying EC2 instance type info. Args: - instance_type: The EC2 instance type (e.g., 'ml.g5.12xlarge') + instance_type: The SageMaker instance type (for example 'ml.g5.12xlarge') Returns: - Number of accelerator devices, or None if cannot be determined + Number of accelerator devices, or None for a CPU instance, an instance type EC2 + does not know, or a failed call """ try: - ec2_client = self.sagemaker_session.boto_session.client("ec2") - ec2_instance_type = instance_type.replace("ml.", "") - response = ec2_client.describe_instance_types(InstanceTypes=[ec2_instance_type]) - if response["InstanceTypes"]: - gpu_info = response["InstanceTypes"][0].get("GpuInfo") - if gpu_info and gpu_info.get("Gpus"): - total_gpus = sum(g.get("Count", 0) for g in gpu_info["Gpus"]) - if total_gpus > 0: - return total_gpus + spec = describe_instance_type( + instance_type, session=self.sagemaker_session.boto_session + ) except Exception as e: logger.warning(f"Could not query GPU count for {instance_type}: {e}.") + return None + if spec.gpu_count > 0: + return spec.gpu_count return None def _is_gpu_instance(self, instance_type: str) -> bool: @@ -3436,7 +3648,7 @@ def _build_single_modelbuilder( raise ValueError( "The trainer passed to ModelBuilder does not have a completed training job. " "Either call trainer.train() first, or manually set " - "trainer._latest_training_job = TrainingJob.get(training_job_name='') " + "trainer.latest_training_job = TrainingJob.get(training_job_name='') " "to attach a previously completed job." ) @@ -3941,9 +4153,11 @@ def _deploy_core_endpoint(self, **kwargs): if "endpoint_name" in kwargs: self.endpoint_name = endpoint_name - tags = kwargs.get("tags", getattr(self, "_tags", None)) - if "tags" in kwargs: - self._tags = tags + # Merge the caller's tags into the builder's list once. Every request below reads + # the merged list, so no second merge is necessary. + if kwargs.get("tags") is not None: + self.add_tags(kwargs["tags"]) + tags = format_tags(getattr(self, "_tags", None)) kms_key = kwargs.get("kms_key", getattr(self, "kms_key", None)) if "kms_key" in kwargs: @@ -3995,6 +4209,8 @@ def _deploy_core_endpoint(self, **kwargs): self._deserializer = deserializer data_capture_config = kwargs.get("data_capture_config", None) + tracing_config_dict = _tracing_config_to_request_dict(kwargs.get("tracing_config")) + fallback_instance_types = kwargs.get("fallback_instance_types", None) volume_size = kwargs.get("volume_size", getattr(self, "volume_size", None)) inference_recommendation_id = kwargs.get("inference_recommendation_id", None) explainer_config = kwargs.get("explainer_config", None) @@ -4005,12 +4221,6 @@ def _deploy_core_endpoint(self, **kwargs): update_endpoint = kwargs.get("update_endpoint", False) wait = kwargs.get("wait", True) - if tags: - self.add_tags(tags) - tags = format_tags(self._tags) - else: - tags = format_tags(self._tags) - routing_config = _resolve_routing_config(routing_config) if ( @@ -4168,6 +4378,11 @@ def _deploy_core_endpoint(self, **kwargs): managed_instance_scaling=managed_instance_scaling_config, routing_config=routing_config, inference_ami_version=inference_ami_version, + instance_pools=( + _instance_pools(instance_type, fallback_instance_types) + if fallback_instance_types + else None + ), ) self.sagemaker_session.endpoint_from_production_variants( @@ -4180,6 +4395,7 @@ def _deploy_core_endpoint(self, **kwargs): role=self.role_arn, live_logging=False, # TODO: enable when IC supports this wait=False, + tracing_config_dict=tracing_config_dict, ) # An IC-based endpoint hosts no model itself, so it never gets an # endpoint log group to stream: wait on its status only. @@ -4254,6 +4470,11 @@ def _deploy_core_endpoint(self, **kwargs): container_startup_health_check_timeout=container_startup_health_check_timeout, routing_config=routing_config, inference_ami_version=inference_ami_version, + instance_pools=( + _instance_pools(instance_type, fallback_instance_types) + if fallback_instance_types and not is_serverless + else None + ), ) if endpoint_name: self.endpoint_name = endpoint_name @@ -4287,22 +4508,14 @@ def _deploy_core_endpoint(self, **kwargs): if update_endpoint: endpoint_config_name = self.sagemaker_session.create_endpoint_config( - name=self.model_name, - model_name=self.model_name, - initial_instance_count=initial_instance_count, - instance_type=instance_type, - accelerator_type=accelerator_type, + name=unique_name_from_base(self.model_name), + production_variants=[production_variant], tags=tags, kms_key=kms_key, data_capture_config_dict=data_capture_config_dict, - volume_size=volume_size, - model_data_download_timeout=model_data_download_timeout, - container_startup_health_check_timeout=container_startup_health_check_timeout, explainer_config_dict=explainer_config_dict, async_inference_config_dict=async_inference_config_dict, - serverless_inference_config_dict=serverless_inference_config_dict, - routing_config=routing_config, - inference_ami_version=inference_ami_version, + tracing_config_dict=tracing_config_dict, ) self.sagemaker_session.update_endpoint(self.endpoint_name, endpoint_config_name) else: @@ -4316,6 +4529,7 @@ def _deploy_core_endpoint(self, **kwargs): explainer_config_dict=explainer_config_dict, async_inference_config_dict=async_inference_config_dict, live_logging=endpoint_logging, + tracing_config_dict=tracing_config_dict, ) self._wait_for_endpoint(endpoint=self.endpoint_name, show_progress=True, wait=wait) @@ -4571,14 +4785,11 @@ def build( # Tag the Model resource that build() creates. Without this the only way to # set _tags was deploy(), so every model built without an endpoint -- a # pipeline ModelStep, batch transform, evaluation, an inference pipeline -- - # was created untagged with no public way to fix it (#6501). - # format_tags() expands the {key: value} mapping form; serialize() turns - # sagemaker.core.shapes.Tag objects into the boto-style {"Key", "Value"} - # dicts the CreateModel request needs, so both representations are accepted. + # was created untagged with no public way to fix it (#6501). add_tags() + # normalizes every accepted representation to the boto-style {"Key", "Value"} + # dicts the CreateModel request needs. if tags is not None: - from sagemaker.core.utils.utils import serialize - - self.add_tags(serialize(format_tags(tags))) + self.add_tags(tags) if hasattr(self, "built_model") and self.built_model is not None: logger.warning( @@ -6084,6 +6295,12 @@ def deploy( role: Optional[str] = None, model_name: Optional[str] = None, endpoint_config_name: Optional[str] = None, + tags: Optional[Tags] = None, + data_capture_config: Optional[DataCaptureConfig] = None, + tracing_config: Optional[Any] = None, + fallback_instance_types: Optional[List[str]] = None, + endpoint_type: Optional[EndpointType] = None, + inference_component_name: Optional[str] = None, **kwargs, ) -> Union[Endpoint, LocalEndpoint, Transformer]: """Deploy the built model to an ``Endpoint``. @@ -6159,12 +6376,50 @@ def deploy( endpoint_config_name (str, optional): Recommendation deploy only. Name for the created EndpointConfig; auto-generated if omitted. Ignored when deploying a normally-built model. + tags (Optional[Tags]): Tags to attach to the endpoint and endpoint config this call + creates, merged with any added through ``add_tags()`` or ``build(tags=...)``. + Accepts a ``{key: value}`` mapping, boto-style ``{"Key": ..., "Value": ...}`` + dicts, ``{"key": ..., "value": ...}`` dicts, or ``sagemaker.core.shapes.Tag`` + objects. The same tags apply on every deploy path: a built model, a fine-tuned + model package, a Nova model, an inference component, and a recommendation. + (Default: None). + data_capture_config (Optional[DataCaptureConfig]): Model Monitor data capture for + the endpoint configuration this call creates. Applies to a single-model + endpoint, a serverless endpoint, an async endpoint, and the fine-tuned and Nova + model-on-variant paths. SageMaker rejects data capture on an inference-component + endpoint, so that path raises ``ValueError``. (Default: None). + tracing_config (optional): A ``sagemaker.core.shapes.TracingConfig`` for the + endpoint configuration this call creates. The installed botocore service model + must know the ``TracingConfig`` member of ``CreateEndpointConfig``, otherwise + ``ValueError`` is raised instead of the member vanishing from the request. + (Default: None). + fallback_instance_types (Optional[List[str]]): Up to 4 instance types, in priority + order, that SageMaker provisions from when ``instance_type`` has no capacity. + The production variant carries ``InstancePools`` with ``instance_type`` at + priority 1 and these at priorities 2 and up. Instance-based endpoints only: + serverless, inference-component, and recommendation deploys raise + ``ValueError``. (Default: None). + endpoint_type (Optional[EndpointType]): ``EndpointType.MODEL_BASED`` places the + model on the production variant, ``EndpointType.INFERENCE_COMPONENT_BASED`` + hosts it as an inference component. The default follows ``inference_config``: + inference-component based for a ``ResourceRequirements``, model based + otherwise. (Default: None). + inference_component_name (Optional[str]): Name of the inference component an + inference-component deploy creates, or that a fine-tuned model deploy creates + or reuses. Auto-generated when omitted. (Default: None). + **kwargs: Options forwarded to the deploy path, one of ``DEPLOY_KWARGS`` + (``accelerator_type``, ``async_inference_config``, ``explainer_config``, + ``kms_key``, ``volume_size``, ``serializer``, ``deserializer``, ...). Raises + ``TypeError`` for a name outside that set. Returns: Union[Endpoint, LocalEndpoint, Transformer]: A ``sagemaker.core.resources.Endpoint`` resource representing the deployed endpoint, a ``LocalEndpoint`` for local mode, or a ``Transformer`` for batch transform inference. Raises: + TypeError: If ``kwargs`` holds a name outside ``DEPLOY_KWARGS``. + ValueError: If an option does not fit the deploy path, for example + ``fallback_instance_types`` on a serverless endpoint or more than 4 of them. UnexpectedStatusException: If ``wait`` is True and the endpoint finishes in a state other than ``InService`` (``CapacityError`` for capacity failures). Model customization and recommendation deployments raise @@ -6184,6 +6439,84 @@ def deploy( ) self._deployed = True + unknown_kwargs = sorted(set(kwargs) - DEPLOY_KWARGS) + if unknown_kwargs: + raise TypeError( + f"deploy() got an unexpected keyword argument {unknown_kwargs[0]!r}. " + f"Options forwarded to the deploy path are: {', '.join(sorted(DEPLOY_KWARGS))}." + ) + endpoint_options = _EndpointOptions( + data_capture_config=data_capture_config, + tracing_config=tracing_config, + fallback_instance_types=fallback_instance_types, + endpoint_type=endpoint_type, + inference_component_name=inference_component_name, + ) + endpoint_options.validate(inference_config) + + # Every deploy path below reads self._tags, so the caller's tags are merged into + # the builder for this call only. The merge is undone on exit, so a second + # deploy() on the same builder (a retry, another endpoint) starts from the tags + # build() and add_tags() set, not from the previous call's. + tags_before_call = format_tags(getattr(self, "_tags", None)) + if tags is not None: + self.add_tags(tags) + try: + return self._deploy_with_options( + endpoint_name=endpoint_name, + initial_instance_count=initial_instance_count, + instance_type=instance_type, + wait=wait, + update_endpoint=update_endpoint, + container_timeout_in_seconds=container_timeout_in_seconds, + inference_config=inference_config, + custom_orchestrator_instance_type=custom_orchestrator_instance_type, + custom_orchestrator_initial_instance_count=custom_orchestrator_initial_instance_count, + reuse_resources=reuse_resources, + use_recommendation=use_recommendation, + recommendation=recommendation, + recommendation_index=recommendation_index, + recommendation_spec_name=recommendation_spec_name, + auto_approve=auto_approve, + role=role, + model_name=model_name, + endpoint_config_name=endpoint_config_name, + endpoint_options=endpoint_options, + **kwargs, + ) + finally: + self._tags = tags_before_call + + def _deploy_with_options( + self, + *, + endpoint_name: Optional[str], + initial_instance_count: Optional[int], + instance_type: Optional[str], + wait: bool, + update_endpoint: Optional[bool], + container_timeout_in_seconds: int, + inference_config: Optional[Any], + custom_orchestrator_instance_type: Optional[str], + custom_orchestrator_initial_instance_count: Optional[int], + reuse_resources: bool, + use_recommendation: Optional[bool], + recommendation: Optional[Any], + recommendation_index: int, + recommendation_spec_name: Optional[str], + auto_approve: bool, + role: Optional[str], + model_name: Optional[str], + endpoint_config_name: Optional[str], + endpoint_options: _EndpointOptions, + **kwargs, + ) -> Union[Endpoint, LocalEndpoint, Transformer]: + """Run the deploy path selection of ``deploy()`` with the validated options. + + ``deploy()`` validates the arguments, merges the per-call tags, and restores them + afterwards. This method holds the path selection so that restore happens on every + return and on every raise. + """ # Recommendation-mode deploy: when a recommendation job is attached # (from a prior generate_deployment_recommendations call, or because the # builder was hydrated via ModelBuilder.from_recommendation_job), deploy @@ -6220,6 +6553,7 @@ def deploy( "or recommendation_spec_name= instead." ) if has_recommendation and use_recommendation is not False: + endpoint_options.reject_for("a recommendation deploy") return self._deploy_recommendation( recommendation_index=recommendation_index, recommendation_spec_name=recommendation_spec_name, @@ -6230,7 +6564,7 @@ def deploy( role=role, instance_type=instance_type, initial_instance_count=initial_instance_count, - tags=kwargs.get("tags"), + tags=format_tags(getattr(self, "_tags", None)), wait=wait, auto_approve=auto_approve, ) @@ -6307,10 +6641,11 @@ def deploy( return self._deploy_model_customization( endpoint_name=endpoint_name, initial_instance_count=initial_instance_count, - inference_component_name=kwargs.pop("inference_component_name", None), + inference_component_name=endpoint_options.inference_component_name, wait=wait, container_timeout_in_seconds=container_timeout_in_seconds, inference_config=inference_config_param, + endpoint_options=endpoint_options, **kwargs, ) @@ -6329,10 +6664,11 @@ def deploy( "update_endpoint": update_endpoint, "container_timeout_in_seconds": container_timeout_in_seconds, "wait": wait, - "endpoint_type": EndpointType.MODEL_BASED, + "endpoint_type": endpoint_options.endpoint_type or EndpointType.MODEL_BASED, } deploy_kwargs.update(kwargs) + deploy_kwargs.update(endpoint_options.as_deploy_kwargs()) return self._deploy(**deploy_kwargs) if isinstance(inference_config, ServerlessInferenceConfig): @@ -6342,10 +6678,11 @@ def deploy( "update_endpoint": update_endpoint, "container_timeout_in_seconds": container_timeout_in_seconds, "wait": wait, - "endpoint_type": EndpointType.MODEL_BASED, + "endpoint_type": endpoint_options.endpoint_type or EndpointType.MODEL_BASED, } deploy_kwargs.update(kwargs) + deploy_kwargs.update(endpoint_options.as_deploy_kwargs()) return self._deploy(**deploy_kwargs) if isinstance(inference_config, AsyncInferenceConfig): @@ -6357,10 +6694,11 @@ def deploy( "update_endpoint": update_endpoint, "container_timeout_in_seconds": container_timeout_in_seconds, "wait": wait, - "endpoint_type": EndpointType.MODEL_BASED, + "endpoint_type": endpoint_options.endpoint_type or EndpointType.MODEL_BASED, } deploy_kwargs.update(kwargs) + deploy_kwargs.update(endpoint_options.as_deploy_kwargs()) return self._deploy(**deploy_kwargs) if isinstance(inference_config, BatchTransformInferenceConfig): @@ -6378,7 +6716,9 @@ def deploy( deploy_kwargs = { "instance_type": self.instance_type, "mode": Mode.SAGEMAKER_ENDPOINT, - "endpoint_type": EndpointType.INFERENCE_COMPONENT_BASED, + "endpoint_type": ( + endpoint_options.endpoint_type or EndpointType.INFERENCE_COMPONENT_BASED + ), "resources": inference_config, "initial_instance_count": initial_instance_count, "role": self.role_arn, @@ -6388,9 +6728,11 @@ def deploy( } deploy_kwargs.update(kwargs) + deploy_kwargs.update(endpoint_options.as_deploy_kwargs()) return self._deploy(**deploy_kwargs) if hasattr(self, "_deployables") and self._deployables: + endpoint_options.reject_for("a modelbuilder_list deploy") endpoints = [] for ic in self._deployables.get("InferenceComponents", []): endpoints.append(self._deploy_for_ic(ic_data=ic, endpoint_name=endpoint_name)) @@ -6443,12 +6785,61 @@ def deploy( raise ValueError("Deployment Options not supported") + def _core_endpoint_tags(self) -> List[Dict[str, str]]: + """Return the builder's accumulated tags in the form ``Endpoint.create`` takes. + + ``add_tags`` stores tags in boto form (``{"Key": ..., "Value": ...}``); the + ``sagemaker.core.resources`` classes take ``{"key": ..., "value": ...}``. + """ + return [ + {"key": tag["Key"], "value": tag["Value"]} + for tag in format_tags(getattr(self, "_tags", None) or []) + ] + + def _create_endpoint_config( + self, + endpoint_config_name: str, + production_variant: ProductionVariant, + endpoint_options: _EndpointOptions, + **create_kwargs, + ) -> None: + """Create the endpoint configuration of a core-resource deploy path. + + Applies the ``data_capture_config`` and ``tracing_config`` of ``deploy()``. The + generated ``EndpointConfig.create`` has no ``tracing_config`` parameter until + sagemaker-core is regenerated from a service model that has the member, so tracing + goes through the session, which checks the installed service model first. + """ + tracing_config_dict = _tracing_config_to_request_dict(endpoint_options.tracing_config) + if tracing_config_dict is not None: + self.sagemaker_session.create_endpoint_config( + name=endpoint_config_name, + production_variants=[serialize(production_variant)], + tags=format_tags(getattr(self, "_tags", None)), + data_capture_config_dict=( + endpoint_options.data_capture_config._to_request_dict() + if endpoint_options.data_capture_config is not None + else None + ), + enable_network_isolation=create_kwargs.get("enable_network_isolation"), + role=create_kwargs.get("execution_role_arn"), + tracing_config_dict=tracing_config_dict, + ) + return + EndpointConfig.create( + endpoint_config_name=endpoint_config_name, + production_variants=[production_variant], + data_capture_config=_data_capture_config_shape(endpoint_options.data_capture_config), + **create_kwargs, + ) + def _deploy_model_customization( self, endpoint_name: str, initial_instance_count: int = 1, inference_component_name: Optional[str] = None, inference_config: Optional[ResourceRequirements] = None, + endpoint_options: Optional[_EndpointOptions] = None, **kwargs, ) -> Endpoint: """Deploy a model customization (fine-tuned) model to an endpoint with inference components. @@ -6464,6 +6855,8 @@ def _deploy_model_customization( inference_component_name (Optional[str]): Name for the inference component inference_config (Optional[ResourceRequirements]): Inference configuration including resource requirements (accelerator count, memory, CPU cores) + endpoint_options (Optional[_EndpointOptions]): The endpoint-configuration options + of ``deploy()`` (data capture, tracing, fallback instance types). **kwargs: Additional deployment parameters Returns: @@ -6471,6 +6864,8 @@ def _deploy_model_customization( """ from sagemaker.core.resources import Tag as CoreTag + endpoint_options = endpoint_options or _EndpointOptions() + # An inference_config of ResourceRequirements requests an inference # component deployment; otherwise the model is placed directly on the # production variant. @@ -6483,10 +6878,12 @@ def _deploy_model_customization( # image, escrow artifacts, and env. is_nova = self._is_nova_model() if is_nova and not is_ic_deploy: + endpoint_options.require_model_based("a Nova model deploy") return self._deploy_nova_model( endpoint_name=endpoint_name, initial_instance_count=initial_instance_count, wait=kwargs.get("wait", True), + endpoint_options=endpoint_options, ) # The model package may be absent (e.g. a Nova CPTTrainer or raw-S3 @@ -6499,21 +6896,24 @@ def _deploy_model_customization( from sagemaker.serve.utils.model_package_utils import is_restricted_model_package if not is_ic_deploy and is_restricted_model_package(model_package): + endpoint_options.require_model_based("a restricted model package deploy") if not endpoint_name: endpoint_name = f"endpoint-{uuid.uuid4().hex[:8]}" - EndpointConfig.create( - endpoint_config_name=endpoint_name, - production_variants=[ - ProductionVariant( - variant_name="AllTraffic", - model_name=self.built_model.model_name, - instance_type=self.instance_type, - initial_instance_count=initial_instance_count or 1, - ) - ], + variant = ProductionVariant( + variant_name="AllTraffic", + model_name=self.built_model.model_name, + instance_type=self.instance_type, + initial_instance_count=initial_instance_count or 1, + ) + self._create_endpoint_config( + endpoint_name, + _with_instance_pools(variant, endpoint_options.fallback_instance_types), + endpoint_options, ) endpoint = Endpoint.create( - endpoint_name=endpoint_name, endpoint_config_name=endpoint_name + endpoint_name=endpoint_name, + endpoint_config_name=endpoint_name, + tags=self._core_endpoint_tags() or None, ) if kwargs.get("wait", True): endpoint.wait_for_status("InService") @@ -6583,35 +6983,32 @@ def _deploy_model_customization( or self._enable_network_isolation ) + endpoint_options.require_inference_component_based("a fine-tuned model deploy") + # Check if endpoint exists is_existing_endpoint = self._does_endpoint_exist(endpoint_name) if not is_existing_endpoint: - EndpointConfig.create( - endpoint_config_name=endpoint_name, - production_variants=[ - ProductionVariant( - variant_name=endpoint_name, - instance_type=self.instance_type, - initial_instance_count=initial_instance_count or 1, - ) - ], + variant = ProductionVariant( + variant_name=endpoint_name, + instance_type=self.instance_type, + initial_instance_count=initial_instance_count or 1, + ) + self._create_endpoint_config( + endpoint_name, + _with_instance_pools(variant, endpoint_options.fallback_instance_types), + endpoint_options, execution_role_arn=self.role_arn, enable_network_isolation=enable_network_isolation, ) logger.info("Endpoint core call starting") - # Apply tags accumulated via add_tags (e.g. the model-source reuse - # tag) to the endpoint so it is discoverable. Stored tags are in - # {"Key":..,"Value":..} form; normalize to the key/value form the - # core resource expects. - endpoint_tags = [ - {"key": tag["Key"], "value": tag["Value"]} - for tag in format_tags(getattr(self, "_tags", None) or []) - ] + # Apply tags accumulated via add_tags (the caller's deploy(tags=...), + # build(tags=...), and the model-source reuse tag) to the endpoint so it + # is discoverable. endpoint = Endpoint.create( endpoint_name=endpoint_name, endpoint_config_name=endpoint_name, - tags=endpoint_tags or None, + tags=self._core_endpoint_tags() or None, ) endpoint.wait_for_status("InService") else: @@ -6800,6 +7197,7 @@ def _deploy_nova_model( endpoint_name: str, initial_instance_count: int = 1, wait: bool = True, + endpoint_options: Optional[_EndpointOptions] = None, ) -> Endpoint: """Deploy a Nova model directly to an endpoint without inference components. @@ -6808,19 +7206,20 @@ def _deploy_nova_model( - No InferenceComponents are created - EnableNetworkIsolation is set on the Model (during build) """ + endpoint_options = endpoint_options or _EndpointOptions() if not endpoint_name: endpoint_name = f"endpoint-{uuid.uuid4().hex[:8]}" - EndpointConfig.create( - endpoint_config_name=endpoint_name, - production_variants=[ - ProductionVariant( - variant_name="AllTraffic", - model_name=self.built_model.model_name, - instance_type=self.instance_type, - initial_instance_count=initial_instance_count, - ) - ], + variant = ProductionVariant( + variant_name="AllTraffic", + model_name=self.built_model.model_name, + instance_type=self.instance_type, + initial_instance_count=initial_instance_count, + ) + self._create_endpoint_config( + endpoint_name, + _with_instance_pools(variant, endpoint_options.fallback_instance_types), + endpoint_options, ) # The jumpstart-model-id tag always applies (resolved from the model @@ -6835,15 +7234,13 @@ def _deploy_nova_model( if base_model is not None and base_model.recipe_name: tags.append({"key": "sagemaker-sdk:recipe-name", "value": base_model.recipe_name}) - # Merge tags accumulated via add_tags (e.g. the model-source reuse tag). - # Those are stored in {"Key": ..., "Value": ...} form, so normalize to the - # {"key": ..., "value": ...} form Endpoint.create expects and de-duplicate. + # Merge tags accumulated via add_tags (the caller's deploy(tags=...), + # build(tags=...), and the model-source reuse tag) and de-duplicate by key. existing_keys = {tag["key"] for tag in tags} - for tag in format_tags(getattr(self, "_tags", None) or []): - key = tag["Key"] - if key not in existing_keys: - tags.append({"key": key, "value": tag["Value"]}) - existing_keys.add(key) + for tag in self._core_endpoint_tags(): + if tag["key"] not in existing_keys: + tags.append(tag) + existing_keys.add(tag["key"]) endpoint = Endpoint.create( endpoint_name=endpoint_name, diff --git a/sagemaker-serve/tests/unit/test_bedrock_deployment.py b/sagemaker-serve/tests/unit/test_bedrock_deployment.py new file mode 100644 index 0000000000..177fb513fa --- /dev/null +++ b/sagemaker-serve/tests/unit/test_bedrock_deployment.py @@ -0,0 +1,427 @@ +# 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. +"""Unit tests for ``sagemaker.serve.bedrock_deployment`` and the builder deployment handle. + +Before this change ``sagemaker.serve`` had no ``bedrock_deployment`` module, so the import +below raised ``ModuleNotFoundError``. ``BedrockModelBuilder`` had no ``deployment`` +attribute and no ``latest_deployment()`` method. Callers kept the response dict and called +``bedrock`` and ``bedrock-runtime`` by hand to list models, to wait out +``ModelNotReadyException``, and to invoke the model. +""" + +import io +import json +from unittest.mock import Mock, call, patch + +import pytest +from botocore.exceptions import ClientError + +from sagemaker.serve.bedrock_deployment import ( + BedrockDeployment, + BedrockFoundationModel, + list_invocable_foundation_models, +) +from .test_bedrock_model_builder import _builder, _make_container, _make_model_package + +MODULE = "sagemaker.serve.bedrock_deployment" +MODEL_ARN = "arn:aws:bedrock:us-west-2:123456789012:imported-model/abcdef123456" +DEPLOYMENT_ARN = "arn:aws:bedrock:us-west-2:123456789012:custom-model-deployment/abcdef123456" +CUSTOM_MODEL_ARN = "arn:aws:bedrock:us-west-2:123456789012:custom-model/nova/abcdef123456" +HAIKU_ID = "anthropic.claude-3-haiku-20240307-v1:0" +NOVA_PRO_ID = "amazon.nova-pro-v1:0" +PROVISIONED_ONLY_ID = "meta.llama3-70b-instruct-v1:0" + + +def _foundation_model_arn(model_id, region="us-west-2"): + return f"arn:aws:bedrock:{region}::foundation-model/{model_id}" + + +def _summary(model_id, inference_types): + return { + "modelArn": _foundation_model_arn(model_id), + "modelId": model_id, + "modelName": model_id.split(".")[1], + "providerName": model_id.split(".")[0].title(), + "inputModalities": ["TEXT"], + "outputModalities": ["TEXT"], + "inferenceTypesSupported": inference_types, + } + + +def _profile(profile_id, model_id, status="ACTIVE", profile_type="SYSTEM_DEFINED"): + return { + "inferenceProfileId": profile_id, + "inferenceProfileArn": f"arn:aws:bedrock:us-west-2:123456789012:inference-profile/{profile_id}", + "status": status, + "type": profile_type, + "models": [ + {"modelArn": _foundation_model_arn(model_id, "us-east-1")}, + {"modelArn": _foundation_model_arn(model_id, "us-west-2")}, + ], + } + + +def _not_ready_error(operation): + return ClientError( + {"Error": {"Code": "ModelNotReadyException", "Message": "Model is not ready"}}, + operation, + ) + + +def _invoke_response(payload): + return { + "body": io.BytesIO(json.dumps(payload).encode("utf-8")), + "contentType": "application/json", + } + + +def _session_with_bedrock(bedrock): + session = Mock() + session.client.return_value = bedrock + return session + + +class TestListInvocableFoundationModels: + def test_two_pages_join_one_profile_and_drop_provisioned_only_models(self): + bedrock = Mock() + bedrock.list_foundation_models.side_effect = [ + { + "modelSummaries": [ + _summary(HAIKU_ID, ["ON_DEMAND"]), + _summary(PROVISIONED_ONLY_ID, ["PROVISIONED"]), + ], + "nextToken": "page-2", + }, + {"modelSummaries": [_summary(NOVA_PRO_ID, ["PROVISIONED"])]}, + ] + bedrock.list_inference_profiles.return_value = { + "inferenceProfileSummaries": [_profile(f"us.{NOVA_PRO_ID}", NOVA_PRO_ID)] + } + session = _session_with_bedrock(bedrock) + + models = list_invocable_foundation_models(region="us-west-2", boto_session=session) + + session.client.assert_called_once_with("bedrock", region_name="us-west-2") + assert bedrock.list_foundation_models.call_args_list == [ + call(byOutputModality="TEXT"), + call(byOutputModality="TEXT", nextToken="page-2"), + ] + bedrock.list_inference_profiles.assert_called_once_with(typeEquals="SYSTEM_DEFINED") + assert [model.model_id for model in models] == [HAIKU_ID, NOVA_PRO_ID] + assert models[0] == BedrockFoundationModel( + model_id=HAIKU_ID, + model_arn=_foundation_model_arn(HAIKU_ID), + model_name="claude-3-haiku-20240307-v1:0", + provider_name="Anthropic", + input_modalities=["TEXT"], + output_modalities=["TEXT"], + inference_types=["ON_DEMAND"], + inference_profile_id=None, + ) + assert models[1].inference_types == ["PROVISIONED", "INFERENCE_PROFILE"] + assert models[1].inference_profile_id == f"us.{NOVA_PRO_ID}" + + def test_output_modality_filter_is_passed_through(self): + bedrock = Mock() + bedrock.list_foundation_models.return_value = {"modelSummaries": []} + bedrock.list_inference_profiles.return_value = {"inferenceProfileSummaries": []} + + models = list_invocable_foundation_models( + boto_session=_session_with_bedrock(bedrock), output_modality="IMAGE" + ) + + assert models == [] + bedrock.list_foundation_models.assert_called_once_with(byOutputModality="IMAGE") + + def test_geographic_profile_wins_over_global_and_inactive_profiles_are_ignored(self): + bedrock = Mock() + bedrock.list_foundation_models.return_value = { + "modelSummaries": [_summary(NOVA_PRO_ID, ["PROVISIONED"])] + } + bedrock.list_inference_profiles.side_effect = [ + { + "inferenceProfileSummaries": [_profile(f"global.{NOVA_PRO_ID}", NOVA_PRO_ID)], + "nextToken": "page-2", + }, + { + "inferenceProfileSummaries": [ + _profile(f"us.{NOVA_PRO_ID}", NOVA_PRO_ID), + _profile(f"eu.{NOVA_PRO_ID}", NOVA_PRO_ID, status="INACTIVE"), + _profile("my-app-profile", NOVA_PRO_ID, profile_type="APPLICATION"), + ] + }, + ] + + models = list_invocable_foundation_models(boto_session=_session_with_bedrock(bedrock)) + + assert bedrock.list_inference_profiles.call_args_list == [ + call(typeEquals="SYSTEM_DEFINED"), + call(typeEquals="SYSTEM_DEFINED", nextToken="page-2"), + ] + assert [model.inference_profile_id for model in models] == [f"us.{NOVA_PRO_ID}"] + + def test_default_session_is_created_when_none_is_given(self): + bedrock = Mock() + bedrock.list_foundation_models.return_value = {"modelSummaries": []} + bedrock.list_inference_profiles.return_value = {"inferenceProfileSummaries": []} + with patch(f"{MODULE}.boto3.Session", return_value=_session_with_bedrock(bedrock)): + assert list_invocable_foundation_models(region="eu-west-1") == [] + bedrock.list_foundation_models.assert_called_once() + + +class TestBedrockDeploymentWaitUntilReady: + def test_imported_model_retries_one_not_ready_then_returns(self): + runtime = Mock() + runtime.invoke_model.side_effect = [ + _not_ready_error("InvokeModel"), + _invoke_response({"generation": "Hi"}), + ] + deployment = BedrockDeployment(MODEL_ARN, _session_with_bedrock(runtime)) + + with patch(f"{MODULE}.time.sleep") as sleep: + deployment.wait_until_ready(timeout=600, poll=30) + + assert runtime.invoke_model.call_count == 2 + sleep.assert_called_once_with(30) + probe_kwargs = runtime.invoke_model.call_args.kwargs + assert probe_kwargs["modelId"] == MODEL_ARN + assert json.loads(probe_kwargs["body"]) == {"prompt": "Hi", "max_tokens": 1} + runtime.converse.assert_not_called() + + def test_custom_model_deployment_probes_converse_with_the_deployment_arn(self): + runtime = Mock() + runtime.converse.side_effect = [_not_ready_error("Converse"), {"output": {}}] + deployment = BedrockDeployment( + CUSTOM_MODEL_ARN, _session_with_bedrock(runtime), deployment_arn=DEPLOYMENT_ARN + ) + + with patch(f"{MODULE}.time.sleep"): + deployment.wait_until_ready(timeout=600, poll=30) + + assert runtime.converse.call_count == 2 + assert runtime.converse.call_args.kwargs["modelId"] == DEPLOYMENT_ARN + assert runtime.converse.call_args.kwargs["inferenceConfig"] == {"maxTokens": 1} + runtime.invoke_model.assert_not_called() + + def test_timeout_raises_with_the_arn(self): + runtime = Mock() + runtime.invoke_model.side_effect = _not_ready_error("InvokeModel") + deployment = BedrockDeployment(MODEL_ARN, _session_with_bedrock(runtime)) + + with patch(f"{MODULE}.time.sleep") as sleep: + with pytest.raises(TimeoutError, match=f"{MODEL_ARN} was not ready after 60s"): + deployment.wait_until_ready(timeout=60, poll=30) + + assert runtime.invoke_model.call_count == 3 + assert sleep.call_count == 2 + + def test_custom_probe_body_replaces_the_completion_default(self): + runtime = Mock() + runtime.invoke_model.return_value = _invoke_response({"embedding": [0.1]}) + deployment = BedrockDeployment(MODEL_ARN, _session_with_bedrock(runtime)) + + deployment.wait_until_ready(timeout=60, poll=30, probe_body={"inputText": "Hi"}) + + runtime.invoke_model.assert_called_once() + assert json.loads(runtime.invoke_model.call_args.kwargs["body"]) == {"inputText": "Hi"} + + def test_other_client_errors_are_raised_at_once(self): + runtime = Mock() + runtime.invoke_model.side_effect = ClientError( + {"Error": {"Code": "ValidationException", "Message": "bad body"}}, "InvokeModel" + ) + deployment = BedrockDeployment(MODEL_ARN, _session_with_bedrock(runtime)) + + with patch(f"{MODULE}.time.sleep") as sleep: + with pytest.raises(ClientError, match="ValidationException"): + deployment.wait_until_ready(timeout=600, poll=30) + + sleep.assert_not_called() + + +class TestBedrockDeploymentInvoke: + def test_dict_body_round_trips_through_json(self): + runtime = Mock() + runtime.invoke_model.return_value = _invoke_response({"outputs": [{"text": "Paris"}]}) + session = _session_with_bedrock(runtime) + deployment = BedrockDeployment(MODEL_ARN, session, region="us-west-2") + + result = deployment.invoke({"prompt": "Capital of France?", "max_tokens": 16}) + + session.client.assert_called_once_with("bedrock-runtime", region_name="us-west-2") + runtime.invoke_model.assert_called_once_with( + modelId=MODEL_ARN, + body=json.dumps({"prompt": "Capital of France?", "max_tokens": 16}), + contentType="application/json", + accept="application/json", + ) + assert result == {"outputs": [{"text": "Paris"}]} + + def test_string_body_and_non_json_accept_return_raw_bytes(self): + runtime = Mock() + runtime.invoke_model.return_value = { + "body": io.BytesIO(b"plain text"), + "contentType": "text/plain", + } + deployment = BedrockDeployment( + CUSTOM_MODEL_ARN, _session_with_bedrock(runtime), deployment_arn=DEPLOYMENT_ARN + ) + + result = deployment.invoke('{"prompt": "x"}', accept="text/plain") + + assert runtime.invoke_model.call_args.kwargs["body"] == '{"prompt": "x"}' + assert runtime.invoke_model.call_args.kwargs["modelId"] == DEPLOYMENT_ARN + assert result == {"body": b"plain text", "contentType": "text/plain"} + + def test_runtime_client_is_created_once(self): + runtime = Mock() + runtime.invoke_model.side_effect = [_invoke_response({}), _invoke_response({})] + session = _session_with_bedrock(runtime) + deployment = BedrockDeployment(MODEL_ARN, session) + + deployment.invoke({"prompt": "a"}) + deployment.invoke({"prompt": "b"}) + + session.client.assert_called_once() + + +class TestBedrockDeploymentConverse: + def test_inference_config_and_system_pass_through(self): + runtime = Mock() + runtime.converse.return_value = {"output": {"message": {"content": [{"text": "Hi"}]}}} + deployment = BedrockDeployment( + CUSTOM_MODEL_ARN, _session_with_bedrock(runtime), deployment_arn=DEPLOYMENT_ARN + ) + messages = [{"role": "user", "content": [{"text": "Hello"}]}] + + result = deployment.converse( + messages, + inference_config={"maxTokens": 100, "temperature": 0.7}, + system=[{"text": "Answer briefly."}], + ) + + runtime.converse.assert_called_once_with( + modelId=DEPLOYMENT_ARN, + messages=messages, + inferenceConfig={"maxTokens": 100, "temperature": 0.7}, + system=[{"text": "Answer briefly."}], + ) + assert result["output"]["message"]["content"][0]["text"] == "Hi" + + def test_optional_maps_are_omitted_when_none(self): + runtime = Mock() + runtime.converse.return_value = {} + deployment = BedrockDeployment(MODEL_ARN, _session_with_bedrock(runtime)) + messages = [{"role": "user", "content": [{"text": "Hello"}]}] + + deployment.converse(messages) + + runtime.converse.assert_called_once_with(modelId=MODEL_ARN, messages=messages) + + +class TestBuilderDeploymentHandle: + def test_latest_deployment_raises_before_deploy(self): + builder = _builder() + assert builder.deployment is None + with pytest.raises(ValueError, match="Call deploy\\(\\) first"): + builder.latest_deployment() + + def test_oss_deploy_sets_the_imported_model_handle(self): + builder = _builder() + builder.model_package = _make_model_package(_make_container(s3_uri="s3://b/model/")) + builder.s3_model_artifacts = "s3://b/model/" + builder.boto_session = Mock() + builder._bedrock_client = Mock() + builder._bedrock_client.create_model_import_job.return_value = {"jobArn": "arn:job"} + builder._bedrock_client.get_model_import_job.return_value = { + "status": "Completed", + "importedModelName": "my-imported-model", + "importedModelArn": MODEL_ARN, + } + runtime = builder.boto_session.client.return_value + runtime.invoke_model.return_value = _invoke_response({"generation": "Paris"}) + + with patch("sagemaker.serve.bedrock_model_builder.time.sleep"): + result = builder.deploy(job_name="j", imported_model_name="m", role_arn="r") + + assert result["importedModelArn"] == MODEL_ARN + assert builder.deployment.model_arn == MODEL_ARN + assert builder.deployment.deployment_arn is None + assert builder.latest_deployment() is builder.deployment + assert builder.deployment.invoke({"prompt": "Capital of France?"}) == { + "generation": "Paris" + } + builder.boto_session.client.assert_any_call("bedrock-runtime", region_name=None) + assert runtime.invoke_model.call_args.kwargs["modelId"] == MODEL_ARN + + def test_nova_deploy_sets_the_custom_model_deployment_handle(self): + builder = _builder() + builder.model_package = _make_model_package(_make_container(recipe_name="nova-micro")) + builder.s3_model_artifacts = "s3://b/ckpt" + builder.boto_session = Mock() + builder._bedrock_client = Mock() + builder._bedrock_client.create_custom_model.return_value = {"modelArn": CUSTOM_MODEL_ARN} + builder._bedrock_client.get_custom_model.return_value = {"modelStatus": "Active"} + builder._bedrock_client.create_custom_model_deployment.return_value = { + "customModelDeploymentArn": DEPLOYMENT_ARN + } + builder._bedrock_client.get_custom_model_deployment.return_value = {"status": "Active"} + + result = builder.deploy(custom_model_name="nova-m", role_arn="r") + + assert result["customModelDeploymentArn"] == DEPLOYMENT_ARN + assert builder.latest_deployment().model_arn == CUSTOM_MODEL_ARN + assert builder.latest_deployment().deployment_arn == DEPLOYMENT_ARN + + def test_reused_nova_model_and_deployment_set_the_handle(self): + builder = _builder() + builder.model_package = _make_model_package(_make_container(recipe_name="nova-micro")) + builder.s3_model_artifacts = "s3://b/ckpt" + builder.boto_session = Mock() + builder._bedrock_client = Mock() + + with ( + patch( + "sagemaker.serve.bedrock_model_builder.find_existing_bedrock_model", + return_value=CUSTOM_MODEL_ARN, + ), + patch( + "sagemaker.serve.bedrock_model_builder.find_active_bedrock_deployment_for_model", + return_value=DEPLOYMENT_ARN, + ), + ): + builder.deploy(custom_model_name="nova-m", role_arn="r", reuse_resources=True) + + builder._bedrock_client.create_custom_model.assert_not_called() + assert builder.latest_deployment().model_arn == CUSTOM_MODEL_ARN + assert builder.latest_deployment().deployment_arn == DEPLOYMENT_ARN + + def test_missing_model_arn_resets_a_stale_handle(self): + builder = _builder() + builder.model_package = _make_model_package(_make_container(s3_uri="s3://b/model/")) + builder.s3_model_artifacts = "s3://b/model/" + builder.boto_session = Mock() + builder.deployment = BedrockDeployment(MODEL_ARN, builder.boto_session) + builder._bedrock_client = Mock() + builder._bedrock_client.create_model_import_job.return_value = {"jobArn": "arn:job"} + builder._bedrock_client.get_model_import_job.return_value = { + "status": "Completed", + "importedModelName": "my-imported-model", + } + + with patch("sagemaker.serve.bedrock_model_builder.time.sleep"): + builder.deploy(job_name="j", imported_model_name="m", role_arn="r") + + assert builder.deployment is None + with pytest.raises(ValueError): + builder.latest_deployment() diff --git a/sagemaker-serve/tests/unit/test_model_builder.py b/sagemaker-serve/tests/unit/test_model_builder.py index 9a68150229..49f6c2e902 100644 --- a/sagemaker-serve/tests/unit/test_model_builder.py +++ b/sagemaker-serve/tests/unit/test_model_builder.py @@ -342,7 +342,13 @@ def test_get_instance_resources(self): """Test getting instance resources from EC2.""" mock_ec2 = Mock() mock_ec2.describe_instance_types.return_value = { - "InstanceTypes": [{"VCpuInfo": {"DefaultVCpus": 8}, "MemoryInfo": {"SizeInMiB": 32768}}] + "InstanceTypes": [ + { + "InstanceType": "g5.2xlarge", + "VCpuInfo": {"DefaultVCpus": 8}, + "MemoryInfo": {"SizeInMiB": 32768}, + } + ] } self.mock_session.boto_session.client.return_value = mock_ec2 @@ -355,6 +361,7 @@ def test_get_instance_resources(self): cpus, memory = builder._get_instance_resources("ml.g5.2xlarge") self.assertEqual(cpus, 8) self.assertEqual(memory, 32768) + mock_ec2.describe_instance_types.assert_called_once_with(InstanceTypes=["g5.2xlarge"]) @patch("sagemaker.serve.model_builder.InferenceComponent") @patch("sagemaker.core.resources.Tag") diff --git a/sagemaker-serve/tests/unit/test_model_builder_deploy_options.py b/sagemaker-serve/tests/unit/test_model_builder_deploy_options.py new file mode 100644 index 0000000000..105edc984c --- /dev/null +++ b/sagemaker-serve/tests/unit/test_model_builder_deploy_options.py @@ -0,0 +1,482 @@ +# 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. +"""``ModelBuilder.deploy`` takes its endpoint options by name and rejects an unknown one. + +Before this change ``data_capture_config``, ``endpoint_type``, and ``inference_component_name`` +travelled 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 not read anywhere, and a misspelled option was a silent no-op that +deployed an endpoint without it. +""" + +from __future__ import absolute_import + +import inspect +import unittest +from unittest.mock import MagicMock, Mock, patch + +import pytest + +from sagemaker.core.enums import EndpointType +from sagemaker.core.helper.session_helper import ( + Session, + _production_variant_instance_types, + production_variant, +) +from sagemaker.core.inference_config import ResourceRequirements, ServerlessInferenceConfig +from sagemaker.core.model_monitor.data_capture_config import DataCaptureConfig +from sagemaker.core.resources import Endpoint +from sagemaker.core.shapes import ProductionVariant +from sagemaker.core.utils.utils import Unassigned +from sagemaker.serve.model_builder import ( + DEPLOY_KWARGS, + ModelBuilder, + _EndpointOptions, + _instance_pools, + _with_instance_pools, +) +from sagemaker.serve.utils.types import ModelServer + +ROLE_ARN = "arn:aws:iam::123456789012:role/TestRole" +REGION = "us-east-2" +IMAGE_URI = "123456789012.dkr.ecr.us-east-2.amazonaws.com/inference:latest" +TRACING = {"EnableTracing": True, "SamplingPercentage": 10} + + +def _session(): + session = Mock() + session.boto_region_name = REGION + session.default_bucket.return_value = "test-bucket" + session.default_bucket_prefix = None + session.config = {} + session.sagemaker_config = {} + session.settings = Mock(include_jumpstart_tags=False) + session.settings._local_download_dir = None + session.sagemaker_client = MagicMock() + session.boto_session = Mock() + session.boto_session.region_name = REGION + session.endpoint_in_service_or_not.return_value = False + return session + + +def _make_builder(): + builder = ModelBuilder( + model=Mock(), + model_server=ModelServer.TORCHSERVE, + image_uri=IMAGE_URI, + s3_model_data_url="s3://test-bucket/model.tar.gz", + role_arn=ROLE_ARN, + instance_type="ml.g5.xlarge", + sagemaker_session=_session(), + ) + builder.region = REGION + builder.built_model = Mock(model_name="built-model") + builder.model_name = "built-model" + builder.model_server = None + return builder + + +def _data_capture_config(): + return DataCaptureConfig( + enable_capture=True, + sampling_percentage=25, + destination_s3_uri="s3://test-bucket/capture", + sagemaker_session=_session(), + ) + + +def _core_path_patches(builder): + return [ + patch.object(ModelBuilder, "_is_model_customization", return_value=False), + patch.object(ModelBuilder, "_resolve_model_source_id", return_value=None), + patch.object(builder, "_wait_for_endpoint"), + patch("sagemaker.serve.model_builder.Endpoint.get", return_value=Mock(spec=Endpoint)), + ] + + +class TestDeploySignature(unittest.TestCase): + """The options are named parameters and an unknown kwarg raises.""" + + def test_deploy_names_the_endpoint_options(self): + parameters = inspect.signature(ModelBuilder.deploy).parameters + + for name in ( + "data_capture_config", + "tracing_config", + "fallback_instance_types", + "endpoint_type", + "inference_component_name", + ): + self.assertIn(name, parameters) + self.assertIsNone(parameters[name].default) + + def test_deploy_rejects_an_unknown_kwarg_before_any_service_call(self): + builder = _make_builder() + + with pytest.raises(TypeError, match="unexpected keyword argument 'fallback_instance_type'"): + builder.deploy(endpoint_name="ep", fallback_instance_type=["ml.g6.xlarge"]) + + builder.sagemaker_session.endpoint_from_production_variants.assert_not_called() + + def test_deploy_kwargs_lists_every_name_a_deploy_path_reads(self): + for name in ("volume_size", "kms_key", "serializer", "async_inference_config", "role"): + self.assertIn(name, DEPLOY_KWARGS) + for name in ("data_capture_config", "endpoint_type", "inference_component_name", "tags"): + self.assertNotIn(name, DEPLOY_KWARGS) + + +class TestEndpointOptionsValidation(unittest.TestCase): + """``_EndpointOptions.validate`` rejects a combination no deploy path can apply.""" + + def test_data_capture_on_an_inference_component_endpoint_raises(self): + options = _EndpointOptions(data_capture_config=_data_capture_config()) + + with pytest.raises(ValueError, match="rejects data capture on an inference-component"): + options.validate(ResourceRequirements(requests={"memory": 1024, "copies": 1})) + + def test_data_capture_config_type_is_checked(self): + with pytest.raises(TypeError, match="DataCaptureConfig"): + _EndpointOptions(data_capture_config={"EnableCapture": True}).validate(None) + + def test_fallback_instance_types_rules(self): + with pytest.raises(ValueError, match="serverless"): + _EndpointOptions(fallback_instance_types=["ml.g6.xlarge"]).validate( + ServerlessInferenceConfig() + ) + with pytest.raises(ValueError, match="at most 4"): + _EndpointOptions( + fallback_instance_types=[f"ml.g6.{size}xlarge" for size in (1, 2, 4, 8, 12)] + ).validate(None) + with pytest.raises(ValueError, match="repeats"): + _EndpointOptions(fallback_instance_types=["ml.g6.xlarge", "ml.g6.xlarge"]).validate( + None + ) + with pytest.raises(TypeError, match="list of instance type names"): + _EndpointOptions(fallback_instance_types="ml.g6.xlarge").validate(None) + + def test_endpoint_type_must_match_the_inference_config(self): + with pytest.raises(ValueError, match="requires inference_config=ResourceRequirements"): + _EndpointOptions(endpoint_type=EndpointType.INFERENCE_COMPONENT_BASED).validate(None) + with pytest.raises(ValueError, match="conflicts with a ResourceRequirements"): + _EndpointOptions(endpoint_type=EndpointType.MODEL_BASED).validate( + ResourceRequirements(requests={"memory": 1024, "copies": 1}) + ) + with pytest.raises(TypeError, match="EndpointType"): + _EndpointOptions(endpoint_type="MODEL_BASED").validate(None) + + def test_as_deploy_kwargs_carries_only_the_set_options(self): + options = _EndpointOptions(fallback_instance_types=["ml.g6.xlarge"], tracing_config=TRACING) + + self.assertEqual( + options.as_deploy_kwargs(), + {"tracing_config": TRACING, "fallback_instance_types": ["ml.g6.xlarge"]}, + ) + self.assertEqual(_EndpointOptions().as_deploy_kwargs(), {}) + + +class TestInstancePools(unittest.TestCase): + """``fallback_instance_types`` becomes ``InstancePools`` with the primary type first.""" + + def test_instance_pools_put_the_primary_type_at_priority_one(self): + self.assertEqual( + _instance_pools("ml.g5.xlarge", ["ml.g6.xlarge", "ml.g5.xlarge", "ml.g6e.xlarge"]), + [ + {"InstanceType": "ml.g5.xlarge", "Priority": 1}, + {"InstanceType": "ml.g6.xlarge", "Priority": 2}, + {"InstanceType": "ml.g6e.xlarge", "Priority": 3}, + ], + ) + with pytest.raises(ValueError, match="requires an instance_type"): + _instance_pools(None, ["ml.g6.xlarge"]) + + def test_production_variant_carries_pools_in_place_of_the_instance_type(self): + pools = _instance_pools("ml.g5.xlarge", ["ml.g6.xlarge"]) + + variant = production_variant( + "built-model", "ml.g5.xlarge", 2, volume_size=64, instance_pools=pools + ) + + self.assertEqual(variant["InstancePools"], pools) + self.assertNotIn("InstanceType", variant) + self.assertEqual(variant["InitialInstanceCount"], 2) + self.assertEqual(variant["VolumeSizeInGB"], 64) + self.assertEqual( + _production_variant_instance_types(variant), ["ml.g5.xlarge", "ml.g6.xlarge"] + ) + plain = production_variant("built-model", "ml.g5.xlarge", 1) + self.assertEqual(plain["InstanceType"], "ml.g5.xlarge") + self.assertNotIn("InstancePools", plain) + + def test_with_instance_pools_moves_a_core_variant_to_pools(self): + variant = ProductionVariant( + variant_name="AllTraffic", + model_name="built-model", + instance_type="ml.g5.xlarge", + initial_instance_count=1, + ) + + pooled = _with_instance_pools(variant, ["ml.g6.xlarge"]) + + self.assertEqual( + [(pool.instance_type, pool.priority) for pool in pooled.instance_pools], + [("ml.g5.xlarge", 1), ("ml.g6.xlarge", 2)], + ) + self.assertIsInstance(pooled.instance_type, Unassigned) + self.assertIs(_with_instance_pools(variant, None), variant) + + +class TestSessionTracingConfig(unittest.TestCase): + """``Session`` sends ``TracingConfig`` only when the installed service model has it.""" + + def _session(self, members): + boto_session = Mock(region_name=REGION) + client = MagicMock() + client.meta.service_model.operation_model.return_value.input_shape.members = members + session = Session(boto_session=boto_session, sagemaker_client=client) + return session, client + + def test_tracing_config_reaches_create_endpoint_config(self): + session, client = self._session({"TracingConfig": object(), "ProductionVariants": object()}) + variants = [production_variant("built-model", "ml.g5.xlarge", 1)] + + with patch.object(session, "_append_sagemaker_config_tags", return_value=[]): + session.create_endpoint_config( + name="cfg", production_variants=variants, tracing_config_dict=TRACING + ) + + request = client.create_endpoint_config.call_args.kwargs + self.assertEqual(request["TracingConfig"], TRACING) + self.assertEqual(request["EndpointConfigName"], "cfg") + + def test_tracing_config_raises_when_the_service_model_lacks_it(self): + session, client = self._session({"ProductionVariants": object()}) + variants = [production_variant("built-model", "ml.g5.xlarge", 1)] + + with patch.object(session, "_append_sagemaker_config_tags", return_value=[]): + with pytest.raises(ValueError, match="TracingConfig is not a CreateEndpointConfig"): + session.create_endpoint_config( + name="cfg", production_variants=variants, tracing_config_dict=TRACING + ) + + client.create_endpoint_config.assert_not_called() + + +class TestCoreEndpointPath(unittest.TestCase): + """The single-model path forwards every option to the endpoint configuration.""" + + def test_options_reach_the_endpoint_config_request(self): + builder = _make_builder() + capture = _data_capture_config() + seen = {} + + def _create(**kwargs): + seen.update(kwargs) + return "ep" + + builder.sagemaker_session.endpoint_from_production_variants = Mock(side_effect=_create) + patches = _core_path_patches(builder) + for active in patches: + active.start() + try: + builder.deploy( + endpoint_name="ep", + wait=False, + data_capture_config=capture, + tracing_config=TRACING, + fallback_instance_types=["ml.g6.xlarge", "ml.g6e.xlarge"], + endpoint_type=EndpointType.MODEL_BASED, + ) + finally: + for active in patches: + active.stop() + + self.assertEqual(seen["data_capture_config_dict"], capture._to_request_dict()) + self.assertEqual(seen["tracing_config_dict"], TRACING) + variant = seen["production_variants"][0] + self.assertEqual( + variant["InstancePools"], + [ + {"InstanceType": "ml.g5.xlarge", "Priority": 1}, + {"InstanceType": "ml.g6.xlarge", "Priority": 2}, + {"InstanceType": "ml.g6e.xlarge", "Priority": 3}, + ], + ) + self.assertNotIn("InstanceType", variant) + + def test_default_path_sends_no_tracing_and_keeps_the_instance_type(self): + builder = _make_builder() + seen = {} + + def _create(**kwargs): + seen.update(kwargs) + return "ep" + + builder.sagemaker_session.endpoint_from_production_variants = Mock(side_effect=_create) + patches = _core_path_patches(builder) + for active in patches: + active.start() + try: + builder.deploy(endpoint_name="ep", wait=False) + finally: + for active in patches: + active.stop() + + self.assertIsNone(seen["tracing_config_dict"]) + self.assertIsNone(seen["data_capture_config_dict"]) + self.assertEqual(seen["production_variants"][0]["InstanceType"], "ml.g5.xlarge") + + def test_inference_component_name_names_the_component(self): + builder = _make_builder() + builder.sagemaker_session.endpoint_from_production_variants = Mock(return_value="ep") + patches = _core_path_patches(builder) + [ + patch( + "sagemaker.serve.model_builder.session_helper.production_variant", + return_value={"VariantName": "AllTraffic"}, + ) + ] + for active in patches: + active.start() + try: + builder.deploy( + endpoint_name="ep", + wait=False, + inference_config=ResourceRequirements(requests={"memory": 1024, "copies": 1}), + inference_component_name="my-component", + ) + finally: + for active in patches: + active.stop() + + created = builder.sagemaker_session.create_inference_component.call_args.kwargs + self.assertEqual(created["inference_component_name"], "my-component") + self.assertEqual(builder.inference_component_name, "my-component") + + def test_update_endpoint_creates_a_config_from_the_production_variant(self): + builder = _make_builder() + builder.sagemaker_session.create_endpoint_config = Mock(return_value="cfg") + patches = _core_path_patches(builder) + for active in patches: + active.start() + try: + builder.deploy( + endpoint_name="existing", update_endpoint=True, wait=False, tracing_config=TRACING + ) + finally: + for active in patches: + active.stop() + + request = builder.sagemaker_session.create_endpoint_config.call_args.kwargs + self.assertEqual(request["production_variants"][0]["InstanceType"], "ml.g5.xlarge") + self.assertEqual(request["tracing_config_dict"], TRACING) + builder.sagemaker_session.update_endpoint.assert_called_once_with("existing", "cfg") + + +class TestCoreResourcePaths(unittest.TestCase): + """The fine-tuned, restricted, and Nova paths apply the options or raise.""" + + def test_nova_deploy_applies_data_capture_and_instance_pools(self): + builder = _make_builder() + capture = _data_capture_config() + seen = {} + + def _create(**kwargs): + seen.update(kwargs) + + with ( + patch.object(ModelBuilder, "_base_model_name", return_value="nova-micro"), + patch.object(ModelBuilder, "_fetch_model_package", return_value=None), + patch("sagemaker.serve.model_builder.EndpointConfig.create", side_effect=_create), + patch( + "sagemaker.serve.model_builder.Endpoint.create", return_value=Mock(spec=Endpoint) + ), + ): + builder._deploy_nova_model( + endpoint_name="ep", + wait=False, + endpoint_options=_EndpointOptions( + data_capture_config=capture, fallback_instance_types=["ml.g6.xlarge"] + ), + ) + + self.assertEqual(seen["data_capture_config"].initial_sampling_percentage, 25) + self.assertEqual(seen["data_capture_config"].destination_s3_uri, "s3://test-bucket/capture") + variant = seen["production_variants"][0] + self.assertEqual( + [(pool.instance_type, pool.priority) for pool in variant.instance_pools], + [("ml.g5.xlarge", 1), ("ml.g6.xlarge", 2)], + ) + + def test_nova_deploy_rejects_an_inference_component_name(self): + builder = _make_builder() + + with patch.object(ModelBuilder, "_is_nova_model", return_value=True): + with pytest.raises(ValueError, match="creates no inference component"): + builder._deploy_model_customization( + endpoint_name="ep", + wait=False, + endpoint_options=_EndpointOptions(inference_component_name="ic"), + ) + + def test_fine_tuned_inference_component_deploy_rejects_data_capture(self): + builder = _make_builder() + + with ( + patch.object(ModelBuilder, "_is_nova_model", return_value=False), + patch.object(ModelBuilder, "_fetch_model_package", return_value=Mock()), + patch( + "sagemaker.serve.utils.model_package_utils.is_restricted_model_package", + return_value=False, + ), + patch.object(ModelBuilder, "_fetch_peft", return_value=None), + patch.object(ModelBuilder, "_prepare_reused_model_customization_deployment_state"), + ): + with pytest.raises(ValueError, match="rejects data capture"): + builder._deploy_model_customization( + endpoint_name="ep", + wait=False, + inference_config=ResourceRequirements(requests={"memory": 1024, "copies": 1}), + endpoint_options=_EndpointOptions(data_capture_config=_data_capture_config()), + ) + + def test_tracing_config_on_a_core_resource_path_goes_through_the_session(self): + builder = _make_builder() + builder.sagemaker_session.create_endpoint_config = Mock(return_value="ep") + + with ( + patch.object(ModelBuilder, "_base_model_name", return_value="nova-micro"), + patch.object(ModelBuilder, "_fetch_model_package", return_value=None), + patch("sagemaker.serve.model_builder.EndpointConfig.create") as core_create, + patch( + "sagemaker.serve.model_builder.Endpoint.create", return_value=Mock(spec=Endpoint) + ), + ): + builder._deploy_nova_model( + endpoint_name="ep", + wait=False, + endpoint_options=_EndpointOptions(tracing_config=TRACING), + ) + + core_create.assert_not_called() + request = builder.sagemaker_session.create_endpoint_config.call_args.kwargs + self.assertEqual(request["tracing_config_dict"], TRACING) + self.assertEqual(request["production_variants"][0]["ModelName"], "built-model") + + def test_recommendation_deploy_rejects_the_endpoint_options(self): + builder = _make_builder() + builder._recommendation_job = Mock() + + with patch.object(ModelBuilder, "_deploy_recommendation") as deploy_recommendation: + with pytest.raises(ValueError, match="cannot be applied by a recommendation deploy"): + builder.deploy(fallback_instance_types=["ml.g6.xlarge"]) + + deploy_recommendation.assert_not_called() diff --git a/sagemaker-serve/tests/unit/test_model_builder_deploy_tags.py b/sagemaker-serve/tests/unit/test_model_builder_deploy_tags.py new file mode 100644 index 0000000000..57f9ba924d --- /dev/null +++ b/sagemaker-serve/tests/unit/test_model_builder_deploy_tags.py @@ -0,0 +1,303 @@ +# 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. +"""``ModelBuilder.deploy(tags=...)`` is one tag contract for every deploy path. + +Before this change ``tags`` reached ``deploy()`` only through ``**kwargs``, so +``inspect.signature`` could not validate it, the fine-tuned (model customization) and +Nova paths ignored it, and ``add_tags`` raised ``KeyError: 'Key'`` on the +``{"key": ..., "value": ...}`` form that the ``sagemaker.core.resources`` classes use. +Every representation now comes out with string keys and values, and the tags a call +passes are merged for that call only, so a second ``deploy()`` on the same builder does +not carry the first call's tags. +""" + +from __future__ import absolute_import + +import inspect +import unittest +from unittest.mock import MagicMock, Mock, patch + +import pytest + +from sagemaker.core.common_utils import _validate_new_tags, format_tags, remove_tag_with_key +from sagemaker.core.enums import EndpointType +from sagemaker.core.inference_config import ResourceRequirements +from sagemaker.core.resources import Endpoint +from sagemaker.core.shapes import Tag +from sagemaker.serve.model_builder import ModelBuilder +from sagemaker.serve.utils.types import ModelServer + +ROLE_ARN = "arn:aws:iam::123456789012:role/TestRole" +REGION = "us-east-2" +IMAGE_URI = "123456789012.dkr.ecr.us-east-2.amazonaws.com/inference:latest" +BOTO_TAG = {"Key": "project", "Value": "example"} +CORE_TAG = {"key": "project", "value": "example"} + + +def _session(): + session = Mock() + session.boto_region_name = REGION + session.default_bucket.return_value = "test-bucket" + session.default_bucket_prefix = None + session.config = {} + session.sagemaker_config = {} + session.settings = Mock(include_jumpstart_tags=False) + session.settings._local_download_dir = None + session.sagemaker_client = MagicMock() + session.boto_session = Mock() + session.boto_session.region_name = REGION + return session + + +def _make_builder(): + builder = ModelBuilder( + model=Mock(), + model_server=ModelServer.TORCHSERVE, + image_uri=IMAGE_URI, + s3_model_data_url="s3://test-bucket/model.tar.gz", + role_arn=ROLE_ARN, + instance_type="ml.m5.large", + sagemaker_session=_session(), + ) + builder.region = REGION + builder.built_model = Mock(model_name="built-model") + builder.model_server = None + return builder + + +class TestTagNormalization(unittest.TestCase): + """``format_tags`` and ``add_tags`` accept every tag representation the SDK uses.""" + + def test_format_tags_accepts_every_representation(self): + expected = [BOTO_TAG] + + self.assertEqual(format_tags({"project": "example"}), expected) + self.assertEqual(format_tags(BOTO_TAG), expected) + self.assertEqual(format_tags([BOTO_TAG]), expected) + self.assertEqual(format_tags([CORE_TAG]), expected) + self.assertEqual(format_tags([Tag(key="project", value="example")]), expected) + self.assertEqual(format_tags([("project", "example")]), expected) + self.assertEqual( + format_tags({"n": 1}), + format_tags([("n", 1)]), + format_tags([{"key": "n", "value": 1}]), + ) + self.assertEqual(format_tags([("n", 1)]), [{"Key": "n", "Value": "1"}]) + self.assertEqual( + format_tags([BOTO_TAG, {"key": "team", "value": "ml"}]), + [BOTO_TAG, {"Key": "team", "Value": "ml"}], + ) + self.assertIsNone(format_tags(None)) + + def test_format_tags_rejects_a_tag_without_a_key(self): + with pytest.raises(TypeError, match="'Key'"): + format_tags([{"Value": "no-key"}]) + with pytest.raises(TypeError, match="Tag object"): + format_tags(["project=example"]) + + def test_validate_new_tags_merges_mixed_casing_by_key(self): + merged = _validate_new_tags([CORE_TAG, {"key": "team", "value": "ml"}], [BOTO_TAG]) + + self.assertEqual(merged, [BOTO_TAG, {"Key": "team", "Value": "ml"}]) + + def test_validate_new_tags_with_tag_objects(self): + merged = _validate_new_tags([Tag(key="team", value="ml")], [CORE_TAG]) + + self.assertEqual(merged, [BOTO_TAG, {"Key": "team", "Value": "ml"}]) + + def test_remove_tag_with_key_accepts_core_casing(self): + self.assertIsNone(remove_tag_with_key("project", [CORE_TAG])) + + def test_add_tags_normalizes_core_casing_to_boto(self): + builder = _make_builder() + + builder.add_tags([CORE_TAG]) + builder.add_tags({"team": "ml"}) + builder.add_tags([Tag(key="env", value="prod")]) + + self.assertEqual( + builder._tags, + [BOTO_TAG, {"Key": "team", "Value": "ml"}, {"Key": "env", "Value": "prod"}], + ) + + +class TestDeployTagsParameter(unittest.TestCase): + """``deploy(tags=...)`` is a named parameter and reaches every deploy path.""" + + def test_deploy_has_a_tags_parameter(self): + parameters = inspect.signature(ModelBuilder.deploy).parameters + + self.assertIn("tags", parameters) + self.assertIsNone(parameters["tags"].default) + + def test_deploy_merges_tags_into_the_builder_before_the_core_path(self): + builder = _make_builder() + builder.add_tags({"team": "ml"}) + seen = {} + + def _capture(**kwargs): + seen["tags"] = list(builder._tags) + seen["kwargs_tags"] = kwargs.get("tags") + return Mock(spec=Endpoint) + + with ( + patch.object(ModelBuilder, "_deploy", side_effect=_capture), + patch.object(ModelBuilder, "_is_model_customization", return_value=False), + patch.object(ModelBuilder, "_resolve_model_source_id", return_value=None), + ): + builder.deploy(endpoint_name="ep", tags=[CORE_TAG]) + + self.assertEqual(seen["tags"], [{"Key": "team", "Value": "ml"}, BOTO_TAG]) + # The core path reads self._tags; no duplicate copy travels through kwargs. + self.assertIsNone(seen["kwargs_tags"]) + + def test_deploy_core_endpoint_normalizes_tags_passed_through_kwargs(self): + builder = _make_builder() + builder.built_model.model_name = "built-model" + builder.model_name = "built-model" + builder.sagemaker_session.endpoint_in_service_or_not.return_value = False + seen = {} + + def _capture(**kwargs): + seen["tags"] = kwargs.get("tags") + return "ep" + + builder.sagemaker_session.endpoint_from_production_variants = Mock(side_effect=_capture) + with ( + patch( + "sagemaker.serve.model_builder.session_helper.production_variant", + return_value={"VariantName": "AllTraffic"}, + ), + patch("sagemaker.serve.model_builder.Endpoint.get", return_value=Mock(spec=Endpoint)), + patch.object(builder, "_wait_for_endpoint"), + ): + builder._deploy_core_endpoint( + instance_type="ml.g5.xlarge", + initial_instance_count=1, + endpoint_name="ep", + endpoint_type=EndpointType.INFERENCE_COMPONENT_BASED, + resources=ResourceRequirements(requests={"memory": 1024, "copies": 1}), + inference_component_name="ic", + tags=[CORE_TAG, Tag(key="team", value="ml")], + ) + + self.assertEqual(seen["tags"], [BOTO_TAG, {"Key": "team", "Value": "ml"}]) + self.assertEqual(builder._tags, [BOTO_TAG, {"Key": "team", "Value": "ml"}]) + + def test_deploy_passes_the_builder_tags_to_a_recommendation_deploy(self): + builder = _make_builder() + builder._recommendation_job = Mock() + seen = {} + + def _capture(**kwargs): + seen["tags"] = kwargs["tags"] + return Mock(spec=Endpoint) + + with patch.object(ModelBuilder, "_deploy_recommendation", side_effect=_capture): + builder.deploy(tags={"project": "example"}) + + self.assertEqual(seen["tags"], [BOTO_TAG]) + + def test_core_endpoint_tags_use_the_core_resource_casing(self): + builder = _make_builder() + builder.add_tags([BOTO_TAG, Tag(key="team", value="ml")]) + + self.assertEqual( + builder._core_endpoint_tags(), + [CORE_TAG, {"key": "team", "value": "ml"}], + ) + + def test_model_customization_restricted_package_endpoint_is_tagged(self): + builder = _make_builder() + builder.instance_type = "ml.g5.xlarge" + created = {} + + def _create_endpoint(**kwargs): + created.update(kwargs) + endpoint = Mock(spec=Endpoint) + endpoint.wait_for_status = Mock() + return endpoint + + with ( + patch.object(ModelBuilder, "_is_nova_model", return_value=False), + patch.object(ModelBuilder, "_fetch_model_package", return_value=Mock()), + patch( + "sagemaker.serve.utils.model_package_utils.is_restricted_model_package", + return_value=True, + ), + patch("sagemaker.serve.model_builder.EndpointConfig.create"), + patch("sagemaker.serve.model_builder.Endpoint.create", side_effect=_create_endpoint), + ): + builder.add_tags([CORE_TAG]) + builder._deploy_model_customization(endpoint_name="ep", wait=False) + + self.assertEqual(created["tags"], [CORE_TAG]) + + def test_nova_endpoint_merges_the_builder_tags(self): + builder = _make_builder() + builder.instance_type = "ml.g5.xlarge" + created = {} + + def _create_endpoint(**kwargs): + created.update(kwargs) + return Mock(spec=Endpoint) + + with ( + patch.object(ModelBuilder, "_base_model_name", return_value="nova-micro"), + patch.object(ModelBuilder, "_fetch_model_package", return_value=None), + patch("sagemaker.serve.model_builder.EndpointConfig.create"), + patch("sagemaker.serve.model_builder.Endpoint.create", side_effect=_create_endpoint), + ): + builder.add_tags([CORE_TAG, {"Key": "sagemaker-sdk:jumpstart-model-id", "Value": "x"}]) + builder._deploy_nova_model(endpoint_name="ep", wait=False) + + self.assertEqual( + created["tags"], + [{"key": "sagemaker-sdk:jumpstart-model-id", "value": "nova-micro"}, CORE_TAG], + ) + + def test_deploy_tags_do_not_outlive_the_call(self): + builder = _make_builder() + builder.add_tags({"team": "ml"}) + seen = {} + + def _capture(**kwargs): + seen["tags"] = list(builder._tags) + return Mock(spec=Endpoint) + + with ( + patch.object(ModelBuilder, "_deploy", side_effect=_capture), + patch.object(ModelBuilder, "_is_model_customization", return_value=False), + patch.object(ModelBuilder, "_resolve_model_source_id", return_value="s3://b/m"), + ): + builder.deploy(endpoint_name="ep", tags={"run": "1"}) + + self.assertEqual( + [tag["Key"] for tag in seen["tags"]], + ["team", "run", "sagemaker.amazonaws.com/model-source"], + ) + self.assertEqual(builder._tags, [{"Key": "team", "Value": "ml"}]) + + def test_deploy_tags_are_restored_when_the_path_raises(self): + builder = _make_builder() + builder.add_tags({"team": "ml"}) + + with ( + patch.object(ModelBuilder, "_deploy", side_effect=RuntimeError("boom")), + patch.object(ModelBuilder, "_is_model_customization", return_value=False), + patch.object(ModelBuilder, "_resolve_model_source_id", return_value=None), + ): + with pytest.raises(RuntimeError, match="boom"): + builder.deploy(endpoint_name="ep", tags={"run": "1"}) + + self.assertEqual(builder._tags, [{"Key": "team", "Value": "ml"}]) diff --git a/sagemaker-serve/tests/unit/test_model_builder_instance_types.py b/sagemaker-serve/tests/unit/test_model_builder_instance_types.py new file mode 100644 index 0000000000..e915d9eecd --- /dev/null +++ b/sagemaker-serve/tests/unit/test_model_builder_instance_types.py @@ -0,0 +1,118 @@ +# 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. +"""``ModelBuilder`` reads instance hardware through ``sagemaker.core.instance_types``. + +Before this change ``_get_instance_resources`` and ``_infer_accelerator_count_from_instance_type`` +each made their own EC2 ``describe_instance_types`` call with its own ``ml.`` strip. +``sagemaker.serve.model_builder`` had no ``describe_instance_type`` to patch, and the two +private copies could drift from each other and from the public module. +""" + +from __future__ import absolute_import + +import unittest +from unittest.mock import Mock, patch + +from botocore.exceptions import ClientError + +from sagemaker.core.instance_types import InstanceSpec +from sagemaker.core.resources import TrainingJob +from sagemaker.serve.model_builder import ModelBuilder + +G5_12XLARGE_INFO = { + "InstanceType": "g5.12xlarge", + "VCpuInfo": {"DefaultVCpus": 48}, + "MemoryInfo": {"SizeInMiB": 196608}, + "GpuInfo": {"Gpus": [{"Name": "A10G", "Count": 4, "MemoryInfo": {"SizeInMiB": 24576}}]}, +} +M5_XLARGE_INFO = { + "InstanceType": "m5.xlarge", + "VCpuInfo": {"DefaultVCpus": 4}, + "MemoryInfo": {"SizeInMiB": 16384}, +} +UNKNOWN_TYPE_ERROR = ClientError( + {"Error": {"Code": "InvalidInstanceType", "Message": "retired.xlarge does not exist"}}, + "DescribeInstanceTypes", +) + + +class ModelBuilderInstanceTypesTest(unittest.TestCase): + """The two hardware helpers delegate to ``describe_instance_type`` and keep their fallbacks.""" + + def setUp(self): + self.mock_session = Mock() + self.mock_session.boto_region_name = "us-east-1" + self.mock_session.boto_session = Mock() + self.mock_session.boto_session.region_name = "us-east-1" + self.mock_session.config = {} + self.mock_session.sagemaker_config = {} + self.ec2_client = Mock() + self.mock_session.boto_session.client.return_value = self.ec2_client + self.builder = ModelBuilder( + model=Mock(spec=TrainingJob), + role_arn="arn:aws:iam::123456789012:role/SageMakerRole", + sagemaker_session=self.mock_session, + ) + + def _ec2_returns(self, *infos): + self.ec2_client.describe_instance_types.return_value = {"InstanceTypes": list(infos)} + + @patch("sagemaker.serve.model_builder.describe_instance_type") + def test_get_instance_resources_passes_the_boto_session_to_the_public_helper(self, describe): + describe.return_value = InstanceSpec("ml.g5.12xlarge", 48, 196608, 4, 98304, "A10G") + + self.assertEqual(self.builder._get_instance_resources("ml.g5.12xlarge"), (48, 196608)) + describe.assert_called_once_with("ml.g5.12xlarge", session=self.mock_session.boto_session) + + def test_get_instance_resources_reads_vcpus_and_memory_from_ec2(self): + self._ec2_returns(G5_12XLARGE_INFO) + + self.assertEqual(self.builder._get_instance_resources("ml.g5.12xlarge"), (48, 196608)) + self.ec2_client.describe_instance_types.assert_called_once_with( + InstanceTypes=["g5.12xlarge"] + ) + + def test_get_instance_resources_falls_back_to_none_for_an_unknown_type(self): + self.ec2_client.describe_instance_types.side_effect = UNKNOWN_TYPE_ERROR + + self.assertEqual(self.builder._get_instance_resources("ml.retired.xlarge"), (None, None)) + + def test_get_instance_resources_falls_back_to_none_when_the_call_fails(self): + self.ec2_client.describe_instance_types.side_effect = ClientError( + {"Error": {"Code": "UnauthorizedOperation", "Message": "not allowed"}}, + "DescribeInstanceTypes", + ) + + self.assertEqual(self.builder._get_instance_resources("ml.g5.12xlarge"), (None, None)) + + def test_infer_accelerator_count_sums_the_gpus(self): + self._ec2_returns(G5_12XLARGE_INFO) + + self.assertEqual( + self.builder._infer_accelerator_count_from_instance_type("ml.g5.12xlarge"), 4 + ) + self.assertTrue(self.builder._is_gpu_instance("ml.g5.12xlarge")) + + def test_infer_accelerator_count_is_none_for_a_cpu_instance(self): + self._ec2_returns(M5_XLARGE_INFO) + + self.assertIsNone(self.builder._infer_accelerator_count_from_instance_type("ml.m5.xlarge")) + self.assertFalse(self.builder._is_gpu_instance("ml.m5.xlarge")) + + def test_infer_accelerator_count_is_none_for_an_unknown_type(self): + self.ec2_client.describe_instance_types.side_effect = UNKNOWN_TYPE_ERROR + + self.assertIsNone( + self.builder._infer_accelerator_count_from_instance_type("ml.retired.xlarge") + ) + self.assertFalse(self.builder._is_gpu_instance("ml.retired.xlarge")) diff --git a/sagemaker-train/src/sagemaker/ai_registry/__init__.py b/sagemaker-train/src/sagemaker/ai_registry/__init__.py index 6549052177..5efb2e32f1 100644 --- a/sagemaker-train/src/sagemaker/ai_registry/__init__.py +++ b/sagemaker-train/src/sagemaker/ai_registry/__init__.py @@ -10,3 +10,16 @@ # 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. +"""AI Registry: datasets, evaluators, and dataset format detection.""" + +from sagemaker.ai_registry.dataset_format_detector import ( # noqa: F401 + FormatDetectionOptions, + FormatDetectionResult, + detect_dataset_format, +) + +__all__ = [ + "FormatDetectionOptions", + "FormatDetectionResult", + "detect_dataset_format", +] diff --git a/sagemaker-train/src/sagemaker/ai_registry/air_constants.py b/sagemaker-train/src/sagemaker/ai_registry/air_constants.py index eb6d1d7619..58a30d0afe 100644 --- a/sagemaker-train/src/sagemaker/ai_registry/air_constants.py +++ b/sagemaker-train/src/sagemaker/ai_registry/air_constants.py @@ -36,7 +36,7 @@ EVALUATOR_SPEC_HUB_CONTENT_SUBTYPE = "AWS/Specification" EVALUATOR_DOCUMENT_SCHEMA_VERSION = "2.0.0" EVALUATOR_DEFAULT_METHOD = "lambda" -EVALUATOR_DEFAULT_RUNTIME = "python3.9" +EVALUATOR_DEFAULT_RUNTIME = "python3.12" EVALUATOR_BYOCODE = "BYOCode" EVALUATOR_BYOLAMBDA = "BYOLambda" diff --git a/sagemaker-train/src/sagemaker/ai_registry/dataset_format_detector.py b/sagemaker-train/src/sagemaker/ai_registry/dataset_format_detector.py index b3ca8d52da..b057414397 100644 --- a/sagemaker-train/src/sagemaker/ai_registry/dataset_format_detector.py +++ b/sagemaker-train/src/sagemaker/ai_registry/dataset_format_detector.py @@ -10,11 +10,124 @@ # 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. -"""Utilities for detecting the format of a customization dataset.""" +"""Utilities for detecting the format of a customization dataset. +``DatasetFormatDetector`` answers whether a local file matches one known format. +``detect_dataset_format`` samples the first bytes of a local file or an ``s3://`` object and +reports the best matching format, its confidence, the other candidates, and the lines that +failed. +""" + +import csv import json -from typing import Dict, Any +import os +from dataclasses import dataclass from pathlib import Path +from typing import Any, Callable, Dict, List, NamedTuple, Optional, Tuple +from urllib.parse import urlparse + +from botocore.exceptions import ClientError +from jsonschema.exceptions import best_match +from jsonschema.validators import validator_for + +from sagemaker.core.helper.session_helper import Session +from sagemaker.core.s3.client import S3Downloader, is_s3_url + +DEFAULT_SAMPLE_SIZE_BYTES = 1_048_576 +RELAXED_MIN_RATIO = 0.5 + +# Names the two agent-skill detectors use for shapes the SDK schemas already describe. +# Lower-case names: SageMakerAIAgentSkills dataset-evaluation/scripts/format_detector.py. +# Capitalized names: the verse dataset-quality/scripts/format_detector.py, aligned with the +# Hub DatasetFileFormat enum. The value is the key in ``FORMAT_MATCHERS``. +SKILL_FORMAT_ALIASES: Dict[str, str] = { + "nova_sft": "converse", + "Bedrock_Converse": "converse", + "nova_dpo": "dpo", + "nova_rlvr": "nova_rlvr", + "Nova_RLVR": "nova_rlvr", + "Nova_PPO": "nova_ppo", + "Nova_CPT": "nova_cpt", + "gpt_oss_sft": "openai_chat", + "OpenAI_Messages": "openai_chat", + "gpt_oss_dpo": "hf_preference", + "open_weights_dpo": "hf_preference", + "HF_Preference": "hf_preference", + "open_weights_sft": "hf_prompt_completion", + "open_weights_sft_conv": "hf_prompt_completion", + "HF_Prompt_Completion": "hf_prompt_completion", + "sagemaker_eval": "genqa", + "SageMaker_Eval": "genqa", + "verl_legacy": "verl", + "Verl": "verl", +} + +NOVA_PPO_REWARD_CATEGORIES = frozenset({"default", "math", "coding", "if", "rag", "rai"}) + +# A bare ``text`` column is a weak signal, so a field of another format rules Nova CPT out. +# The list matches the verse skill and the Data Quality service. +NON_CPT_FIELDS = frozenset( + { + "prompt", + "completion", + "messages", + "query", + "chosen", + "rejected", + "label", + "label_text", + "labels", + "intent", + "category", + "answer", + } +) + +# ``None`` means the record matches. A string names the first mismatch. +MismatchReason = Optional[str] +Matcher = Callable[[Any], MismatchReason] +LineError = Tuple[int, str] + + +@dataclass(frozen=True) +class FormatDetectionResult: + """Outcome of ``detect_dataset_format``. + + Attributes: + format: The detected format name, ``None`` when nothing matched. + confidence: Matched lines divided by sampled lines for ``format``, 0.0 to 1.0. + lines_sampled: Non-empty complete lines in the sample. + line_errors: ``": "`` for each line that failed to parse or + failed ``format``. Without a detected format the reasons refer to the best candidate. + candidates: Every format that matched at least one line, with its ratio. + """ + + format: Optional[str] + confidence: float + lines_sampled: int + line_errors: List[str] + candidates: Dict[str, float] + + +@dataclass +class FormatDetectionOptions: + """Options for ``detect_dataset_format``. + + Attributes: + sample_size_bytes: Bytes read from the start of the source. Defaults to 1 MiB. + strict: When ``True`` every sampled line must match the format. When ``False`` the + best candidate above a 0.5 ratio wins. + sagemaker_session: Session for ``s3://`` sources. ``None`` creates a default session. + """ + + sample_size_bytes: int = DEFAULT_SAMPLE_SIZE_BYTES + strict: bool = False + sagemaker_session: Optional[Session] = None + + +class _Sample(NamedTuple): + data: bytes + complete: bool class DatasetFormatDetector: @@ -145,3 +258,257 @@ def _is_rft_format(data: Dict[str, Any]) -> bool: return False return True + + +def _json_path(error) -> str: + parts = [f"[{part}]" if isinstance(part, int) else f".{part}" for part in error.absolute_path] + return "$" + "".join(parts) + + +def _schema_matcher(format_name: str) -> Matcher: + """Return a matcher that validates a record against the SDK JSON schema of ``format_name``.""" + schema = DatasetFormatDetector._load_schema(format_name) + if not schema: + raise ValueError( + f"No JSON schema named {format_name!r} in {DatasetFormatDetector.SCHEMA_DIR}" + ) + validator = validator_for(schema)(schema) + + def match(record: Any) -> MismatchReason: + error = best_match(validator.iter_errors(record)) + if error is None: + return None + return f"{_json_path(error)}: {error.message}" + + return match + + +_match_openai_chat = _schema_matcher("openai_chat") + + +def _not_an_object(record: Any) -> str: + return f"$: {type(record).__name__} is not an object" + + +def _match_nova_rlvr(record: Any) -> MismatchReason: + """Nova RLVR: OpenAI-style messages plus a ``reference_answer`` object.""" + if not isinstance(record, dict): + return _not_an_object(record) + if not isinstance(record.get("reference_answer"), dict): + return "$.reference_answer: missing or not an object" + return _match_openai_chat(record) + + +def _is_nova_ppo_turn(turn: Any) -> bool: + return isinstance(turn, str) and turn.startswith(("User:", "Bot:")) + + +def _match_nova_ppo(record: Any) -> MismatchReason: + """Nova PPO: ``turns``, ``turns_to_mask``, and a known ``reward_category``.""" + if not isinstance(record, dict): + return _not_an_object(record) + turns = record.get("turns") + if not isinstance(turns, list) or not turns or not all(map(_is_nova_ppo_turn, turns)): + return "$.turns: expected a non-empty list of strings that start with 'User:' or 'Bot:'" + turns_to_mask = record.get("turns_to_mask") + if not isinstance(turns_to_mask, list) or not all(isinstance(i, int) for i in turns_to_mask): + return "$.turns_to_mask: expected a list of integers" + if record.get("reward_category") not in NOVA_PPO_REWARD_CATEGORIES: + return f"$.reward_category: expected one of {sorted(NOVA_PPO_REWARD_CATEGORIES)}" + return None + + +def _match_nova_cpt(record: Any) -> MismatchReason: + """Nova CPT: a ``text`` string and no field of another format.""" + if not isinstance(record, dict): + return _not_an_object(record) + if not isinstance(record.get("text"), str): + return "$.text: missing or not a string" + other_fields = sorted(NON_CPT_FIELDS.intersection(record)) + if other_fields: + return f"$: fields {other_fields} belong to another format" + return None + + +# Priority order: the most specific shape first. A tie in match ratio resolves to the earlier +# entry, so an SFT converse line (which the dpo schema also accepts) reports ``converse`` and a +# prompt/completion line (which the verl schema also accepts) reports ``hf_prompt_completion``. +FORMAT_MATCHERS: Dict[str, Matcher] = { + "nova_rlvr": _match_nova_rlvr, + "nova_ppo": _match_nova_ppo, + "converse": _schema_matcher("converse"), + "dpo": _schema_matcher("dpo"), + "openai_chat": _match_openai_chat, + "hf_preference": _schema_matcher("hf_preference"), + "genqa": _schema_matcher("genqa"), + "hf_prompt_completion": _schema_matcher("hf_prompt_completion"), + "nova_cpt": _match_nova_cpt, + "verl": _schema_matcher("verl"), +} + + +def _match_parquet(sample: _Sample) -> MismatchReason: + if sample.data[:4] == b"PAR1": + return None + return "missing the Parquet magic bytes PAR1" + + +def _match_csv(sample: _Sample) -> MismatchReason: + first_line = sample.data.split(b"\n", 1)[0].decode("utf-8", errors="replace") + header = next(csv.reader([first_line]), None) + if header: + return None + return "missing a CSV header row" + + +def _match_json(sample: _Sample) -> MismatchReason: + text = sample.data.decode("utf-8", errors="replace") + if not sample.complete: + return None if text.lstrip().startswith("[") else "expected a JSON array" + try: + parsed = json.loads(text) + except json.JSONDecodeError as error: + return f"invalid JSON: {error}" + if isinstance(parsed, list) and parsed: + return None + return "expected a non-empty JSON array" + + +# Formats that a file extension names. The matcher reads the sampled bytes, not JSONL lines. +EXTENSION_MATCHERS: Dict[str, Callable[[_Sample], MismatchReason]] = { + ".parquet": _match_parquet, + ".csv": _match_csv, + ".json": _match_json, +} + + +def _read_s3_sample(source: str, options: FormatDetectionOptions) -> bytes: + byte_range = (0, options.sample_size_bytes - 1) + try: + return S3Downloader.read_bytes(source, options.sagemaker_session, byte_range=byte_range) + except ClientError as error: + # GetObject with a Range header fails with 416 InvalidRange on a zero-byte object. + if error.response.get("Error", {}).get("Code") == "InvalidRange": + return b"" + raise + + +def _read_sample(source: str, options: FormatDetectionOptions) -> _Sample: + if options.sample_size_bytes < 1: + raise ValueError(f"sample_size_bytes must be at least 1, got {options.sample_size_bytes}.") + if is_s3_url(source): + data = _read_s3_sample(source, options) + else: + with open(source, "rb") as handle: + data = handle.read(options.sample_size_bytes) + return _Sample(data=data, complete=len(data) < options.sample_size_bytes) + + +def _complete_lines(sample: _Sample) -> List[Tuple[int, str]]: + """Return ``(line number, text)`` for each non-empty complete line of the sample.""" + data = sample.data + if not sample.complete: + # Trim at the last newline byte before decoding. A byte range can cut a multibyte + # character, and 0x0A is never a UTF-8 continuation byte. + data = data[: data.rfind(b"\n") + 1] + text = data.decode("utf-8", errors="replace") + return [(number, line) for number, line in enumerate(text.split("\n"), start=1) if line.strip()] + + +def _parse_lines(lines: List[Tuple[int, str]]) -> Tuple[List[LineError], List[Tuple[int, Any]]]: + parse_errors: List[LineError] = [] + records: List[Tuple[int, Any]] = [] + for number, line in lines: + try: + records.append((number, json.loads(line))) + except json.JSONDecodeError as error: + parse_errors.append((number, f"invalid JSON: {error}")) + return parse_errors, records + + +def _mismatches(matcher: Matcher, records: List[Tuple[int, Any]]) -> List[LineError]: + mismatches: List[LineError] = [] + for number, record in records: + reason = matcher(record) + if reason is not None: + mismatches.append((number, reason)) + return mismatches + + +def _select_format(best_name: Optional[str], ratio: float, strict: bool) -> Optional[str]: + if best_name is None: + return None + if strict: + return best_name if ratio >= 1.0 else None + return best_name if ratio > RELAXED_MIN_RATIO else None + + +def _extension_result(format_name: str, reason: MismatchReason) -> FormatDetectionResult: + if reason is None: + return FormatDetectionResult( + format=format_name, + confidence=1.0, + lines_sampled=0, + line_errors=[], + candidates={format_name: 1.0}, + ) + return FormatDetectionResult( + format=None, confidence=0.0, lines_sampled=0, line_errors=[f"1: {reason}"], candidates={} + ) + + +def _detect_jsonl(lines: List[Tuple[int, str]], strict: bool) -> FormatDetectionResult: + if not lines: + return FormatDetectionResult( + format=None, confidence=0.0, lines_sampled=0, line_errors=[], candidates={} + ) + parse_errors, records = _parse_lines(lines) + mismatches_by_format = { + name: _mismatches(matcher, records) for name, matcher in FORMAT_MATCHERS.items() + } + candidates = { + name: (len(records) - len(mismatches)) / len(lines) + for name, mismatches in mismatches_by_format.items() + if len(mismatches) < len(records) + } + best_name = max(candidates, key=candidates.get, default=None) + detected = _select_format(best_name, candidates.get(best_name, 0.0), strict) + reported = detected or best_name + line_errors = list(parse_errors) + if reported is not None: + line_errors.extend(mismatches_by_format[reported]) + return FormatDetectionResult( + format=detected, + confidence=candidates[detected] if detected else 0.0, + lines_sampled=len(lines), + line_errors=[f"{number}: {reason}" for number, reason in sorted(line_errors)], + candidates=candidates, + ) + + +def detect_dataset_format( + source: str, options: Optional[FormatDetectionOptions] = None +) -> FormatDetectionResult: + """Detect the format of a dataset from the first bytes of a local file or an S3 object. + + An ``s3://`` source is read with one ranged ``GetObject`` through + ``S3Downloader.read_bytes``. A ``.parquet``, ``.csv``, or ``.json`` extension selects a + whole-file matcher. Every other source is JSONL: each sampled line is matched against + ``FORMAT_MATCHERS``, the SDK schema names plus the Nova shapes the agent skills detect. + + Args: + source: A local path or an ``s3://bucket/key`` URI. + options: Sample size, strict or relaxed selection, and the session for S3 reads. + + Returns: + FormatDetectionResult: The detected format, its confidence, the sampled line count, + the lines that failed, and every candidate with its match ratio. + """ + options = options or FormatDetectionOptions() + sample = _read_sample(source, options) + path = urlparse(source).path if is_s3_url(source) else source + extension = os.path.splitext(path)[1].lower() + extension_matcher = EXTENSION_MATCHERS.get(extension) + if extension_matcher is not None: + return _extension_result(extension.lstrip("."), extension_matcher(sample)) + return _detect_jsonl(_complete_lines(sample), options.strict) diff --git a/sagemaker-train/src/sagemaker/ai_registry/evaluator.py b/sagemaker-train/src/sagemaker/ai_registry/evaluator.py index 6d197ea117..bc8be4d543 100644 --- a/sagemaker-train/src/sagemaker/ai_registry/evaluator.py +++ b/sagemaker-train/src/sagemaker/ai_registry/evaluator.py @@ -246,6 +246,7 @@ def create( role: Optional[str] = None, domain_id: Optional[str] = None, sagemaker_session: Optional[Session] = None, + runtime: Optional[str] = None, ) -> "Evaluator": """Create a new Evaluator entity in the AI Registry. @@ -259,6 +260,8 @@ def create( is visible in Studio. If not provided, it is auto-detected from the Studio environment; supply it explicitly when creating evaluators outside Studio (e.g. from a laptop or CI) so they still appear in the target domain. + runtime: Lambda runtime of the function created from a local file. Defaults to + ``EVALUATOR_DEFAULT_RUNTIME`` (``python3.12``). Ignored for a Lambda ARN. Returns: Evaluator: Newly created Evaluator instance @@ -284,7 +287,7 @@ def create( if type == REWARD_PROMPT: reference = cls._handle_reward_prompt(name, source) elif type == REWARD_FUNCTION: - method, reference = cls._handle_reward_function(name, source, role) + method, reference = cls._handle_reward_function(name, source, role, runtime) else: raise ValueError(f"Unsupported evaluator type: {type}") @@ -381,13 +384,15 @@ def _handle_reward_prompt(cls, name: str, source: Optional[str]) -> str: @classmethod def _handle_reward_function( - cls, name: str, source: Optional[str], role: Optional[str] + cls, name: str, source: Optional[str], role: Optional[str], runtime: Optional[str] = None ) -> tuple[EvaluatorMethod, str]: """Handle creation of reward function evaluator. Args: name: Name of the evaluator source: Lambda ARN or local file path + role: IAM role ARN of the Lambda function created from a local file + runtime: Lambda runtime of the function created from a local file Returns: Tuple of (method, reference) @@ -400,17 +405,19 @@ def _handle_reward_function( return EvaluatorMethod.LAMBDA, source else: # BYOC - create Lambda function from local file - return cls._create_lambda_function(name, source, role) + return cls._create_lambda_function(name, source, role, runtime) @classmethod def _create_lambda_function( - cls, name: str, source_file: str, role: Optional[str] + cls, name: str, source_file: str, role: Optional[str], runtime: Optional[str] = None ) -> tuple[EvaluatorMethod, str]: """Create Lambda function from local Python file. Args: name: Name of the evaluator source_file: Path to local Python file + role: IAM role ARN of the Lambda function + runtime: Lambda runtime. Defaults to ``EVALUATOR_DEFAULT_RUNTIME``. Returns: Tuple of (EvaluatorMethod.BYOC, lambda_arn) @@ -434,7 +441,7 @@ def _create_lambda_function( try: lambda_response = lambda_client.create_function( FunctionName=function_name, - Runtime=EVALUATOR_DEFAULT_RUNTIME, + Runtime=runtime or EVALUATOR_DEFAULT_RUNTIME, Role=role, Handler=handler_name, Code={"ZipFile": zip_buffer.read()}, diff --git a/sagemaker-train/src/sagemaker/train/__init__.py b/sagemaker-train/src/sagemaker/train/__init__.py index e32812f15e..9f222722da 100644 --- a/sagemaker-train/src/sagemaker/train/__init__.py +++ b/sagemaker-train/src/sagemaker/train/__init__.py @@ -53,6 +53,10 @@ def __getattr__(name): from sagemaker.train.cpt_trainer import CPTTrainer return CPTTrainer + elif name in ("FineTuningSweep", "SweepConfig", "SweepTrial"): + from sagemaker.train import fine_tuning_sweep + + return getattr(fine_tuning_sweep, name) elif name == "DataMixingConfig": from sagemaker.train.data_mixing_config import DataMixingConfig @@ -114,22 +118,26 @@ def __getattr__(name): from sagemaker.train.custom_agent_lambda import CustomAgentLambda return CustomAgentLambda - elif name == "plot_training_metrics": - from sagemaker.train.common_utils.metrics_visualizer import plot_training_metrics + elif name == "verify_reward_function": + from sagemaker.train.common_utils.rlvr_reward_verifier import verify_reward_function - return plot_training_metrics - elif name == "get_available_metrics": - from sagemaker.train.common_utils.metrics_visualizer import get_available_metrics + return verify_reward_function + elif name == "RewardVerificationOptions": + from sagemaker.train.common_utils.rlvr_reward_verifier import RewardVerificationOptions - return get_available_metrics - elif name == "get_studio_url": - from sagemaker.train.common_utils.metrics_visualizer import get_studio_url + return RewardVerificationOptions + elif name == "RewardFunctionOutput": + from sagemaker.train.common_utils.rlvr_reward_verifier import RewardFunctionOutput - return get_studio_url - elif name == "get_mlflow_url": - from sagemaker.train.common_utils.trainer_wait import get_mlflow_url + return RewardFunctionOutput + elif name == "RewardSample": + from sagemaker.train.common_utils.rlvr_reward_verifier import RewardSample - return get_mlflow_url + return RewardSample + elif name == "EvalRewardSample": + from sagemaker.train.common_utils.rlvr_reward_verifier import EvalRewardSample + + return EvalRewardSample elif name == "plot_training_metrics": from sagemaker.train.common_utils.metrics_visualizer import plot_training_metrics @@ -138,6 +146,14 @@ def __getattr__(name): from sagemaker.train.common_utils.metrics_visualizer import get_available_metrics return get_available_metrics + elif name == "get_metric_history": + from sagemaker.train.common_utils.metrics_visualizer import get_metric_history + + return get_metric_history + elif name == "MetricPoint": + from sagemaker.train.common_utils.metrics_visualizer import MetricPoint + + return MetricPoint elif name == "get_studio_url": from sagemaker.train.common_utils.metrics_visualizer import get_studio_url diff --git a/sagemaker-train/src/sagemaker/train/base_trainer.py b/sagemaker-train/src/sagemaker/train/base_trainer.py index a60f5ba52b..c0af10b1e7 100644 --- a/sagemaker-train/src/sagemaker/train/base_trainer.py +++ b/sagemaker-train/src/sagemaker/train/base_trainer.py @@ -10,7 +10,7 @@ import time from abc import ABC, abstractmethod from datetime import datetime as _datetime -from typing import Optional, Dict, Any, List, Union +from typing import Optional, Dict, Any, Callable, List, Union from urllib.parse import urlparse import boto3 @@ -120,7 +120,49 @@ class BaseTrainer(ABC): input_data_config: Optional[List[Union[Channel, InputData]]] = None environment: Optional[Dict[str, str]] = None training_image: Optional[str] = None - latest_training_job: Optional[TrainingJob] = None + + @property + def latest_training_job(self) -> Optional[TrainingJob]: + """The ``TrainingJob`` created by the most recent ``train()`` call. + + ``None`` until ``train()`` has submitted a job (a ``dry_run`` or a + ``PipelineSession`` capture submits nothing). The resource is bound to the + trainer's session and region, so ``refresh()``, ``wait()``, ``stop()`` and the + ``training_job_name`` / ``training_job_arn`` attributes work without a second + lookup. ``show_metrics()`` and ``stream_logs()`` read this same job. + + Assign a ``TrainingJob`` to re-attach a trainer to a job it did not create:: + + trainer.latest_training_job = TrainingJob.get("my-job") + + Returns: + Optional[TrainingJob]: The latest created training job, or ``None``. + """ + return getattr(self, "_latest_training_job", None) + + @latest_training_job.setter + def latest_training_job(self, training_job: Optional[TrainingJob]) -> None: + if training_job is not None and not isinstance(training_job, TrainingJob): + raise TypeError( + "latest_training_job must be a sagemaker.core.resources.TrainingJob or None, " + f"got {type(training_job).__name__}. Use TrainingJob.get() to attach " + "an existing job." + ) + self._latest_training_job = training_job + + def _record_submitted_job( + self, + training_job: TrainingJob, + on_submit: Optional[Callable[[TrainingJob], None]] = None, + ) -> None: + """Store the job the service accepted and hand it to ``on_submit`` before any wait. + + A ``PipelineSession`` capture submits no job, so the handle is ``None`` and the + callback is skipped. + """ + self._latest_training_job = training_job + if on_submit is not None and training_job is not None: + on_submit(training_job) @classmethod @_telemetry_emitter( @@ -186,6 +228,7 @@ def __init__( self.disable_output_compression = disable_output_compression self.notification_rule_arn = None self.source_code = None + self._latest_training_job: Optional[TrainingJob] = None # Set up notifications if configured if notifications: @@ -518,17 +561,17 @@ def show_metrics( ValueError: If no training job has been run yet, no logs/metrics are found, or MLflow is not configured for OSS models. """ - # Resolve the job reference. Prefer _latest_training_job (CreateTrainingJob), + # Resolve the job reference. Prefer latest_training_job (CreateTrainingJob), # fall back to _latest_job (generic CreateJob API used by MTRL). - resolved_job = getattr(self, "_latest_training_job", None) + resolved_job = self.latest_training_job if resolved_job is None: latest_job = getattr(self, "_latest_job", None) if latest_job is None: raise ValueError( "No training job found. Call .train() first, then call .show_metrics() " - "to view training metrics. If training has already completed, set the " - "job name directly via trainer._latest_training_job = '' or " - "trainer._latest_job = ''." + "to view training metrics. If training has already completed, attach " + "the job with trainer.latest_training_job = TrainingJob.get('') " + "or trainer._latest_job = ''." ) # For AgentRFTJob (MTRL), delegate to its own get_training_metrics() # which fetches metrics from MLflow via the Job API (not TrainingJob). @@ -787,17 +830,17 @@ def stream_logs( Raises: ValueError: If no training job has been run yet. """ - # Resolve the job reference. Prefer _latest_training_job (CreateTrainingJob), + # Resolve the job reference. Prefer latest_training_job (CreateTrainingJob), # fall back to _latest_job (generic CreateJob API used by MTRL). - resolved_job = getattr(self, "_latest_training_job", None) + resolved_job = self.latest_training_job if resolved_job is None: latest_job = getattr(self, "_latest_job", None) if latest_job is None: raise ValueError( "No training job found. Call .train(wait=False) first, " "then call .stream_logs() to stream logs in real-time. " - "If training has already completed, set the job name directly via " - "trainer._latest_training_job = '' or " + "If training has already completed, attach the job with " + "trainer.latest_training_job = TrainingJob.get('') or " "trainer._latest_job = ''." ) # For AgentRFTJob (MTRL), delegate to its own stream_logs which @@ -964,8 +1007,14 @@ def train( logs: bool = True, wait_timeout: Optional[int] = None, dry_run: bool = False, + on_submit: Optional[Callable[[TrainingJob], None]] = None, ): - """Common training method that calls the specific implementation.""" + """Common training method that calls the specific implementation. + + ``on_submit`` is called once with the created ``TrainingJob`` as soon as the + service accepts the job and before any wait, so a caller can record the job + name or ARN of a long run without ``wait=False`` plus a second lookup. + """ def _get_extra_smtj_hyperparameters(self) -> Dict[str, Any]: """Return extra hyperparameters to inject for SMTJ training. @@ -986,6 +1035,7 @@ def _train_serverful_smtj( wait_timeout=None, poll=5, dry_run=False, + on_submit=None, ): """Execute training on serverful SageMaker Training Job (SMTJ) compute. @@ -1373,14 +1423,12 @@ def _yaml_safe_default(value): logger.info("Dry-run validation passed. No job submitted.") return None - # Execute training - model_trainer.train( - wait=wait, - logs=wait, - ) - - # Store latest training job reference - self._latest_training_job = model_trainer._latest_training_job + # Execute training. The job is submitted first and waited on afterwards, so the + # callback sees it before the wait starts. + model_trainer.train(wait=False, logs=False) + self._record_submitted_job(model_trainer.latest_training_job, on_submit) + if wait and self._latest_training_job is not None: + self._latest_training_job.wait(logs=True) if wait: job_name = None @@ -1496,6 +1544,7 @@ def _train_hyperpod( wait_timeout=None, poll=5, dry_run=False, + on_submit=None, ): """Execute training on a SageMaker HyperPod cluster. @@ -1696,6 +1745,7 @@ def _train_hyperpod( logger.info(f"HyperPod job submitted: {job_name}") training_job = TrainingJob(training_job_name=job_name) + self._record_submitted_job(training_job, on_submit) if wait: try: checkpoint_path = self._resolve_checkpoint_from_manifest( @@ -1709,6 +1759,5 @@ def _train_hyperpod( ) except Exception as e: logger.warning("Could not resolve checkpoint from manifest for %s: %s", job_name, e) - self._latest_training_job = training_job return job_name diff --git a/sagemaker-train/src/sagemaker/train/common_utils/metrics_visualizer.py b/sagemaker-train/src/sagemaker/train/common_utils/metrics_visualizer.py index 5749716039..a4cc1a80e5 100644 --- a/sagemaker-train/src/sagemaker/train/common_utils/metrics_visualizer.py +++ b/sagemaker-train/src/sagemaker/train/common_utils/metrics_visualizer.py @@ -2,13 +2,93 @@ import io import base64 +import importlib.util import logging -from typing import Optional, List +from dataclasses import dataclass +from typing import Dict, List, Mapping, Optional, Sequence, Set, Union from sagemaker.core.resources import TrainingJob logger = logging.getLogger(__name__) +DEFAULT_METRIC_ALIASES: Mapping[str, Sequence[str]] = { + "loss": ("total_loss", "train/loss", "loss"), + "val_loss": ("val_eval_total_loss", "val/loss", "eval_loss"), +} + + +@dataclass(frozen=True) +class MetricPoint: + """One logged value of a training metric. + + Attributes: + step: Training step the value was logged at. + value: Logged metric value. + timestamp: Log time in epoch milliseconds, as MLflow returns it. + """ + + step: int + value: float + timestamp: int + + +def _resolve_training_job(training_job: Union[str, TrainingJob]) -> TrainingJob: + """Return the TrainingJob for a job object or a job name.""" + if isinstance(training_job, str): + return TrainingJob.get(training_job_name=training_job) + return training_job + + +def _connect_mlflow_run(training_job: TrainingJob) -> Optional[str]: + """Point MLflow at the job's tracking server and return the job's run id. + + Returns: + The MLflow run id, or None when the job has no MLflow config or no run id. + """ + import mlflow + + mlflow_config = getattr(training_job, "mlflow_config", None) + mlflow_details = getattr(training_job, "mlflow_details", None) + if not mlflow_config or not mlflow_details or not mlflow_details.mlflow_run_id: + return None + mlflow.set_tracking_uri(mlflow_config.mlflow_resource_arn) + return mlflow_details.mlflow_run_id + + +def _logged_metric_names(run_id: str) -> List[str]: + """Return the names of the metrics logged on an MLflow run.""" + import mlflow + + run = mlflow.get_run(run_id) + return list(run.data.metrics.keys()) + + +def _resolve_metric_names( + requested: Sequence[str], logged: Set[str], aliases: Mapping[str, Sequence[str]] +) -> Dict[str, str]: + """Map each requested name to the first alias candidate the run logged. + + The requested name itself is the last candidate. + + Raises: + ValueError: If a requested name matches no logged metric. + """ + resolved: Dict[str, str] = {} + missing: List[str] = [] + for name in requested: + candidates = (*aliases.get(name, ()), name) + matches = [candidate for candidate in candidates if candidate in logged] + if matches: + resolved[name] = matches[0] + else: + missing.append(f"{name} (candidates: {' / '.join(candidates)})") + if missing: + raise ValueError( + f"Metrics not logged on the MLflow run: {', '.join(missing)}. " + f"Available metrics: {', '.join(sorted(logged))}" + ) + return resolved + def _is_in_studio() -> bool: """Check if running inside SageMaker Studio.""" @@ -203,7 +283,9 @@ def display_job_links_html(rows: list, as_html: bool = False): def plot_training_metrics( - training_job: TrainingJob, metrics: Optional[List[str]] = None, figsize: tuple = (12, 6) + training_job: Union[str, TrainingJob], + metrics: Optional[List[str]] = None, + figsize: tuple = (12, 6), ) -> None: """Plot training metrics from MLflow for a completed training job. @@ -213,23 +295,21 @@ def plot_training_metrics( figsize: Figure size as (width, height) """ import matplotlib.pyplot as plt - import mlflow from mlflow.tracking import MlflowClient from IPython.display import display logging.getLogger("botocore.credentials").setLevel(logging.WARNING) - if isinstance(training_job, str): - training_job = TrainingJob.get(training_job_name=training_job) - - run_id = training_job.mlflow_details.mlflow_run_id + training_job = _resolve_training_job(training_job) + run_id = _connect_mlflow_run(training_job) + if run_id is None: + logger.warning( + "Training job %s has no MLflow run. Nothing to plot.", training_job.training_job_name + ) + return - mlflow.set_tracking_uri(training_job.mlflow_config.mlflow_resource_arn) client = MlflowClient() - - run = mlflow.get_run(run_id) - available_metrics = list(run.data.metrics.keys()) - metrics_to_plot = metrics if metrics else available_metrics + metrics_to_plot = metrics if metrics else _logged_metric_names(run_id) # Fetch metric histories metric_data = {} @@ -291,7 +371,7 @@ def plot_training_metrics( ) -def get_available_metrics(training_job: TrainingJob) -> List[str]: +def get_available_metrics(training_job: Union[str, TrainingJob]) -> List[str]: """Get list of available metrics for a training job. Args: @@ -300,24 +380,64 @@ def get_available_metrics(training_job: TrainingJob) -> List[str]: Returns: List of metric names """ - try: - import mlflow - except ImportError: + if importlib.util.find_spec("mlflow") is None: logger.error("mlflow package not installed") return [] - # Handle string input - if isinstance(training_job, str): - training_job = TrainingJob.get(training_job_name=training_job) - - if not hasattr(training_job, "mlflow_config") or not training_job.mlflow_config: + training_job = _resolve_training_job(training_job) + run_id = _connect_mlflow_run(training_job) + if run_id is None: return [] + return _logged_metric_names(run_id) - mlflow_details = training_job.mlflow_details - if not mlflow_details or not mlflow_details.mlflow_run_id: - return [] - mlflow.set_tracking_uri(training_job.mlflow_config.mlflow_resource_arn) - run = mlflow.get_run(mlflow_details.mlflow_run_id) +def get_metric_history( + training_job: Union[str, TrainingJob], + metrics: Optional[Sequence[str]] = None, + aliases: Optional[Mapping[str, Sequence[str]]] = None, +) -> Dict[str, List[MetricPoint]]: + """Get the MLflow metric history of a training job as data. - return list(run.data.metrics.keys()) + Args: + training_job: SageMaker TrainingJob object or job name string + metrics: Names to return. A name is a key of ``aliases`` or a metric name the run + logged. If None, returns every metric the run logged under its own name. + aliases: Caller name to candidate run metric names, first logged candidate wins. + Defaults to ``DEFAULT_METRIC_ALIASES``, which maps ``loss`` and ``val_loss`` to + the LLMFT and Verl recipe metric names. + + Returns: + Dict keyed by the requested name. Each value is the list of MetricPoint sorted by step. + + Raises: + ValueError: If the job has no MLflow run, or a requested metric matches no metric + the run logged. The message lists the metrics the run has. + + Example: + >>> from sagemaker.train import get_metric_history + >>> history = get_metric_history("my-training-job", metrics=["loss"]) + >>> history["loss"][-1].value + """ + from mlflow.tracking import MlflowClient + + training_job = _resolve_training_job(training_job) + run_id = _connect_mlflow_run(training_job) + if run_id is None: + raise ValueError(f"Training job {training_job.training_job_name} has no MLflow run.") + + logged = set(_logged_metric_names(run_id)) + if metrics is None: + resolved = {name: name for name in sorted(logged)} + else: + alias_map = DEFAULT_METRIC_ALIASES if aliases is None else aliases + resolved = _resolve_metric_names(list(metrics), logged, alias_map) + + client = MlflowClient() + history: Dict[str, List[MetricPoint]] = {} + for name, run_metric in resolved.items(): + points = [ + MetricPoint(step=entry.step, value=entry.value, timestamp=entry.timestamp) + for entry in client.get_metric_history(run_id, run_metric) + ] + history[name] = sorted(points, key=lambda point: point.step) + return history diff --git a/sagemaker-train/src/sagemaker/train/common_utils/rlvr_reward_verifier.py b/sagemaker-train/src/sagemaker/train/common_utils/rlvr_reward_verifier.py index 2877594a14..3b6a9bd5cf 100644 --- a/sagemaker-train/src/sagemaker/train/common_utils/rlvr_reward_verifier.py +++ b/sagemaker-train/src/sagemaker/train/common_utils/rlvr_reward_verifier.py @@ -16,7 +16,8 @@ import json import logging import re -from typing import Any, Dict, List, Literal, Optional, Union +from dataclasses import dataclass, fields +from typing import Any, Callable, Dict, List, Literal, Optional, Type, Union import boto3 from botocore.exceptions import ClientError @@ -30,6 +31,85 @@ r"^arn:aws[a-zA-Z-]*:lambda:[a-z0-9-]+:\d{12}:function:[A-Za-z0-9-_]+$" ) +RewardEnvelope = Literal["oss", "nova", "eval_oss", "eval_nova"] + + +class RewardSample(BaseModel): + """One training sample sent to an RLVR reward function. + + The ``oss`` and ``nova`` envelopes send a list of these in one invocation. + ``reference_answer`` is a plain string or a ``{"text": ...}`` object, as the + dataset provides it. + """ + + id: str + messages: List[Dict[str, Any]] + reference_answer: Optional[Union[str, Dict[str, Any]]] = None + + +class EvalRewardSample(BaseModel): + """One evaluation sample sent to a custom-scorer reward function. + + The ``eval_oss`` and ``eval_nova`` envelopes send one invocation per sample, each + with a list that holds that one sample. + """ + + id: str + model_response: str + query: Optional[str] = None + response: Optional[str] = None + reference_answer: Optional[Union[str, Dict[str, Any]]] = None + + +@dataclass(frozen=True) +class RewardVerificationOptions: + """Options of ``verify_reward_function``. + + Attributes: + validate_format: Validate each result against ``RewardFunctionOutput``. + compute: Compute of the training job. A ``HyperPodCompute`` with a Nova model + makes the Lambda function name contain ``SageMaker``. + is_nova: Response format when ``envelope`` is None. True expects a bare list, + False expects ``{statusCode: 200, body: json.dumps(results)}``. + envelope: Request and response contract. ``oss`` and ``nova`` send every sample + in one invocation and validate the samples as ``RewardSample``. ``eval_oss`` + and ``eval_nova`` send one invocation per sample and validate the samples as + ``EvalRewardSample``. ``nova`` expects a bare list back, the other three + expect the ``statusCode`` and ``body`` envelope. None keeps the ``is_nova`` + behavior and does not validate the samples. + """ + + validate_format: bool = True + compute: Optional[Union[TrainingJobCompute, HyperPodCompute]] = None + is_nova: bool = True + envelope: Optional[RewardEnvelope] = None + + +@dataclass(frozen=True) +class _EnvelopeSpec: + """How one envelope shapes the request, the response, and the sample model.""" + + sample_model: Optional[Type[BaseModel]] + per_sample: bool + nova_model: bool + bare_list_response: bool + + +_ENVELOPE_SPECS: Dict[str, _EnvelopeSpec] = { + "oss": _EnvelopeSpec( + sample_model=RewardSample, per_sample=False, nova_model=False, bare_list_response=False + ), + "nova": _EnvelopeSpec( + sample_model=RewardSample, per_sample=False, nova_model=True, bare_list_response=True + ), + "eval_oss": _EnvelopeSpec( + sample_model=EvalRewardSample, per_sample=True, nova_model=False, bare_list_response=False + ), + "eval_nova": _EnvelopeSpec( + sample_model=EvalRewardSample, per_sample=True, nova_model=True, bare_list_response=False + ), +} + class RewardMetric(BaseModel): """A single metric or reward entry from a reward function output.""" @@ -100,12 +180,92 @@ def _unwrap_response(payload: Any, is_nova: bool) -> Any: return body +def _resolve_options( + options: Optional[RewardVerificationOptions], option_fields: Dict[str, Any] +) -> RewardVerificationOptions: + """Build the options from the ``options`` object or from its fields passed as keywords.""" + if options is not None and option_fields: + known = ", ".join(field.name for field in fields(RewardVerificationOptions)) + raise ValueError(f"Pass 'options' or the keyword arguments {known}, not the two together.") + if options is not None: + return options + return RewardVerificationOptions(**option_fields) + + +def _resolve_envelope(options: RewardVerificationOptions) -> _EnvelopeSpec: + """Return the envelope spec, with ``is_nova`` as the fallback when no envelope is set.""" + if options.envelope is None: + return _EnvelopeSpec( + sample_model=None, + per_sample=False, + nova_model=options.is_nova, + bare_list_response=options.is_nova, + ) + return _ENVELOPE_SPECS[options.envelope] + + +def _validate_samples(sample_data: List[Dict[str, Any]], sample_model: Type[BaseModel]) -> None: + """Validate every sample against the envelope's sample model. + + Raises: + ValueError: If a sample does not match, with one line per field error. + """ + errors: List[str] = [] + for idx, sample in enumerate(sample_data): + try: + sample_model.model_validate(sample) + except ValidationError as e: + errors.extend( + f"Sample {idx}: {'.'.join(str(p) for p in err['loc'])} - {err['msg']}" + for err in e.errors() + ) + if errors: + raise ValueError( + f"Sample data does not match the {sample_model.__name__} contract:\n - " + + "\n - ".join(errors) + ) + + +def _invoke_batches( + invoke: Callable[[List[Dict[str, Any]]], Any], + sample_data: List[Dict[str, Any]], + spec: _EnvelopeSpec, +) -> List[Any]: + """Return one raw response per invocation. + + A batch envelope sends every sample in one invocation. A per-sample envelope sends + one invocation per sample. + """ + if not spec.per_sample: + return [invoke(sample_data)] + return [invoke([sample]) for sample in sample_data] + + +def _unwrap_batches(raw_responses: List[Any], spec: _EnvelopeSpec) -> Any: + """Unwrap each raw response. A per-sample envelope joins the per-sample lists. + + Raises: + ValueError: If a response does not match the envelope. + """ + unwrapped = [_unwrap_response(raw, spec.bare_list_response) for raw in raw_responses] + if not spec.per_sample: + return unwrapped[0] + results: List[Any] = [] + for item in unwrapped: + if not isinstance(item, list): + raise ValueError( + "Reward function must return a list with one result for each sample, " + f"got {type(item).__name__}" + ) + results.extend(item) + return results + + def verify_reward_function( reward_function: str, sample_data: List[Dict[str, Any]], - validate_format: bool = True, - compute: Optional[Union[TrainingJobCompute, HyperPodCompute]] = None, - is_nova: bool = True, + options: Optional[RewardVerificationOptions] = None, + **option_fields: Any, ) -> Dict[str, Any]: """Verify a reward function with sample data before using it in RLVR training or evaluation. @@ -117,15 +277,17 @@ def verify_reward_function( or a path to a local Python file containing the reward function. sample_data: List of conversation samples to test. Each sample should be a dict with 'id', 'messages', and optionally 'reference_answer' keys. - validate_format: If True, validates that the output matches expected format - (default: True). - compute: Compute configuration (TrainingJobCompute or HyperPodCompute). Required - when using Lambda ARN with Nova models. When set to HyperPodCompute, - validates that Lambda ARN contains 'SageMaker' in the function name as - required by SageMaker HyperPod. Optional for local files. - is_nova: If True (default), expects Nova format where the reward function returns - a list of results directly. If False, expects OSS format where the function - returns {statusCode: 200, headers: {...}, body: json.dumps(results)}. + options: ``RewardVerificationOptions`` with ``validate_format``, ``compute``, + ``is_nova``, and ``envelope``. + **option_fields: The fields of ``RewardVerificationOptions`` as keyword arguments, + as an alternative to ``options``. ``validate_format`` (default True) + validates each result against ``RewardFunctionOutput``. ``compute`` + (``TrainingJobCompute`` or ``HyperPodCompute``) with a Nova model makes + the Lambda function name contain 'SageMaker', as SageMaker HyperPod + demands. ``is_nova`` (default True) expects a bare list back, False + expects {statusCode: 200, headers: {...}, body: json.dumps(results)}. + ``envelope`` selects the request and response contract, see + ``RewardVerificationOptions``. Returns: Dict containing: @@ -137,7 +299,8 @@ def verify_reward_function( Raises: ValueError: If any validation errors are encountered, with a detailed error message - listing all issues found. + listing all issues found, or if a sample does not match the envelope's + sample model. Example: >>> # Test training reward function with Lambda ARN @@ -177,12 +340,26 @@ def verify_reward_function( ... sample_data=[...], ... compute=HyperPodCompute(cluster_name="my-cluster") ... ) + + >>> # Test an evaluation custom scorer for an OSS model, one invocation per sample + >>> result = verify_reward_function( + ... reward_function="./my_scorer.py", + ... sample_data=[{"id": "s1", "model_response": "4", "query": "2+2?"}], + ... envelope="eval_oss", + ... ) """ + resolved_options = _resolve_options(options, option_fields) + validate_format = resolved_options.validate_format + compute = resolved_options.compute + spec = _resolve_envelope(resolved_options) + if spec.sample_model is not None: + _validate_samples(sample_data, spec.sample_model) + # Determine if it's a Lambda ARN or local file is_lambda = bool(LAMBDA_ARN_REGEX.match(reward_function)) # Validate Lambda ARN format for HyperPod compute - if is_nova and isinstance(compute, HyperPodCompute): + if spec.nova_model and isinstance(compute, HyperPodCompute): if is_lambda: # Extract function name from ARN: arn:aws[-*]:lambda:region:account:function:function-name function_name_match = re.search( @@ -227,19 +404,19 @@ def verify_reward_function( logger.info(f"Invoking Lambda: {reward_function}") logger.info(f"Number of samples: {len(sample_data)}") - # Invoke Lambda with all samples at once (Lambda expects array) - response = lambda_client.invoke( - FunctionName=reward_function, - InvocationType="RequestResponse", - Payload=json.dumps(sample_data), - ) + def invoke_lambda(batch: List[Dict[str, Any]]) -> Any: + response = lambda_client.invoke( + FunctionName=reward_function, + InvocationType="RequestResponse", + Payload=json.dumps(batch), + ) + return json.loads(response["Payload"].read()) - # Parse response - payload = json.loads(response["Payload"].read()) + raw_responses = _invoke_batches(invoke_lambda, sample_data, spec) # Unwrap based on output format (Nova vs OSS) try: - payload = _unwrap_response(payload, is_nova) + payload = _unwrap_batches(raw_responses, spec) except ValueError as e: # Unwrap failure means the response format is wrong error_msg = str(e) @@ -361,13 +538,13 @@ def verify_reward_function( handler = namespace["lambda_handler"] - # Call handler with all samples at once (matches Lambda behavior) + # Call handler per the envelope (matches Lambda behavior) try: - result = handler(sample_data, {}) + raw_responses = _invoke_batches(lambda batch: handler(batch, {}), sample_data, spec) # Unwrap based on output format (Nova vs OSS) try: - result = _unwrap_response(result, is_nova) + result = _unwrap_batches(raw_responses, spec) except ValueError as e: error_msg = str(e) for idx in range(len(sample_data)): diff --git a/sagemaker-train/src/sagemaker/train/cpt_trainer.py b/sagemaker-train/src/sagemaker/train/cpt_trainer.py index b8283dea5a..7b9349b7b9 100644 --- a/sagemaker-train/src/sagemaker/train/cpt_trainer.py +++ b/sagemaker-train/src/sagemaker/train/cpt_trainer.py @@ -12,12 +12,12 @@ # language governing permissions and limitations under the License. """CPTTrainer — Continued Pre-Training on foundation models using SageMaker HyperPod.""" -from typing import Any, Dict, Optional, Union +from typing import Any, Callable, Dict, Optional, Union import logging from sagemaker.train.base_trainer import BaseTrainer from sagemaker.train.common import TrainingType, CustomizationTechnique -from sagemaker.core.resources import ModelPackageGroup, ModelPackage +from sagemaker.core.resources import ModelPackageGroup, ModelPackage, TrainingJob from sagemaker.core.shapes import VpcConfig from sagemaker.ai_registry.dataset import DataSet from sagemaker.train.configs import StoppingCondition @@ -209,6 +209,7 @@ def train( wait_timeout: Optional[int] = None, poll: int = 5, dry_run: bool = False, + on_submit: Optional[Callable[[TrainingJob], None]] = None, ): """Execute the CPT training job on HyperPod. @@ -227,6 +228,9 @@ def train( If True, runs validation without submitting a job. Returns None on success, raises on validation failure. Defaults to False. + on_submit (Optional[Callable[[TrainingJob], None]]): + Called once with a ``TrainingJob`` handle for the submitted HyperPod job, + before any wait. Not called for a ``dry_run``. Returns: str: The HyperPod job name, or None if dry_run=True. @@ -280,4 +284,5 @@ def train( wait_timeout=wait_timeout, poll=poll, dry_run=dry_run, + on_submit=on_submit, ) diff --git a/sagemaker-train/src/sagemaker/train/custom_agent_lambda.py b/sagemaker-train/src/sagemaker/train/custom_agent_lambda.py index c142d60b9b..b292a308e5 100644 --- a/sagemaker-train/src/sagemaker/train/custom_agent_lambda.py +++ b/sagemaker-train/src/sagemaker/train/custom_agent_lambda.py @@ -23,6 +23,8 @@ import boto3 +from sagemaker.ai_registry.air_constants import EVALUATOR_DEFAULT_RUNTIME + S3_URI_PATTERN = re.compile(r"^s3://[^/]+(/.*)?$") @@ -50,7 +52,7 @@ def create( source: str, function_name: Optional[str] = None, role: Optional[str] = None, - runtime: str = "python3.12", + runtime: str = EVALUATOR_DEFAULT_RUNTIME, handler: str = "lambda_function.handler", timeout: int = 900, memory_size: int = 256, @@ -72,7 +74,7 @@ def create( function_name: Lambda function name. If not provided, a unique name is generated automatically. role: IAM role ARN for the Lambda execution role. - runtime: Lambda runtime (default: ``"python3.12"``). + runtime: Lambda runtime (default: ``EVALUATOR_DEFAULT_RUNTIME``, ``"python3.12"``). handler: Lambda handler (default: ``"lambda_function.handler"``). timeout: Lambda timeout in seconds (default: 900). memory_size: Lambda memory in MB (default: 256). diff --git a/sagemaker-train/src/sagemaker/train/dpo_trainer.py b/sagemaker-train/src/sagemaker/train/dpo_trainer.py index b0c0cf6871..4fcc8861f1 100644 --- a/sagemaker-train/src/sagemaker/train/dpo_trainer.py +++ b/sagemaker-train/src/sagemaker/train/dpo_trainer.py @@ -1,6 +1,6 @@ """DPO (Direct Preference Optimization) trainer for SageMaker fine-tuning.""" -from typing import Any, Dict, Optional, Union +from typing import Any, Callable, Dict, Optional, Union import logging from sagemaker.ai_registry.dataset import DataSet from sagemaker.train.base_trainer import BaseTrainer @@ -258,6 +258,7 @@ def train( wait_timeout: Optional[int] = None, poll: int = 5, dry_run: bool = False, + on_submit: Optional[Callable[[TrainingJob], None]] = None, ): """Execute the DPO training job. @@ -279,6 +280,10 @@ def train( If True, runs all validation (IAM, hyperparameters, infrastructure, data paths) without submitting a job. Returns None on success, raises on validation failure. Defaults to False. + on_submit (Optional[Callable[[TrainingJob], None]]): + Called once with the created ``TrainingJob`` as soon as the service accepts + it, before any wait. Use it to record the job name or ARN of a long run. + Not called for a ``dry_run``. Returns: TrainingJob: The SageMaker training job object, or None if dry_run=True. @@ -292,6 +297,7 @@ def train( wait_timeout=wait_timeout, poll=poll, dry_run=dry_run, + on_submit=on_submit, ) if isinstance(self.compute, TrainingJobCompute): return self._train_serverful_smtj( @@ -301,6 +307,7 @@ def train( wait_timeout=wait_timeout, poll=poll, dry_run=dry_run, + on_submit=on_submit, ) # Default: serverless compute (None) @@ -419,6 +426,7 @@ def train( except Exception as e: logger.error("Error: %s", e) raise e + self._record_submitted_job(training_job, on_submit) if wait: from sagemaker.train.common_utils.trainer_wait import wait as _wait @@ -433,5 +441,4 @@ def train( except TimeoutExceededError as e: logger.error("Error: %s", e) - self.latest_training_job = training_job return training_job diff --git a/sagemaker-train/src/sagemaker/train/evaluate/llm_as_judge_evaluator.py b/sagemaker-train/src/sagemaker/train/evaluate/llm_as_judge_evaluator.py index 9e03432acb..7305e930bc 100644 --- a/sagemaker-train/src/sagemaker/train/evaluate/llm_as_judge_evaluator.py +++ b/sagemaker-train/src/sagemaker/train/evaluate/llm_as_judge_evaluator.py @@ -35,6 +35,13 @@ "evaluation-judge.html#evaluation-judge-supported" ) +# Limits of the Bedrock LLM-as-Judge custom-metric format: +# https://docs.aws.amazon.com/bedrock/latest/userguide/model-evaluation-custom-metrics-prompt-formats.html +CUSTOM_METRICS_MAX_COUNT = 10 +CUSTOM_METRIC_INSTRUCTIONS_MAX_CHARS = 5000 +CUSTOM_METRIC_DEFINITION_MAX_CHARS = 100 +CUSTOM_METRIC_PLACEHOLDERS = ("{{prompt}}", "{{prediction}}") + # S3 key of the service-maintained supported-judge-models list, mirrored per region # under the JumpStart cache bucket. This file is the source of truth for which # Bedrock models are supported as LLM-as-Judge evaluators (kept current by the @@ -134,6 +141,82 @@ def _resolve_bedrock_model_id(base_model_name: str, region: str) -> Optional[str return f"{prefix}.{base_id}" +def _rating_value_type(value: Any) -> Optional[str]: + """Return ``float`` or ``string`` for a rating value that sets exactly one of the two keys.""" + if not isinstance(value, dict): + return None + has_float = value.get("floatValue") is not None + has_string = value.get("stringValue") is not None + if has_float == has_string: + return None + return "float" if has_float else "string" + + +def _validate_rating_scale(prefix: str, rating_scale: Any) -> List[str]: + """Return the violations of one ``ratingScale`` list.""" + if not isinstance(rating_scale, list): + return [f"{prefix}: ratingScale must be a list."] + messages: List[str] = [] + value_types = set() + for index, entry in enumerate(rating_scale): + entry_prefix = f"{prefix}: ratingScale[{index}]" + if not isinstance(entry, dict): + messages.append(f"{entry_prefix} must be an object.") + continue + definition = entry.get("definition") + if not isinstance(definition, str): + messages.append(f"{entry_prefix}.definition must be a string.") + elif len(definition) > CUSTOM_METRIC_DEFINITION_MAX_CHARS: + messages.append( + f"{entry_prefix}.definition is {len(definition)} characters, " + f"the limit is {CUSTOM_METRIC_DEFINITION_MAX_CHARS}." + ) + value_type = _rating_value_type(entry.get("value")) + if value_type is None: + messages.append( + f"{entry_prefix}.value must set exactly one of floatValue or stringValue." + ) + else: + value_types.add(value_type) + if len(value_types) > 1: + messages.append(f"{prefix}: ratingScale mixes floatValue and stringValue. Use one type.") + return messages + + +def _validate_instructions(prefix: str, instructions: Any) -> List[str]: + """Return the violations of one ``instructions`` string.""" + if not isinstance(instructions, str): + return [f"{prefix}: instructions must be a string."] + messages: List[str] = [] + if len(instructions) > CUSTOM_METRIC_INSTRUCTIONS_MAX_CHARS: + messages.append( + f"{prefix}: instructions is {len(instructions)} characters, " + f"the limit is {CUSTOM_METRIC_INSTRUCTIONS_MAX_CHARS}." + ) + if not any(placeholder in instructions for placeholder in CUSTOM_METRIC_PLACEHOLDERS): + messages.append( + f"{prefix}: instructions must contain {' or '.join(CUSTOM_METRIC_PLACEHOLDERS)}." + ) + return messages + + +def _validate_custom_metric(prefix: str, item: Any) -> List[str]: + """Return the violations of one ``{"customMetricDefinition": {...}}`` item.""" + if not isinstance(item, dict) or "customMetricDefinition" not in item: + return [f"{prefix} must be an object with a customMetricDefinition key."] + definition = item["customMetricDefinition"] + if not isinstance(definition, dict): + return [f"{prefix}: customMetricDefinition must be an object."] + messages: List[str] = [] + if not isinstance(definition.get("name"), str): + messages.append(f"{prefix}: name must be a string.") + messages.extend(_validate_instructions(prefix, definition.get("instructions"))) + rating_scale = definition.get("ratingScale") + if rating_scale is not None: + messages.extend(_validate_rating_scale(prefix, rating_scale)) + return messages + + class LLMAsJudgeEvaluator(BaseEvaluator): r"""LLM-as-judge evaluation job. @@ -160,7 +243,8 @@ class LLMAsJudgeEvaluator(BaseEvaluator): removed if present. Examples: ['Correctness', 'Faithfulness'] or ['Builtin.Correctness', 'Builtin.Faithfulness']. Optional. custom_metrics (Optional[str]): JSON string containing array of custom metric definitions. - Optional. For format details, see: + Optional. ``validate_custom_metrics`` checks the format before the job starts. + For format details, see: https://docs.aws.amazon.com/bedrock/latest/userguide/model-evaluation-custom-metrics-prompt-formats.html mlflow_resource_arn (Optional[str]): ARN of the MLflow tracking server for experiment tracking. Optional. If not provided, the system will attempt to resolve it using the default @@ -456,6 +540,41 @@ def _process_builtin_metrics(self, metrics: Optional[List[str]]) -> List[str]: return processed_metrics + @staticmethod + def validate_custom_metrics(custom_metrics: Any) -> List[str]: + """Return every violation of the Bedrock LLM-as-Judge custom-metric format. + + The rules: an array of 1 to 10 metrics, each an object with a ``customMetricDefinition`` + that has a string ``name`` and string ``instructions`` of at most 5000 characters that + contain ``{{prompt}}`` or ``{{prediction}}``. An optional ``ratingScale`` is a list of + entries with a ``definition`` of at most 100 characters and a ``value`` that sets exactly + one of ``floatValue`` or ``stringValue``, with one value type across the scale. + + Args: + custom_metrics: The metric definitions as a list, or as the JSON string of that list. + + Returns: + List[str]: One message per violation. An empty list means the metrics are valid. + """ + if isinstance(custom_metrics, str): + try: + custom_metrics = json.loads(custom_metrics) + except json.JSONDecodeError as e: + return [f"Invalid JSON: {e}"] + if not isinstance(custom_metrics, list): + return ["custom_metrics must be a JSON array of metric definitions."] + if len(custom_metrics) == 0: + return ["custom_metrics is empty. Define at least one metric."] + if len(custom_metrics) > CUSTOM_METRICS_MAX_COUNT: + return [ + f"custom_metrics has {len(custom_metrics)} metrics, " + f"the maximum is {CUSTOM_METRICS_MAX_COUNT}." + ] + messages: List[str] = [] + for index, item in enumerate(custom_metrics): + messages.extend(_validate_custom_metric(f"Metric [{index}]", item)) + return messages + def _validate_custom_metrics_json(self, custom_metrics_json: Optional[str]) -> Optional[str]: """Validate custom metrics JSON string if provided. @@ -466,17 +585,21 @@ def _validate_custom_metrics_json(self, custom_metrics_json: Optional[str]) -> O Optional[str]: Validated JSON string or None Raises: - ValueError: If JSON is invalid + ValueError: If the JSON is invalid or the metrics violate the custom-metric format """ if not custom_metrics_json: return None try: - json.loads(custom_metrics_json) # Validate JSON - return custom_metrics_json + parsed_metrics = json.loads(custom_metrics_json) except json.JSONDecodeError as e: raise ValueError(f"Invalid JSON in custom_metrics: {e}") + violations = self.validate_custom_metrics(parsed_metrics) + if violations: + raise ValueError("Invalid custom_metrics:\n - " + "\n - ".join(violations)) + return custom_metrics_json + def _resolve_llmaj_proxy_model_arn(self, region: str) -> str: """Resolve a non-Nova model ARN for the LLMAJEvaluation judging step. diff --git a/sagemaker-train/src/sagemaker/train/fine_tuning_sweep.py b/sagemaker-train/src/sagemaker/train/fine_tuning_sweep.py new file mode 100644 index 0000000000..d9f56503a5 --- /dev/null +++ b/sagemaker-train/src/sagemaker/train/fine_tuning_sweep.py @@ -0,0 +1,515 @@ +"""Run serverless fine-tuning sweeps without an AMT training definition.""" + +import copy +import itertools +import json +import logging +import math +import random +import threading +import time +from concurrent.futures import Future, ThreadPoolExecutor +from dataclasses import asdict, dataclass, field, replace +from typing import Any, Callable, Dict, List, Optional, Sequence + +import mlflow +from mlflow.tracking import MlflowClient + +from sagemaker.core.helper.session_helper import Session +from sagemaker.core.resources import TrainingJob +from sagemaker.core.s3 import S3Downloader, S3Uploader, parse_s3_url +from sagemaker.train.base_trainer import BaseTrainer + +logger = logging.getLogger(__name__) + +_POLL_SECONDS = 5 +_MLFLOW_LOCK = threading.Lock() +_TERMINAL_STATUSES = frozenset(("Completed", "Stopped", "Failed")) +_TRIAL_STATUSES = frozenset(("pending", "running", "completed", "stopped", "failed")) + + +@dataclass +class SweepTrial: + """Store a trial's parameters, remote identity, and result.""" + + trial_id: int + hyperparameters: Dict[str, Any] + training_job_name: Optional[str] = None + status: str = "pending" + objective_value: Optional[float] = None + rung_values: List[float] = field(default_factory=list) + + +@dataclass(frozen=True) +class SweepConfig: + """Configure a sweep with finite parameter choices. + + Grid and ASHA use the product order of ``search_space``. The random scheduler + samples with replacement. ASHA compares each concurrent batch at the given + steps and retains at least one trial per rung, so a batch of one never stops. + ``rung_maximize`` ranks the rung metric. It defaults to ``maximize``, so set it + when a loss gates an accuracy objective. A resumed ASHA sweep batches the + unfinished trials again in order. With a manifest, every parameter value must + be JSON serializable. + """ + + search_space: Dict[str, Sequence[Any]] + max_trials: int + max_parallel: int = 1 + scheduler: str = "grid" + maximize: bool = False + seed: Optional[int] = None + asha_rungs: Optional[Sequence[int]] = None + asha_reduction_factor: int = 3 + rung_metric: Optional[str] = None + rung_maximize: Optional[bool] = None + manifest_s3_uri: Optional[str] = None + + def __post_init__(self): + """Reject settings that cannot schedule a trial.""" + for name in ("max_trials", "max_parallel", "asha_reduction_factor"): + if type(getattr(self, name)) is not int or getattr(self, name) < 1: + raise ValueError(f"{name} must be a positive integer") + if self.scheduler not in ("grid", "random", "asha"): + raise ValueError("scheduler must be grid, random, or asha") + self._validate_search_space() + if self.scheduler == "asha": + self._validate_asha() + if self.manifest_s3_uri is not None: + bucket, key = parse_s3_url(self.manifest_s3_uri) + if not bucket or not key: + raise ValueError("manifest_s3_uri must identify an S3 object") + + def _validate_search_space(self): + if not self.search_space: + raise ValueError("search_space cannot be empty") + for name, choices in self.search_space.items(): + if not isinstance(name, str) or not name or not choices: + raise ValueError("Each parameter must have a name and at least one choice") + if isinstance(choices, (str, bytes)) or not isinstance(choices, Sequence): + raise ValueError("Parameter choices must be a sequence, not a string") + + def _validate_asha(self): + if self.asha_reduction_factor < 2: + raise ValueError("ASHA reduction factor must be at least 2") + if not self.rung_metric or not self.asha_rungs: + raise ValueError("ASHA must have a rung_metric and asha_rungs") + if any(type(step) is not int or step < 1 for step in self.asha_rungs): + raise ValueError("ASHA rungs must be positive integer steps") + if list(self.asha_rungs) != sorted(set(self.asha_rungs)): + raise ValueError("ASHA rungs must increase without duplicates") + + +@dataclass(frozen=True) +class _ActiveTrial: + trial_id: int + training_job: TrainingJob + + +def _get_rung_value(training_job: TrainingJob, metric: str, step: int) -> Optional[float]: + """Return the metric at the largest recorded step at or past the rung.""" + tracking_uri = getattr( + getattr(training_job, "mlflow_config", None), "mlflow_resource_arn", None + ) + run_id = getattr(getattr(training_job, "mlflow_details", None), "mlflow_run_id", None) + if not tracking_uri or not run_id: + return None + # sagemaker-mlflow signs each request with the ARN in the global tracking URI, + # so one lock holds the global setting and the read together. + with _MLFLOW_LOCK: + mlflow.set_tracking_uri(tracking_uri) + readings = MlflowClient(tracking_uri=tracking_uri).get_metric_history(run_id, metric) + eligible = [reading for reading in readings if reading.step >= step] + if not eligible: + return None + value = float(max(eligible, key=lambda reading: reading.step).value) + if not math.isfinite(value): + raise ValueError(f"MLflow metric {metric} has a non-finite value") + return value + + +def _refresh_status(training_job: TrainingJob) -> Optional[str]: + """Return the job status, or None when the describe call fails.""" + try: + training_job.refresh() + except Exception: + logger.exception("Could not describe training job %s", training_job.training_job_name) + return None + return training_job.training_job_status + + +def _stop_quietly(training_job: TrainingJob, reason: str, expected: bool = False) -> bool: + """Stop the job and log the reason. Return False when the stop call fails.""" + level = logging.INFO if expected else logging.ERROR + logger.log(level, "Stopping %s: %s", training_job.training_job_name, reason) + try: + training_job.stop() + except Exception: + logger.exception("Could not stop training job %s", training_job.training_job_name) + return False + return True + + +def _poll_until_terminal(training_job: TrainingJob) -> Optional[str]: + """Poll the job status without the rich display that ``wait()`` opens. + + Two ``TrainingJob.wait()`` calls in parallel threads share one console, and + rich 13 raises ``LiveError`` for the second live display. + """ + while True: + status = _refresh_status(training_job) + if status is None or status in _TERMINAL_STATUSES: + return status + time.sleep(_POLL_SECONDS) + + +def _make_trials(config: SweepConfig) -> List[SweepTrial]: + names = list(config.search_space) + choices = list(config.search_space.values()) + if config.scheduler == "random": + generator = random.Random(config.seed) + points = ( + tuple(generator.choice(values) for values in choices) for _ in range(config.max_trials) + ) + else: + points = itertools.islice(itertools.product(*choices), config.max_trials) + return [SweepTrial(index, dict(zip(names, point))) for index, point in enumerate(points)] + + +def _validate_trial(trial: SweepTrial) -> None: + if trial.status not in _TRIAL_STATUSES: + raise ValueError(f"Unknown trial status: {trial.status}") + if trial.status == "running" and not trial.training_job_name: + raise ValueError("A running trial must have a training job name") + if trial.status == "completed" and trial.objective_value is None: + raise ValueError("A completed trial must have an objective value") + + +def _load_trials(records: List[Dict[str, Any]], config: SweepConfig) -> List[SweepTrial]: + try: + trials = [SweepTrial(**record) for record in records] + except TypeError as error: + raise ValueError(f"The manifest has a malformed trial record: {error}") from error + if [trial.trial_id for trial in trials] != list(range(len(trials))): + raise ValueError("The manifest must contain ordered, consecutive trial identifiers") + if len(trials) > config.max_trials: + raise ValueError("The manifest has more trials than max_trials") + for trial in trials: + _validate_trial(trial) + return trials + + +class FineTuningSweep: + """Run independent trainers with a bounded number of concurrent jobs. + + Args: + trainer_factory: Build a configured trainer for each parameter dictionary. + objective: Read a finite score from a completed training job. + config: Select parameters, concurrency, and optional S3 checkpoints. + sagemaker_session: Bind S3 calls and resumed jobs to this session. + """ + + def __init__( + self, + trainer_factory: Callable[[Dict[str, Any]], BaseTrainer], + objective: Callable[[TrainingJob], float], + config: SweepConfig, + sagemaker_session: Optional[Session] = None, + ): + self._trainer_factory = trainer_factory + self._objective = objective + self._config = copy.deepcopy(config) + self._session = sagemaker_session + self._trials = _make_trials(self._config) + self._lock = threading.Lock() + self._upload_lock = threading.Lock() + self._execution: Optional[Future] = None + + @property + def trials(self) -> List[SweepTrial]: + """Return snapshots that callers can read without a lock.""" + with self._lock: + return copy.deepcopy(self._trials) + + @property + def best_trial(self) -> Optional[SweepTrial]: + """Return the best completed trial, or None before the first result.""" + completed = [trial for trial in self.trials if trial.status == "completed"] + if not completed: + return None + ordered = sorted( + completed, key=lambda trial: trial.objective_value, reverse=self._config.maximize + ) + return ordered[0] + + def _trial(self, trial_id): + with self._lock: + return copy.deepcopy(self._trials[trial_id]) + + def _checkpoint(self): + """Write the manifest. The upload lock keeps the newest snapshot last.""" + if self._config.manifest_s3_uri is None: + return + with self._upload_lock: + with self._lock: + manifest = { + "config": asdict(self._config), + "trials": [asdict(trial) for trial in self._trials], + } + S3Uploader.upload_string_as_file_body( + body=json.dumps(manifest, allow_nan=False), + desired_s3_uri=self._config.manifest_s3_uri, + sagemaker_session=self._session, + ) + + def _update(self, trial_id, **changes): + """Apply a trial change, then write the manifest. A stopped trial stays stopped.""" + with self._lock: + current = self._trials[trial_id] + if current.status == "stopped": + changes.update(status="stopped", objective_value=None) + self._trials[trial_id] = replace(current, **changes) + updated = copy.deepcopy(self._trials[trial_id]) + self._checkpoint() + return updated + + def _mark_stopped(self, trial_id): + """Stop a pending or running trial. A finished trial keeps its result.""" + with self._lock: + current = self._trials[trial_id] + if current.status not in ("pending", "running"): + return + self._trials[trial_id] = replace(current, status="stopped", objective_value=None) + self._checkpoint() + + def _attach(self, job_name): + if self._session is None: + return TrainingJob.get(training_job_name=job_name) + return TrainingJob.get( + training_job_name=job_name, + session=self._session.boto_session, + region=self._session.boto_region_name, + ) + + def _submit(self, trial): + """Attach to a trial's recorded job, or submit a new one and return its handle.""" + if trial.training_job_name is not None: + return self._attach(trial.training_job_name) + trainer = self._trainer_factory(copy.deepcopy(trial.hyperparameters)) + training_job = trainer.train(wait=False) + if isinstance(training_job, TrainingJob): + return training_job + # A HyperPod trainer returns the job name and keeps the handle on the trainer. + training_job = trainer.latest_training_job + if isinstance(training_job, TrainingJob): + return training_job + raise TypeError("The trainer must return a TrainingJob or set latest_training_job") + + def _start_trial(self, trial): + """Submit or attach the trial. A recorded job that cannot be attached stays running.""" + if self._trial(trial.trial_id).status == "stopped": + return None + try: + training_job = self._submit(trial) + except Exception: + logger.exception("Trial %s did not start", trial.trial_id) + if trial.training_job_name is None: + self._update(trial.trial_id, status="failed") + return None + try: + current = self._update( + trial.trial_id, training_job_name=training_job.training_job_name, status="running" + ) + except Exception: + if trial.training_job_name is None: + # The manifest does not name the new job, so a resume would submit it again. + _stop_quietly(training_job, "the manifest write failed") + raise + if current.status == "stopped": + _stop_quietly(training_job, "the sweep stopped before the job started") + return _ActiveTrial(trial.trial_id, training_job) + + def _finish_trial(self, active): + """Score the job. A trial whose job cannot be described stays running.""" + remote_status = _poll_until_terminal(active.training_job) + if remote_status is None: + return + if remote_status == "Stopped" or self._trial(active.trial_id).status == "stopped": + self._update(active.trial_id, status="stopped") + return + if remote_status != "Completed": + logger.error("Trial %s ended with status %s", active.trial_id, remote_status) + self._update(active.trial_id, status="failed") + return + try: + objective_value = float(self._objective(active.training_job)) + if not math.isfinite(objective_value): + raise ValueError("The objective must return a finite value") + except Exception: + logger.exception("Trial %s failed", active.trial_id) + self._update(active.trial_id, status="failed") + return + self._update(active.trial_id, status="completed", objective_value=objective_value) + + def _poll_rung(self, active, rung_index): + """Record the rung value. + + A job that ends before the rung leaves the comparison. A non-finite reading + stops the trial, because it cannot rank against the others. + """ + current = self._trial(active.trial_id) + if current.status == "stopped" or len(current.rung_values) > rung_index: + return True + status = _refresh_status(active.training_job) + if status is None: + return True + try: + value = _get_rung_value( + active.training_job, self._config.rung_metric, self._config.asha_rungs[rung_index] + ) + except ValueError as error: + if _stop_quietly(active.training_job, str(error)): + self._update(active.trial_id, status="stopped") + return True + if value is not None: + self._update(active.trial_id, rung_values=current.rung_values + [value]) + return True + return status in _TERMINAL_STATUSES + + def _prune_rung(self, active_trials, rung_index): + """Rank the jobs still in progress and stop the bottom share. + + A job that cannot be described keeps its last known status. + """ + snapshots = self.trials + candidates = [] + for active in active_trials: + trial = snapshots[active.trial_id] + if trial.status != "running" or len(trial.rung_values) <= rung_index: + continue + status = _refresh_status(active.training_job) or active.training_job.training_job_status + if status == "InProgress": + candidates.append(active) + rung_maximize = self._config.rung_maximize + if rung_maximize is None: + rung_maximize = self._config.maximize + candidates.sort( + key=lambda active: snapshots[active.trial_id].rung_values[rung_index], + reverse=rung_maximize, + ) + retained = max(1, len(candidates) // self._config.asha_reduction_factor) + for active in candidates[retained:]: + reason = f"ranked below rung {rung_index}" + if _stop_quietly(active.training_job, reason, expected=True): + self._update(active.trial_id, status="stopped") + + def _run_rungs(self, active_trials): + for rung_index in range(len(self._config.asha_rungs)): + pending = list(active_trials) + while pending: + pending = [active for active in pending if not self._poll_rung(active, rung_index)] + if pending: + time.sleep(_POLL_SECONDS) + self._prune_rung(active_trials, rung_index) + + def _run_trial(self, trial): + active = self._start_trial(trial) + if active is not None: + self._finish_trial(active) + + def _run_batch(self, trials, executor): + submissions = [executor.submit(self._start_trial, trial) for trial in trials] + active_trials = [ + active for future in submissions if (active := future.result()) is not None + ] + self._run_rungs(active_trials) + completions = [executor.submit(self._finish_trial, active) for active in active_trials] + for future in completions: + future.result() + + def _execute(self): + try: + self._checkpoint() + self._run_trials() + except Exception: + logger.exception("The sweep stopped before every trial ended") + raise + + def _run_trials(self): + trials = [trial for trial in self.trials if trial.status in ("pending", "running")] + batch_size = self._config.max_parallel + executor = ThreadPoolExecutor(max_workers=batch_size) + try: + if self._config.scheduler != "asha": + for future in [executor.submit(self._run_trial, trial) for trial in trials]: + future.result() + return + # ASHA compares the trials of one batch at each rung. + for start in range(0, len(trials), batch_size): + self._run_batch(trials[start : start + batch_size], executor) + finally: + # A sweep error leaves the unsubmitted trials pending for a later resume. + executor.shutdown(wait=True, cancel_futures=True) + + def run(self, wait: bool = True) -> List[SweepTrial]: + """Start the sweep once, or wait for an earlier asynchronous run. + + Args: + wait: Wait for all trials. False returns the current snapshots. + + Returns: + Trial snapshots. A trial that did not submit or score has status ``failed``. + A trial whose job the sweep could not describe keeps status ``running`` for + ``resume()``. A manifest write or MLflow read failure ends the sweep, leaves + the submitted jobs running, and raises from the next ``run()`` that waits. + """ + with self._lock: + if self._execution is None: + coordinator = ThreadPoolExecutor(max_workers=1, thread_name_prefix="sweep") + self._execution = coordinator.submit(self._execute) + coordinator.shutdown(wait=False) + execution = self._execution + if wait: + execution.result() + return self.trials + + def stop(self) -> None: + """Cancel pending trials and stop each running job. + + A trial whose remote stop fails stays running, so ``resume()`` attaches to it. + """ + failures = [] + for trial in self.trials: + if trial.status not in ("pending", "running"): + continue + if trial.training_job_name is not None: + try: + training_job = self._attach(trial.training_job_name) + if training_job.training_job_status not in _TERMINAL_STATUSES: + training_job.stop() + except Exception as error: + failures.append(f"{trial.training_job_name}: {error}") + continue + self._mark_stopped(trial.trial_id) + if failures: + raise RuntimeError("Could not stop training jobs: " + ", ".join(failures)) + + @classmethod + def resume(cls, manifest_s3_uri, trainer_factory, objective, sagemaker_session=None): + """Rebuild a sweep from a manifest. The call submits no job. + + Args: + manifest_s3_uri: Read the saved state from this S3 object. + trainer_factory: Build trainers for pending trials. + objective: Read a score from completed training jobs. + sagemaker_session: Bind S3 calls and resumed jobs to this session. + + Returns: + A sweep whose ``run()`` skips terminal trials and attaches to running jobs. + """ + body = S3Downloader.read_file(manifest_s3_uri, sagemaker_session=sagemaker_session) + manifest = json.loads(body) + config = SweepConfig(**dict(manifest["config"], manifest_s3_uri=manifest_s3_uri)) + sweep = cls(trainer_factory, objective, config, sagemaker_session) + sweep._trials = _load_trials(manifest["trials"], config) + return sweep diff --git a/sagemaker-train/src/sagemaker/train/model_trainer.py b/sagemaker-train/src/sagemaker/train/model_trainer.py index 4ccc6247da..108ef1eb39 100644 --- a/sagemaker-train/src/sagemaker/train/model_trainer.py +++ b/sagemaker-train/src/sagemaker/train/model_trainer.py @@ -21,7 +21,7 @@ import shlex import shutil from tempfile import TemporaryDirectory -from typing import Optional, List, Union, Dict, Any, ClassVar +from typing import Optional, List, Union, Dict, Any, Callable, ClassVar import yaml from botocore.exceptions import ClientError, NoCredentialsError, NoRegionError @@ -352,6 +352,35 @@ class ModelTrainer(BaseModel): config_mgr: SageMakerConfig = SageMakerConfig() + @property + def latest_training_job(self) -> Optional[resources.TrainingJob]: + """The ``TrainingJob`` created by the most recent ``train()`` call. + + ``None`` until ``train()`` has submitted a job (a ``dry_run`` or a + ``PipelineSession`` capture submits nothing). The resource is bound to the + session and region the job was created with, so ``refresh()``, ``wait()``, + ``stop()`` and the ``training_job_name`` / ``training_job_arn`` attributes + work without a second lookup. + + Assign a ``TrainingJob`` to re-attach a trainer to a job it did not create:: + + trainer.latest_training_job = TrainingJob.get("my-job") + + Returns: + Optional[TrainingJob]: The latest created training job, or ``None``. + """ + return self._latest_training_job + + @latest_training_job.setter + def latest_training_job(self, training_job: Optional[resources.TrainingJob]) -> None: + if training_job is not None and not isinstance(training_job, resources.TrainingJob): + raise TypeError( + "latest_training_job must be a sagemaker.core.resources.TrainingJob or None, " + f"got {type(training_job).__name__}. Use TrainingJob.get() to attach " + "an existing job." + ) + self._latest_training_job = training_job + @property def output_data(self) -> str: """The S3 URI of the training job's ``output.tar.gz`` archive. @@ -923,7 +952,8 @@ def train( wait: Optional[bool] = True, logs: Optional[bool] = True, dry_run: bool = False, - ): + on_submit: Optional[Callable[[resources.TrainingJob], None]] = None, + ) -> Optional[resources.TrainingJob]: """Train a model using AWS SageMaker. Args: @@ -942,6 +972,17 @@ def train( If True, runs all validation (configuration, input resolution, hyperparameters) without submitting a job. Returns None on success, raises on validation failure. Defaults to False. + on_submit (Optional[Callable[[TrainingJob], None]]): + Called once with the created ``TrainingJob`` as soon as the service accepts + it, before any wait. Use it to record the job name or ARN of a long run. Not + called for a ``dry_run``, a ``PipelineSession`` capture, or local mode. + + Returns: + Optional[TrainingJob]: The created ``sagemaker.core.resources.TrainingJob`` for a + SageMaker training job, bound to the trainer's session and region. ``None`` for + a ``dry_run``, a ``PipelineSession`` capture (the step runs the job later), and + local-mode training. The same object is available afterwards as + ``latest_training_job``. """ training_request = self._create_training_job_args(input_data_config=input_data_config) @@ -983,12 +1024,14 @@ def train( ) raise self._latest_training_job = training_job + if on_submit is not None: + on_submit(training_job) if wait: training_job.wait(logs=logs) if logs and not wait: logger.warning( - "Not displaing the training container logs as 'wait' is set to False." + "Not displaying the training container logs as 'wait' is set to False." ) else: @@ -1015,6 +1058,9 @@ def train( local_container.train(wait) if self._temp_code_dir is not None: self._temp_code_dir.cleanup() + if self.training_mode == Mode.SAGEMAKER_TRAINING_JOB: + return self._latest_training_job + return None def _resolve_staging_bucket(self) -> tuple[str, str]: """Resolve the S3 bucket and key prefix for staging training artifacts. diff --git a/sagemaker-train/src/sagemaker/train/rlaif_trainer.py b/sagemaker-train/src/sagemaker/train/rlaif_trainer.py index 6bb676f8c4..caef07a6e5 100644 --- a/sagemaker-train/src/sagemaker/train/rlaif_trainer.py +++ b/sagemaker-train/src/sagemaker/train/rlaif_trainer.py @@ -1,6 +1,6 @@ """RLAIF (Reinforcement Learning from AI Feedback) trainer for SageMaker fine-tuning.""" -from typing import Any, Dict, Optional, Union +from typing import Any, Callable, Dict, Optional, Union import logging from sagemaker.train.base_trainer import BaseTrainer from sagemaker.train.common import TrainingType, CustomizationTechnique, JOB_TYPE @@ -254,6 +254,7 @@ def train( wait_timeout: Optional[int] = None, poll: int = 5, dry_run: bool = False, + on_submit: Optional[Callable[[TrainingJob], None]] = None, ): """Execute the RLAIF training job. @@ -275,6 +276,10 @@ def train( If True, runs all validation (IAM, hyperparameters, infrastructure, data paths) without submitting a job. Returns None on success, raises on validation failure. Defaults to False. + on_submit (Optional[Callable[[TrainingJob], None]]): + Called once with the created ``TrainingJob`` as soon as the service accepts + it, before any wait. Use it to record the job name or ARN of a long run. + Not called for a ``dry_run``. Returns: TrainingJob: The SageMaker training job object, or None if dry_run=True. @@ -393,6 +398,7 @@ def train( except Exception as e: logger.error("Error: %s", e) raise e + self._record_submitted_job(training_job, on_submit) if wait: from sagemaker.train.common_utils.trainer_wait import wait as _wait @@ -407,7 +413,6 @@ def train( except TimeoutExceededError as e: logger.error("Error: %s", e) - self._latest_training_job = training_job return training_job def _process_hyperparameters(self): diff --git a/sagemaker-train/src/sagemaker/train/rlvr_trainer.py b/sagemaker-train/src/sagemaker/train/rlvr_trainer.py index b3bfe5854e..3782bb492f 100644 --- a/sagemaker-train/src/sagemaker/train/rlvr_trainer.py +++ b/sagemaker-train/src/sagemaker/train/rlvr_trainer.py @@ -2,7 +2,7 @@ import json import logging -from typing import Any, Dict, List, Optional, Union +from typing import Any, Callable, Dict, List, Optional, Union from sagemaker.train.base_trainer import BaseTrainer from sagemaker.train.common import TrainingType, CustomizationTechnique, JOB_TYPE @@ -419,6 +419,7 @@ def train( wait_timeout: Optional[int] = None, poll: int = 5, dry_run: bool = False, + on_submit: Optional[Callable[[TrainingJob], None]] = None, ): """Execute the RLVR training job. @@ -440,6 +441,10 @@ def train( If True, runs all validation (IAM, hyperparameters, infrastructure, data paths) without submitting a job. Returns None on success, raises on validation failure. Defaults to False. + on_submit (Optional[Callable[[TrainingJob], None]]): + Called once with the created ``TrainingJob`` as soon as the service accepts + it, before any wait. Use it to record the job name or ARN of a long run. + Not called for a ``dry_run``. Returns: TrainingJob: The SageMaker training job object, or None if dry_run=True. @@ -468,6 +473,7 @@ def train( wait_timeout=wait_timeout, poll=poll, dry_run=dry_run, + on_submit=on_submit, ) if isinstance(self.compute, TrainingJobCompute): return self._train_serverful_smtj( @@ -477,6 +483,7 @@ def train( wait_timeout=wait_timeout, poll=poll, dry_run=dry_run, + on_submit=on_submit, ) # Default: serverless compute (None) @@ -627,6 +634,7 @@ def train( except Exception as e: logger.error("Error: %s", e) raise e + self._record_submitted_job(training_job, on_submit) if wait: from sagemaker.train.common_utils.trainer_wait import wait as _wait @@ -641,7 +649,6 @@ def train( except TimeoutExceededError as e: logger.error("Error: %s", e) - self._latest_training_job = training_job return training_job def _get_extra_smtj_hyperparameters(self) -> Dict[str, Any]: diff --git a/sagemaker-train/src/sagemaker/train/sft_trainer.py b/sagemaker-train/src/sagemaker/train/sft_trainer.py index c92184fa8f..e82dcae317 100644 --- a/sagemaker-train/src/sagemaker/train/sft_trainer.py +++ b/sagemaker-train/src/sagemaker/train/sft_trainer.py @@ -1,6 +1,6 @@ """SFT (Supervised Fine-Tuning) trainer for SageMaker fine-tuning.""" -from typing import Any, Dict, Optional, Union +from typing import Any, Callable, Dict, Optional, Union import logging from sagemaker.train.base_trainer import BaseTrainer from sagemaker.train.common import TrainingType, CustomizationTechnique, JOB_TYPE @@ -290,6 +290,7 @@ def train( wait_timeout: Optional[int] = None, poll: int = 5, dry_run: bool = False, + on_submit: Optional[Callable[[TrainingJob], None]] = None, ): """Execute the SFT training job. @@ -311,6 +312,10 @@ def train( If True, runs all validation (IAM, hyperparameters, infrastructure, data paths) without submitting a job. Returns None on success, raises on validation failure. Defaults to False. + on_submit (Optional[Callable[[TrainingJob], None]]): + Called once with the created ``TrainingJob`` as soon as the service accepts + it, before any wait. Use it to record the job name or ARN of a long run. + Not called for a ``dry_run``. Returns: TrainingJob: The SageMaker training job object, or None if dry_run=True. @@ -349,6 +354,7 @@ def train( wait_timeout=wait_timeout, poll=poll, dry_run=dry_run, + on_submit=on_submit, ) if isinstance(self.compute, TrainingJobCompute): if self.data_mixing_config is not None: @@ -360,6 +366,7 @@ def train( wait_timeout=wait_timeout, poll=poll, dry_run=dry_run, + on_submit=on_submit, ) # Default: serverless compute (None) @@ -491,6 +498,7 @@ def train( except Exception as e: logger.error("Error: %s", e) raise e + self._record_submitted_job(training_job, on_submit) if wait: from sagemaker.train.common_utils.trainer_wait import wait as _wait @@ -505,5 +513,4 @@ def train( except TimeoutExceededError as e: logger.error("Error: %s", e) - self._latest_training_job = training_job return training_job diff --git a/sagemaker-train/tests/unit/ai_registry/data/format_detection/valid_gpt_oss_sft.jsonl b/sagemaker-train/tests/unit/ai_registry/data/format_detection/valid_gpt_oss_sft.jsonl new file mode 100644 index 0000000000..005b532e6d --- /dev/null +++ b/sagemaker-train/tests/unit/ai_registry/data/format_detection/valid_gpt_oss_sft.jsonl @@ -0,0 +1,4 @@ +{"messages": [{"role": "user", "content": "What is machine learning?"}, {"role": "assistant", "content": "Machine learning is a subset of artificial intelligence that enables systems to learn and improve from experience without being explicitly programmed."}]} +{"messages": [{"role": "user", "content": "How does a neural network work?"}, {"role": "assistant", "content": "A neural network consists of interconnected layers of nodes (neurons) that process information by adjusting weights based on training data to recognize patterns."}]} +{"messages": [{"role": "user", "content": "What is the difference between supervised and unsupervised learning?"}, {"role": "assistant", "content": "Supervised learning uses labeled data to train models, while unsupervised learning finds patterns in unlabeled data without predefined categories."}]} +{"messages": [{"role": "user", "content": "Explain gradient descent."}, {"role": "assistant", "content": "Gradient descent is an optimization algorithm that iteratively adjusts model parameters to minimize the loss function by moving in the direction of steepest descent."}]} diff --git a/sagemaker-train/tests/unit/ai_registry/data/format_detection/valid_nova_dpo.jsonl b/sagemaker-train/tests/unit/ai_registry/data/format_detection/valid_nova_dpo.jsonl new file mode 100644 index 0000000000..0a89ed8b9f --- /dev/null +++ b/sagemaker-train/tests/unit/ai_registry/data/format_detection/valid_nova_dpo.jsonl @@ -0,0 +1,3 @@ +{"messages": [{"role": "user", "content": [{"text": "Should I learn Python or JavaScript first?"}]}, {"role": "assistant", "candidates": [{"content": [{"text": "Python is generally recommended for beginners due to its simpler syntax and readability."}], "preferenceLabel": "preferred"}, {"content": [{"text": "JavaScript is better because it runs in browsers."}], "preferenceLabel": "non-preferred"}]}]} +{"messages": [{"role": "user", "content": [{"text": "What's better for weight loss: cardio or strength training?"}]}, {"role": "assistant", "candidates": [{"content": [{"text": "A combination of both is most effective, as cardio burns calories while strength training builds muscle."}], "preferenceLabel": "preferred"}, {"content": [{"text": "Just do cardio, it burns more calories."}], "preferenceLabel": "non-preferred"}]}]} +{"messages": [{"role": "user", "content": [{"text": "How do I prepare for a technical interview?"}]}, {"role": "assistant", "candidates": [{"content": [{"text": "Practice coding problems, review data structures and algorithms, and research the company."}], "preferenceLabel": "preferred"}, {"content": [{"text": "Just wing it and be confident."}], "preferenceLabel": "non-preferred"}]}]} diff --git a/sagemaker-train/tests/unit/ai_registry/data/format_detection/valid_nova_rlvr.jsonl b/sagemaker-train/tests/unit/ai_registry/data/format_detection/valid_nova_rlvr.jsonl new file mode 100644 index 0000000000..43bd2b4001 --- /dev/null +++ b/sagemaker-train/tests/unit/ai_registry/data/format_detection/valid_nova_rlvr.jsonl @@ -0,0 +1,3 @@ +{"messages": [{"role": "system", "content": "You are a helpful assistant."}, {"role": "user", "content": "What is 7*7?"}], "reference_answer": {"solution": "49"}} +{"messages": [{"role": "system", "content": "You are a math tutor."}, {"role": "user", "content": "What is 12+15?"}], "reference_answer": {"solution": "27"}} +{"messages": [{"role": "system", "content": "You are a helpful assistant."}, {"role": "user", "content": "What is the square root of 144?"}], "reference_answer": {"solution": "12"}} diff --git a/sagemaker-train/tests/unit/ai_registry/data/format_detection/valid_open_weights_sft.jsonl b/sagemaker-train/tests/unit/ai_registry/data/format_detection/valid_open_weights_sft.jsonl new file mode 100644 index 0000000000..377311c55e --- /dev/null +++ b/sagemaker-train/tests/unit/ai_registry/data/format_detection/valid_open_weights_sft.jsonl @@ -0,0 +1,5 @@ +{"prompt": "What is cloud computing?", "completion": "Cloud computing is the delivery of computing services including servers, storage, databases, networking, software, and analytics over the internet (the cloud) to offer faster innovation and flexible resources."} +{"prompt": "Explain the difference between IaaS, PaaS, and SaaS.", "completion": "IaaS (Infrastructure as a Service) provides virtualized computing resources, PaaS (Platform as a Service) offers a platform for developing and deploying applications, and SaaS (Software as a Service) delivers software applications over the internet."} +{"prompt": "What are the benefits of microservices architecture?", "completion": "Microservices architecture offers benefits including independent deployment, technology flexibility, fault isolation, scalability of individual services, and easier maintenance of smaller codebases."} +{"prompt": "How does containerization work?", "completion": "Containerization packages an application and its dependencies into a container that can run consistently across different computing environments, isolating the application from the underlying infrastructure."} +{"prompt": "What is continuous integration and continuous deployment?", "completion": "CI/CD is a practice where code changes are automatically built, tested, and deployed to production, enabling faster and more reliable software delivery through automation."} diff --git a/sagemaker-train/tests/unit/ai_registry/data/format_detection/valid_sagemaker_eval.jsonl b/sagemaker-train/tests/unit/ai_registry/data/format_detection/valid_sagemaker_eval.jsonl new file mode 100644 index 0000000000..e070f26b15 --- /dev/null +++ b/sagemaker-train/tests/unit/ai_registry/data/format_detection/valid_sagemaker_eval.jsonl @@ -0,0 +1,3 @@ +{"system": "You are an English major.", "query": "What symbol ends a question?", "response": "?", "category": "Grammar"} +{"system": "You are a science teacher.", "query": "What planet is closest to the sun?", "response": "Mercury is the closest planet to the sun.", "category": "Astronomy"} +{"query": "What is the capital of France?", "response": "Paris is the capital of France."} diff --git a/sagemaker-train/tests/unit/ai_registry/data/format_detection/valid_verl.jsonl b/sagemaker-train/tests/unit/ai_registry/data/format_detection/valid_verl.jsonl new file mode 100644 index 0000000000..bdd8f50c30 --- /dev/null +++ b/sagemaker-train/tests/unit/ai_registry/data/format_detection/valid_verl.jsonl @@ -0,0 +1,3 @@ +{"data_source": "openai/gsm8k", "prompt": [{"role": "system", "content": "You are a math tutor."}, {"role": "user", "content": "How many clips did Natalia sell?"}], "ability": "math", "reward_model": {"ground_truth": "72", "style": "rule"}, "extra_info": {"answer": "72", "index": 0}} +{"data_source": "openai/gsm8k", "prompt": [{"role": "system", "content": "You are a math tutor."}, {"role": "user", "content": "If a train travels 60 mph for 2 hours, how far does it go?"}], "ability": "math", "reward_model": {"ground_truth": "120", "style": "rule"}, "extra_info": {"answer": "120", "index": 1}} +{"data_source": "openai/gsm8k", "prompt": [{"role": "system", "content": "You are a math tutor."}, {"role": "user", "content": "What is 15% of 200?"}], "ability": "math", "reward_model": {"ground_truth": "30", "style": "rule"}, "extra_info": {"answer": "30", "index": 2}} diff --git a/sagemaker-train/tests/unit/ai_registry/test_detect_dataset_format.py b/sagemaker-train/tests/unit/ai_registry/test_detect_dataset_format.py new file mode 100644 index 0000000000..7750985bee --- /dev/null +++ b/sagemaker-train/tests/unit/ai_registry/test_detect_dataset_format.py @@ -0,0 +1,280 @@ +# 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. +"""``detect_dataset_format`` names the format of a local or S3 dataset sample. + +Before this change ``sagemaker.ai_registry`` exported nothing and +``DatasetFormatDetector.validate_dataset`` returned only ``True`` or ``False`` for a local +path: it stopped at the first line, named no format, read no ``s3://`` source, and did not +know the Nova RLVR, PPO, or CPT shapes. Two SageMaker agent skill scripts each carried their +own detector with 9 and 11 format names. +""" + +from __future__ import absolute_import + +import json +from pathlib import Path +from unittest.mock import Mock, patch + +import pytest +from botocore.exceptions import ClientError + +from sagemaker.ai_registry import ( + FormatDetectionOptions, + FormatDetectionResult, + detect_dataset_format, +) +from sagemaker.ai_registry.dataset_format_detector import ( + FORMAT_MATCHERS, + SKILL_FORMAT_ALIASES, + DatasetFormatDetector, +) + +DATA_DIR = Path(__file__).parent / "data" / "format_detection" +READ_BYTES = "sagemaker.ai_registry.dataset_format_detector.S3Downloader.read_bytes" + +EXPECTED_FORMATS = { + "valid_nova_dpo.jsonl": "dpo", + "valid_gpt_oss_sft.jsonl": "openai_chat", + "valid_open_weights_sft.jsonl": "hf_prompt_completion", + "valid_sagemaker_eval.jsonl": "genqa", + "valid_verl.jsonl": "verl", + "valid_nova_rlvr.jsonl": "nova_rlvr", +} + +NOVA_SFT_LINE = json.dumps( + { + "messages": [ + {"role": "user", "content": [{"text": "What is the capital of France?"}]}, + {"role": "assistant", "content": [{"text": "Paris."}]}, + ] + } +) +GPT_OSS_LINE = json.dumps( + { + "messages": [ + {"role": "user", "content": "What is machine learning?"}, + {"role": "assistant", "content": "A field of AI."}, + ] + } +) + + +def _write(tmp_path, name, text): + path = tmp_path / name + path.write_text(text) + return str(path) + + +@pytest.mark.parametrize("fixture_name, expected", sorted(EXPECTED_FORMATS.items())) +def test_local_fixture_resolves_to_its_format(fixture_name, expected): + result = detect_dataset_format(str(DATA_DIR / fixture_name)) + + assert result.format == expected + assert result.confidence == 1.0 + assert result.lines_sampled == len((DATA_DIR / fixture_name).read_text().splitlines()) + assert result.line_errors == [] + assert result.candidates[expected] == 1.0 + + +@pytest.mark.parametrize("fixture_name, expected", sorted(EXPECTED_FORMATS.items())) +def test_strict_mode_accepts_every_fixture(fixture_name, expected): + result = detect_dataset_format( + str(DATA_DIR / fixture_name), FormatDetectionOptions(strict=True) + ) + + assert result.format == expected + + +def test_s3_source_reads_one_byte_range_through_s3_downloader(): + payload = (NOVA_SFT_LINE + "\n") * 3 + NOVA_SFT_LINE[:20] + session = Mock(name="sagemaker_session") + with patch(READ_BYTES, return_value=payload.encode("utf-8")) as read_bytes: + result = detect_dataset_format( + "s3://bucket/train/data.jsonl", + FormatDetectionOptions(sample_size_bytes=len(payload), sagemaker_session=session), + ) + + read_bytes.assert_called_once_with( + "s3://bucket/train/data.jsonl", session, byte_range=(0, len(payload) - 1) + ) + assert result.format == "converse" + assert result.lines_sampled == 3 + assert result.line_errors == [] + + +def test_s3_sample_drops_the_last_partial_line_only_when_the_sample_is_full(): + payload = (GPT_OSS_LINE + "\n") * 2 + GPT_OSS_LINE + with patch(READ_BYTES, return_value=payload.encode("utf-8")): + full_sample = detect_dataset_format( + "s3://bucket/data.jsonl", FormatDetectionOptions(sample_size_bytes=len(payload)) + ) + short_file = detect_dataset_format( + "s3://bucket/data.jsonl", FormatDetectionOptions(sample_size_bytes=len(payload) + 1) + ) + + assert full_sample.lines_sampled == 2 + assert short_file.lines_sampled == 3 + assert short_file.format == "openai_chat" + + +def test_s3_zero_byte_object_returns_no_format(): + error = ClientError({"Error": {"Code": "InvalidRange", "Message": "x"}}, "GetObject") + with patch(READ_BYTES, side_effect=error): + result = detect_dataset_format("s3://bucket/empty.jsonl") + + assert result == FormatDetectionResult( + format=None, confidence=0.0, lines_sampled=0, line_errors=[], candidates={} + ) + + +def test_other_s3_errors_propagate(): + error = ClientError({"Error": {"Code": "NoSuchKey", "Message": "x"}}, "GetObject") + with patch(READ_BYTES, side_effect=error): + with pytest.raises(ClientError): + detect_dataset_format("s3://bucket/missing.jsonl") + + +def test_strict_mode_reports_the_first_mismatching_line(tmp_path): + path = _write(tmp_path, "mixed.jsonl", "\n".join([NOVA_SFT_LINE, NOVA_SFT_LINE, GPT_OSS_LINE])) + + relaxed = detect_dataset_format(path) + strict = detect_dataset_format(path, FormatDetectionOptions(strict=True)) + + assert relaxed.format == "converse" + assert relaxed.confidence == pytest.approx(2 / 3) + # The dpo schema also accepts a plain converse line. The tie resolves to converse by order. + assert relaxed.candidates == { + "converse": pytest.approx(2 / 3), + "dpo": pytest.approx(2 / 3), + "openai_chat": pytest.approx(1 / 3), + } + assert strict.format is None + assert strict.confidence == 0.0 + assert strict.lines_sampled == 3 + assert len(strict.line_errors) == 1 + assert strict.line_errors[0].startswith("3: $.messages[") + assert strict.line_errors[0].endswith("is not of type 'array'") + + +def test_relaxed_mode_rejects_a_candidate_at_half_or_below(tmp_path): + path = _write(tmp_path, "half.jsonl", "\n".join([NOVA_SFT_LINE, GPT_OSS_LINE])) + + result = detect_dataset_format(path) + + assert result.format is None + assert result.candidates == {"converse": 0.5, "dpo": 0.5, "openai_chat": 0.5} + assert len(result.line_errors) == 1 + assert result.line_errors[0].startswith("2: $.messages[") + assert result.line_errors[0].endswith("is not of type 'array'") + + +def test_invalid_json_lines_count_as_errors_and_lower_confidence(tmp_path): + path = _write(tmp_path, "broken.jsonl", "\n".join([GPT_OSS_LINE, "{not json", GPT_OSS_LINE])) + + result = detect_dataset_format(path) + + assert result.format == "openai_chat" + assert result.confidence == pytest.approx(2 / 3) + assert result.lines_sampled == 3 + assert len(result.line_errors) == 1 + assert result.line_errors[0].startswith("2: invalid JSON: ") + + +def test_empty_file_returns_no_format(tmp_path): + path = _write(tmp_path, "empty.jsonl", "") + + result = detect_dataset_format(path) + + assert result == FormatDetectionResult( + format=None, confidence=0.0, lines_sampled=0, line_errors=[], candidates={} + ) + + +def test_unknown_shape_returns_no_format_and_no_candidates(tmp_path): + path = _write(tmp_path, "unknown.jsonl", json.dumps({"foo": 1}) + "\n") + + result = detect_dataset_format(path) + + assert result.format is None + assert result.candidates == {} + assert result.line_errors == [] + + +def test_local_sample_size_limits_the_lines_read(tmp_path): + path = _write(tmp_path, "many.jsonl", (GPT_OSS_LINE + "\n") * 10) + + result = detect_dataset_format( + path, FormatDetectionOptions(sample_size_bytes=len(GPT_OSS_LINE) * 3 + 2) + ) + + assert result.lines_sampled == 2 + assert result.format == "openai_chat" + + +def test_nova_cpt_text_only_lines(tmp_path): + path = _write(tmp_path, "cpt.jsonl", json.dumps({"text": "a paragraph"}) + "\n") + mixed_path = _write( + tmp_path, "cpt_mixed.jsonl", json.dumps({"text": "a", "prompt": "b"}) + "\n" + ) + + assert detect_dataset_format(path).format == "nova_cpt" + assert "nova_cpt" not in detect_dataset_format(mixed_path).candidates + + +def test_nova_ppo_lines(tmp_path): + record = {"turns": ["User: hi", "Bot: hello"], "turns_to_mask": [0], "reward_category": "math"} + path = _write(tmp_path, "ppo.jsonl", json.dumps(record) + "\n") + bad_record = dict(record, reward_category="unknown") + bad_path = _write(tmp_path, "ppo_bad.jsonl", json.dumps(bad_record) + "\n") + + assert detect_dataset_format(path).format == "nova_ppo" + assert "nova_ppo" not in detect_dataset_format(bad_path).candidates + + +def test_parquet_csv_and_json_extensions(tmp_path): + parquet_path = tmp_path / "data.parquet" + parquet_path.write_bytes(b"PAR1" + b"\x00" * 16) + csv_path = _write(tmp_path, "data.csv", "prompt,completion\na,b\n") + json_path = _write(tmp_path, "data.json", json.dumps([{"prompt": "a", "completion": "b"}])) + bad_parquet_path = tmp_path / "bad.parquet" + bad_parquet_path.write_bytes(b"notparquet") + + assert detect_dataset_format(str(parquet_path)).format == "parquet" + assert detect_dataset_format(csv_path).format == "csv" + assert detect_dataset_format(json_path).format == "json" + bad = detect_dataset_format(str(bad_parquet_path)) + assert bad.format is None + assert bad.line_errors == ["1: missing the Parquet magic bytes PAR1"] + + +def test_extension_of_an_s3_key_selects_the_whole_file_matcher(): + with patch(READ_BYTES, return_value=b"PAR1" + b"\x00" * 16): + result = detect_dataset_format("s3://bucket/dir/part-0.parquet") + + assert result.format == "parquet" + assert result.candidates == {"parquet": 1.0} + + +def test_sample_size_below_one_is_rejected(tmp_path): + path = _write(tmp_path, "x.jsonl", GPT_OSS_LINE) + + with pytest.raises(ValueError): + detect_dataset_format(path, FormatDetectionOptions(sample_size_bytes=0)) + + +def test_matchers_cover_the_sdk_schemas_and_the_skill_aliases(): + schema_names = {path.stem for path in DatasetFormatDetector.SCHEMA_DIR.glob("*.json")} + + assert schema_names <= set(FORMAT_MATCHERS) + assert set(SKILL_FORMAT_ALIASES.values()) <= set(FORMAT_MATCHERS) + assert {"nova_rlvr", "nova_ppo", "nova_cpt"} <= set(FORMAT_MATCHERS) - schema_names diff --git a/sagemaker-train/tests/unit/ai_registry/test_evaluator_runtime.py b/sagemaker-train/tests/unit/ai_registry/test_evaluator_runtime.py new file mode 100644 index 0000000000..cc55e10a01 --- /dev/null +++ b/sagemaker-train/tests/unit/ai_registry/test_evaluator_runtime.py @@ -0,0 +1,110 @@ +# 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. +"""Tests for the Lambda runtime of ``Evaluator.create``. + +Before this change ``EVALUATOR_DEFAULT_RUNTIME`` was ``python3.9``, ``Evaluator.create`` had +no ``runtime`` parameter, and ``CustomAgentLambda.create`` carried its own ``python3.12`` +literal, so the two Lambda paths of the SDK deployed different runtimes. +""" + +import pytest +from unittest.mock import patch, MagicMock + +from sagemaker.ai_registry.evaluator import Evaluator, EvaluatorMethod +from sagemaker.ai_registry.air_constants import ( + EVALUATOR_DEFAULT_RUNTIME, + RESPONSE_KEY_HUB_CONTENT_VERSION, + RESPONSE_KEY_HUB_CONTENT_ARN, + RESPONSE_KEY_CREATION_TIME, + RESPONSE_KEY_LAST_MODIFIED_TIME, + REWARD_FUNCTION, +) +from sagemaker.train.custom_agent_lambda import CustomAgentLambda + +DUMMY_ROLE = "arn:aws:iam::123456789012:role/SageMakerRole" + + +def test_default_runtime_is_python312(): + assert EVALUATOR_DEFAULT_RUNTIME == "python3.12" + + +def test_custom_agent_lambda_shares_the_evaluator_runtime(): + defaults = CustomAgentLambda.create.__func__.__defaults__ + assert EVALUATOR_DEFAULT_RUNTIME in defaults + + +class TestEvaluatorRuntime: + @pytest.fixture(autouse=True) + def stub_aws_resolution(self): + """Keep the tests offline, same stubs as ``test_evaluator.py``.""" + session = MagicMock() + with ( + patch("sagemaker.ai_registry.evaluator.Session", return_value=session), + patch("sagemaker.ai_registry.evaluator._get_current_domain_id", return_value=None), + patch( + "sagemaker.train.defaults.TrainDefaults.get_sagemaker_session", + return_value=session, + ), + patch("sagemaker.train.defaults.TrainDefaults.get_role", return_value=DUMMY_ROLE), + patch( + "sagemaker.ai_registry.air_hub_entity.AIRHub.get_hub_name", + return_value="test-hub", + ), + ): + yield + + @pytest.fixture + def lambda_client(self): + with ( + patch("sagemaker.ai_registry.evaluator._get_default_bucket", return_value="bucket"), + patch("sagemaker.ai_registry.evaluator.boto3") as mock_boto3, + patch("sagemaker.ai_registry.evaluator.AIRHub") as mock_air_hub, + patch("zipfile.ZipFile"), + ): + client = MagicMock() + mock_boto3.client.return_value = client + client.create_function.return_value = {"FunctionArn": "lambda-arn"} + mock_air_hub.upload_to_s3.return_value = "s3://bucket/path" + mock_air_hub.import_hub_content.return_value = {"HubContentArn": "test-arn"} + mock_air_hub.describe_hub_content.return_value = { + RESPONSE_KEY_HUB_CONTENT_VERSION: "1.0.0", + RESPONSE_KEY_HUB_CONTENT_ARN: "test-arn", + RESPONSE_KEY_CREATION_TIME: "2024-01-01", + RESPONSE_KEY_LAST_MODIFIED_TIME: "2024-01-01", + } + yield client + + def test_create_from_file_uses_python312_by_default(self, lambda_client): + evaluator = Evaluator.create( + name="test-evaluator", + source="/local/path/function.py", + type=REWARD_FUNCTION, + wait=False, + ) + + assert evaluator.method == EvaluatorMethod.BYOC + call_kwargs = lambda_client.create_function.call_args.kwargs + assert call_kwargs["Runtime"] == "python3.12" + + def test_runtime_reaches_the_lambda_create_call(self, lambda_client): + Evaluator.create( + name="test-evaluator", + source="/local/path/function.py", + type=REWARD_FUNCTION, + wait=False, + runtime="python3.13", + ) + + call_kwargs = lambda_client.create_function.call_args.kwargs + assert call_kwargs["Runtime"] == "python3.13" + assert call_kwargs["Handler"] == "lambda_function.lambda_handler" diff --git a/sagemaker-train/tests/unit/train/common_utils/test_get_metric_history.py b/sagemaker-train/tests/unit/train/common_utils/test_get_metric_history.py new file mode 100644 index 0000000000..89db637cf7 --- /dev/null +++ b/sagemaker-train/tests/unit/train/common_utils/test_get_metric_history.py @@ -0,0 +1,197 @@ +# 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. +"""Tests for ``get_metric_history``. + +Before this change the SDK had no public function that returned MLflow metric history as +data. ``plot_training_metrics`` rendered a PNG and ``get_available_metrics`` returned only +names, so callers set the tracking URI and called ``MlflowClient().get_metric_history`` +themselves, with recipe-specific metric-name aliases copied into each caller. +""" + +from __future__ import absolute_import + +from unittest.mock import Mock, patch + +import pytest +from mlflow.entities import Metric + +from sagemaker.train.common_utils.metrics_visualizer import ( + DEFAULT_METRIC_ALIASES, + MetricPoint, + get_metric_history, +) + +TRACKING_ARN = "arn:aws:sagemaker:us-west-2:123456789012:mlflow-tracking-server/abc" +RUN_ID = "run-123" + + +def _training_job(): + job = Mock() + job.training_job_name = "my-job" + job.mlflow_config.mlflow_resource_arn = TRACKING_ARN + job.mlflow_details.mlflow_run_id = RUN_ID + return job + + +def _metric(name, step, value): + return Metric(key=name, value=value, timestamp=1_700_000_000_000 + step, step=step) + + +def _history_side_effect(histories): + def get_metric_history(run_id, metric_name): + assert run_id == RUN_ID + return histories[metric_name] + + return get_metric_history + + +@pytest.fixture +def mlflow_run(): + with ( + patch("mlflow.set_tracking_uri") as set_uri, + patch("mlflow.get_run") as get_run, + patch("mlflow.tracking.MlflowClient") as client_cls, + ): + yield set_uri, get_run, client_cls.return_value + + +def _logged(get_run, names): + get_run.return_value.data.metrics = {name: 0.0 for name in names} + + +def test_alias_resolves_to_train_loss_when_total_loss_is_absent(mlflow_run): + set_uri, get_run, client = mlflow_run + _logged(get_run, ["train/loss", "val/loss", "lr"]) + client.get_metric_history.side_effect = _history_side_effect( + {"train/loss": [_metric("train/loss", 1, 2.0), _metric("train/loss", 2, 1.5)]} + ) + + history = get_metric_history(_training_job(), metrics=["loss"]) + + set_uri.assert_called_once_with(TRACKING_ARN) + get_run.assert_called_once_with(RUN_ID) + client.get_metric_history.assert_called_once_with(RUN_ID, "train/loss") + assert list(history) == ["loss"] + assert history["loss"] == [ + MetricPoint(step=1, value=2.0, timestamp=1_700_000_000_001), + MetricPoint(step=2, value=1.5, timestamp=1_700_000_000_002), + ] + + +def test_alias_prefers_total_loss_when_present(mlflow_run): + _, get_run, client = mlflow_run + _logged(get_run, ["total_loss", "train/loss"]) + client.get_metric_history.side_effect = _history_side_effect( + {"total_loss": [_metric("total_loss", 1, 3.0)]} + ) + + history = get_metric_history(_training_job(), metrics=["loss"]) + + client.get_metric_history.assert_called_once_with(RUN_ID, "total_loss") + assert history["loss"][0].value == 3.0 + + +def test_explicit_metric_list_returns_only_those_metrics(mlflow_run): + _, get_run, client = mlflow_run + _logged(get_run, ["lr", "grad_norm", "total_loss"]) + client.get_metric_history.side_effect = _history_side_effect( + {"lr": [_metric("lr", 1, 0.1)], "grad_norm": [_metric("grad_norm", 1, 0.5)]} + ) + + history = get_metric_history(_training_job(), metrics=["lr", "grad_norm"]) + + assert set(history) == {"lr", "grad_norm"} + assert history["lr"] == [MetricPoint(step=1, value=0.1, timestamp=1_700_000_000_001)] + assert client.get_metric_history.call_count == 2 + + +def test_no_metric_list_returns_every_logged_metric(mlflow_run): + _, get_run, client = mlflow_run + _logged(get_run, ["lr", "total_loss"]) + client.get_metric_history.side_effect = _history_side_effect( + {"lr": [_metric("lr", 1, 0.1)], "total_loss": [_metric("total_loss", 1, 3.0)]} + ) + + history = get_metric_history(_training_job()) + + assert list(history) == ["lr", "total_loss"] + + +@patch("sagemaker.train.common_utils.metrics_visualizer.TrainingJob") +def test_job_name_string_goes_through_training_job_get(training_job_cls, mlflow_run): + _, get_run, client = mlflow_run + training_job_cls.get.return_value = _training_job() + _logged(get_run, ["lr"]) + client.get_metric_history.side_effect = _history_side_effect({"lr": [_metric("lr", 1, 0.1)]}) + + history = get_metric_history("my-job", metrics=["lr"]) + + training_job_cls.get.assert_called_once_with(training_job_name="my-job") + assert history["lr"][0].step == 1 + + +def test_unknown_metric_raises_with_available_names(mlflow_run): + _, get_run, client = mlflow_run + _logged(get_run, ["lr", "total_loss"]) + + with pytest.raises(ValueError) as excinfo: + get_metric_history(_training_job(), metrics=["val_loss"]) + + message = str(excinfo.value) + assert "val_loss" in message + assert "val_eval_total_loss / val/loss / eval_loss / val_loss" in message + assert "Available metrics: lr, total_loss" in message + client.get_metric_history.assert_not_called() + + +def test_points_are_sorted_by_step(mlflow_run): + _, get_run, client = mlflow_run + _logged(get_run, ["lr"]) + client.get_metric_history.side_effect = _history_side_effect( + {"lr": [_metric("lr", 3, 0.3), _metric("lr", 1, 0.1), _metric("lr", 2, 0.2)]} + ) + + history = get_metric_history(_training_job(), metrics=["lr"]) + + assert [point.step for point in history["lr"]] == [1, 2, 3] + assert [point.value for point in history["lr"]] == [0.1, 0.2, 0.3] + + +def test_custom_aliases_replace_the_defaults(mlflow_run): + _, get_run, client = mlflow_run + _logged(get_run, ["my/reward"]) + client.get_metric_history.side_effect = _history_side_effect( + {"my/reward": [_metric("my/reward", 1, 0.9)]} + ) + + history = get_metric_history( + _training_job(), metrics=["reward"], aliases={"reward": ("my/reward",)} + ) + + assert history["reward"][0].value == 0.9 + assert "loss" in DEFAULT_METRIC_ALIASES + + +def test_job_without_mlflow_run_raises(mlflow_run): + job = _training_job() + job.mlflow_details.mlflow_run_id = None + + with pytest.raises(ValueError, match="has no MLflow run"): + get_metric_history(job, metrics=["lr"]) + + +def test_export_resolves_from_sagemaker_train(): + import sagemaker.train as train_module + + assert train_module.get_metric_history is get_metric_history + assert train_module.MetricPoint is MetricPoint diff --git a/sagemaker-train/tests/unit/train/common_utils/test_reward_function_contract.py b/sagemaker-train/tests/unit/train/common_utils/test_reward_function_contract.py new file mode 100644 index 0000000000..e4193e5bcf --- /dev/null +++ b/sagemaker-train/tests/unit/train/common_utils/test_reward_function_contract.py @@ -0,0 +1,261 @@ +# 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. +"""Tests for the typed reward-function contract of ``verify_reward_function``. + +Before this change the verifier had no typed input contract: it accepted any dict as a +sample, sent every sample in one invocation, and selected the response format only with +``is_nova``. ``verify_reward_function`` and ``RewardFunctionOutput`` were importable only +from ``sagemaker.train.common_utils.rlvr_reward_verifier``. +""" + +from __future__ import absolute_import + +import json +import tempfile +from pathlib import Path + +import pytest +from pydantic import ValidationError + +import sagemaker.train as train_module +from sagemaker.train.common_utils.rlvr_reward_verifier import ( + EvalRewardSample, + RewardFunctionOutput, + RewardSample, + RewardVerificationOptions, + verify_reward_function, +) + +TRAIN_SAMPLES = [ + { + "id": "sample_1", + "messages": [ + {"role": "user", "content": "What is 2+2?"}, + {"role": "assistant", "content": "4"}, + ], + "reference_answer": "4", + }, + { + "id": "sample_2", + "messages": [ + {"role": "user", "content": "Capital of France?"}, + {"role": "assistant", "content": "Paris"}, + ], + "reference_answer": {"text": "Paris"}, + }, +] + +EVAL_SAMPLES = [ + { + "id": "eval_1", + "model_response": "4", + "query": "What is 2+2?", + "response": "4", + "reference_answer": {"text": "4"}, + }, + {"id": "eval_2", "model_response": "Paris", "query": "Capital of France?"}, +] + +RECORDING_HANDLER = """ +import json + +CALLS = [] + +def lambda_handler(event, context): + CALLS.append(event) + results = [ + {"id": sample["id"], "aggregate_reward_score": 1.0} for sample in event + ] + with open(RECORD_PATH, "a") as record: + record.write(json.dumps(event) + "\\n") + return ENVELOPE(results) +""" + +OSS_RETURN = """ +def ENVELOPE(results): + import json + return {"statusCode": 200, "headers": {}, "body": json.dumps(results)} +""" + +NOVA_RETURN = """ +def ENVELOPE(results): + return results +""" + + +@pytest.fixture +def recorder(tmp_path): + """Write a handler that records each event it gets and returns the selected envelope.""" + record_path = tmp_path / "events.jsonl" + + def write(return_code): + code = f'RECORD_PATH = "{record_path}"\n' + return_code + RECORDING_HANDLER + with tempfile.NamedTemporaryFile(mode="w", suffix=".py", delete=False) as handle: + handle.write(code) + return handle.name + + def events(): + if not record_path.exists(): + return [] + lines = record_path.read_text().splitlines() + return [json.loads(line) for line in lines] + + return write, events + + +def _run(reward_file, sample_data, **options): + try: + return verify_reward_function( + reward_function=reward_file, sample_data=sample_data, **options + ) + finally: + Path(reward_file).unlink() + + +def test_oss_envelope_sends_all_samples_in_one_invocation(recorder): + write, events = recorder + + result = _run(write(OSS_RETURN), TRAIN_SAMPLES, envelope="oss") + + assert result["successful_samples"] == 2 + assert events() == [TRAIN_SAMPLES] + + +def test_nova_envelope_accepts_a_bare_list_response(recorder): + write, events = recorder + + result = _run(write(NOVA_RETURN), TRAIN_SAMPLES, envelope="nova") + + assert result["successful_samples"] == 2 + assert events() == [TRAIN_SAMPLES] + + +def test_oss_envelope_rejects_a_bare_list_response(recorder): + write, _ = recorder + + with pytest.raises(ValueError) as excinfo: + _run(write(NOVA_RETURN), TRAIN_SAMPLES, envelope="oss") + + message = str(excinfo.value) + assert "2/2 sample(s) failed validation" in message + assert "OSS reward function must return a dict with 'statusCode' and 'body'" in message + + +def test_envelope_wins_over_is_nova(recorder): + write, _ = recorder + + result = _run(write(OSS_RETURN), TRAIN_SAMPLES, envelope="oss", is_nova=True) + + assert result["successful_samples"] == 2 + + +def test_eval_oss_envelope_sends_one_invocation_per_sample(recorder): + write, events = recorder + + result = _run(write(OSS_RETURN), EVAL_SAMPLES, envelope="eval_oss") + + assert result["successful_samples"] == 2 + assert events() == [[EVAL_SAMPLES[0]], [EVAL_SAMPLES[1]]] + assert [item["output"]["id"] for item in result["results"]] == ["eval_1", "eval_2"] + + +def test_eval_nova_envelope_expects_the_status_code_body_shape(recorder): + write, events = recorder + + result = _run(write(OSS_RETURN), EVAL_SAMPLES, envelope="eval_nova") + + assert result["successful_samples"] == 2 + assert events() == [[EVAL_SAMPLES[0]], [EVAL_SAMPLES[1]]] + + +def test_no_envelope_keeps_is_nova_behavior(recorder): + write, events = recorder + untyped_samples = [{"messages": [{"role": "user", "content": "hi"}]}] + code = 'RECORD_PATH = "unused"\n' + NOVA_RETURN + """ +def lambda_handler(event, context): + return ENVELOPE([{"id": "x", "aggregate_reward_score": 0.5} for _ in event]) +""" + with tempfile.NamedTemporaryFile(mode="w", suffix=".py", delete=False) as handle: + handle.write(code) + reward_file = handle.name + + result = _run(reward_file, untyped_samples) + + assert result["successful_samples"] == 1 + + +def test_typed_samples_reject_a_missing_id(recorder): + write, events = recorder + sample_without_id = {"messages": [{"role": "user", "content": "hi"}]} + + with pytest.raises(ValueError) as excinfo: + _run(write(NOVA_RETURN), [sample_without_id], envelope="nova") + + message = str(excinfo.value) + assert "RewardSample" in message + assert "Sample 0: id - Field required" in message + assert events() == [] + + +def test_eval_envelope_rejects_a_training_sample(recorder): + write, _ = recorder + + with pytest.raises(ValueError) as excinfo: + _run(write(OSS_RETURN), TRAIN_SAMPLES, envelope="eval_oss") + + message = str(excinfo.value) + assert "EvalRewardSample" in message + assert "model_response - Field required" in message + + +def test_options_object_selects_the_envelope(recorder): + write, events = recorder + options = RewardVerificationOptions(envelope="eval_oss") + + result = _run(write(OSS_RETURN), EVAL_SAMPLES, options=options) + + assert result["successful_samples"] == 2 + assert len(events()) == 2 + + +def test_options_object_and_keywords_together_raise(recorder): + write, _ = recorder + + with pytest.raises(ValueError, match="not the two together"): + _run(write(OSS_RETURN), EVAL_SAMPLES, options=RewardVerificationOptions(), envelope="oss") + + +def test_reward_sample_model_requires_id_and_messages(): + with pytest.raises(ValidationError) as excinfo: + RewardSample.model_validate({"messages": []}) + assert excinfo.value.errors()[0]["loc"] == ("id",) + + sample = RewardSample.model_validate(TRAIN_SAMPLES[1]) + assert sample.reference_answer == {"text": "Paris"} + + +def test_eval_reward_sample_model_requires_id_and_model_response(): + with pytest.raises(ValidationError) as excinfo: + EvalRewardSample.model_validate({"query": "q"}) + assert {error["loc"] for error in excinfo.value.errors()} == {("id",), ("model_response",)} + + sample = EvalRewardSample.model_validate(EVAL_SAMPLES[1]) + assert sample.response is None + + +def test_exports_resolve_from_sagemaker_train(): + assert train_module.verify_reward_function is verify_reward_function + assert train_module.RewardFunctionOutput is RewardFunctionOutput + assert train_module.RewardSample is RewardSample + assert train_module.EvalRewardSample is EvalRewardSample + assert train_module.RewardVerificationOptions is RewardVerificationOptions diff --git a/sagemaker-train/tests/unit/train/evaluate/test_llm_as_judge_evaluator.py b/sagemaker-train/tests/unit/train/evaluate/test_llm_as_judge_evaluator.py index 6de1512b18..390b0f9f0c 100644 --- a/sagemaker-train/tests/unit/train/evaluate/test_llm_as_judge_evaluator.py +++ b/sagemaker-train/tests/unit/train/evaluate/test_llm_as_judge_evaluator.py @@ -84,6 +84,16 @@ def _patch_supported_models(model_ids=None, side_effect=None): DEFAULT_BASE_MODEL_ARN = "arn:aws:sagemaker:us-west-2:aws:hub-content/SageMakerPublicHub/Model/llama3-2-1b-instruct/1.0.0" # noqa: E501 DEFAULT_ARTIFACT_ARN = "arn:aws:sagemaker:us-west-2:123456789012:artifact/test-artifact" DEFAULT_EVALUATOR_MODEL = "anthropic.claude-sonnet-4-5-20250929-v1:0" +VALID_CUSTOM_METRIC = { + "customMetricDefinition": { + "name": "TestMetric", + "instructions": "Rate the response. Prompt: {{prompt}} Response: {{prediction}}", + "ratingScale": [ + {"definition": "Good", "value": {"floatValue": 1.0}}, + {"definition": "Poor", "value": {"floatValue": 0.0}}, + ], + } +} @patch("sagemaker.train.common_utils.model_resolution._resolve_base_model") @@ -434,7 +444,7 @@ def test_llm_as_judge_evaluator_validate_custom_metrics_json_valid(mock_artifact sagemaker_session=mock_session, ) - valid_json = json.dumps([{"name": "test"}]) + valid_json = json.dumps([VALID_CUSTOM_METRIC]) result = evaluator._validate_custom_metrics_json(valid_json) assert result == valid_json @@ -502,7 +512,7 @@ def test_llm_as_judge_evaluator_get_llmaj_template_additions( mock_session.get_caller_identity_arn.return_value = DEFAULT_ROLE builtin_metrics = ["Builtin.Correctness", "Helpfulness"] - custom_metrics = json.dumps([{"name": "test"}]) + custom_metrics = json.dumps([VALID_CUSTOM_METRIC]) evaluator = LLMAsJudgeEvaluator( evaluator_model=DEFAULT_EVALUATOR_MODEL, @@ -641,7 +651,7 @@ def test_llm_as_judge_evaluator_custom_metrics_only_no_builtin( mock_session.boto_session = Mock() mock_session.get_caller_identity_arn.return_value = DEFAULT_ROLE - custom_metrics_json = json.dumps([{"customMetricDefinition": {"name": "TestMetric"}}]) + custom_metrics_json = json.dumps([VALID_CUSTOM_METRIC]) evaluator = LLMAsJudgeEvaluator( evaluator_model=DEFAULT_EVALUATOR_MODEL, diff --git a/sagemaker-train/tests/unit/train/evaluate/test_validate_custom_metrics.py b/sagemaker-train/tests/unit/train/evaluate/test_validate_custom_metrics.py new file mode 100644 index 0000000000..70598078ea --- /dev/null +++ b/sagemaker-train/tests/unit/train/evaluate/test_validate_custom_metrics.py @@ -0,0 +1,219 @@ +# 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. +"""Tests for ``LLMAsJudgeEvaluator.validate_custom_metrics``. + +Before this change ``_validate_custom_metrics_json`` only parsed the JSON. A metric with no +``{{prompt}}`` or ``{{prediction}}`` placeholder, 11 metrics, a 101-character rating +definition, or a rating value with neither ``floatValue`` nor ``stringValue`` reached the +evaluation job and failed there. ``LLMAsJudgeEvaluator`` had no ``validate_custom_metrics``. +""" + +from __future__ import absolute_import + +import copy +import json +from unittest.mock import Mock, patch + +import pytest + +from sagemaker.train.evaluate.llm_as_judge_evaluator import LLMAsJudgeEvaluator + +DEFAULT_REGION = "us-west-2" +DEFAULT_ROLE = "arn:aws:iam::123456789012:role/test-role" +DEFAULT_MODEL = "llama3-2-1b-instruct" +DEFAULT_DATASET = "s3://test-bucket/dataset.jsonl" +DEFAULT_S3_OUTPUT = "s3://test-bucket/outputs/" +DEFAULT_MLFLOW_ARN = "arn:aws:sagemaker:us-west-2:123456789012:mlflow-tracking-server/test-server" +DEFAULT_MODEL_PACKAGE_GROUP_ARN = ( + "arn:aws:sagemaker:us-west-2:123456789012:model-package-group/test-group" +) +DEFAULT_BASE_MODEL_ARN = "arn:aws:sagemaker:us-west-2:aws:hub-content/SageMakerPublicHub/Model/llama3-2-1b-instruct/1.0.0" # noqa: E501 +DEFAULT_ARTIFACT_ARN = "arn:aws:sagemaker:us-west-2:123456789012:artifact/test-artifact" +DEFAULT_EVALUATOR_MODEL = "anthropic.claude-sonnet-4-5-20250929-v1:0" + +VALID_METRIC = { + "customMetricDefinition": { + "name": "PositiveSentiment", + "instructions": "Rate the sentiment. Prompt: {{prompt}} Response: {{prediction}}", + "ratingScale": [ + {"definition": "Positive", "value": {"floatValue": 1.0}}, + {"definition": "Negative", "value": {"floatValue": 0.0}}, + ], + } +} + + +def _metric(**definition_changes): + metric = copy.deepcopy(VALID_METRIC) + metric["customMetricDefinition"].update(definition_changes) + return metric + + +def _validate(metrics): + return LLMAsJudgeEvaluator.validate_custom_metrics(metrics) + + +def test_valid_payload_returns_no_violation(): + string_scale = _metric( + name="Tone", + ratingScale=[ + {"definition": "Formal", "value": {"stringValue": "formal"}}, + {"definition": "Casual", "value": {"stringValue": "casual"}}, + ], + ) + no_scale = _metric(name="Free") + del no_scale["customMetricDefinition"]["ratingScale"] + + assert _validate([VALID_METRIC, string_scale, no_scale]) == [] + assert _validate(json.dumps([VALID_METRIC])) == [] + + +def test_invalid_json_string_is_one_violation(): + violations = _validate("not valid json {") + + assert len(violations) == 1 + assert violations[0].startswith("Invalid JSON:") + + +def test_non_array_is_rejected(): + assert _validate(VALID_METRIC) == ["custom_metrics must be a JSON array of metric definitions."] + + +def test_empty_array_is_rejected(): + assert _validate([]) == ["custom_metrics is empty. Define at least one metric."] + + +def test_more_than_ten_metrics_is_rejected(): + violations = _validate([_metric(name=f"m{index}") for index in range(11)]) + + assert violations == ["custom_metrics has 11 metrics, the maximum is 10."] + assert _validate([_metric(name=f"m{index}") for index in range(10)]) == [] + + +def test_item_without_custom_metric_definition_is_rejected(): + violations = _validate([{"name": "test"}]) + + assert violations == ["Metric [0] must be an object with a customMetricDefinition key."] + + +def test_name_must_be_a_string(): + violations = _validate([_metric(name=7)]) + + assert violations == ["Metric [0]: name must be a string."] + + +def test_instructions_longer_than_5000_characters_are_rejected(): + long_instructions = "{{prompt}}" + "x" * (5000 - len("{{prompt}}")) + + assert len(long_instructions) == 5000 + assert _validate([_metric(instructions=long_instructions)]) == [] + violations = _validate([_metric(instructions=long_instructions + "y")]) + assert violations == ["Metric [0]: instructions is 5001 characters, the limit is 5000."] + + +def test_instructions_without_a_placeholder_are_rejected(): + violations = _validate([_metric(instructions="Rate the answer from 1 to 5.")]) + + assert violations == ["Metric [0]: instructions must contain {{prompt}} or {{prediction}}."] + + +def test_instructions_with_only_prediction_placeholder_are_valid(): + assert _validate([_metric(instructions="Rate {{prediction}}.")]) == [] + + +def test_rating_definition_longer_than_100_characters_is_rejected(): + scale = [ + {"definition": "d" * 100, "value": {"floatValue": 1.0}}, + {"definition": "d" * 101, "value": {"floatValue": 0.0}}, + ] + + violations = _validate([_metric(ratingScale=scale)]) + + assert violations == [ + "Metric [0]: ratingScale[1].definition is 101 characters, the limit is 100." + ] + + +def test_rating_value_with_neither_or_two_values_is_rejected(): + scale = [ + {"definition": "Neither", "value": {}}, + {"definition": "Two", "value": {"floatValue": 1.0, "stringValue": "one"}}, + ] + + violations = _validate([_metric(ratingScale=scale)]) + + assert violations == [ + "Metric [0]: ratingScale[0].value must set exactly one of floatValue or stringValue.", + "Metric [0]: ratingScale[1].value must set exactly one of floatValue or stringValue.", + ] + + +def test_rating_scale_with_mixed_value_types_is_rejected(): + scale = [ + {"definition": "Good", "value": {"floatValue": 1.0}}, + {"definition": "Bad", "value": {"stringValue": "bad"}}, + ] + + violations = _validate([_metric(ratingScale=scale)]) + + assert violations == ["Metric [0]: ratingScale mixes floatValue and stringValue. Use one type."] + + +def test_violations_of_every_metric_are_reported(): + violations = _validate([_metric(instructions="no placeholder"), _metric(name=None)]) + + assert violations == [ + "Metric [0]: instructions must contain {{prompt}} or {{prediction}}.", + "Metric [1]: name must be a string.", + ] + + +@patch("sagemaker.train.common_utils.model_resolution._resolve_base_model") +@patch("sagemaker.core.resources.Artifact") +def test_validate_custom_metrics_json_raises_with_the_joined_messages(mock_artifact, mock_resolve): + mock_info = Mock() + mock_info.base_model_name = DEFAULT_MODEL + mock_info.base_model_arn = DEFAULT_BASE_MODEL_ARN + mock_info.source_model_package_arn = None + mock_resolve.return_value = mock_info + + mock_artifact.get_all.return_value = iter([]) + mock_artifact_instance = Mock() + mock_artifact_instance.artifact_arn = DEFAULT_ARTIFACT_ARN + mock_artifact.create.return_value = mock_artifact_instance + + mock_session = Mock() + mock_session.boto_region_name = DEFAULT_REGION + mock_session.boto_session = Mock() + mock_session.get_caller_identity_arn.return_value = DEFAULT_ROLE + + evaluator = LLMAsJudgeEvaluator( + evaluator_model=DEFAULT_EVALUATOR_MODEL, + dataset=DEFAULT_DATASET, + model=DEFAULT_MODEL, + s3_output_path=DEFAULT_S3_OUTPUT, + mlflow_resource_arn=DEFAULT_MLFLOW_ARN, + model_package_group=DEFAULT_MODEL_PACKAGE_GROUP_ARN, + sagemaker_session=mock_session, + ) + + valid_json = json.dumps([VALID_METRIC]) + assert evaluator._validate_custom_metrics_json(valid_json) == valid_json + + with pytest.raises(ValueError) as excinfo: + evaluator._validate_custom_metrics_json(json.dumps([{"name": "test"}])) + + assert str(excinfo.value) == ( + "Invalid custom_metrics:\n" + " - Metric [0] must be an object with a customMetricDefinition key." + ) diff --git a/sagemaker-train/tests/unit/train/test_fine_tuning_sweep.py b/sagemaker-train/tests/unit/train/test_fine_tuning_sweep.py new file mode 100644 index 0000000000..f96bd3f86b --- /dev/null +++ b/sagemaker-train/tests/unit/train/test_fine_tuning_sweep.py @@ -0,0 +1,679 @@ +"""Serverless trainers could not use AMT to schedule and resume parameter sweeps. + +``HyperparameterTuner`` builds a ``HyperParameterTrainingJobDefinition`` without a +``serverless_job_config``, so ``SFTTrainer`` and the other serverless trainers had no sweep. +""" + +import json +import threading +import time +from dataclasses import asdict +from types import SimpleNamespace +from typing import Optional +from unittest.mock import Mock, patch + +import pytest + +from sagemaker.core.resources import TrainingJob +from sagemaker.train import FineTuningSweep, SweepConfig, SweepTrial +from sagemaker.train.fine_tuning_sweep import _get_rung_value + +MANIFEST_URI = "s3://sweep-state/manifest.json" + + +class FakeTrainer: + """Return a job that completes after ``in_progress_polls`` refreshes.""" + + def __init__(self, hyperparameters, in_progress_polls: Optional[int] = None): + self.hyperparameters = hyperparameters + self.in_progress_polls = in_progress_polls or 0 + self.job = Mock(spec=TrainingJob) + self.job.training_job_name = "trial-" + "-".join( + str(value) for value in hyperparameters.values() + ) + self.job.training_job_status = "InProgress" + self.job.score = sum(hyperparameters.values()) + self.job.refresh.side_effect = self._refresh + self.job.stop.side_effect = self._stop + + def _refresh(self): + if self.job.training_job_status != "InProgress": + return + if self.in_progress_polls > 0: + self.in_progress_polls -= 1 + return + self.job.training_job_status = "Completed" + + def _stop(self): + self.job.training_job_status = "Stopped" + + def train(self, wait=False, on_submit=None): + """Return the job without a service call.""" + assert wait is False + if on_submit is not None: + on_submit(self.job) + return self.job + + +def objective(training_job): + return training_job.score + + +def blocking_trainer(entered, release): + """Build a trainer whose refresh blocks until the test sets ``release``.""" + + class BlockingTrainer(FakeTrainer): + def _refresh(self): + entered.set() + assert release.wait(timeout=5) + super()._refresh() + + return BlockingTrainer + + +def test_grid_enumerates_product_and_selects_minimum(): + config = SweepConfig({"rate": [2, 1], "epochs": [4, 3]}, max_trials=10) + sweep = FineTuningSweep(FakeTrainer, objective, config) + + assert sweep.best_trial is None + trials = sweep.run() + + assert [trial.hyperparameters for trial in trials] == [ + {"rate": 2, "epochs": 4}, + {"rate": 2, "epochs": 3}, + {"rate": 1, "epochs": 4}, + {"rate": 1, "epochs": 3}, + ] + assert [trial.status for trial in trials] == ["completed"] * 4 + assert [trial.training_job_name for trial in trials] == [ + "trial-2-4", + "trial-2-3", + "trial-1-4", + "trial-1-3", + ] + assert sweep.best_trial.hyperparameters == {"rate": 1, "epochs": 3} + assert sweep.best_trial.objective_value == 4 + + +def test_maximize_selects_largest_objective(): + config = SweepConfig({"rate": [1, 5, 3]}, max_trials=3, maximize=True) + sweep = FineTuningSweep(FakeTrainer, objective, config) + + sweep.run() + + assert sweep.best_trial.objective_value == 5 + + +def test_grid_respects_max_trials(): + config = SweepConfig({"rate": [1, 2, 3]}, max_trials=2) + sweep = FineTuningSweep(FakeTrainer, objective, config) + + assert len(sweep.run()) == 2 + + +def test_random_seed_reproduces_samples_and_trial_limit(): + config = SweepConfig( + {"rate": [1, 2, 3], "epochs": [4, 5]}, max_trials=9, scheduler="random", seed=12 + ) + first = FineTuningSweep(FakeTrainer, objective, config) + second = FineTuningSweep(FakeTrainer, objective, config) + + assert len(first.run()) == 9 + assert [trial.hyperparameters for trial in first.trials] == [ + trial.hyperparameters for trial in second.run() + ] + assert all(trial.hyperparameters["rate"] in (1, 2, 3) for trial in first.trials) + + +@pytest.mark.parametrize("maximize,kept", [(False, {1, 2}), (True, {5, 6})]) +def test_asha_stops_bottom_share_at_rung(maximize, kept): + config = SweepConfig( + {"rate": [1, 2, 3, 4, 5, 6]}, + max_trials=6, + max_parallel=6, + scheduler="asha", + maximize=maximize, + asha_rungs=[10], + rung_metric="total_loss", + ) + sweep = FineTuningSweep( + lambda hyperparameters: FakeTrainer(hyperparameters, in_progress_polls=2), + objective, + config, + ) + with patch( + "sagemaker.train.fine_tuning_sweep._get_rung_value", + side_effect=lambda job, metric, step: job.score, + ) as rung_value: + trials = sweep.run() + + completed = {trial.hyperparameters["rate"] for trial in trials if trial.status == "completed"} + assert completed == kept + assert sum(trial.status == "stopped" for trial in trials) == 4 + assert all(len(trial.rung_values) == 1 for trial in trials) + assert all(trial.objective_value is None for trial in trials if trial.status == "stopped") + assert {call.args[1:] for call in rung_value.call_args_list} == {("total_loss", 10)} + + +def test_asha_rung_direction_can_differ_from_objective_direction(): + config = SweepConfig( + {"rate": [1, 2, 3, 4, 5, 6]}, + max_trials=6, + max_parallel=6, + scheduler="asha", + maximize=True, + rung_maximize=False, + asha_rungs=[10], + rung_metric="total_loss", + ) + sweep = FineTuningSweep( + lambda hyperparameters: FakeTrainer(hyperparameters, in_progress_polls=2), + objective, + config, + ) + with patch( + "sagemaker.train.fine_tuning_sweep._get_rung_value", + side_effect=lambda job, metric, step: job.score, + ): + trials = sweep.run() + + completed = {trial.hyperparameters["rate"] for trial in trials if trial.status == "completed"} + assert completed == {1, 2} + assert sweep.best_trial.hyperparameters == {"rate": 2} + + +def test_asha_nonfinite_rung_reading_stops_only_that_trial(): + config = SweepConfig( + {"rate": [1, 2, 3]}, + max_trials=3, + max_parallel=3, + scheduler="asha", + asha_rungs=[10], + rung_metric="total_loss", + ) + trainers = {} + + def factory(hyperparameters): + trainers[hyperparameters["rate"]] = FakeTrainer(hyperparameters, in_progress_polls=2) + return trainers[hyperparameters["rate"]] + + def rung_value(job, metric, step): + if job.score == 2: + raise ValueError("MLflow metric total_loss has a non-finite value") + return job.score + + sweep = FineTuningSweep(factory, objective, config) + with patch("sagemaker.train.fine_tuning_sweep._get_rung_value", side_effect=rung_value): + trials = sweep.run() + + assert [trial.status for trial in trials] == ["completed", "stopped", "stopped"] + trainers[2].job.stop.assert_called_once_with() + assert trials[1].rung_values == [] + + +def test_asha_retains_one_trial_at_second_rung(): + config = SweepConfig( + {"rate": [1, 2, 3, 4, 5, 6]}, + max_trials=6, + max_parallel=6, + scheduler="asha", + asha_rungs=[10, 20], + rung_metric="total_loss", + ) + sweep = FineTuningSweep( + lambda hyperparameters: FakeTrainer(hyperparameters, in_progress_polls=4), + objective, + config, + ) + with patch( + "sagemaker.train.fine_tuning_sweep._get_rung_value", + side_effect=lambda job, metric, step: job.score, + ): + trials = sweep.run() + + assert sum(trial.status == "completed" for trial in trials) == 1 + assert sweep.best_trial.hyperparameters == {"rate": 1} + assert sweep.best_trial.rung_values == [1, 1] + + +def test_max_parallel_runs_two_trainers_at_once(): + barrier = threading.Barrier(2) + lock = threading.Lock() + counts = {"active": 0, "maximum": 0} + + class ConcurrentTrainer(FakeTrainer): + def train(self, wait=False, on_submit=None): + with lock: + counts["active"] += 1 + counts["maximum"] = max(counts.values()) + barrier.wait(timeout=5) + with lock: + counts["active"] -= 1 + return super().train(wait=wait, on_submit=on_submit) + + config = SweepConfig({"rate": [1, 2, 3, 4]}, max_trials=4, max_parallel=2) + sweep = FineTuningSweep(ConcurrentTrainer, objective, config) + + assert all(trial.status == "completed" for trial in sweep.run()) + assert counts == {"active": 0, "maximum": 2} + + +def test_rolling_window_starts_next_trial_when_one_slot_frees(): + entered = threading.Event() + release = threading.Event() + started = [] + lock = threading.Lock() + + def factory(hyperparameters): + with lock: + started.append(hyperparameters["rate"]) + if hyperparameters["rate"] == 1: + return blocking_trainer(entered, release)(hyperparameters) + return FakeTrainer(hyperparameters) + + config = SweepConfig({"rate": [1, 2, 3]}, max_trials=3, max_parallel=2) + sweep = FineTuningSweep(factory, objective, config) + try: + sweep.run(wait=False) + assert entered.wait(timeout=5) + deadline = time.monotonic() + 5 + while time.monotonic() < deadline and sweep.trials[2].status != "completed": + time.sleep(0.01) + assert [trial.status for trial in sweep.trials] == ["running", "completed", "completed"] + finally: + release.set() + + assert [trial.status for trial in sweep.run()] == ["completed"] * 3 + assert started == [1, 2, 3] + + +def test_manifest_records_initial_and_each_trial_state(): + config = SweepConfig({"rate": [1, 2]}, max_trials=2, manifest_s3_uri=MANIFEST_URI) + session = Mock() + sweep = FineTuningSweep(FakeTrainer, objective, config, session) + + with patch("sagemaker.train.fine_tuning_sweep.S3Uploader.upload_string_as_file_body") as upload: + sweep.run() + + manifests = [json.loads(call.kwargs["body"]) for call in upload.call_args_list] + assert [[trial["status"] for trial in manifest["trials"]] for manifest in manifests] == [ + ["pending", "pending"], + ["running", "pending"], + ["completed", "pending"], + ["completed", "running"], + ["completed", "completed"], + ] + assert manifests[-1]["config"] == asdict(config) + assert manifests[-1]["trials"][1] == asdict(sweep.trials[1]) + assert all(call.kwargs["sagemaker_session"] is session for call in upload.call_args_list) + assert all(call.kwargs["desired_s3_uri"] == MANIFEST_URI for call in upload.call_args_list) + + +def test_resume_skips_completed_trial_and_attaches_running_trial(): + config = SweepConfig({"rate": [1, 2, 3]}, max_trials=3) + trials = [ + SweepTrial(0, {"rate": 1}, "done", "completed", 1), + SweepTrial(1, {"rate": 2}, "in-progress", "running"), + SweepTrial(2, {"rate": 3}), + ] + manifest = json.dumps({"config": asdict(config), "trials": [asdict(trial) for trial in trials]}) + resumed_job = FakeTrainer({"rate": 2}).job + resumed_job.training_job_name = "in-progress" + session = Mock() + factory = Mock(side_effect=FakeTrainer) + with ( + patch( + "sagemaker.train.fine_tuning_sweep.S3Downloader.read_file", return_value=manifest + ) as read_file, + patch("sagemaker.train.fine_tuning_sweep.S3Uploader.upload_string_as_file_body") as upload, + patch.object(TrainingJob, "get", return_value=resumed_job) as get_job, + ): + sweep = FineTuningSweep.resume(MANIFEST_URI, factory, objective, session) + result = sweep.run() + + read_file.assert_called_once_with(MANIFEST_URI, sagemaker_session=session) + factory.assert_called_once_with({"rate": 3}) + get_job.assert_called_once_with( + training_job_name="in-progress", + session=session.boto_session, + region=session.boto_region_name, + ) + assert [trial.status for trial in result] == ["completed"] * 3 + assert [trial.training_job_name for trial in result] == ["done", "in-progress", "trial-3"] + assert upload.call_args.kwargs["desired_s3_uri"] == MANIFEST_URI + + +def test_rung_helper_uses_largest_step_at_or_past_rung(): + training_job = FakeTrainer({"rate": 1}).job + training_job.mlflow_config = SimpleNamespace(mlflow_resource_arn="tracking-resource") + training_job.mlflow_details = SimpleNamespace(mlflow_run_id="run-id") + readings = [ + SimpleNamespace(step=3, value=6), + SimpleNamespace(step=20, value=2), + SimpleNamespace(step=10, value=4), + ] + with ( + patch("sagemaker.train.fine_tuning_sweep.mlflow.set_tracking_uri") as set_uri, + patch("sagemaker.train.fine_tuning_sweep.MlflowClient") as client, + ): + client.return_value.get_metric_history.return_value = readings + assert _get_rung_value(training_job, "loss", 10) == 2 + assert _get_rung_value(training_job, "loss", 30) is None + + set_uri.assert_called_with("tracking-resource") + client.assert_called_with(tracking_uri="tracking-resource") + client.return_value.get_metric_history.assert_called_with("run-id", "loss") + + +def test_rung_helper_waits_for_mlflow_identity(): + training_job = FakeTrainer({"rate": 1}).job + + assert _get_rung_value(training_job, "loss", 10) is None + + +def test_model_trainer_style_return_falls_back_to_latest_training_job(): + trainer = FakeTrainer({"rate": 1}) + trainer.latest_training_job = trainer.job + trainer.train = Mock(return_value=None) + sweep = FineTuningSweep( + lambda hyperparameters: trainer, objective, SweepConfig({"rate": [1]}, 1) + ) + + assert sweep.run()[0].status == "completed" + trainer.train.assert_called_once_with(wait=False) + + +def test_trainer_without_training_job_fails_trial(): + trainer = FakeTrainer({"rate": 1}) + trainer.latest_training_job = None + trainer.train = Mock(return_value=None) + sweep = FineTuningSweep( + lambda hyperparameters: trainer, objective, SweepConfig({"rate": [1]}, 1) + ) + + assert sweep.run()[0].status == "failed" + + +def test_hyperpod_style_job_name_return_uses_latest_training_job(): + trainer = FakeTrainer({"rate": 1}) + trainer.latest_training_job = trainer.job + trainer.train = Mock(return_value="trial-1") + sweep = FineTuningSweep( + lambda hyperparameters: trainer, objective, SweepConfig({"rate": [1]}, 1) + ) + + trials = sweep.run() + + assert trials[0].status == "completed" + assert trials[0].training_job_name == "trial-1" + + +def test_manifest_failure_after_attach_keeps_the_resumed_job(): + config = SweepConfig({"rate": [1]}, max_trials=1, manifest_s3_uri=MANIFEST_URI) + trials = [SweepTrial(0, {"rate": 1}, "in-progress", "running")] + manifest = json.dumps({"config": asdict(config), "trials": [asdict(trials[0])]}) + resumed_job = FakeTrainer({"rate": 1}).job + upload = Mock(side_effect=[None, RuntimeError("S3 denied")]) + with ( + patch("sagemaker.train.fine_tuning_sweep.S3Downloader.read_file", return_value=manifest), + patch("sagemaker.train.fine_tuning_sweep.S3Uploader.upload_string_as_file_body", upload), + patch.object(TrainingJob, "get", return_value=resumed_job), + ): + sweep = FineTuningSweep.resume(MANIFEST_URI, FakeTrainer, objective) + with pytest.raises(RuntimeError, match="S3 denied"): + sweep.run() + + resumed_job.stop.assert_not_called() + + +def test_manifest_failure_after_submission_stops_the_job(): + config = SweepConfig({"rate": [1]}, max_trials=1, manifest_s3_uri=MANIFEST_URI) + trainer = FakeTrainer({"rate": 1}) + sweep = FineTuningSweep(lambda hyperparameters: trainer, objective, config) + upload = Mock(side_effect=[None, RuntimeError("S3 denied")]) + with patch("sagemaker.train.fine_tuning_sweep.S3Uploader.upload_string_as_file_body", upload): + with pytest.raises(RuntimeError, match="S3 denied"): + sweep.run() + + trainer.job.stop.assert_called_once_with() + assert sweep.trials[0].status == "running" + + +def test_manifest_failure_still_raises_when_the_stop_fails(): + config = SweepConfig({"rate": [1]}, max_trials=1, manifest_s3_uri=MANIFEST_URI) + trainer = FakeTrainer({"rate": 1}) + trainer.job.stop.side_effect = RuntimeError("AccessDenied") + sweep = FineTuningSweep(lambda hyperparameters: trainer, objective, config) + upload = Mock(side_effect=[None, RuntimeError("S3 denied")]) + with patch("sagemaker.train.fine_tuning_sweep.S3Uploader.upload_string_as_file_body", upload): + with pytest.raises(RuntimeError, match="S3 denied"): + sweep.run() + + trainer.job.stop.assert_called_once_with() + + +def test_resume_keeps_running_trial_when_attach_fails(): + config = SweepConfig({"rate": [1]}, max_trials=1) + trials = [SweepTrial(0, {"rate": 1}, "in-progress", "running")] + manifest = json.dumps({"config": asdict(config), "trials": [asdict(trials[0])]}) + with ( + patch("sagemaker.train.fine_tuning_sweep.S3Downloader.read_file", return_value=manifest), + patch("sagemaker.train.fine_tuning_sweep.S3Uploader.upload_string_as_file_body"), + patch.object(TrainingJob, "get", side_effect=RuntimeError("AccessDenied")), + ): + sweep = FineTuningSweep.resume(MANIFEST_URI, FakeTrainer, objective) + result = sweep.run() + + assert result[0].status == "running" + assert result[0].training_job_name == "in-progress" + + +def test_lost_contact_leaves_trial_running_for_resume(): + trainer = FakeTrainer({"rate": 1}) + trainer.job.refresh.side_effect = RuntimeError("ThrottlingException") + sweep = FineTuningSweep( + lambda hyperparameters: trainer, objective, SweepConfig({"rate": [1]}, 1) + ) + + trials = sweep.run() + + assert trials[0].status == "running" + assert trials[0].training_job_name == "trial-1" + trainer.job.stop.assert_not_called() + assert sweep.best_trial is None + + +def test_individual_failure_does_not_block_other_trials(): + def factory(hyperparameters): + if hyperparameters["rate"] == 1: + raise RuntimeError("Rejected training job") + return FakeTrainer(hyperparameters) + + config = SweepConfig({"rate": [1, 2]}, max_trials=2) + sweep = FineTuningSweep(factory, objective, config) + + assert [trial.status for trial in sweep.run()] == ["failed", "completed"] + assert sweep.best_trial.hyperparameters == {"rate": 2} + + +def test_failed_training_job_fails_trial(): + trainer = FakeTrainer({"rate": 1}) + trainer.job.refresh.side_effect = lambda: setattr(trainer.job, "training_job_status", "Failed") + sweep = FineTuningSweep( + lambda hyperparameters: trainer, objective, SweepConfig({"rate": [1]}, 1) + ) + + assert sweep.run()[0].status == "failed" + + +@pytest.mark.parametrize("value", [float("nan"), float("inf"), float("-inf")]) +def test_nonfinite_objective_fails_trial(value): + config = SweepConfig({"rate": [1]}, max_trials=1) + sweep = FineTuningSweep(FakeTrainer, lambda job: value, config) + + assert sweep.run()[0].status == "failed" + assert sweep.best_trial is None + + +def test_trial_snapshots_do_not_mutate_state(): + config = SweepConfig({"rate": [1]}, max_trials=1) + sweep = FineTuningSweep(FakeTrainer, objective, config) + snapshot = sweep.trials[0] + snapshot.hyperparameters["rate"] = 99 + snapshot.status = "completed" + + assert sweep.trials[0].hyperparameters == {"rate": 1} + assert sweep.trials[0].status == "pending" + + +def test_async_run_returns_before_training_and_can_wait_again(): + entered = threading.Event() + release = threading.Event() + config = SweepConfig({"rate": [1]}, max_trials=1) + sweep = FineTuningSweep(blocking_trainer(entered, release), objective, config) + try: + sweep.run(wait=False) + assert entered.wait(timeout=5) + assert sweep.trials[0].status == "running" + finally: + release.set() + + assert sweep.run()[0].status == "completed" + + +def test_stop_ends_running_trial(): + entered = threading.Event() + release = threading.Event() + trainers = [] + + def factory(hyperparameters): + trainers.append(blocking_trainer(entered, release)(hyperparameters)) + return trainers[-1] + + sweep = FineTuningSweep(factory, objective, SweepConfig({"rate": [1]}, max_trials=1)) + try: + sweep.run(wait=False) + assert entered.wait(timeout=5) + with patch.object(TrainingJob, "get", side_effect=lambda **kwargs: trainers[0].job): + sweep.stop() + finally: + release.set() + + assert sweep.run()[0].status == "stopped" + trainers[0].job.stop.assert_called_once_with() + + +def test_stop_keeps_trial_running_when_remote_stop_fails(): + entered = threading.Event() + release = threading.Event() + trainers = [] + + def factory(hyperparameters): + trainers.append(blocking_trainer(entered, release)(hyperparameters)) + return trainers[-1] + + sweep = FineTuningSweep(factory, objective, SweepConfig({"rate": [1]}, max_trials=1)) + try: + sweep.run(wait=False) + assert entered.wait(timeout=5) + trainers[0].job.stop.side_effect = RuntimeError("AccessDenied") + with patch.object(TrainingJob, "get", side_effect=lambda **kwargs: trainers[0].job): + with pytest.raises(RuntimeError, match="trial-1: AccessDenied"): + sweep.stop() + assert sweep.trials[0].status == "running" + finally: + release.set() + + assert sweep.run()[0].status == "completed" + + +def test_stop_wins_over_a_later_refresh_failure(): + entered = threading.Event() + release = threading.Event() + trainers = [] + + def factory(hyperparameters): + trainers.append(blocking_trainer(entered, release)(hyperparameters)) + return trainers[-1] + + sweep = FineTuningSweep(factory, objective, SweepConfig({"rate": [1]}, max_trials=1)) + try: + sweep.run(wait=False) + assert entered.wait(timeout=5) + with patch.object(TrainingJob, "get", side_effect=lambda **kwargs: trainers[0].job): + sweep.stop() + trainers[0].job.refresh.side_effect = RuntimeError("ThrottlingException") + finally: + release.set() + + assert sweep.run()[0].status == "stopped" + + +def test_stop_cancels_pending_trials_without_starting_trainers(): + factory = Mock(side_effect=FakeTrainer) + sweep = FineTuningSweep(factory, objective, SweepConfig({"rate": [1, 2]}, max_trials=2)) + + sweep.stop() + + assert [trial.status for trial in sweep.run()] == ["stopped", "stopped"] + factory.assert_not_called() + + +def test_checkpoint_failure_does_not_start_training(): + config = SweepConfig({"rate": [1]}, max_trials=1, manifest_s3_uri=MANIFEST_URI) + factory = Mock(side_effect=FakeTrainer) + sweep = FineTuningSweep(factory, objective, config) + with patch( + "sagemaker.train.fine_tuning_sweep.S3Uploader.upload_string_as_file_body", + side_effect=RuntimeError("S3 denied"), + ): + with pytest.raises(RuntimeError, match="S3 denied"): + sweep.run() + + factory.assert_not_called() + + +@pytest.mark.parametrize( + "settings", + [ + {"max_trials": 0}, + {"max_parallel": 0}, + {"scheduler": "unknown"}, + {"search_space": {}}, + {"search_space": {"rate": []}}, + {"search_space": {"rate": "123"}}, + {"scheduler": "asha"}, + {"scheduler": "asha", "rung_metric": "loss", "asha_rungs": [2, 1]}, + {"scheduler": "asha", "rung_metric": "loss", "asha_rungs": [1, 1]}, + {"scheduler": "asha", "rung_metric": "loss", "asha_rungs": [1], "asha_reduction_factor": 1}, + {"manifest_s3_uri": "s3://bucket-only"}, + ], +) +def test_invalid_config_rejects_before_submission(settings): + arguments = {"search_space": {"rate": [1]}, "max_trials": 1} + arguments.update(settings) + + with pytest.raises(ValueError): + SweepConfig(**arguments) + + +@pytest.mark.parametrize( + "records", + [ + [{"trial_id": 1, "hyperparameters": {"rate": 1}}], + [{"trial_id": 0, "hyperparameters": {"rate": 1}, "status": "running"}], + [{"trial_id": 0, "hyperparameters": {"rate": 1}, "status": "completed"}], + [{"trial_id": 0, "hyperparameters": {"rate": 1}, "status": "unknown"}], + [{"trial_id": 0}], + [{"trial_id": 0, "hyperparameters": {"rate": 1}, "job": "extra-key"}], + ], +) +def test_resume_rejects_inconsistent_manifest(records): + config = SweepConfig({"rate": [1]}, max_trials=1) + manifest = json.dumps({"config": asdict(config), "trials": records}) + + with ( + patch("sagemaker.train.fine_tuning_sweep.S3Downloader.read_file", return_value=manifest), + pytest.raises(ValueError), + ): + FineTuningSweep.resume(MANIFEST_URI, FakeTrainer, objective) diff --git a/sagemaker-train/tests/unit/train/test_latest_training_job.py b/sagemaker-train/tests/unit/train/test_latest_training_job.py new file mode 100644 index 0000000000..1dd26574a3 --- /dev/null +++ b/sagemaker-train/tests/unit/train/test_latest_training_job.py @@ -0,0 +1,350 @@ +# 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. +"""``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. +""" + +from __future__ import absolute_import + +import json +from types import SimpleNamespace +from unittest.mock import MagicMock, Mock, patch + +import pytest + +from sagemaker.core.helper.session_helper import Session +from sagemaker.core.resources import TrainingJob +from sagemaker.core.shapes import OutputDataConfig +from sagemaker.core.training.configs import Compute +from sagemaker.train.base_trainer import BaseTrainer +from sagemaker.train.dpo_trainer import DPOTrainer +from sagemaker.train.model_trainer import ModelTrainer +from sagemaker.train.sft_trainer import SFTTrainer +from sagemaker.train.common import TrainingType + +ROLE = "arn:aws:iam::123456789012:role/SageMakerRole" +IMAGE = "123456789012.dkr.ecr.us-west-2.amazonaws.com/my-algo:latest" +BUCKET = "my-bucket" + + +@pytest.fixture(scope="module", autouse=True) +def modules_session(): + with ( + patch("sagemaker.train.Session", spec=Session) as session_mock, + patch("sagemaker.train.defaults.resolve_and_validate_role", return_value=ROLE), + ): + session_instance = session_mock.return_value + session_instance.default_bucket.return_value = BUCKET + session_instance.get_caller_identity_arn.return_value = ROLE + session_instance.default_bucket_prefix = None + session_instance.boto_session = MagicMock(spec="boto3.session.Session") + session_instance.boto_region_name = "us-west-2" + yield session_instance + + +@pytest.fixture +def model_trainer(modules_session): + return ModelTrainer( + training_image=IMAGE, + role=ROLE, + sagemaker_session=modules_session, + compute=Compute(instance_type="ml.m5.xlarge", instance_count=1), + output_data_config=OutputDataConfig(s3_output_path=f"s3://{BUCKET}/output"), + ) + + +@patch("sagemaker.train.model_trainer.TrainingJob") +def test_model_trainer_train_returns_the_created_job(mock_training_job, model_trainer): + created = TrainingJob(training_job_name="created-job") + mock_training_job.create.return_value = created + + returned = model_trainer.train(wait=False) + + assert returned is created + assert model_trainer.latest_training_job is created + assert model_trainer.latest_training_job.training_job_name == "created-job" + + +@patch("sagemaker.train.model_trainer.TrainingJob") +def test_model_trainer_train_with_wait_returns_the_job_after_waiting( + mock_training_job, model_trainer +): + created = Mock(spec=TrainingJob) + mock_training_job.create.return_value = created + + returned = model_trainer.train(wait=True, logs=False) + + created.wait.assert_called_once_with(logs=False) + assert returned is created + + +def test_model_trainer_latest_training_job_is_none_before_train(model_trainer): + assert model_trainer.latest_training_job is None + + +@patch("sagemaker.train.model_trainer.TrainingJob") +def test_model_trainer_dry_run_returns_none_and_creates_nothing(mock_training_job, model_trainer): + assert model_trainer.train(wait=False, dry_run=True) is None + mock_training_job.create.assert_not_called() + assert model_trainer.latest_training_job is None + + +def test_model_trainer_latest_training_job_setter_attaches_an_existing_job(model_trainer): + existing = TrainingJob(training_job_name="existing-job") + + model_trainer.latest_training_job = existing + + assert model_trainer.latest_training_job is existing + assert model_trainer.output_data.endswith("/existing-job/output/output.tar.gz") + + +def test_model_trainer_latest_training_job_setter_rejects_a_job_name(model_trainer): + with pytest.raises(TypeError, match="TrainingJob.get"): + model_trainer.latest_training_job = "existing-job" + + +class _StubTrainer(BaseTrainer): + """Minimal concrete trainer: BaseTrainer is abstract on ``train``.""" + + def train( + self, + input_data_config=None, + wait=True, + logs=True, + wait_timeout=None, + dry_run=False, + on_submit=None, + ): + return None + + +@patch("sagemaker.train.base_trainer.TrainDefaults.get_sagemaker_session") +def test_base_trainer_latest_training_job_is_none_before_train(mock_get_session): + mock_get_session.return_value = Mock(sagemaker_config={}) + + trainer = _StubTrainer() + + assert trainer.latest_training_job is None + + +@patch("sagemaker.train.base_trainer.TrainDefaults.get_sagemaker_session") +def test_base_trainer_latest_training_job_setter_round_trips(mock_get_session): + mock_get_session.return_value = Mock(sagemaker_config={}) + trainer = _StubTrainer() + existing = TrainingJob(training_job_name="existing-job") + + trainer.latest_training_job = existing + + assert trainer.latest_training_job is existing + with pytest.raises(TypeError, match="TrainingJob.get"): + trainer.latest_training_job = "existing-job" + + +@patch("sagemaker.train.base_trainer.TrainDefaults.get_sagemaker_session") +def test_base_trainer_show_metrics_reads_the_public_handle(mock_get_session): + mock_get_session.return_value = Mock(sagemaker_config={}) + trainer = _StubTrainer() + + with pytest.raises(ValueError, match="trainer.latest_training_job = TrainingJob.get"): + trainer.show_metrics() + with pytest.raises(ValueError, match="trainer.latest_training_job = TrainingJob.get"): + trainer.stream_logs() + + +def _serverless_trainer_patches(module): + return [ + patch("sagemaker.train.common_utils.finetune_utils._get_beta_session"), + patch(f"{module}.validate_data_path_exists"), + patch(f"{module}._resolve_model_and_name", return_value=("test-model", "test-model")), + patch(f"{module}.TrainDefaults.get_role", return_value=ROLE), + patch( + f"{module}.TrainDefaults.get_sagemaker_session", + return_value=Mock(sagemaker_config={}), + ), + patch(f"{module}._get_unique_name", return_value="test-job-name"), + patch(f"{module}._validate_and_resolve_model_package_group", return_value="test-group"), + patch(f"{module}._create_input_data_config", return_value=[Mock()]), + patch(f"{module}._convert_input_data_to_channels", return_value=[Mock()]), + patch(f"{module}._create_output_config", return_value=Mock()), + patch(f"{module}._create_serverless_config", return_value=Mock()), + patch(f"{module}._create_mlflow_config", return_value=Mock()), + patch(f"{module}._create_model_package_config", return_value=Mock()), + ] + + +@pytest.mark.parametrize( + "trainer_cls, module", + [ + (SFTTrainer, "sagemaker.train.sft_trainer"), + (DPOTrainer, "sagemaker.train.dpo_trainer"), + ], +) +def test_serverless_trainers_expose_the_created_job(trainer_cls, module): + options = Mock() + options.to_dict.return_value = {"learning_rate": "0.001"} + created = Mock(spec=TrainingJob) + created.training_job_name = "test-job-name" + patches = _serverless_trainer_patches(module) + [ + patch( + f"{module}._get_fine_tuning_options_and_model_arn", return_value=(options, "arn", False) + ), + patch("sagemaker.core.resources.TrainingJob.create", return_value=created), + ] + for active in patches: + active.start() + try: + trainer = trainer_cls( + model="test-model", + training_type=TrainingType.LORA, + model_package_group="test-group", + training_dataset="s3://bucket/train", + ) + assert trainer.latest_training_job is None + + returned = trainer.train(wait=False) + + assert returned is created + assert trainer.latest_training_job is created + assert trainer.latest_training_job.training_job_name == "test-job-name" + finally: + for active in patches: + active.stop() + + +@patch("sagemaker.train.model_trainer.TrainingJob") +def test_model_trainer_on_submit_sees_the_job_before_the_wait(mock_training_job, model_trainer): + created = Mock(spec=TrainingJob) + created.training_job_name = "created-job" + mock_training_job.create.return_value = created + order = [] + created.wait.side_effect = lambda **_: order.append("wait") + + model_trainer.train(wait=True, logs=False, on_submit=lambda job: order.append(job)) + + assert order == [created, "wait"] + + +@patch("sagemaker.train.model_trainer.TrainingJob") +def test_model_trainer_dry_run_skips_on_submit(mock_training_job, model_trainer): + callback = Mock() + + model_trainer.train(wait=False, dry_run=True, on_submit=callback) + + callback.assert_not_called() + mock_training_job.create.assert_not_called() + + +@pytest.mark.parametrize( + "trainer_cls, module", + [ + (SFTTrainer, "sagemaker.train.sft_trainer"), + (DPOTrainer, "sagemaker.train.dpo_trainer"), + ], +) +def test_serverless_trainers_call_on_submit_before_the_wait(trainer_cls, module): + options = Mock() + options.to_dict.return_value = {"learning_rate": "0.001"} + created = Mock(spec=TrainingJob) + created.training_job_name = "test-job-name" + order = [] + patches = _serverless_trainer_patches(module) + [ + patch( + f"{module}._get_fine_tuning_options_and_model_arn", return_value=(options, "arn", False) + ), + patch("sagemaker.core.resources.TrainingJob.create", return_value=created), + patch( + "sagemaker.train.common_utils.trainer_wait.wait", + side_effect=lambda job, **_: order.append("wait"), + ), + ] + for active in patches: + active.start() + try: + trainer = trainer_cls( + model="test-model", + training_type=TrainingType.LORA, + model_package_group="test-group", + training_dataset="s3://bucket/train", + ) + + returned = trainer.train(wait=True, on_submit=lambda job: order.append(job)) + + assert order == [created, "wait"] + assert returned is created + assert trainer.latest_training_job is created + finally: + for active in patches: + active.stop() + + +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 + + +@patch("sagemaker.train.base_trainer.subprocess") +@patch("sagemaker.train.base_trainer.TrainDefaults.verify_hyperpod_caller_permissions") +@patch("sagemaker.train.base_trainer.validate_hyperpod_compute") +@patch("sagemaker.train.base_trainer.TrainDefaults.get_sagemaker_session") +@patch("sagemaker.train.base_trainer.get_hyperpod_recipe_path", return_value="recipes/test") +@patch("sagemaker.train.base_trainer.flatten_resolved_recipe", return_value={}) +def test_hyperpod_train_calls_on_submit_with_the_submitted_job( + mock_flatten, + mock_get_recipe_path, + mock_get_session, + mock_validate, + mock_verify, + mock_subprocess, +): + mock_get_session.return_value = MagicMock() + mock_subprocess.run.return_value = SimpleNamespace(stdout="NAME: my-job-123\n", stderr="") + trainer = _hyperpod_sft_trainer() + submitted = [] + + with ( + patch("sagemaker.train.common_utils.finetune_utils.get_training_image", return_value=None), + patch.object(trainer, "get_resolved_recipe", return_value={"training_config": {}}), + ): + job_name = trainer._train_hyperpod(wait=False, on_submit=submitted.append) + + assert job_name == "my-job-123" + 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])