From bca9deebc57cc297f8335484a09cf5f65ad628df Mon Sep 17 00:00:00 2001 From: Lisa Ni Date: Tue, 7 Jul 2026 00:30:35 +0000 Subject: [PATCH 1/2] feat(serve): add opt-in model source tag-based resource reuse Add reuse_resources to ModelBuilder.build/deploy and BedrockModelBuilder.deploy. On a hit, discover an existing resource by the model-source tag and return it instead of creating a duplicate (warn, do not raise). Honored per call. - New sagemaker/serve/model_reuse.py: tag helpers + service-client discovery - Consolidate Nova manifest/checkpoint reading into sagemaker/core/training/utils.py - SageMaker: build() skips Model creation on reuse (sets built_model to the existing Model); deploy() reuses the endpoint after validating env vars/image/ instance type (PrimaryContainer with Containers[0] fallback for Nova) - Reuse gates are skipped for inference-component builds/deploys so IC create/update (via _deploy_for_ic) is never silently intercepted - Bedrock: reuse custom model + active deployment; response includes modelArn - Reuse discovery uses the cached session/bedrock clients - Support raw S3 URI model input via model_metadata BASE_MODEL_NAME - Unit tests + notebook examples --- .../src/sagemaker/core/training/utils.py | 136 ++++++ .../tests/unit/test_training_utils.py | 146 ++++++ .../sagemaker/serve/bedrock_model_builder.py | 246 +++++++--- .../src/sagemaker/serve/model_builder.py | 422 ++++++++++++++++-- .../src/sagemaker/serve/model_reuse.py | 303 +++++++++++++ ...est_nova_model_customization_deployment.py | 16 + .../tests/unit/test_bedrock_model_builder.py | 252 ++++++++++- .../tests/unit/test_model_builder.py | 333 ++++++++++++++ .../tests/unit/test_model_reuse.py | 297 ++++++++++++ .../bedrock-modelbuilder-deployment.ipynb | 32 ++ .../model_builder_deployment_notebook.ipynb | 234 +++++++--- 11 files changed, 2255 insertions(+), 162 deletions(-) create mode 100644 sagemaker-core/tests/unit/test_training_utils.py create mode 100644 sagemaker-serve/src/sagemaker/serve/model_reuse.py create mode 100644 sagemaker-serve/tests/unit/test_model_reuse.py diff --git a/sagemaker-core/src/sagemaker/core/training/utils.py b/sagemaker-core/src/sagemaker/core/training/utils.py index 67009c9131..2ee2afdee7 100644 --- a/sagemaker-core/src/sagemaker/core/training/utils.py +++ b/sagemaker-core/src/sagemaker/core/training/utils.py @@ -13,8 +13,12 @@ """Training utilities.""" from __future__ import absolute_import +import io +import json import os +import tarfile from typing import Any, Literal +from urllib.parse import urlparse from sagemaker.core.utils.utils import Unassigned @@ -75,3 +79,135 @@ def _is_valid_s3_uri(path: str, path_type: Literal["File", "Directory", "Any"] = return path.endswith("/") return path_type == "Any" + + +_MANIFEST_CHECKPOINT_KEY = "checkpoint_s3_bucket" + + +def build_nova_manifest_s3_uri(s3_output_path: str, training_job_name: str) -> str: + """Build the manifest.json S3 URI for a Nova training job. + + Args: + s3_output_path: The training job's ``output_data_config.s3_output_path``. + training_job_name: The training job name. + + Returns: + Fully-qualified S3 URI to the job's manifest.json. + """ + output_path = s3_output_path.rstrip("/") + return f"{output_path}/{training_job_name}/output/output/manifest.json" + + +def build_nova_output_tar_gz_s3_uri(s3_output_path: str, training_job_name: str) -> str: + """Build the output.tar.gz S3 URI for a Nova training job. + + Args: + s3_output_path: The training job's ``output_data_config.s3_output_path``. + training_job_name: The training job name. + + Returns: + Fully-qualified S3 URI to the job's output.tar.gz. + """ + output_path = s3_output_path.rstrip("/") + return f"{output_path}/{training_job_name}/output/output.tar.gz" + + +def _split_s3_uri(s3_uri: str) -> tuple: + """Split an S3 URI into (bucket, key).""" + parsed = urlparse(s3_uri) + return parsed.netloc, parsed.path.lstrip("/") + + +def read_nova_checkpoint_uri_from_manifest(s3_client, s3_uri: str) -> str: + """Read the checkpoint URI from a raw manifest.json object in S3. + + Args: + s3_client: A boto3 S3 client. + s3_uri: S3 URI of the manifest.json object. + + Returns: + The ``checkpoint_s3_bucket`` value from the manifest. + + Raises: + ValueError: If the object is missing, unparseable, or lacks the key. + """ + bucket, key = _split_s3_uri(s3_uri) + try: + response = s3_client.get_object(Bucket=bucket, Key=key) + manifest = json.loads(response["Body"].read().decode("utf-8")) + except s3_client.exceptions.NoSuchKey: + raise ValueError(f"manifest.json not found at s3://{bucket}/{key}") + except json.JSONDecodeError as e: + raise ValueError(f"Failed to parse manifest.json: {e}") + + checkpoint_uri = manifest.get(_MANIFEST_CHECKPOINT_KEY) + if not checkpoint_uri: + raise ValueError( + f"'{_MANIFEST_CHECKPOINT_KEY}' not found in manifest.json. " + f"Available keys: {list(manifest.keys())}" + ) + return checkpoint_uri + + +def _read_checkpoint_uri_from_tar_gz(s3_client, s3_uri: str) -> str: + """Read the checkpoint URI from a manifest.json inside an output.tar.gz in S3. + + Args: + s3_client: A boto3 S3 client. + s3_uri: S3 URI of the output.tar.gz object. + + Returns: + The ``checkpoint_s3_bucket`` value from the embedded manifest. + + Raises: + ValueError: If the archive or manifest is missing or lacks the key. + """ + bucket, key = _split_s3_uri(s3_uri) + response = s3_client.get_object(Bucket=bucket, Key=key) + body = response["Body"].read() + with tarfile.open(fileobj=io.BytesIO(body), mode="r:gz") as tar: + for member in tar.getmembers(): + if not member.name.endswith("manifest.json"): + continue + extracted = tar.extractfile(member) + if extracted is None: + continue + manifest = json.loads(extracted.read().decode("utf-8")) + checkpoint_uri = manifest.get(_MANIFEST_CHECKPOINT_KEY) + if checkpoint_uri: + return checkpoint_uri + + raise ValueError( + f"'{_MANIFEST_CHECKPOINT_KEY}' not found in manifest.json within " + f"s3://{bucket}/{key}" + ) + + +def resolve_nova_checkpoint_uri( + s3_client, + s3_output_path: str, + training_job_name: str, +) -> str: + """Resolve the Nova checkpoint (escrow) URI from a training job's output. + + Reads ``checkpoint_s3_bucket`` from the job's manifest.json. The manifest is + first looked up as a raw object, and if that fails, it falls back to the copy + packaged inside ``output.tar.gz``. + + Args: + s3_client: A boto3 S3 client. + s3_output_path: The training job's ``output_data_config.s3_output_path``. + training_job_name: The training job name. + + Returns: + The checkpoint URI recorded in the manifest. + + Raises: + ValueError: If the checkpoint URI cannot be resolved from either source. + """ + manifest_uri = build_nova_manifest_s3_uri(s3_output_path, training_job_name) + try: + return read_nova_checkpoint_uri_from_manifest(s3_client, manifest_uri) + except Exception: + tar_gz_uri = build_nova_output_tar_gz_s3_uri(s3_output_path, training_job_name) + return _read_checkpoint_uri_from_tar_gz(s3_client, tar_gz_uri) diff --git a/sagemaker-core/tests/unit/test_training_utils.py b/sagemaker-core/tests/unit/test_training_utils.py new file mode 100644 index 0000000000..0cd41d62bc --- /dev/null +++ b/sagemaker-core/tests/unit/test_training_utils.py @@ -0,0 +1,146 @@ +# 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 Nova manifest/checkpoint helpers in training/utils.py.""" +import io +import json +import tarfile + +import pytest +from unittest.mock import Mock + +from sagemaker.core.training.utils import ( + build_nova_manifest_s3_uri, + build_nova_output_tar_gz_s3_uri, + read_nova_checkpoint_uri_from_manifest, + resolve_nova_checkpoint_uri, +) + +CHECKPOINT_URI = "s3://bucket/ckpt/step_100" + + +def test_build_nova_manifest_s3_uri(): + result = build_nova_manifest_s3_uri("s3://bucket/output/", "my-job") + assert result == "s3://bucket/output/my-job/output/output/manifest.json" + + +def test_build_nova_manifest_s3_uri_strips_trailing_slash(): + assert build_nova_manifest_s3_uri( + "s3://bucket/output//", "my-job" + ) == "s3://bucket/output/my-job/output/output/manifest.json" + + +def test_build_nova_output_tar_gz_s3_uri(): + result = build_nova_output_tar_gz_s3_uri("s3://bucket/output", "my-job") + assert result == "s3://bucket/output/my-job/output/output.tar.gz" + + +def _s3_client_returning(body_bytes): + client = Mock() + client.exceptions = Mock() + client.exceptions.NoSuchKey = type("NoSuchKey", (Exception,), {}) + body = Mock() + body.read.return_value = body_bytes + client.get_object.return_value = {"Body": body} + return client + + +def test_read_manifest_returns_checkpoint_uri(): + client = _s3_client_returning( + json.dumps({"checkpoint_s3_bucket": CHECKPOINT_URI}).encode("utf-8") + ) + result = read_nova_checkpoint_uri_from_manifest( + client, "s3://bucket/output/my-job/output/output/manifest.json" + ) + assert result == CHECKPOINT_URI + client.get_object.assert_called_once_with( + Bucket="bucket", Key="output/my-job/output/output/manifest.json" + ) + + +def test_read_manifest_missing_key_raises(): + client = _s3_client_returning(json.dumps({"other": "value"}).encode("utf-8")) + with pytest.raises(ValueError, match="checkpoint_s3_bucket"): + read_nova_checkpoint_uri_from_manifest(client, "s3://bucket/manifest.json") + + +def test_read_manifest_not_found_raises(): + client = Mock() + client.exceptions = Mock() + client.exceptions.NoSuchKey = type("NoSuchKey", (Exception,), {}) + client.get_object.side_effect = client.exceptions.NoSuchKey() + with pytest.raises(ValueError, match="manifest.json not found"): + read_nova_checkpoint_uri_from_manifest(client, "s3://bucket/manifest.json") + + +def test_read_manifest_invalid_json_raises(): + client = _s3_client_returning(b"not-json") + with pytest.raises(ValueError, match="Failed to parse manifest.json"): + read_nova_checkpoint_uri_from_manifest(client, "s3://bucket/manifest.json") + + +def _make_tar_gz_with_manifest(manifest_dict): + content = json.dumps(manifest_dict).encode("utf-8") + buf = io.BytesIO() + with tarfile.open(fileobj=buf, mode="w:gz") as tar: + info = tarfile.TarInfo(name="manifest.json") + info.size = len(content) + tar.addfile(info, io.BytesIO(content)) + return buf.getvalue() + + +def test_resolve_checkpoint_uri_from_raw_manifest(): + client = _s3_client_returning( + json.dumps({"checkpoint_s3_bucket": CHECKPOINT_URI}).encode("utf-8") + ) + result = resolve_nova_checkpoint_uri(client, "s3://bucket/output/", "my-job") + assert result == CHECKPOINT_URI + + +def test_resolve_checkpoint_uri_falls_back_to_tar_gz(): + client = Mock() + client.exceptions = Mock() + no_such_key = type("NoSuchKey", (Exception,), {}) + client.exceptions.NoSuchKey = no_such_key + tar_bytes = _make_tar_gz_with_manifest({"checkpoint_s3_bucket": CHECKPOINT_URI}) + + def get_object(Bucket, Key): + if Key.endswith("manifest.json"): + raise no_such_key() + body = Mock() + body.read.return_value = tar_bytes + return {"Body": body} + + client.get_object.side_effect = get_object + + result = resolve_nova_checkpoint_uri(client, "s3://bucket/output/", "my-job") + assert result == CHECKPOINT_URI + + +def test_resolve_checkpoint_uri_raises_when_both_sources_fail(): + client = Mock() + client.exceptions = Mock() + no_such_key = type("NoSuchKey", (Exception,), {}) + client.exceptions.NoSuchKey = no_such_key + tar_bytes = _make_tar_gz_with_manifest({"other": "value"}) + + def get_object(Bucket, Key): + if Key.endswith("manifest.json"): + raise no_such_key() + body = Mock() + body.read.return_value = tar_bytes + return {"Body": body} + + client.get_object.side_effect = get_object + + with pytest.raises(ValueError, match="checkpoint_s3_bucket"): + resolve_nova_checkpoint_uri(client, "s3://bucket/output/", "my-job") diff --git a/sagemaker-serve/src/sagemaker/serve/bedrock_model_builder.py b/sagemaker-serve/src/sagemaker/serve/bedrock_model_builder.py index d84cfa3818..aedd622269 100644 --- a/sagemaker-serve/src/sagemaker/serve/bedrock_model_builder.py +++ b/sagemaker-serve/src/sagemaker/serve/bedrock_model_builder.py @@ -18,11 +18,20 @@ import time import logging from datetime import datetime, timezone - -from sagemaker.serve.utils.model_package_utils import is_restricted_model_package from typing import Optional, Dict, Any, Union from urllib.parse import urlparse +from sagemaker.serve.utils.model_package_utils import is_restricted_model_package +from sagemaker.serve.model_reuse import ( + build_source_tag, + find_active_bedrock_deployment_for_model, + find_existing_bedrock_model, +) +from sagemaker.core.training.utils import ( + build_nova_manifest_s3_uri, + read_nova_checkpoint_uri_from_manifest, + resolve_nova_checkpoint_uri, +) from sagemaker.core.helper.session_helper import Session from sagemaker.core.helper.iam_role_resolver import resolve_and_validate_role from sagemaker.core.resources import TrainingJob, ModelPackage @@ -85,7 +94,18 @@ class BedrockModelBuilder: This class provides functionality to deploy SageMaker models to Bedrock using either model import jobs or custom model creation, depending on - the model type (Nova models vs. other models). + the model type. Nova models use the ``create_custom_model`` + + ``create_custom_model_deployment`` path; other (OSS) models use the + ``create_model_import_job`` path. + + Resource reuse: + ``deploy()`` accepts ``reuse_resources`` (default False). When True, it + tags each custom model it creates with the model source and, on a + subsequent deploy of the same source, reuses the existing custom model + (and its active deployment if present) instead of creating duplicates, + logging a warning rather than raising. This matters because Bedrock + custom-model import is slow and consumes the limited imported-model + quota per account/region. Args: model: The model to deploy. Can be a ModelTrainer, MultiTurnRLTrainer, @@ -145,6 +165,80 @@ def _get_sagemaker_client(self): self._sagemaker_client = self.boto_session.client("sagemaker") return self._sagemaker_client + def _resolve_model_source_id(self) -> Optional[str]: + """Determine the model source identifier for reuse lookups. + + Resolution order: + 1. Checkpoint URI from training job manifest (Nova models) + 2. Model package ARN (RMP models) + 3. Training job / trainer S3 model artifacts + 4. Direct S3 artifact path + + Returns: + Source identifier string, or None if cannot be determined. + """ + try: + # Nova models resolve to the checkpoint URI read from the manifest. + if isinstance(self.model, TrainingJob) and self.model_package: + spec = getattr(self.model_package, "inference_specification", None) + containers = getattr(spec, "containers", None) if spec else None + container = containers[0] if containers else None + if container and _is_nova_model(container): + return self._get_checkpoint_uri_from_manifest_safe() + + # Restricted model packages resolve to their ARN. + if self._is_rmp and self.model_package: + return self.model_package.model_package_arn + + # TrainingJob and ModelTrainer/BaseTrainer both expose a training job + # (directly or via _latest_training_job) carrying the artifacts. + training_job = self._resolve_training_job() + if training_job is not None: + mp_arn = getattr(training_job, "output_model_package_arn", None) + if isinstance(mp_arn, str) and mp_arn: + return mp_arn + s3_path = self._s3_artifacts_from_training_job(training_job) + if s3_path: + return s3_path + + if self.s3_model_artifacts: + return self.s3_model_artifacts + + except Exception as e: + logger.warning("Could not resolve model source identifier: %s", e) + + return None + + @staticmethod + def _s3_artifacts_from_training_job(training_job) -> Optional[str]: + """Return s3_model_artifacts from a training job's model_artifacts, if set.""" + artifacts = getattr(training_job, "model_artifacts", None) + if artifacts and not isinstance(artifacts, Unassigned): + s3_path = getattr(artifacts, "s3_model_artifacts", None) + if isinstance(s3_path, str) and s3_path: + return s3_path + return None + + def _get_checkpoint_uri_from_manifest_safe(self) -> Optional[str]: + """Attempt to read checkpoint URI from manifest, returning None on failure.""" + if not isinstance(self.model, TrainingJob): + return None + + output_data_config = getattr(self.model, "output_data_config", None) + s3_output_path = getattr(output_data_config, "s3_output_path", None) + if not s3_output_path: + return None + + try: + return resolve_nova_checkpoint_uri( + self.boto_session.client("s3"), + s3_output_path, + self.model.training_job_name, + ) + except Exception as e: + logger.warning("Could not read checkpoint URI from manifest: %s", e) + return None + def _is_nova_model_for_telemetry(self) -> bool: """Check if the model is a Nova model for telemetry tracking.""" try: @@ -168,6 +262,7 @@ def deploy( client_request_token: Optional[str] = None, imported_model_kms_key_id: Optional[str] = None, deployment_name: Optional[str] = None, + reuse_resources: bool = False, ) -> Dict[str, Any]: """Deploy the model to Bedrock. @@ -197,9 +292,19 @@ def deploy( imported_model_kms_key_id: KMS key ID for encryption (OSS models only). deployment_name: Name for the deployment (Nova models only). If not provided, defaults to custom_model_name suffixed with '-deployment'. + reuse_resources: If False (default), always creates new resources. If + True, checks for an existing custom model (and active deployment) + with the same source tag and reuses them instead of creating + duplicates. Newly created models are always tagged for future + discovery regardless of this flag. Returns: - For Nova models: the create_custom_model_deployment response. + For Nova models: the create_custom_model_deployment response. The + response always includes a ``modelArn`` key identifying the custom + model that was created or reused. When ``reuse_resources=True`` and a + match is found, returns ``{"modelArn": ..., "customModelDeploymentArn": + ...}`` for the reused model and its existing active deployment (or a + newly created deployment on the reused model if none exists). For OSS models: the completed get_model_import_job response. Raises: @@ -235,6 +340,45 @@ def deploy( sagemaker_session=self.sagemaker_session, ) + source_id = self._resolve_model_source_id() + + if source_id and reuse_resources: + existing_arn = find_existing_bedrock_model( + self._get_bedrock_client(), + source_id, + ) + if existing_arn: + model_arn = existing_arn + # Reuse an existing active deployment on the model if present; + # otherwise create a new deployment on the reused model. + existing_deployment = find_active_bedrock_deployment_for_model( + self._get_bedrock_client(), model_arn + ) + if existing_deployment: + logger.warning( + "Reusing existing custom model %s and deployment %s " + "(matched model-source tag). No new resources were created. " + "Pass reuse_resources=False to force new resources.", + model_arn, + existing_deployment, + ) + return { + "modelArn": model_arn, + "customModelDeploymentArn": existing_deployment, + } + logger.warning( + "Reusing existing custom model %s (matched model-source tag); " + "creating a new deployment on it. Pass reuse_resources=False to " + "force a new model.", + model_arn, + ) + 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) + return response + if self._is_rmp: params = { "modelName": custom_model_name, @@ -251,8 +395,17 @@ def deploy( "modelSourceConfig": {"s3DataSource": {"s3Uri": self.s3_model_artifacts}}, "roleArn": role_arn, } - if model_tags: - params["modelTags"] = model_tags + + merged_tags = list(model_tags) if model_tags else [] + if source_id: + source_tag = build_source_tag(source_id) + merged_tags = [ + t for t in merged_tags if t.get("key") != source_tag["key"] + ] + merged_tags.append(source_tag) + if merged_tags: + params["modelTags"] = merged_tags + params = {k: v for k, v in params.items() if v is not None} logger.info("Creating custom model %s for Nova deployment", custom_model_name) @@ -261,7 +414,9 @@ def deploy( model_arn = create_response.get("modelArn") deploy_name = deployment_name or f"{custom_model_name}-deployment" - return self.create_deployment(model_arn=model_arn, deployment_name=deploy_name) + response = self.create_deployment(model_arn=model_arn, deployment_name=deploy_name) + response.setdefault("modelArn", model_arn) + return response else: # Resolve and validate the Bedrock role: the provided role_arn if given, # otherwise the caller's own identity role. A RoleValidationError @@ -657,33 +812,27 @@ def _get_s3_artifacts(self) -> Optional[str]: return self._resolve_hf_model_path(s3_uri) return None - # No model_package — resolve from model_artifacts directly. - if isinstance(self.model, TrainingJob): - artifacts = getattr(self.model, 'model_artifacts', None) - if artifacts and not isinstance(artifacts, Unassigned): - s3_path = getattr(artifacts, 's3_model_artifacts', None) - if s3_path and isinstance(s3_path, str): - logger.info( - "Resolved S3 artifacts from TrainingJob model_artifacts: %s", s3_path - ) - return s3_path - return None + # No model_package — resolve from the training job's model_artifacts, + # whether the model is a TrainingJob or a trainer wrapping one. + training_job = self._resolve_training_job() + if training_job is not None: + s3_path = self._s3_artifacts_from_training_job(training_job) + if s3_path: + logger.info("Resolved S3 artifacts from training job: %s", s3_path) + return s3_path - # ModelTrainer or BaseTrainer — resolve from _latest_training_job.model_artifacts. - if isinstance(self.model, (ModelTrainer, BaseTrainer)): - training_job = getattr(self.model, '_latest_training_job', None) - if not training_job: - return None - artifacts = getattr(training_job, 'model_artifacts', None) - if artifacts and not isinstance(artifacts, Unassigned): - s3_path = getattr(artifacts, 's3_model_artifacts', None) - if s3_path and isinstance(s3_path, str): - logger.info( - "Resolved S3 artifacts from trainer's training job: %s", s3_path - ) - return s3_path - return None + return None + + def _resolve_training_job(self): + """Return the underlying TrainingJob for the model, if any. + Handles a direct ``TrainingJob`` as well as ``ModelTrainer``/``BaseTrainer`` + instances that expose one via ``_latest_training_job``. + """ + if isinstance(self.model, TrainingJob): + return self.model + if isinstance(self.model, (ModelTrainer, BaseTrainer)): + return getattr(self.model, "_latest_training_job", None) return None def _resolve_hf_model_path(self, s3_uri: str) -> str: @@ -826,38 +975,13 @@ def _get_checkpoint_uri_from_manifest(self) -> Optional[str]: if not s3_output_path: raise ValueError("No S3 output path found in training job output_data_config") - output_path = s3_output_path.rstrip("/") - manifest_path = f"{output_path}/{self.model.training_job_name}/output/output/manifest.json" - - logger.info("Manifest path: %s", manifest_path) - - parsed = urlparse(manifest_path) - bucket = parsed.netloc - manifest_key = parsed.path.lstrip("/") - - logger.info("Looking for manifest at s3://%s/%s", bucket, manifest_key) + manifest_uri = build_nova_manifest_s3_uri(s3_output_path, self.model.training_job_name) + logger.info("Looking for manifest at %s", manifest_uri) - s3_client = self.boto_session.client("s3") try: - response = s3_client.get_object(Bucket=bucket, Key=manifest_key) - manifest = json.loads(response["Body"].read().decode("utf-8")) - logger.info("Manifest content: %s", manifest) - - checkpoint_uri = manifest.get("checkpoint_s3_bucket") - if not checkpoint_uri: - raise ValueError( - "'checkpoint_s3_bucket' not found in manifest. " - "Available keys: %s" % list(manifest.keys()) - ) - - logger.info("Checkpoint URI: %s", checkpoint_uri) - return checkpoint_uri - except s3_client.exceptions.NoSuchKey: - raise ValueError( - "manifest.json not found at s3://%s/%s" % (bucket, manifest_key) + return read_nova_checkpoint_uri_from_manifest( + self.boto_session.client("s3"), manifest_uri ) - except json.JSONDecodeError as e: - raise ValueError("Failed to parse manifest.json: %s" % e) except ValueError: raise except Exception as e: diff --git a/sagemaker-serve/src/sagemaker/serve/model_builder.py b/sagemaker-serve/src/sagemaker/serve/model_builder.py index b7dc98b768..1ad0a9d96d 100644 --- a/sagemaker-serve/src/sagemaker/serve/model_builder.py +++ b/sagemaker-serve/src/sagemaker/serve/model_builder.py @@ -151,6 +151,12 @@ from sagemaker.train.base_trainer import BaseTrainer from sagemaker.core.telemetry.telemetry_logging import _telemetry_emitter from sagemaker.core.telemetry.constants import Feature +from sagemaker.serve.model_reuse import ( + build_source_tag, + find_existing_sagemaker_endpoint, + MODEL_SOURCE_TAG_KEY, +) +from sagemaker.core.training.utils import resolve_nova_checkpoint_uri _LOWEST_MMS_VERSION = "1.2" SCRIPT_PARAM_NAME = "sagemaker_program" @@ -176,6 +182,21 @@ class ModelBuilder(_InferenceRecommenderMixin, _ModelBuilderServers, _ModelBuild 2. Call build() to create a deployable Model resource 3. Call deploy() to create an Endpoint resource for inference + Resource reuse: + Both build() and deploy() accept ``reuse_resources`` (default False), and it + must be set on each call you want to reuse — the flag is honored per call, not + inherited. When True, ModelBuilder tags each resource it creates with the model + source and, on a subsequent call for the same source, discovers the existing + endpoint instead of creating a duplicate (logging a warning; it does not raise). + This is useful because endpoint creation is slow and constrained by + accelerated-instance capacity. + + - build(reuse_resources=True): on a hit, creates no new Model/EndpointConfig/ + Endpoint and sets ``built_model`` to the existing Model backing the reused + endpoint. + - deploy(reuse_resources=True): on a hit, returns the existing endpoint instead + of creating a new one. + Example: >>> from sagemaker.serve.model_builder import ModelBuilder >>> from sagemaker.serve.mode.function_pointers import Mode @@ -190,6 +211,10 @@ class ModelBuilder(_InferenceRecommenderMixin, _ModelBuilderServers, _ModelBuild >>> # Build the model (creates SageMaker Model resource) >>> model = model_builder.build() >>> + >>> # Reuse an existing endpoint if one was already built from this source + >>> model_builder.build(reuse_resources=True) + >>> endpoint = model_builder.deploy(reuse_resources=True) + >>> >>> # Deploy to endpoint (creates SageMaker Endpoint resource) >>> endpoint = model_builder.deploy(endpoint_name="my-endpoint") >>> @@ -198,7 +223,10 @@ class ModelBuilder(_InferenceRecommenderMixin, _ModelBuilderServers, _ModelBuild Args: model: The model to deploy. Can be a trained model object, ModelTrainer, TrainingJob, - ModelPackage, or JumpStart model ID string. Either model or inference_spec is required. + ModelPackage, JumpStart model ID string, or a raw S3 URI string pointing to model + artifacts. For a raw S3 URI that targets a Nova base model, supply the base model + via ``model_metadata={"BASE_MODEL_NAME": "..."}``. Either model or inference_spec + is required. model_path: Local directory path where model artifacts are stored or will be downloaded. inference_spec: Custom inference specification with load() and invoke() functions. schema_builder: Defines input/output schema for serialization and deserialization. @@ -323,7 +351,9 @@ class ModelBuilder(_InferenceRecommenderMixin, _ModelBuilderServers, _ModelBuild "help": "Dictionary to override model metadata. Supported keys: HF_TASK (for HuggingFace " "models without task metadata), MLFLOW_MODEL_PATH (local or S3 path to MLflow artifacts), " "FINE_TUNING_MODEL_PATH (S3 path to fine-tuned model), FINE_TUNING_JOB_NAME (fine-tuning " - "job name), and CUSTOM_MODEL_PATH (local or S3 path to custom model artifacts). " + "job name), CUSTOM_MODEL_PATH (local or S3 path to custom model artifacts), and " + "BASE_MODEL_NAME (base model identifier for a raw S3 URI model, e.g. a Nova checkpoint " + "deployed without a ModelPackage). " "FINE_TUNING_MODEL_PATH and FINE_TUNING_JOB_NAME are mutually exclusive." }, ) @@ -1071,11 +1101,39 @@ def _fetch_and_cache_recipe_config(self): ], } + def _base_model_name(self) -> Optional[str]: + """Resolve the base model identifier once. + + Comes from the model package's ``base_model.hub_content_name`` when a + package is available; otherwise from the trainer's ``base_model_name`` + (e.g. a BaseTrainer built from an S3 checkpoint). + + Returns: + The base model name string, or None if it cannot be determined. + """ + model_package = self._fetch_model_package() + if model_package is not None: + base_model = model_package.inference_specification.containers[0].base_model + return getattr(base_model, "hub_content_name", None) if base_model else None + if isinstance(self.model, BaseTrainer): + return self.model.base_model_name + # Raw S3 checkpoint: base model identity is supplied via model_metadata. + if self.model_metadata: + return self.model_metadata.get("BASE_MODEL_NAME") + return None + + def _is_raw_s3_model(self) -> bool: + """Return True if the model was provided as a raw S3 URI string.""" + return isinstance(self.model, str) and self.model.startswith("s3://") + def _is_nova_model(self) -> bool: """Check if the model is a Nova model based on recipe name or hub content name.""" model_package = self._fetch_model_package() if not model_package: - return False + # No model package (e.g. a trainer built from an S3 checkpoint). Fall + # back to the base model name carried by the trainer. + base_model_name = self._base_model_name() + return bool(base_model_name) and "nova" in base_model_name.lower() containers = getattr(model_package.inference_specification, "containers", None) if not containers: return False @@ -1696,6 +1754,11 @@ def _is_model_customization(self) -> bool: ): return True + # Raw S3 checkpoint with a Nova base model supplied via model_metadata + # (e.g. a HyperPod/SMTJ Nova checkpoint deployed without a ModelPackage). + if self._is_raw_s3_model() and self._is_nova_model(): + return True + # AgentRFTJob from MultiTurnRLTrainer.attach() if isinstance(self.model, AgentRFTJob): return True @@ -1704,6 +1767,15 @@ def _is_model_customization(self) -> bool: if isinstance(self.model, MultiTurnRLTrainer): return True if isinstance(self.model, BaseTrainer) and hasattr(self.model, "_latest_training_job"): + # Trainer built from an S3 checkpoint (e.g. Serverful SMTJ): no model + # package, but the completed training job has an S3 output path that + # holds the customized artifacts. + output_data_config = getattr( + self.model._latest_training_job, "output_data_config", None + ) + s3_output_path = getattr(output_data_config, "s3_output_path", None) + if s3_output_path and not isinstance(s3_output_path, Unassigned): + return True # Check model_package_config first (new location) if ( hasattr(self.model._latest_training_job, "model_package_config") @@ -1826,6 +1898,166 @@ def _fetch_model_package(self) -> Optional[ModelPackage]: return ModelPackage.get(arn) return None + def _resolve_model_source_id(self) -> Optional[str]: + """Determine the model source identifier for reuse lookups. + + Resolution order: + 1. Model package ARN (if model_package available) + 2. Nova escrow checkpoint URI (for Nova model customization) + 3. Raw S3 URI passed as the model + 4. S3 model artifact URI + 5. JumpStart model ID + + Returns: + Source identifier string, or None if cannot be determined. + """ + model_package_arn = self._fetch_model_package_arn() + if model_package_arn: + return model_package_arn + + # For Nova model customization, the stable identifier is the escrow + # checkpoint URI from the training job manifest. Resolve it before falling + # back to s3_model_data_url, which may be an auto-generated per-deploy + # "model-builder//" upload path that is useless as a reuse key. + if self._is_model_customization() and self._is_nova_model(): + try: + escrow_uri = self._resolve_nova_escrow_uri() + if escrow_uri: + return escrow_uri + except Exception as e: + logger.warning("Could not resolve Nova escrow URI for reuse: %s", e) + + if isinstance(self.model, str): + # A model passed as an S3 URI (e.g. a Nova escrow checkpoint path) is + # itself the source identifier for reuse. + if self.model.startswith("s3://"): + return self.model + if self._is_jumpstart_model_id(): + return self.model + + if self.s3_model_data_url and isinstance(self.s3_model_data_url, str): + return self.s3_model_data_url + + return None + + def _get_model_for_endpoint(self, endpoint_name: str) -> Optional[Model]: + """Return the Model resource backing an existing endpoint, if resolvable. + + Reads the endpoint's config production variant to find the model name, + then fetches that Model. Returns None if it cannot be resolved. + """ + sagemaker_client = self.sagemaker_session.sagemaker_client + try: + endpoint_desc = sagemaker_client.describe_endpoint(EndpointName=endpoint_name) + config_desc = sagemaker_client.describe_endpoint_config( + EndpointConfigName=endpoint_desc["EndpointConfigName"] + ) + variants = config_desc.get("ProductionVariants", []) + if not variants: + return None + model_name = variants[0]["ModelName"] + return Model.get(model_name=model_name, region=self.region) + except Exception as e: + logger.warning( + "Could not resolve the Model backing endpoint %s: %s.", + endpoint_name, + e, + ) + return None + + def _find_reusable_endpoint(self, instance_type: Optional[str] = None) -> Optional[str]: + """Return the name of an existing endpoint that can be reused, if any. + + A candidate must carry the same model-source tag and match the requested + deployment configuration (env vars, image URI, instance type). + + Args: + instance_type: The instance type requested for this deploy, if any. + + Returns: + The reusable endpoint name, or None if there is no suitable match. + """ + source_id = self._resolve_model_source_id() + if not source_id: + return None + + existing_arn = find_existing_sagemaker_endpoint( + self.sagemaker_session.sagemaker_client, + source_id, + ) + if not existing_arn: + return None + + existing_name = existing_arn.rsplit("/", 1)[-1] + if self._reused_endpoint_matches_config( + existing_name, instance_type=instance_type or self.instance_type + ): + return existing_name + + logger.info( + "Existing endpoint %s matches the model source but has different " + "deployment configuration; creating a new endpoint.", + existing_name, + ) + return None + + def _reused_endpoint_matches_config( + self, endpoint_name: str, instance_type: Optional[str] = None + ) -> bool: + """Check that a reuse candidate endpoint matches the requested deploy config. + + A source-tag match only proves the endpoint was built from the same model + artifacts. Before reusing it, confirm the runtime configuration the caller + requested (container environment variables, instance type, and image URI) + also matches, so a differently-configured request does not silently get an + endpoint that contradicts it. + + Args: + endpoint_name: Name of the candidate endpoint to inspect. + instance_type: The instance type requested for this deploy, if any. + + Returns: + True if the candidate matches (or the config cannot be read and reuse + should be attempted), False if a definite mismatch is detected. + """ + sagemaker_client = self.sagemaker_session.sagemaker_client + try: + endpoint_desc = sagemaker_client.describe_endpoint(EndpointName=endpoint_name) + config_desc = sagemaker_client.describe_endpoint_config( + EndpointConfigName=endpoint_desc["EndpointConfigName"] + ) + variants = config_desc.get("ProductionVariants", []) + if not variants: + return True + variant = variants[0] + model_desc = sagemaker_client.describe_model(ModelName=variant["ModelName"]) + # Check PrimaryContainer and fallback to Containers list if PrimaryContainer is empty + container = model_desc.get("PrimaryContainer") + if not container: + containers = model_desc.get("Containers") or [] + container = containers[0] if containers else {} + except Exception as e: + logger.warning( + "Could not read configuration of existing endpoint %s: %s. " + "Proceeding with reuse.", + endpoint_name, + e, + ) + return True + + existing_env = container.get("Environment") or {} + requested_env = self.env_vars or {} + if requested_env and requested_env != existing_env: + return False + + if self.image_uri and container.get("Image") and self.image_uri != container["Image"]: + return False + + if instance_type and variant.get("InstanceType") and instance_type != variant["InstanceType"]: + return False + + return True + def _convert_model_data_source_to_local(self, model_data_source): """Convert Core ModelDataSource to Local dictionary format.""" if not model_data_source: @@ -2595,15 +2827,19 @@ def _build_single_modelbuilder( self.built_model = Model.create(**create_kwargs) return self.built_model - # Fetch recipe config first to set image_uri, instance_type, env_vars, and s3_upload_path - base_model = model_package.inference_specification.containers[0].base_model - if base_model is not None: - self._fetch_and_cache_recipe_config() + # Fetch recipe config first to set image_uri, instance_type, env_vars, + # and s3_upload_path. Only possible when a model package is available; + # trainers built from an S3 checkpoint carry no package, so the caller + # must supply image_uri/instance_type/env_vars directly. + if model_package is not None: + base_model = model_package.inference_specification.containers[0].base_model + if base_model is not None: + self._fetch_and_cache_recipe_config() # Nova models use a completely different deployment architecture if self._is_nova_model(): escrow_uri = self._resolve_nova_escrow_uri() - base_model = model_package.inference_specification.containers[0].base_model + base_model_name = self._base_model_name() container_def = ContainerDefinition( image=self.image_uri, @@ -2617,15 +2853,23 @@ def _build_single_modelbuilder( }, ) model_name = self.model_name or f"model-{uuid.uuid4().hex[:10]}" + nova_tags = [ + {"key": "sagemaker-sdk:jumpstart-model-id", "value": base_model_name}, + ] + # Tag the Model with the model source so it is discoverable and + # trackable, mirroring the endpoint tagging done at deploy time. + source_id = self._resolve_model_source_id() + if source_id: + source_tag = build_source_tag(source_id) + nova_tags.append( + {"key": source_tag["key"], "value": source_tag["value"]} + ) self.built_model = Model.create( execution_role_arn=self.role_arn, model_name=model_name, containers=[container_def], enable_network_isolation=True, - tags=[ - {"key": "sagemaker-sdk:jumpstart-model-id", - "value": base_model.hub_content_name}, - ], + tags=nova_tags, ) return self.built_model @@ -3545,6 +3789,7 @@ def build( role_arn: Optional[str] = None, sagemaker_session: Optional[Session] = None, region: Optional[str] = None, + reuse_resources: bool = False, ) -> Union[Model, "ModelBuilder", None]: """Build a deployable ``Model`` instance with ``ModelBuilder``. @@ -3568,6 +3813,11 @@ def build( configuration chain. (Default: None). region (str, optional): The AWS region for deployment. If specified and different from the current region, a new session will be created. (Default: None). + reuse_resources (bool, optional): If True, checks for an existing endpoint built + from the same model source (with matching deployment configuration) before + creating anything. On a match, build() creates no new resources and sets + ``built_model`` to the existing Model backing that endpoint; the subsequent + deploy() returns the existing endpoint. (Default: False). Returns: Union[Model, ModelBuilder, None]: A ``sagemaker.core.resources.Model`` resource @@ -3622,6 +3872,34 @@ def build( self.accept_eula = getattr(self, "accept_eula", None) self.container_log_level = getattr(self, "container_log_level", None) + # Inference-component builds (modelbuilder_list or a custom orchestrator + # inference spec) populate self._deployables and manage their own reuse by + # IC name at deploy time. The endpoint-return reuse short-circuit only + # applies to single-model builds, so those IC builds still run normally. + is_inference_component_build = bool(self.modelbuilder_list) or isinstance( + self.inference_spec, (CustomOrchestrator, AsyncCustomOrchestrator) + ) + + # Resource reuse: if an existing endpoint built from the same model source + # (with matching config) is found, skip creating any new resources — no + # Model, EndpointConfig, or Endpoint. The subsequent deploy() returns the + # existing endpoint. + if reuse_resources and not is_inference_component_build: + self.serve_settings = self._get_serve_setting() + reusable_endpoint = self._find_reusable_endpoint() + if reusable_endpoint: + logger.warning( + "Reusing existing endpoint %r (matched model-source tag and " + "deployment configuration). No new resources will be created. " + "Pass reuse_resources=False to force new resources.", + reusable_endpoint, + ) + self._reused_endpoint_name = reusable_endpoint + # Point built_model at the existing Model backing the reused + # endpoint, so deploy() has a valid built_model to return. + self.built_model = self._get_model_for_endpoint(reusable_endpoint) + return self.built_model + deployables = {} if not self.modelbuilder_list and not isinstance( @@ -4418,6 +4696,7 @@ def deploy( ] = None, custom_orchestrator_instance_type: str = None, custom_orchestrator_initial_instance_count: int = None, + reuse_resources: bool = False, **kwargs, ) -> Union[Endpoint, LocalEndpoint, Transformer]: """Deploy the built model to an ``Endpoint``. @@ -4451,6 +4730,22 @@ def deploy( orchestrator deployment. (Default: None). custom_orchestrator_initial_instance_count (int, optional): Initial instance count for custom orchestrator deployment. (Default: None). + reuse_resources (bool): If False (default), always creates a new endpoint. + If True, checks for an existing endpoint created from the same model + source (with matching deployment configuration) and returns it instead + of creating a duplicate. New endpoints are always tagged for future + discovery regardless of this flag. + + Note: this flag must be set on ``deploy()`` for it to reuse an endpoint; + reuse is not inherited from ``build()``. Passing ``reuse_resources=True`` + here only avoids creating a new *endpoint* — the ``Model`` is created by + ``build()``, which runs first. To also avoid creating a new Model on a + reuse hit, pass ``reuse_resources=True`` to ``build()`` as well (build + then sets ``built_model`` to the existing Model backing the endpoint). + Inference-component deployments (``inference_config`` is a + ``ResourceRequirements``, or a ``modelbuilder_list`` build) are not + intercepted by this flag — they manage their own reuse by inference + component name (create vs. in-place update). Returns: Union[Endpoint, LocalEndpoint, Transformer]: A ``sagemaker.core.resources.Endpoint`` resource representing the deployed endpoint, a ``LocalEndpoint`` for local mode, @@ -4473,6 +4768,54 @@ def deploy( if not hasattr(self, "built_model") and not hasattr(self, "_deployables"): raise ValueError("Model needs to be built before deploying") + # Inference component deployments manage their own reuse by IC name + # (create vs. in-place update in _deploy_for_ic). The endpoint-return + # reuse gate must not intercept them, or an intended IC create/update + # would be silently skipped. + is_inference_component_deploy = isinstance( + inference_config, ResourceRequirements + ) or bool(getattr(self, "_deployables", None)) + + # Resource reuse is opt-in per call. When requested, reuse an existing + # endpoint built from the same model source (with matching config). If + # build() already resolved one (build's reuse gate), use that; otherwise + # discover it here. + if reuse_resources and not is_inference_component_deploy: + reusable_endpoint = getattr( + self, "_reused_endpoint_name", None + ) or self._find_reusable_endpoint( + instance_type=instance_type or self.instance_type + ) + if reusable_endpoint: + if endpoint_name and endpoint_name != reusable_endpoint: + logger.warning( + "Requested endpoint name %r is ignored; reusing existing " + "endpoint %r which matches the model source and deployment " + "configuration.", + endpoint_name, + reusable_endpoint, + ) + logger.warning( + "Reusing existing endpoint %r (matched model-source tag and " + "deployment configuration). No new resources were created. " + "Pass reuse_resources=False to force a new endpoint.", + reusable_endpoint, + ) + return Endpoint.get( + endpoint_name=reusable_endpoint, + session=self.sagemaker_session.boto_session, + region=self.region, + ) + + source_id = self._resolve_model_source_id() + + if source_id: + tag = build_source_tag(source_id) + # Pass as a single-element list; a bare {"Key":..., "Value":...} dict + # is ambiguous and gets expanded by format_tags into two junk tags + # keyed "Key" and "Value". + self.add_tags([{"Key": MODEL_SOURCE_TAG_KEY, "Value": tag["value"]}]) + # Handle model customization deployment if self._is_model_customization(): logger.info("Deploying Model Customization model") @@ -4903,8 +5246,9 @@ def _resolve_nova_escrow_uri(self) -> str: Nova training jobs write artifacts to an escrow S3 bucket. The location is recorded in manifest.json in the training job output directory. """ - import json - from urllib.parse import urlparse + # Raw S3 checkpoint: the provided URI is itself the escrow location. + if self._is_raw_s3_model(): + return self.model.rstrip("/") if isinstance(self.model, TrainingJob): training_job = self.model @@ -4916,24 +5260,13 @@ def _resolve_nova_escrow_uri(self) -> str: else: raise ValueError("Nova escrow URI resolution requires a TrainingJob or ModelTrainer") - output_path = training_job.output_data_config.s3_output_path.rstrip("/") - manifest_s3 = f"{output_path}/{training_job.training_job_name}/output/output/manifest.json" - - parsed = urlparse(manifest_s3) - bucket = parsed.netloc - key = parsed.path.lstrip("/") - - s3_client = self.sagemaker_session.boto_session.client("s3") - resp = s3_client.get_object(Bucket=bucket, Key=key) - manifest = json.loads(resp["Body"].read().decode()) - - escrow_uri = manifest.get("checkpoint_s3_bucket") - if not escrow_uri: - raise ValueError( - f"'checkpoint_s3_bucket' not found in manifest.json. " - f"Available keys: {list(manifest.keys())}" - ) - return escrow_uri + # Resolve the checkpoint URI from the job's manifest.json, which may be a + # raw object or packaged inside output.tar.gz. + return resolve_nova_checkpoint_uri( + self.sagemaker_session.boto_session.client("s3"), + training_job.output_data_config.s3_output_path, + training_job.training_job_name, + ) def _deploy_nova_model( self, @@ -4950,9 +5283,6 @@ def _deploy_nova_model( """ from sagemaker.core.shapes import ProductionVariant - model_package = self._fetch_model_package() - base_model = model_package.inference_specification.containers[0].base_model - if not endpoint_name: endpoint_name = f"endpoint-{uuid.uuid4().hex[:8]}" @@ -4968,11 +5298,27 @@ def _deploy_nova_model( ], ) + # The jumpstart-model-id tag always applies (resolved from the model + # package or the trainer's base_model_name). The recipe-name tag is only + # available when a model package is present. tags = [ - {"key": "sagemaker-sdk:jumpstart-model-id", "value": base_model.hub_content_name}, + {"key": "sagemaker-sdk:jumpstart-model-id", "value": self._base_model_name()}, ] - if base_model.recipe_name: - tags.append({"key": "sagemaker-sdk:recipe-name", "value": base_model.recipe_name}) + model_package = self._fetch_model_package() + if model_package is not None: + base_model = model_package.inference_specification.containers[0].base_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. + 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) endpoint = Endpoint.create( endpoint_name=endpoint_name, diff --git a/sagemaker-serve/src/sagemaker/serve/model_reuse.py b/sagemaker-serve/src/sagemaker/serve/model_reuse.py new file mode 100644 index 0000000000..1f72c1402c --- /dev/null +++ b/sagemaker-serve/src/sagemaker/serve/model_reuse.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. +"""Model source tag-based resource reuse utilities.""" +from __future__ import annotations + +import hashlib +import logging +import time +from typing import Callable, Optional + +logger = logging.getLogger(__name__) + +MODEL_SOURCE_TAG_KEY = "sagemaker.amazonaws.com/model-source" + +_TAG_VALUE_MAX_LENGTH = 256 +_TAG_TRUNCATE_PREFIX_LENGTH = 224 +_TAG_HASH_SUFFIX_LENGTH = 31 + +_ACTIVE_STATUSES = {"Active", "InService"} +_CREATING_STATUSES = {"Creating"} +_FAILED_STATUSES = {"Failed"} + + +def normalize_tag_value(value: str) -> str: + """Normalize a tag value to fit within the 256-character AWS tag limit. + + If the value is <= 256 chars, returns as-is. + Otherwise, truncates to 224 chars + "-" + 31 hex chars of SHA-256. + """ + if len(value) <= _TAG_VALUE_MAX_LENGTH: + return value + hash_suffix = hashlib.sha256(value.encode()).hexdigest()[:_TAG_HASH_SUFFIX_LENGTH] + return f"{value[:_TAG_TRUNCATE_PREFIX_LENGTH]}-{hash_suffix}" + + +def find_existing_bedrock_model( + bedrock_client, + source_id: str, + poll_interval: int = 30, + max_wait: int = 900, +) -> Optional[str]: + """Find an existing Bedrock custom model tagged with a matching source id. + + Enumerates custom models and matches on the + ``sagemaker.amazonaws.com/model-source`` tag, then validates the model + status before returning it for reuse. + + Args: + bedrock_client: A boto3 Bedrock client. + source_id: Raw source identifier (will be normalized). + poll_interval: Seconds between status polls for "Creating" resources. + max_wait: Maximum wait time for "Creating" resources. + + Returns: + Model ARN if an active/ready model is found, None otherwise. + + Raises: + TimeoutError: If a creating model doesn't become ready within max_wait. + """ + tag_value = normalize_tag_value(source_id) + try: + resource_arn = _find_bedrock_model_arn_by_tag(bedrock_client, tag_value) + except Exception as e: + logger.warning("Could not list Bedrock custom models: %s. Proceeding without.", e) + return None + + if not resource_arn: + return None + + return _resolve_ready_arn( + bedrock_client, resource_arn, check_bedrock_model_status, poll_interval, max_wait + ) + + +def find_active_bedrock_deployment_for_model(bedrock_client, model_arn: str) -> Optional[str]: + """Find an existing active custom model deployment for a Bedrock model. + + Args: + bedrock_client: A boto3 Bedrock client. + model_arn: ARN of the custom model whose deployment to reuse. + + Returns: + The ARN of an existing Active deployment on the model, or None. + """ + try: + next_token = None + while True: + kwargs = {"nextToken": next_token} if next_token else {} + response = bedrock_client.list_custom_model_deployments(**kwargs) + for summary in response.get("modelDeploymentSummaries", []): + if summary.get("modelArn") != model_arn: + continue + if summary.get("status") in _ACTIVE_STATUSES: + return summary.get("customModelDeploymentArn") + next_token = response.get("nextToken") + if not next_token: + return None + except Exception as e: + logger.warning( + "Could not list Bedrock custom model deployments: %s. Proceeding without.", e + ) + return None + + +def find_existing_sagemaker_endpoint( + sagemaker_client, + source_id: str, + poll_interval: int = 30, + max_wait: int = 900, +) -> Optional[str]: + """Find an existing SageMaker endpoint tagged with a matching source id. + + Enumerates endpoints and matches on the + ``sagemaker.amazonaws.com/model-source`` tag, then validates the endpoint + status before returning it for reuse. + + Args: + sagemaker_client: A boto3 SageMaker client. + source_id: Raw source identifier (will be normalized). + poll_interval: Seconds between status polls for "Creating" resources. + max_wait: Maximum wait time for "Creating" resources. + + Returns: + Endpoint ARN if an in-service/ready endpoint is found, None otherwise. + + Raises: + TimeoutError: If a creating endpoint doesn't become ready within max_wait. + """ + tag_value = normalize_tag_value(source_id) + try: + resource_arn = _find_sagemaker_endpoint_arn_by_tag(sagemaker_client, tag_value) + except Exception as e: + logger.warning("Could not list SageMaker endpoints: %s. Proceeding without.", e) + return None + + if not resource_arn: + return None + + return _resolve_ready_arn( + sagemaker_client, resource_arn, check_sagemaker_endpoint_status, poll_interval, max_wait + ) + + +def _find_bedrock_model_arn_by_tag(bedrock_client, tag_value: str) -> Optional[str]: + """Return the ARN of the first Bedrock custom model carrying the source tag.""" + next_token = None + while True: + kwargs = {"nextToken": next_token} if next_token else {} + response = bedrock_client.list_custom_models(**kwargs) + for summary in response.get("modelSummaries", []): + arn = summary.get("modelArn") + if arn and _bedrock_resource_has_tag(bedrock_client, arn, tag_value): + return arn + next_token = response.get("nextToken") + if not next_token: + return None + + +def _bedrock_resource_has_tag(bedrock_client, resource_arn: str, tag_value: str) -> bool: + """Return True if the Bedrock resource carries the source tag with tag_value.""" + tags = bedrock_client.list_tags_for_resource(resourceARN=resource_arn).get("tags", []) + return any( + tag.get("key") == MODEL_SOURCE_TAG_KEY and tag.get("value") == tag_value + for tag in tags + ) + + +def _find_sagemaker_endpoint_arn_by_tag(sagemaker_client, tag_value: str) -> Optional[str]: + """Return the ARN of the first SageMaker endpoint carrying the source tag.""" + next_token = None + while True: + kwargs = {"NextToken": next_token} if next_token else {} + response = sagemaker_client.list_endpoints(**kwargs) + for endpoint in response.get("Endpoints", []): + arn = endpoint.get("EndpointArn") + if arn and _sagemaker_resource_has_tag(sagemaker_client, arn, tag_value): + return arn + next_token = response.get("NextToken") + if not next_token: + return None + + +def _sagemaker_resource_has_tag(sagemaker_client, resource_arn: str, tag_value: str) -> bool: + """Return True if the SageMaker resource carries the source tag with tag_value.""" + tags = sagemaker_client.list_tags(ResourceArn=resource_arn).get("Tags", []) + return any( + tag.get("Key") == MODEL_SOURCE_TAG_KEY and tag.get("Value") == tag_value + for tag in tags + ) + + +def _resolve_ready_arn( + client, + resource_arn: str, + status_checker: Callable, + poll_interval: int, + max_wait: int, +) -> Optional[str]: + """Validate a resource's status and return its ARN only when ready. + + Returns the ARN for active resources, polls creating resources until ready, + and returns None for failed or unexpected statuses. + """ + try: + status = status_checker(client, resource_arn) + except Exception as e: + logger.warning("Could not check resource status: %s. Proceeding without.", e) + return None + + if status in _ACTIVE_STATUSES: + return resource_arn + + if status in _FAILED_STATUSES: + logger.warning("Found resource %s in Failed status. Proceeding to create new.", resource_arn) + return None + + if status in _CREATING_STATUSES: + return _poll_until_ready(client, resource_arn, status_checker, poll_interval, max_wait) + + logger.warning("Resource %s has unexpected status '%s'. Proceeding to create new.", resource_arn, status) + return None + + +def _poll_until_ready( + client, + resource_arn: str, + status_checker: Callable, + poll_interval: int, + max_wait: int, +) -> Optional[str]: + """Poll a resource in Creating status until it becomes ready or times out.""" + elapsed = 0 + while elapsed < max_wait: + time.sleep(poll_interval) + elapsed += poll_interval + + try: + status = status_checker(client, resource_arn) + except Exception as e: + logger.warning("Could not check resource status during poll: %s. Proceeding without.", e) + return None + + if status in _ACTIVE_STATUSES: + return resource_arn + + if status in _FAILED_STATUSES: + logger.warning( + "Resource %s transitioned to Failed during poll. Proceeding to create new.", + resource_arn, + ) + return None + + if status not in _CREATING_STATUSES: + logger.warning( + "Resource %s has unexpected status '%s' during poll. Proceeding to create new.", + resource_arn, + status, + ) + return None + + raise TimeoutError( + f"Resource {resource_arn} did not become ready within {max_wait} seconds." + ) + + +def build_source_tag(source_id: str) -> dict: + """Build a tag dict for the model source.""" + return {"key": MODEL_SOURCE_TAG_KEY, "value": normalize_tag_value(source_id)} + + +def check_bedrock_model_status(bedrock_client, model_arn: str) -> str: + """Return the status of a Bedrock custom model.""" + try: + response = bedrock_client.get_custom_model(modelIdentifier=model_arn) + return response["modelStatus"] + except Exception as e: + logger.warning("Could not get Bedrock model status: %s. Proceeding without.", e) + raise + + +def check_sagemaker_endpoint_status(sagemaker_client, endpoint_arn: str) -> str: + """Return the status of a SageMaker endpoint.""" + try: + response = sagemaker_client.describe_endpoint(EndpointName=_arn_to_name(endpoint_arn)) + return response["EndpointStatus"] + except Exception as e: + logger.warning("Could not get endpoint status: %s. Proceeding without.", e) + raise + + +def _arn_to_name(arn: str) -> str: + """Extract the resource name from an ARN (last segment after '/').""" + return arn.rsplit("/", 1)[-1] diff --git a/sagemaker-serve/tests/integ/test_nova_model_customization_deployment.py b/sagemaker-serve/tests/integ/test_nova_model_customization_deployment.py index d4247774c4..09c9e069eb 100644 --- a/sagemaker-serve/tests/integ/test_nova_model_customization_deployment.py +++ b/sagemaker-serve/tests/integ/test_nova_model_customization_deployment.py @@ -26,6 +26,7 @@ import pytest import random from sagemaker.serve import ModelBuilder +from sagemaker.serve.model_reuse import MODEL_SOURCE_TAG_KEY from sagemaker.core.resources import TrainingJob logger = logging.getLogger(__name__) @@ -226,6 +227,13 @@ def test_deploy_from_training_job(self, training_job_name, endpoint_name, cleanu assert endpoint.endpoint_arn is not None assert endpoint.endpoint_status == "InService" + # The endpoint should carry the model-source tag that powers resource reuse. + sm_client = boto3.client("sagemaker", region_name=AWS_REGION) + endpoint_tags = sm_client.list_tags(ResourceArn=endpoint.endpoint_arn).get("Tags", []) + assert MODEL_SOURCE_TAG_KEY in {t["Key"] for t in endpoint_tags}, ( + f"Endpoint {endpoint.endpoint_arn} missing model-source tag for reuse" + ) + time.sleep(10) # brief buffer for inference component readiness invoke_response = endpoint.invoke( @@ -491,6 +499,14 @@ def test_nova_bedrock_deployment_active(self, deployed_nova_model, bedrock_clien ) assert deployment.get("status") == "Active" + def test_nova_bedrock_custom_model_tagged_for_reuse(self, deployed_nova_model, bedrock_client): + """The Nova custom model should carry the model-source tag that powers reuse.""" + model_arn = deployed_nova_model["model_arn"] + tags = bedrock_client.list_tags_for_resource(resourceARN=model_arn).get("tags", []) + assert MODEL_SOURCE_TAG_KEY in {t["key"] for t in tags}, ( + f"Custom model {model_arn} missing model-source tag for reuse" + ) + @pytest.mark.slow def test_nova_bedrock_invoke(self, deployed_nova_model, bedrock_runtime): """Invoke the deployed Nova model on Bedrock end-to-end.""" diff --git a/sagemaker-serve/tests/unit/test_bedrock_model_builder.py b/sagemaker-serve/tests/unit/test_bedrock_model_builder.py index 9b504ae5b2..bfbc1a7594 100644 --- a/sagemaker-serve/tests/unit/test_bedrock_model_builder.py +++ b/sagemaker-serve/tests/unit/test_bedrock_model_builder.py @@ -491,6 +491,12 @@ def _stub_role_validation(self): ): yield + @pytest.fixture(autouse=True) + def _stub_find_existing_bedrock_model(self): + """Patch find_existing_bedrock_model to return None by default for non-reuse tests.""" + with patch(f"{MODULE}.find_existing_bedrock_model", return_value=None): + yield + def test_oss_waits_for_import_and_returns_job_details(self): """OSS deploy: import job → wait → return job details.""" c = _make_container(s3_uri="s3://b/m.tar.gz") @@ -605,7 +611,9 @@ def test_nova_with_tags(self): tags = [{"Key": "env", "Value": "test"}] b.deploy(custom_model_name="m", role_arn="r", model_tags=tags) kw = b._bedrock_client.create_custom_model.call_args[1] - assert kw["modelTags"] == tags + assert {"Key": "env", "Value": "test"} in kw["modelTags"] + source_tag = {"key": "sagemaker.amazonaws.com/model-source", "value": "s3://b/k"} + assert source_tag in kw["modelTags"] def test_no_model_package_raises(self): b = _builder() @@ -1150,3 +1158,245 @@ def test_model_trainer_no_latest_training_job(self): b = BedrockModelBuilder(model=mock_trainer) assert b.s3_model_artifacts is None + + +class TestResolveModelSourceId: + def test_training_job_manifest_json(self): + b = _builder() + mock_job = Mock(spec=TrainingJob) + mock_job.output_data_config = Mock() + mock_job.output_data_config.s3_output_path = "s3://bucket/output/" + mock_job.training_job_name = "my-job" + b.model = mock_job + + nova_container = _make_container(recipe_name="nova-micro") + b.model_package = _make_model_package(nova_container) + b._is_rmp = False + b.s3_model_artifacts = "s3://bucket/ckpt" + + manifest = {"checkpoint_s3_bucket": "s3://bucket/ckpt/step_100"} + body = Mock() + body.read.return_value = json.dumps(manifest).encode() + mock_s3 = Mock() + mock_s3.get_object.return_value = {"Body": body} + mock_s3.exceptions = Mock() + mock_s3.exceptions.NoSuchKey = ClientError + session = Mock() + session.client.return_value = mock_s3 + b.boto_session = session + + with patch(f"{MODULE}.TrainingJob", type(mock_job)): + result = b._resolve_model_source_id() + + assert result == "s3://bucket/ckpt/step_100" + + def test_training_job_output_tar_gz_fallback(self): + b = _builder() + mock_job = Mock(spec=TrainingJob) + mock_job.output_data_config = Mock() + mock_job.output_data_config.s3_output_path = "s3://bucket/output/" + mock_job.training_job_name = "my-job" + b.model = mock_job + + nova_container = _make_container(recipe_name="nova-micro") + b.model_package = _make_model_package(nova_container) + b._is_rmp = False + b.s3_model_artifacts = "s3://bucket/ckpt" + + import tarfile + import io + + manifest_content = json.dumps({"checkpoint_s3_bucket": "s3://bucket/ckpt/step_50"}).encode() + buf = io.BytesIO() + with tarfile.open(fileobj=buf, mode="w:gz") as tar: + info = tarfile.TarInfo(name="manifest.json") + info.size = len(manifest_content) + tar.addfile(info, io.BytesIO(manifest_content)) + tar_bytes = buf.getvalue() + + mock_s3 = Mock() + manifest_err = ClientError({"Error": {"Code": "NoSuchKey"}}, "GetObject") + mock_s3.exceptions = Mock() + mock_s3.exceptions.NoSuchKey = ClientError + + def get_object_side_effect(Bucket, Key): + if Key.endswith("manifest.json"): + raise manifest_err + body = Mock() + body.read.return_value = tar_bytes + return {"Body": body} + + mock_s3.get_object.side_effect = get_object_side_effect + session = Mock() + session.client.return_value = mock_s3 + b.boto_session = session + + with patch(f"{MODULE}.TrainingJob", type(mock_job)): + result = b._resolve_model_source_id() + + assert result == "s3://bucket/ckpt/step_50" + + def test_model_package_arn_for_rmp(self): + b = _builder() + b.model = Mock() + b._is_rmp = True + b.model_package = Mock() + b.model_package.model_package_arn = "arn:aws:sagemaker:us-west-2:123456789012:model-package/my-pkg" + b.s3_model_artifacts = None + + with patch(f"{MODULE}.TrainingJob", _SentinelA), \ + patch(f"{MODULE}.ModelTrainer", _SentinelB), \ + patch(f"{MODULE}.BaseTrainer", _SentinelC): + result = b._resolve_model_source_id() + + assert result == "arn:aws:sagemaker:us-west-2:123456789012:model-package/my-pkg" + + def test_s3_model_artifacts_direct(self): + b = _builder() + b.model = "not-a-known-type" + b._is_rmp = False + b.model_package = None + b.s3_model_artifacts = "s3://my-bucket/checkpoints/" + + with patch(f"{MODULE}.TrainingJob", _SentinelA), \ + patch(f"{MODULE}.ModelTrainer", _SentinelB), \ + patch(f"{MODULE}.BaseTrainer", _SentinelC): + result = b._resolve_model_source_id() + + assert result == "s3://my-bucket/checkpoints/" + + def test_returns_none_when_no_source(self): + b = _builder() + b.model = None + b._is_rmp = False + b.model_package = None + b.s3_model_artifacts = None + + result = b._resolve_model_source_id() + + assert result is None + + +class TestModelReuseDeploy: + @pytest.fixture(autouse=True) + def _stub_role_validation(self): + with patch( + f"{MODULE}.resolve_and_validate_role", + side_effect=lambda provided_role, **kwargs: provided_role or "auto-role", + ): + yield + + def test_deploy_existing_model_skips_create(self): + c = _make_container(recipe_name="nova-micro") + b = _builder() + b.model_package = _make_model_package(c) + b.s3_model_artifacts = "s3://b/ckpt" + b._is_rmp = False + b._bedrock_client = Mock() + b._bedrock_client.get_custom_model.return_value = {"modelStatus": "Active"} + b._bedrock_client.create_custom_model_deployment.return_value = { + "customModelDeploymentArn": "arn:dep" + } + b._bedrock_client.get_custom_model_deployment.return_value = {"status": "Active"} + + with ( + patch(f"{MODULE}.find_existing_bedrock_model", return_value="arn:existing-model"), + patch(f"{MODULE}.find_active_bedrock_deployment_for_model", return_value=None), + ): + result = b.deploy(custom_model_name="m", role_arn="r", reuse_resources=True) + + b._bedrock_client.create_custom_model.assert_not_called() + assert result["customModelDeploymentArn"] == "arn:dep" + assert result["modelArn"] == "arn:existing-model" + + def test_deploy_existing_model_and_deployment_reused(self): + c = _make_container(recipe_name="nova-micro") + b = _builder() + b.model_package = _make_model_package(c) + b.s3_model_artifacts = "s3://b/ckpt" + b._is_rmp = False + b._bedrock_client = Mock() + + with ( + patch(f"{MODULE}.find_existing_bedrock_model", return_value="arn:existing-model"), + patch( + f"{MODULE}.find_active_bedrock_deployment_for_model", + return_value="arn:existing-dep", + ), + ): + result = b.deploy(custom_model_name="m", role_arn="r", reuse_resources=True) + + # Neither a new model nor a new deployment should be created. + b._bedrock_client.create_custom_model.assert_not_called() + b._bedrock_client.create_custom_model_deployment.assert_not_called() + assert result["modelArn"] == "arn:existing-model" + assert result["customModelDeploymentArn"] == "arn:existing-dep" + + def test_deploy_no_existing_model_creates_and_tags(self): + c = _make_container(recipe_name="nova-micro") + b = _builder() + b.model_package = _make_model_package(c) + b.s3_model_artifacts = "s3://b/ckpt" + b._is_rmp = False + b._bedrock_client = Mock() + b._bedrock_client.create_custom_model.return_value = {"modelArn": "arn:new-model"} + b._bedrock_client.get_custom_model.return_value = {"modelStatus": "Active"} + b._bedrock_client.create_custom_model_deployment.return_value = { + "customModelDeploymentArn": "arn:dep" + } + b._bedrock_client.get_custom_model_deployment.return_value = {"status": "Active"} + + with patch(f"{MODULE}.find_existing_bedrock_model", return_value=None): + result = b.deploy(custom_model_name="m", role_arn="r", reuse_resources=True) + + b._bedrock_client.create_custom_model.assert_called_once() + kw = b._bedrock_client.create_custom_model.call_args[1] + source_tag = {"key": "sagemaker.amazonaws.com/model-source", "value": "s3://b/ckpt"} + assert source_tag in kw["modelTags"] + assert result["customModelDeploymentArn"] == "arn:dep" + + def test_deploy_default_skips_lookup_but_tags(self): + c = _make_container(recipe_name="nova-micro") + b = _builder() + b.model_package = _make_model_package(c) + b.s3_model_artifacts = "s3://b/ckpt" + b._is_rmp = False + b._bedrock_client = Mock() + b._bedrock_client.create_custom_model.return_value = {"modelArn": "arn:new-model"} + b._bedrock_client.get_custom_model.return_value = {"modelStatus": "Active"} + b._bedrock_client.create_custom_model_deployment.return_value = { + "customModelDeploymentArn": "arn:dep" + } + b._bedrock_client.get_custom_model_deployment.return_value = {"status": "Active"} + + # Default is reuse_resources=False: no lookup, but new model is still tagged. + with patch(f"{MODULE}.find_existing_bedrock_model") as mock_find: + result = b.deploy(custom_model_name="m", role_arn="r") + + mock_find.assert_not_called() + b._bedrock_client.create_custom_model.assert_called_once() + kw = b._bedrock_client.create_custom_model.call_args[1] + source_tag = {"key": "sagemaker.amazonaws.com/model-source", "value": "s3://b/ckpt"} + assert source_tag in kw["modelTags"] + + def test_deploy_without_target_still_applies_source_tag(self): + b = BedrockModelBuilder(model="s3://bucket/my-artifacts/") + b._bedrock_client = Mock() + b._bedrock_client.create_custom_model.return_value = {"modelArn": "arn:m"} + b._bedrock_client.get_custom_model.return_value = {"modelStatus": "Active"} + b._bedrock_client.create_custom_model_deployment.return_value = { + "customModelDeploymentArn": "arn:dep" + } + b._bedrock_client.get_custom_model_deployment.return_value = {"status": "Active"} + + with patch(f"{MODULE}.find_existing_bedrock_model", return_value=None): + result = b.deploy(custom_model_name="my-model", role_arn="r") + + b._bedrock_client.create_custom_model.assert_called_once() + kw = b._bedrock_client.create_custom_model.call_args[1] + source_tag = { + "key": "sagemaker.amazonaws.com/model-source", + "value": "s3://bucket/my-artifacts/", + } + assert source_tag in kw["modelTags"] + assert result["customModelDeploymentArn"] == "arn:dep" diff --git a/sagemaker-serve/tests/unit/test_model_builder.py b/sagemaker-serve/tests/unit/test_model_builder.py index ae792c3133..2622b81976 100644 --- a/sagemaker-serve/tests/unit/test_model_builder.py +++ b/sagemaker-serve/tests/unit/test_model_builder.py @@ -889,3 +889,336 @@ def test_lora_build_passes_accept_eula_true(self, mock_model, mock_container_def finally: for p in patches: p.stop() + + +class TestModelReuse(unittest.TestCase): + """Test ModelBuilder model reuse integration.""" + + def setUp(self): + self.mock_session = Mock() + self.mock_session.boto_region_name = "us-west-2" + self.mock_session.default_bucket.return_value = "test-bucket" + self.mock_session.default_bucket_prefix = "test-prefix" + self.mock_session.boto_session = Mock() + self.mock_session.boto_session.region_name = "us-west-2" + self.mock_session.config = {} + self.mock_session.sagemaker_config = {} + self.mock_session.settings = Mock() + self.mock_session.settings.include_jumpstart_tags = False + self.mock_session.settings._local_download_dir = None + + def _make_builder(self, **overrides): + defaults = dict( + model=Mock(), + model_server=ModelServer.TORCHSERVE, + role_arn="arn:aws:iam::123456789012:role/SageMakerRole", + sagemaker_session=self.mock_session, + ) + defaults.update(overrides) + return ModelBuilder(**defaults) + + def test_resolve_model_source_id_returns_model_package_arn(self): + model_package = Mock(spec=["model_package_arn"]) + model_package.model_package_arn = "arn:aws:sagemaker:us-west-2:123456789012:model-package/my-pkg/1" + + from sagemaker.core.resources import ModelPackage as CoreModelPackage + + with patch.object(ModelBuilder, "_fetch_model_package_arn") as mock_fetch: + mock_fetch.return_value = "arn:aws:sagemaker:us-west-2:123456789012:model-package/my-pkg/1" + builder = self._make_builder() + result = builder._resolve_model_source_id() + + assert result == "arn:aws:sagemaker:us-west-2:123456789012:model-package/my-pkg/1" + + def test_resolve_model_source_id_returns_s3_artifact_uri(self): + with patch.object(ModelBuilder, "_fetch_model_package_arn", return_value=None): + builder = self._make_builder(s3_model_data_url="s3://my-bucket/artifacts/model.tar.gz") + result = builder._resolve_model_source_id() + + assert result == "s3://my-bucket/artifacts/model.tar.gz" + + def test_resolve_model_source_id_returns_none_when_no_source(self): + with patch.object(ModelBuilder, "_fetch_model_package_arn", return_value=None): + builder = self._make_builder(s3_model_data_url=None) + builder.model = 12345 + result = builder._resolve_model_source_id() + + assert result is None + + def test_resolve_model_source_id_returns_raw_s3_model(self): + with patch.object(ModelBuilder, "_fetch_model_package_arn", return_value=None): + builder = self._make_builder(model="s3://bucket/checkpoint/", s3_model_data_url=None) + result = builder._resolve_model_source_id() + + assert result == "s3://bucket/checkpoint/" + + def test_is_raw_s3_model(self): + builder = self._make_builder(model="s3://bucket/checkpoint/") + assert builder._is_raw_s3_model() is True + + builder = self._make_builder(model="nova-textgeneration-lite") + assert builder._is_raw_s3_model() is False + + def test_base_model_name_from_model_metadata(self): + with patch.object(ModelBuilder, "_fetch_model_package", return_value=None): + builder = self._make_builder( + model="s3://bucket/checkpoint/", + model_metadata={"BASE_MODEL_NAME": "nova-textgeneration-lite"}, + ) + assert builder._base_model_name() == "nova-textgeneration-lite" + + def test_is_model_customization_raw_s3_nova(self): + with patch.object(ModelBuilder, "_fetch_model_package", return_value=None): + builder = self._make_builder( + model="s3://bucket/checkpoint/", + model_metadata={"BASE_MODEL_NAME": "nova-textgeneration-lite"}, + ) + assert builder._is_model_customization() is True + + def test_is_model_customization_raw_s3_non_nova_is_false(self): + with patch.object(ModelBuilder, "_fetch_model_package", return_value=None): + builder = self._make_builder( + model="s3://bucket/checkpoint/", + model_metadata={"BASE_MODEL_NAME": "llama-3-8b"}, + ) + assert builder._is_model_customization() is False + + def test_is_model_customization_raw_s3_without_base_model_is_false(self): + with patch.object(ModelBuilder, "_fetch_model_package", return_value=None): + builder = self._make_builder(model="s3://bucket/checkpoint/", model_metadata=None) + assert builder._is_model_customization() is False + + def test_resolve_nova_escrow_uri_raw_s3(self): + with patch.object(ModelBuilder, "_fetch_model_package", return_value=None): + builder = self._make_builder( + model="s3://bucket/checkpoint/", + model_metadata={"BASE_MODEL_NAME": "nova-textgeneration-lite"}, + ) + assert builder._resolve_nova_escrow_uri() == "s3://bucket/checkpoint" + + @patch("sagemaker.serve.model_builder.Endpoint.get") + @patch("sagemaker.serve.model_builder.find_existing_sagemaker_endpoint") + def test_deploy_with_existing_endpoint_returns_without_creating( + self, mock_find, mock_endpoint_get + ): + existing_arn = "arn:aws:sagemaker:us-west-2:123456789012:endpoint/existing-ep" + mock_find.return_value = existing_arn + mock_endpoint = Mock() + mock_endpoint_get.return_value = mock_endpoint + + with ( + patch.object(ModelBuilder, "_resolve_model_source_id", return_value="s3://bucket/model"), + patch.object(ModelBuilder, "_reused_endpoint_matches_config", return_value=True), + ): + builder = self._make_builder() + builder.built_model = Mock() + builder.region = "us-west-2" + + result = builder.deploy(endpoint_name="existing-ep", reuse_resources=True) + + # deploy() discovers the reusable endpoint via _find_reusable_endpoint, + # which queries find_existing_sagemaker_endpoint with the session's + # sagemaker_client and the resolved source id. + mock_find.assert_called_once_with( + self.mock_session.sagemaker_client, + "s3://bucket/model", + ) + mock_endpoint_get.assert_called_once_with( + endpoint_name="existing-ep", + session=self.mock_session.boto_session, + region="us-west-2", + ) + assert result == mock_endpoint + + @patch("sagemaker.serve.model_builder.find_existing_sagemaker_endpoint") + def test_deploy_with_no_existing_endpoint_creates_and_tags(self, mock_find): + mock_find.return_value = None + + with ( + patch.object(ModelBuilder, "_resolve_model_source_id", return_value="s3://bucket/model"), + patch.object(ModelBuilder, "_is_model_customization", return_value=False), + patch.object(ModelBuilder, "_get_deploy_wrapper") as mock_get_wrapper, + patch.object(ModelBuilder, "add_tags") as mock_add_tags, + ): + mock_deploy = Mock(return_value=Mock()) + mock_get_wrapper.return_value = mock_deploy + + builder = self._make_builder() + builder.built_model = Mock() + builder.region = "us-west-2" + + builder.deploy(endpoint_name="new-ep", reuse_resources=True) + + mock_find.assert_called_once() + mock_add_tags.assert_called_once() + tag_arg = mock_add_tags.call_args[0][0][0] + assert tag_arg["Key"] == "sagemaker.amazonaws.com/model-source" + assert tag_arg["Value"] == "s3://bucket/model" + + @patch("sagemaker.serve.model_builder.find_existing_sagemaker_endpoint") + def test_deploy_default_skips_lookup_but_tags(self, mock_find): + with ( + patch.object(ModelBuilder, "_resolve_model_source_id", return_value="s3://bucket/model"), + patch.object(ModelBuilder, "_is_model_customization", return_value=False), + patch.object(ModelBuilder, "_get_deploy_wrapper") as mock_get_wrapper, + patch.object(ModelBuilder, "add_tags") as mock_add_tags, + ): + mock_deploy = Mock(return_value=Mock()) + mock_get_wrapper.return_value = mock_deploy + + builder = self._make_builder() + builder.built_model = Mock() + builder.region = "us-west-2" + + # Default is reuse_resources=False: no lookup, but new endpoint is tagged. + builder.deploy(endpoint_name="forced-ep") + + mock_find.assert_not_called() + mock_add_tags.assert_called_once() + tag_arg = mock_add_tags.call_args[0][0][0] + assert tag_arg["Key"] == "sagemaker.amazonaws.com/model-source" + + @patch("sagemaker.serve.model_builder.Endpoint.get") + @patch("sagemaker.serve.model_builder.find_existing_sagemaker_endpoint") + def test_deploy_skips_reuse_when_config_mismatch(self, mock_find, mock_endpoint_get): + mock_find.return_value = ( + "arn:aws:sagemaker:us-west-2:123456789012:endpoint/existing-ep" + ) + + with ( + patch.object(ModelBuilder, "_resolve_model_source_id", return_value="s3://bucket/model"), + patch.object(ModelBuilder, "_reused_endpoint_matches_config", return_value=False), + patch.object(ModelBuilder, "_is_model_customization", return_value=False), + patch.object(ModelBuilder, "_get_deploy_wrapper") as mock_get_wrapper, + patch.object(ModelBuilder, "add_tags"), + ): + mock_deploy = Mock(return_value=Mock()) + mock_get_wrapper.return_value = mock_deploy + + builder = self._make_builder() + builder.built_model = Mock() + builder.region = "us-west-2" + + builder.deploy(endpoint_name="new-ep", reuse_resources=True) + + # A config mismatch must not return the existing endpoint; a new one is created. + mock_endpoint_get.assert_not_called() + mock_deploy.assert_called_once() + + @patch("sagemaker.serve.model_builder.Endpoint.get") + @patch("sagemaker.serve.model_builder.find_existing_sagemaker_endpoint") + def test_reuse_does_not_intercept_inference_component_deploy( + self, mock_find, mock_endpoint_get + ): + # reuse_resources=True must NOT short-circuit an IC deploy; the IC path + # (inference_config is a ResourceRequirements) manages its own reuse. + from sagemaker.core.compute_resource_requirements.resource_requirements import ( + ResourceRequirements, + ) + + with ( + patch.object(ModelBuilder, "_resolve_model_source_id", return_value="s3://bucket/model"), + patch.object(ModelBuilder, "_is_model_customization", return_value=False), + patch.object(ModelBuilder, "_find_reusable_endpoint") as mock_find_reusable, + patch.object(ModelBuilder, "_deploy") as mock_deploy, + patch.object(ModelBuilder, "add_tags"), + ): + mock_deploy.return_value = Mock() + + builder = self._make_builder() + builder.built_model = Mock() + builder.region = "us-west-2" + + builder.deploy( + endpoint_name="ic-ep", + inference_config=ResourceRequirements( + requests={"num_accelerators": 1, "memory": 1024, "copies": 1} + ), + reuse_resources=True, + ) + + # The reuse gate must be bypassed entirely for IC deployments. + mock_find_reusable.assert_not_called() + mock_endpoint_get.assert_not_called() + mock_deploy.assert_called_once() + + +class TestReusedEndpointMatchesConfig(unittest.TestCase): + """Tests for ModelBuilder._reused_endpoint_matches_config.""" + + def _make_builder(self, **overrides): + session = Mock() + session.boto_session = Mock() + defaults = dict( + model=Mock(), + model_server=ModelServer.TORCHSERVE, + role_arn="arn:aws:iam::123456789012:role/SageMakerRole", + sagemaker_session=session, + ) + defaults.update(overrides) + builder = ModelBuilder(**defaults) + return builder + + def _stub_sagemaker_client(self, builder, env=None, image=None, instance_type=None): + client = Mock() + client.describe_endpoint.return_value = {"EndpointConfigName": "cfg"} + client.describe_endpoint_config.return_value = { + "ProductionVariants": [ + {"ModelName": "m", "InstanceType": instance_type or "ml.g5.xlarge"} + ] + } + client.describe_model.return_value = { + "PrimaryContainer": { + "Environment": env or {}, + "Image": image or "img:1", + } + } + builder.sagemaker_session.sagemaker_client = client + return client + + def test_matches_when_env_and_image_and_instance_match(self): + builder = self._make_builder(env_vars={"A": "1"}, image_uri="img:1") + self._stub_sagemaker_client( + builder, env={"A": "1"}, image="img:1", instance_type="ml.g5.xlarge" + ) + assert builder._reused_endpoint_matches_config("ep", instance_type="ml.g5.xlarge") is True + + def test_matches_nova_model_using_containers_list(self): + # Nova / model-customization models expose config via Containers, not + # PrimaryContainer. The match must read from Containers[0] in that case. + builder = self._make_builder(env_vars={"A": "1"}, image_uri="img:1") + client = Mock() + client.describe_endpoint.return_value = {"EndpointConfigName": "cfg"} + client.describe_endpoint_config.return_value = { + "ProductionVariants": [{"ModelName": "m", "InstanceType": "ml.p4d.24xlarge"}] + } + client.describe_model.return_value = { + "PrimaryContainer": None, + "Containers": [{"Environment": {"A": "1"}, "Image": "img:1"}], + } + builder.sagemaker_session.sagemaker_client = client + assert builder._reused_endpoint_matches_config("ep", instance_type="ml.p4d.24xlarge") is True + + def test_mismatch_on_env_vars(self): + builder = self._make_builder(env_vars={"A": "2"}, image_uri="img:1") + self._stub_sagemaker_client(builder, env={"A": "1"}, image="img:1") + assert builder._reused_endpoint_matches_config("ep") is False + + def test_mismatch_on_image(self): + builder = self._make_builder(env_vars={"A": "1"}, image_uri="img:2") + self._stub_sagemaker_client(builder, env={"A": "1"}, image="img:1") + assert builder._reused_endpoint_matches_config("ep") is False + + def test_mismatch_on_instance_type(self): + builder = self._make_builder(env_vars={"A": "1"}, image_uri="img:1") + self._stub_sagemaker_client( + builder, env={"A": "1"}, image="img:1", instance_type="ml.g5.xlarge" + ) + assert builder._reused_endpoint_matches_config("ep", instance_type="ml.p4d.24xlarge") is False + + def test_matches_when_describe_fails(self): + builder = self._make_builder(env_vars={"A": "1"}) + client = Mock() + client.describe_endpoint.side_effect = Exception("boom") + builder.sagemaker_session.sagemaker_client = client + assert builder._reused_endpoint_matches_config("ep") is True diff --git a/sagemaker-serve/tests/unit/test_model_reuse.py b/sagemaker-serve/tests/unit/test_model_reuse.py new file mode 100644 index 0000000000..46dda673e9 --- /dev/null +++ b/sagemaker-serve/tests/unit/test_model_reuse.py @@ -0,0 +1,297 @@ +# 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 model_reuse.py.""" + +import hashlib +import pytest +from unittest.mock import Mock, patch, call + +from sagemaker.serve.model_reuse import ( + MODEL_SOURCE_TAG_KEY, + normalize_tag_value, + find_existing_bedrock_model, + find_existing_sagemaker_endpoint, + build_source_tag, + check_bedrock_model_status, + check_sagemaker_endpoint_status, + _arn_to_name, +) + + +@pytest.fixture +def boto_session(): + return Mock() + + +@pytest.fixture +def bedrock_client(boto_session): + client = Mock() + boto_session.client.return_value = client + return client + + +@pytest.fixture +def sagemaker_client(boto_session): + client = Mock() + boto_session.client.return_value = client + return client + + +SAMPLE_ARN = "arn:aws:bedrock:us-east-1:123456789012:custom-model/my-model" +ENDPOINT_ARN = "arn:aws:sagemaker:us-east-1:123456789012:endpoint/my-endpoint" + + +def _bedrock_with_tagged_model(bedrock_client, arn, tag_value): + """Configure a bedrock client mock to return one model carrying the source tag.""" + bedrock_client.list_custom_models.return_value = {"modelSummaries": [{"modelArn": arn}]} + bedrock_client.list_tags_for_resource.return_value = { + "tags": [{"key": MODEL_SOURCE_TAG_KEY, "value": tag_value}] + } + + +def _sagemaker_with_tagged_endpoint(sagemaker_client, arn, tag_value): + """Configure a sagemaker client mock to return one endpoint carrying the source tag.""" + sagemaker_client.list_endpoints.return_value = {"Endpoints": [{"EndpointArn": arn}]} + sagemaker_client.list_tags.return_value = { + "Tags": [{"Key": MODEL_SOURCE_TAG_KEY, "Value": tag_value}] + } + + +@pytest.mark.parametrize( + "length", + [0, 100, 256], + ids=["empty", "short", "at_limit"], +) +def test_normalize_tag_value_within_limit(length): + value = "a" * length + assert normalize_tag_value(value) == value + + +@pytest.mark.parametrize( + "length", + [257, 512], + ids=["one_over", "long"], +) +def test_normalize_tag_value_exceeds_limit(length): + value = "x" * length + result = normalize_tag_value(value) + assert len(result) == 256 + assert result[:224] == value[:224] + assert result[224] == "-" + + +def test_normalize_tag_value_sha256_suffix(): + value = "s3://my-bucket/" + "a" * 300 + result = normalize_tag_value(value) + expected_hash = hashlib.sha256(value.encode()).hexdigest()[:31] + assert result.endswith(expected_hash) + assert result == f"{value[:224]}-{expected_hash}" + + +def test_find_existing_bedrock_model_returns_arn_when_active(boto_session, bedrock_client): + _bedrock_with_tagged_model(bedrock_client, SAMPLE_ARN, "source-id") + bedrock_client.get_custom_model.return_value = {"modelStatus": "Active"} + + result = find_existing_bedrock_model(bedrock_client, "source-id") + + assert result == SAMPLE_ARN + bedrock_client.list_tags_for_resource.assert_called_once_with(resourceARN=SAMPLE_ARN) + + +def test_find_existing_bedrock_model_paginates(boto_session, bedrock_client): + other_arn = "arn:aws:bedrock:us-east-1:123456789012:custom-model/other" + bedrock_client.list_custom_models.side_effect = [ + {"modelSummaries": [{"modelArn": other_arn}], "nextToken": "page-2"}, + {"modelSummaries": [{"modelArn": SAMPLE_ARN}]}, + ] + bedrock_client.list_tags_for_resource.side_effect = [ + {"tags": [{"key": MODEL_SOURCE_TAG_KEY, "value": "different"}]}, + {"tags": [{"key": MODEL_SOURCE_TAG_KEY, "value": "source-id"}]}, + ] + bedrock_client.get_custom_model.return_value = {"modelStatus": "Active"} + + result = find_existing_bedrock_model(bedrock_client, "source-id") + + assert result == SAMPLE_ARN + assert bedrock_client.list_custom_models.call_count == 2 + + +@patch("sagemaker.serve.model_reuse.time.sleep") +def test_find_existing_bedrock_model_polls_creating_until_ready(mock_sleep, boto_session, bedrock_client): + _bedrock_with_tagged_model(bedrock_client, SAMPLE_ARN, "source-id") + bedrock_client.get_custom_model.side_effect = [ + {"modelStatus": "Creating"}, + {"modelStatus": "Creating"}, + {"modelStatus": "Active"}, + ] + + result = find_existing_bedrock_model( + bedrock_client, "source-id", poll_interval=5, max_wait=900 + ) + + assert result == SAMPLE_ARN + assert mock_sleep.call_count == 2 + mock_sleep.assert_called_with(5) + + +@patch("sagemaker.serve.model_reuse.time.sleep") +def test_find_existing_bedrock_model_raises_timeout_on_creating(mock_sleep, boto_session, bedrock_client): + _bedrock_with_tagged_model(bedrock_client, SAMPLE_ARN, "source-id") + bedrock_client.get_custom_model.return_value = {"modelStatus": "Creating"} + + with pytest.raises(TimeoutError, match="did not become ready"): + find_existing_bedrock_model( + bedrock_client, "source-id", poll_interval=5, max_wait=10 + ) + + +def test_find_existing_bedrock_model_returns_none_on_failed(boto_session, bedrock_client): + _bedrock_with_tagged_model(bedrock_client, SAMPLE_ARN, "source-id") + bedrock_client.get_custom_model.return_value = {"modelStatus": "Failed"} + + result = find_existing_bedrock_model(bedrock_client, "source-id") + + assert result is None + + +def test_find_existing_bedrock_model_returns_none_on_list_failure(boto_session, bedrock_client): + bedrock_client.list_custom_models.side_effect = Exception("Access denied") + + result = find_existing_bedrock_model(bedrock_client, "source-id") + + assert result is None + + +def test_find_existing_bedrock_model_returns_none_when_no_match(boto_session, bedrock_client): + bedrock_client.list_custom_models.return_value = { + "modelSummaries": [{"modelArn": SAMPLE_ARN}] + } + bedrock_client.list_tags_for_resource.return_value = { + "tags": [{"key": MODEL_SOURCE_TAG_KEY, "value": "different"}] + } + + result = find_existing_bedrock_model(bedrock_client, "source-id") + + assert result is None + + +def test_find_existing_bedrock_model_returns_none_when_no_models(boto_session, bedrock_client): + bedrock_client.list_custom_models.return_value = {"modelSummaries": []} + + result = find_existing_bedrock_model(bedrock_client, "source-id") + + assert result is None + + +def test_find_existing_sagemaker_endpoint_returns_arn_when_in_service(boto_session, sagemaker_client): + _sagemaker_with_tagged_endpoint(sagemaker_client, ENDPOINT_ARN, "source-id") + sagemaker_client.describe_endpoint.return_value = {"EndpointStatus": "InService"} + + result = find_existing_sagemaker_endpoint(sagemaker_client, "source-id") + + assert result == ENDPOINT_ARN + sagemaker_client.list_tags.assert_called_once_with(ResourceArn=ENDPOINT_ARN) + + +def test_find_existing_sagemaker_endpoint_paginates(boto_session, sagemaker_client): + other_arn = "arn:aws:sagemaker:us-east-1:123456789012:endpoint/other" + sagemaker_client.list_endpoints.side_effect = [ + {"Endpoints": [{"EndpointArn": other_arn}], "NextToken": "page-2"}, + {"Endpoints": [{"EndpointArn": ENDPOINT_ARN}]}, + ] + sagemaker_client.list_tags.side_effect = [ + {"Tags": [{"Key": MODEL_SOURCE_TAG_KEY, "Value": "different"}]}, + {"Tags": [{"Key": MODEL_SOURCE_TAG_KEY, "Value": "source-id"}]}, + ] + sagemaker_client.describe_endpoint.return_value = {"EndpointStatus": "InService"} + + result = find_existing_sagemaker_endpoint(sagemaker_client, "source-id") + + assert result == ENDPOINT_ARN + assert sagemaker_client.list_endpoints.call_count == 2 + + +def test_find_existing_sagemaker_endpoint_returns_none_on_failed(boto_session, sagemaker_client): + _sagemaker_with_tagged_endpoint(sagemaker_client, ENDPOINT_ARN, "source-id") + sagemaker_client.describe_endpoint.return_value = {"EndpointStatus": "Failed"} + + result = find_existing_sagemaker_endpoint(sagemaker_client, "source-id") + + assert result is None + + +def test_find_existing_sagemaker_endpoint_returns_none_on_list_failure(boto_session, sagemaker_client): + sagemaker_client.list_endpoints.side_effect = Exception("Access denied") + + result = find_existing_sagemaker_endpoint(sagemaker_client, "source-id") + + assert result is None + + +def test_find_existing_sagemaker_endpoint_returns_none_when_no_endpoints(boto_session, sagemaker_client): + sagemaker_client.list_endpoints.return_value = {"Endpoints": []} + + result = find_existing_sagemaker_endpoint(sagemaker_client, "source-id") + + assert result is None + + +def test_build_source_tag_returns_correct_dict(): + source_id = "s3://bucket/path/to/model" + tag = build_source_tag(source_id) + + assert tag == {"key": MODEL_SOURCE_TAG_KEY, "value": source_id} + + +def test_build_source_tag_normalizes_long_value(): + source_id = "s3://bucket/" + "a" * 300 + tag = build_source_tag(source_id) + + assert tag["key"] == MODEL_SOURCE_TAG_KEY + assert len(tag["value"]) == 256 + + +def test_check_bedrock_model_status_returns_model_status(): + bedrock_client = Mock() + bedrock_client.get_custom_model.return_value = {"modelStatus": "Active"} + + result = check_bedrock_model_status(bedrock_client, SAMPLE_ARN) + + assert result == "Active" + bedrock_client.get_custom_model.assert_called_once_with(modelIdentifier=SAMPLE_ARN) + + +def test_check_bedrock_model_status_raises_on_failure(): + bedrock_client = Mock() + bedrock_client.get_custom_model.side_effect = Exception("Not found") + + with pytest.raises(Exception, match="Not found"): + check_bedrock_model_status(bedrock_client, SAMPLE_ARN) + + +def test_check_sagemaker_endpoint_status_returns_endpoint_status(): + sm_client = Mock() + sm_client.describe_endpoint.return_value = {"EndpointStatus": "InService"} + + result = check_sagemaker_endpoint_status(sm_client, ENDPOINT_ARN) + + assert result == "InService" + sm_client.describe_endpoint.assert_called_once_with(EndpointName="my-endpoint") + + +def test_check_sagemaker_endpoint_status_raises_on_failure(): + sm_client = Mock() + sm_client.describe_endpoint.side_effect = Exception("Endpoint not found") + + with pytest.raises(Exception, match="Endpoint not found"): + check_sagemaker_endpoint_status(sm_client, ENDPOINT_ARN) diff --git a/v3-examples/model-customization-examples/bedrock-modelbuilder-deployment.ipynb b/v3-examples/model-customization-examples/bedrock-modelbuilder-deployment.ipynb index b0340bf086..9598dd53e0 100644 --- a/v3-examples/model-customization-examples/bedrock-modelbuilder-deployment.ipynb +++ b/v3-examples/model-customization-examples/bedrock-modelbuilder-deployment.ipynb @@ -268,6 +268,38 @@ "output = json.loads(response[\"body\"].read().decode())\n", "print(output)" ] + }, + { + "cell_type": "markdown", + "metadata": {}, + "source": [ + "## Reusing an Existing Custom Model\n", + "\n", + "Importing a custom model into Bedrock is slow and consumes the limited imported-model quota per account/region. If you redeploy the *same* model artifacts, you can avoid a duplicate import by opting into resource reuse with `reuse_resources=True`.\n", + "\n", + "When enabled, `BedrockModelBuilder` tags each custom model it creates with the model source and, on a later deploy of the same source, reuses the existing custom model (and its active deployment if one exists) instead of creating duplicates. It logs a warning rather than raising, and the returned response includes the `modelArn` that was reused.\n", + "\n", + "Reuse is off by default (`reuse_resources=False`)." + ] + }, + { + "cell_type": "code", + "execution_count": null, + "metadata": {}, + "outputs": [], + "source": [ + "# Reuse an existing custom model built from the same source, if one exists.\n", + "reuse_result = builder.deploy(\n", + " custom_model_name=\"my-nova-model\",\n", + " role_arn=ROLE_ARN,\n", + " reuse_resources=True,\n", + ")\n", + "\n", + "# On a reuse hit you'll see a warning noting no new resources were created, and:\n", + "# reuse_result[\"modelArn\"] -> the reused custom model\n", + "# reuse_result[\"customModelDeploymentArn\"] -> the reused/created deployment\n", + "print(reuse_result)" + ] } ], "metadata": { diff --git a/v3-examples/model-customization-examples/model_builder_deployment_notebook.ipynb b/v3-examples/model-customization-examples/model_builder_deployment_notebook.ipynb index 7b9641fc93..2bba6c7130 100644 --- a/v3-examples/model-customization-examples/model_builder_deployment_notebook.ipynb +++ b/v3-examples/model-customization-examples/model_builder_deployment_notebook.ipynb @@ -2,8 +2,10 @@ "cells": [ { "cell_type": "code", + "execution_count": null, "id": "777b47454f7d860b_setup", "metadata": {}, + "outputs": [], "source": [ "from pprint import pprint\n", "\n", @@ -11,14 +13,14 @@ "from sagemaker.core.utils.utils import Unassigned\n", "! ada credentials update --provider=isengard --account=<> --role=Admin --profile=default --once\n", "! aws configure set region us-west-2" - ], - "outputs": [], - "execution_count": null + ] }, { "cell_type": "code", + "execution_count": null, "id": "da22762d06751e9b", "metadata": {}, + "outputs": [], "source": [ "from sagemaker.core.resources import Endpoint\n", "\n", @@ -26,14 +28,14 @@ "for endpoint in Endpoint.get_all():\n", " if endpoint.endpoint_name.startswith('e2e-'):\n", " endpoint.delete()\n" - ], - "outputs": [], - "execution_count": null + ] }, { "cell_type": "code", + "execution_count": null, "id": "95367703", "metadata": {}, + "outputs": [], "source": [ "from sagemaker.core.resources import TrainingJob, HubContent, InferenceComponent, ModelPackage\n", "from sagemaker.core.utils.utils import Unassigned\n", @@ -48,13 +50,14 @@ " print(model_package.inference_specification.containers[0].image)\n", " except:\n", " pass\n" - ], - "outputs": [], - "execution_count": null + ] }, { - "metadata": {}, "cell_type": "code", + "execution_count": null, + "id": "2415b1cb715a304c", + "metadata": {}, + "outputs": [], "source": [ "from sagemaker.core.resources import TrainingJob\n", "import random\n", @@ -67,22 +70,24 @@ "print(model.model_arn)\n", "import random\n", "#endpoint = model_builder.deploy(endpoint_name=name)" - ], - "id": "2415b1cb715a304c", - "outputs": [], - "execution_count": null + ] }, { - "metadata": {}, "cell_type": "code", - "source": "endpoint = model_builder.deploy(endpoint_name=name)", + "execution_count": null, "id": "8b8bc9eb4299ecba", + "metadata": {}, "outputs": [], - "execution_count": null + "source": [ + "endpoint = model_builder.deploy(endpoint_name=name)" + ] }, { - "metadata": {}, "cell_type": "code", + "execution_count": null, + "id": "58b5d5995791bd96", + "metadata": {}, + "outputs": [], "source": [ "from sagemaker.core.resources import InferenceComponent, Tag\n", "from pprint import pprint\n", @@ -92,50 +97,47 @@ " for tag in Tag.get_all(resource_arn=inference_component.inference_component_arn):\n", " pprint(tag)\n", "\n" - ], - "id": "58b5d5995791bd96", - "outputs": [], - "execution_count": null + ] }, { "cell_type": "code", + "execution_count": null, "id": "2833eab06285f075", "metadata": {}, + "outputs": [], "source": [ "import json\n", "# Note this is expected to fail since Endpoint invoke is only available for authorized users. The Invoke call here is the sagemaker-core Endpoint.invoke call .\n", "print(endpoint.endpoint_arn)\n", "endpoint.invoke(body=json.dumps({\"inputs\": \"What is the capital of France?\", \"parameters\": {\"max_new_tokens\": 50}}))" - ], - "outputs": [], - "execution_count": null + ] }, { "cell_type": "code", + "execution_count": null, "id": "695a83cf38e46cea", "metadata": {}, + "outputs": [], "source": [ "from sagemaker.core.resources import TrainingJob\n", "from sagemaker.serve import ModelBuilder\n", "\n", "model_builder = ModelBuilder(model=TrainingJob.get(training_job_name=\"meta-textgeneration-llama-3-2-1b-instruct-sft-20251123162832\"))\n", "model_builder.fetch_endpoint_names_for_base_model()" - ], - "outputs": [], - "execution_count": null + ] }, { "cell_type": "code", + "execution_count": null, "id": "92e0da7904ffb743", "metadata": {}, + "outputs": [], "source": [ "name = f\"e2e-{random.randint(100, 10000)}\"\n", "model_builder.name = name\n", "endpoint = model_builder.deploy(endpoint_name=name, inference_component_name=f\"{name}-adapter\")\n", "sda" - ], - "outputs": [], - "execution_count": null + ] }, { "cell_type": "markdown", @@ -165,24 +167,130 @@ { "cell_type": "markdown", "id": "a9aa326f", - "source": "## Configuring CONTEXT_LENGTH and MAX_CONCURRENCY for Nova Models\n\nWhen deploying Nova models to SageMaker endpoints, you can configure inference performance\nby setting `CONTEXT_LENGTH` and `MAX_CONCURRENCY` via the `env_vars` parameter. These control\nthe maximum input context window and concurrent request capacity of the endpoint.\n\n**Supported configurations vary by model and instance type.** Each (model, instance) combination\nsupports specific tier bounds — higher context lengths reduce the maximum concurrency allowed.\n\nFor example, `nova-textgeneration-micro` on `ml.p5.48xlarge` supports:\n| CONTEXT_LENGTH | MAX_CONCURRENCY |\n|---|---|\n| ≤ 16,000 | up to 128 |\n| ≤ 64,000 | up to 32 |\n| ≤ 128,000 | up to 8 |\n\nIf invalid values are provided, `ModelBuilder` will raise a `ValueError` at build time\nwith a clear message indicating the supported limits — rather than failing at container startup.", - "metadata": {} + "metadata": {}, + "source": [ + "## Configuring CONTEXT_LENGTH and MAX_CONCURRENCY for Nova Models\n", + "\n", + "When deploying Nova models to SageMaker endpoints, you can configure inference performance\n", + "by setting `CONTEXT_LENGTH` and `MAX_CONCURRENCY` via the `env_vars` parameter. These control\n", + "the maximum input context window and concurrent request capacity of the endpoint.\n", + "\n", + "**Supported configurations vary by model and instance type.** Each (model, instance) combination\n", + "supports specific tier bounds — higher context lengths reduce the maximum concurrency allowed.\n", + "\n", + "For example, `nova-textgeneration-micro` on `ml.p5.48xlarge` supports:\n", + "| CONTEXT_LENGTH | MAX_CONCURRENCY |\n", + "|---|---|\n", + "| ≤ 16,000 | up to 128 |\n", + "| ≤ 64,000 | up to 32 |\n", + "| ≤ 128,000 | up to 8 |\n", + "\n", + "If invalid values are provided, `ModelBuilder` will raise a `ValueError` at build time\n", + "with a clear message indicating the supported limits — rather than failing at container startup." + ] }, { "cell_type": "code", + "execution_count": null, "id": "f743a286", - "source": "# Deploy a Nova model with custom CONTEXT_LENGTH and MAX_CONCURRENCY\nfrom sagemaker.core.resources import TrainingJob\nfrom sagemaker.serve import ModelBuilder\n\ntraining_job = TrainingJob.get(training_job_name=\"\")\n\n# Configure for high-throughput: lower context, higher concurrency\nmodel_builder = ModelBuilder(\n model=training_job,\n role_arn=\"arn:aws:iam:::role/\",\n instance_type=\"ml.p5.48xlarge\",\n env_vars={\n \"CONTEXT_LENGTH\": \"16000\",\n \"MAX_CONCURRENCY\": \"128\",\n },\n)\n\nmodel = model_builder.build(model_name=\"nova-high-throughput\")\nendpoint = model_builder.deploy(endpoint_name=\"nova-high-throughput\")", "metadata": {}, - "execution_count": null, - "outputs": [] + "outputs": [], + "source": [ + "# Deploy a Nova model with custom CONTEXT_LENGTH and MAX_CONCURRENCY\n", + "from sagemaker.core.resources import TrainingJob\n", + "from sagemaker.serve import ModelBuilder\n", + "\n", + "training_job = TrainingJob.get(training_job_name=\"\")\n", + "\n", + "# Configure for high-throughput: lower context, higher concurrency\n", + "model_builder = ModelBuilder(\n", + " model=training_job,\n", + " role_arn=\"arn:aws:iam:::role/\",\n", + " instance_type=\"ml.p5.48xlarge\",\n", + " env_vars={\n", + " \"CONTEXT_LENGTH\": \"16000\",\n", + " \"MAX_CONCURRENCY\": \"128\",\n", + " },\n", + ")\n", + "\n", + "model = model_builder.build(model_name=\"nova-high-throughput\")\n", + "endpoint = model_builder.deploy(endpoint_name=\"nova-high-throughput\")" + ] }, { "cell_type": "code", + "execution_count": null, "id": "910ab04f", - "source": "# Example: what happens with invalid configuration\n# This will raise a ValueError before deployment starts:\n#\n# model_builder = ModelBuilder(\n# model=training_job,\n# instance_type=\"ml.p5.48xlarge\",\n# env_vars={\n# \"CONTEXT_LENGTH\": \"64000\",\n# \"MAX_CONCURRENCY\": \"50\", # Exceeds limit of 32 at this context length\n# },\n# )\n# model_builder.build()\n# >>> ValueError: MAX_CONCURRENCY=50 exceeds maximum supported value of 32\n# >>> for 'nova-textgeneration-micro' on ml.p5.48xlarge at CONTEXT_LENGTH<=64000.", "metadata": {}, + "outputs": [], + "source": [ + "# Example: what happens with invalid configuration\n", + "# This will raise a ValueError before deployment starts:\n", + "#\n", + "# model_builder = ModelBuilder(\n", + "# model=training_job,\n", + "# instance_type=\"ml.p5.48xlarge\",\n", + "# env_vars={\n", + "# \"CONTEXT_LENGTH\": \"64000\",\n", + "# \"MAX_CONCURRENCY\": \"50\", # Exceeds limit of 32 at this context length\n", + "# },\n", + "# )\n", + "# model_builder.build()\n", + "# >>> ValueError: MAX_CONCURRENCY=50 exceeds maximum supported value of 32\n", + "# >>> for 'nova-textgeneration-micro' on ml.p5.48xlarge at CONTEXT_LENGTH<=64000." + ] + }, + { + "cell_type": "markdown", + "id": "reuse_resources_md", + "metadata": {}, + "source": [ + "## Reusing an Existing Endpoint\n", + "\n", + "If you redeploy the *same* model artifacts (e.g. re-running this notebook, CI, or sharing a checkpoint), you can opt into resource reuse with `reuse_resources=True`.\n", + "\n", + "When enabled, `ModelBuilder` tags each endpoint it creates with the model source and, on a later deploy of the same source, discovers and returns the existing endpoint instead of creating a duplicate. It logs a warning (it does not raise) and returns the existing endpoint as the normal return value.\n", + "\n", + "**Where to pass the flag** (it is honored per call, not inherited — set it on each call you want to reuse):\n", + "- `build(reuse_resources=True)`: on a hit, creates no new Model/EndpointConfig/Endpoint and sets `built_model` to the existing Model backing the reused endpoint.\n", + "- `deploy(reuse_resources=True)`: on a hit, returns the existing endpoint instead of creating a new one.\n", + "- Pass it to **both** `build()` and `deploy()` to reuse end to end and create nothing new.\n", + "\n", + "Reuse is off by default (`reuse_resources=False`). A reused endpoint is only returned when its deployment configuration (env vars, image URI, instance type) matches the request; otherwise a new endpoint is created." + ] + }, + { + "cell_type": "code", "execution_count": null, - "outputs": [] + "id": "reuse_resources_code", + "metadata": {}, + "outputs": [], + "source": [ + "# Reuse an existing endpoint built from the same model source, if one exists.\n", + "from sagemaker.core.resources import TrainingJob\n", + "from sagemaker.serve import ModelBuilder\n", + "\n", + "training_job = TrainingJob.get(training_job_name=\"\")\n", + "\n", + "model_builder = ModelBuilder(\n", + " model=training_job,\n", + " role_arn=\"arn:aws:iam:::role/\",\n", + " instance_type=\"ml.p5.48xlarge\",\n", + ")\n", + "\n", + "# Pass reuse_resources=True to build() so no new resources are created on a hit.\n", + "model_builder.build(model_name=\"reuse-demo\", reuse_resources=True)\n", + "endpoint = model_builder.deploy(endpoint_name=\"reuse-demo\", reuse_resources=True)\n", + "\n", + "# On a reuse hit you'll see a warning like:\n", + "# Reusing existing endpoint 'reuse-demo' (matched model-source tag and\n", + "# deployment configuration). No new resources will be created.\n", + "result = endpoint.invoke(\n", + " body='{\"messages\": [{\"role\": \"user\", \"content\": [{\"type\": \"text\", \"text\": \"Hello\"}]}]}',\n", + " content_type=\"application/json\",\n", + " accept=\"application/json\",\n", + ")" + ] }, { "cell_type": "markdown", @@ -231,8 +339,10 @@ }, { "cell_type": "code", + "execution_count": null, "id": "778be153d0a87d13", "metadata": {}, + "outputs": [], "source": [ "import random\n", "from sagemaker.serve import ModelBuilder\n", @@ -243,9 +353,7 @@ "model_package = ModelPackage.get(model_package_name=\"arn:aws:sagemaker:us-west-2:<>:model-package/test-finetuned-models-gamma/68\")\n", "model_builder = ModelBuilder(model=model_package)\n", "model_builder.build()" - ], - "outputs": [], - "execution_count": null + ] }, { "cell_type": "markdown", @@ -273,11 +381,13 @@ }, { "cell_type": "code", + "execution_count": null, "id": "ef3384c868dd58d5", "metadata": {}, - "source": "endpoint = model_builder.deploy( endpoint_name=name)\n", "outputs": [], - "execution_count": null + "source": [ + "endpoint = model_builder.deploy( endpoint_name=name)\n" + ] }, { "cell_type": "markdown", @@ -289,8 +399,10 @@ }, { "cell_type": "code", + "execution_count": null, "id": "d17136303e9b7c9e", "metadata": {}, + "outputs": [], "source": [ "import boto3\n", "import json\n", @@ -332,14 +444,14 @@ ")\n", "\n", "print(\"config.json uploaded successfully\")\n" - ], - "outputs": [], - "execution_count": null + ] }, { "cell_type": "code", + "execution_count": null, "id": "865777e899016a07", "metadata": {}, + "outputs": [], "source": [ "import boto3\n", "import json\n", @@ -347,24 +459,24 @@ "s3 = boto3.client('s3', region_name='us-west-2')\n", "config = {\"add_bos_token\": True, \"add_eos_token\": False, \"bos_token\": \"<|begin_of_text|>\", \"eos_token\": \"<|end_of_text|>\", \"pad_token\": \"<|end_of_text|>\", \"model_max_length\": 131072, \"tokenizer_class\": \"LlamaTokenizer\"}\n", "s3.put_object(Bucket=\"open-models-testing-pdx\", Key=\"output/meta-textgeneration-llama-3-2-1b-instruct-sft-20251114104310/output/model/tokenizer_config.json\", Body=json.dumps(config))\n" - ], - "outputs": [], - "execution_count": null + ] }, { "cell_type": "code", + "execution_count": null, "id": "533d0f1022d169eb", "metadata": {}, + "outputs": [], "source": [ "! ada credentials update --provider=isengard --account=<> --role=Admin --profile=default --once\n" - ], - "outputs": [], - "execution_count": null + ] }, { "cell_type": "code", + "execution_count": null, "id": "798f5b8668305f43", "metadata": {}, + "outputs": [], "source": [ "from sagemaker.core.resources import TrainingJob\n", "import random\n", @@ -375,14 +487,14 @@ "\n", "# bedrock_builder = BedrockModelBuilder(model=training_job)\n", "# bedrock_builder.deploy(job_name=name, imported_model_name=name, role_arn=\"arn:aws:iam::<>:role/Admin\")" - ], - "outputs": [], - "execution_count": null + ] }, { "cell_type": "code", + "execution_count": null, "id": "6fdd61406713c8c9", "metadata": {}, + "outputs": [], "source": [ "# Assuming you previously did something like:\n", "# bedrock_builder = BedrockModelBuilder(model_trainer)\n", @@ -400,9 +512,7 @@ " }\n", " })\n", ")\n" - ], - "outputs": [], - "execution_count": null + ] }, { "cell_type": "markdown", @@ -438,8 +548,10 @@ }, { "cell_type": "code", + "execution_count": null, "id": "aefa0ec7cd360d5c", "metadata": {}, + "outputs": [], "source": [ "import boto3\n", "\n", @@ -460,9 +572,7 @@ " model_arn = model['modelArn']\n", " print(f\"Deleting imported model: {model_arn}\")\n", " bedrock.delete_imported_model(modelIdentifier=model_arn)\n" - ], - "outputs": [], - "execution_count": null + ] } ], "metadata": { @@ -486,4 +596,4 @@ }, "nbformat": 4, "nbformat_minor": 5 -} \ No newline at end of file +} From 075863019331dd03840e5c15206cc1da80587975 Mon Sep 17 00:00:00 2001 From: Lisa Ni Date: Tue, 7 Jul 2026 22:09:34 +0000 Subject: [PATCH 2/2] feat(serve): support Nova inference-component deployment and harden reuse Route Nova model-customization deploys through the shared single-inference -component path when a ResourceRequirements inference_config is supplied, so each Nova checkpoint (full-rank or LoRA-merged) is hosted as one inference component referencing the built Model. Nova without an inference_config keeps the direct model-on-variant path. - Broaden _is_nova_model to identify Nova from a package-less source (raw S3 checkpoint or trainer) via base_model_name, in addition to the model package recipe/hub-content name. - Set EnableNetworkIsolation on the IC endpoint config to match the built Model (always True for Nova), fixing CreateInferenceComponent rejection on mismatched network isolation. - Guard model-package-dependent logic (restricted-package path, PEFT/recipe metadata, lineage tracking) so package-less Nova checkpoints deploy cleanly. - Apply accumulated tags (including the model-source reuse tag) to endpoints created on the shared IC path so they remain discoverable. - build(reuse_resources=True) only short-circuits when the backing Model can be resolved; IC endpoints and stale/deleted configs fall through and build a real Model, preventing a None built_model on later IC deploys. - deploy() warns that reuse_resources has no effect for inference-component deployments, which manage their own reuse by component name. - Surface both the manifest.json and output.tar.gz errors when Nova checkpoint URI resolution fails, instead of masking the primary failure. Add unit tests covering the Nova IC path (routing, network isolation, IC spec) and the model-on-variant fallback. --- .../src/sagemaker/core/training/utils.py | 16 +- .../src/sagemaker/serve/model_builder.py | 153 ++++++++++++------ .../tests/unit/test_model_builder.py | 84 +++++++++- 3 files changed, 203 insertions(+), 50 deletions(-) diff --git a/sagemaker-core/src/sagemaker/core/training/utils.py b/sagemaker-core/src/sagemaker/core/training/utils.py index 2ee2afdee7..061cf4c746 100644 --- a/sagemaker-core/src/sagemaker/core/training/utils.py +++ b/sagemaker-core/src/sagemaker/core/training/utils.py @@ -208,6 +208,18 @@ def resolve_nova_checkpoint_uri( manifest_uri = build_nova_manifest_s3_uri(s3_output_path, training_job_name) try: return read_nova_checkpoint_uri_from_manifest(s3_client, manifest_uri) - except Exception: + except Exception as manifest_error: + # The raw manifest.json is the primary source (e.g. HyperPod). Only fall + # back to the copy inside output.tar.gz (e.g. some SMTJ jobs) if it exists. + # If the fallback also fails, surface BOTH errors so the raw-manifest + # failure isn't masked by a misleading "not found in tar.gz" message. tar_gz_uri = build_nova_output_tar_gz_s3_uri(s3_output_path, training_job_name) - return _read_checkpoint_uri_from_tar_gz(s3_client, tar_gz_uri) + try: + return _read_checkpoint_uri_from_tar_gz(s3_client, tar_gz_uri) + except Exception as tar_gz_error: + raise ValueError( + "Could not resolve the Nova checkpoint URI. " + f"Reading manifest.json at {manifest_uri} failed: {manifest_error}. " + f"Falling back to output.tar.gz at {tar_gz_uri} also failed: " + f"{tar_gz_error}." + ) from manifest_error diff --git a/sagemaker-serve/src/sagemaker/serve/model_builder.py b/sagemaker-serve/src/sagemaker/serve/model_builder.py index 1ad0a9d96d..66da1fdda0 100644 --- a/sagemaker-serve/src/sagemaker/serve/model_builder.py +++ b/sagemaker-serve/src/sagemaker/serve/model_builder.py @@ -45,6 +45,10 @@ ModelLifeCycle, DriftCheckBaselines, InferenceComponentComputeResourceRequirements, + InferenceComponentSpecification, + InferenceComponentContainerSpecification, + InferenceComponentRuntimeConfig, + ProductionVariant, ) from sagemaker.core.resources import ( ModelPackage, @@ -1127,22 +1131,28 @@ def _is_raw_s3_model(self) -> bool: return isinstance(self.model, str) and self.model.startswith("s3://") def _is_nova_model(self) -> bool: - """Check if the model is a Nova model based on recipe name or hub content name.""" + """Check if the model is a Nova model. + + Recognizes Nova from the model package's recipe/hub-content name, and also + from a package-less source (raw S3 checkpoint or trainer) via the resolved + ``base_model_name``. All supported Nova checkpoints — full-rank custom + models and LoRA-merged models — are identified here. + """ model_package = self._fetch_model_package() - if not model_package: - # No model package (e.g. a trainer built from an S3 checkpoint). Fall - # back to the base model name carried by the trainer. - base_model_name = self._base_model_name() - return bool(base_model_name) and "nova" in base_model_name.lower() - containers = getattr(model_package.inference_specification, "containers", None) - if not containers: - return False - base_model = getattr(containers[0], "base_model", None) - if not base_model: - return False - recipe_name = getattr(base_model, "recipe_name", "") or "" - hub_content_name = getattr(base_model, "hub_content_name", "") or "" - return "nova" in recipe_name.lower() or "nova" in hub_content_name.lower() + if model_package: + containers = getattr(model_package.inference_specification, "containers", None) + if containers: + base_model = getattr(containers[0], "base_model", None) + if base_model: + recipe_name = getattr(base_model, "recipe_name", "") or "" + hub_content_name = getattr(base_model, "hub_content_name", "") or "" + if "nova" in recipe_name.lower() or "nova" in hub_content_name.lower(): + return True + + # No (or non-Nova) model package: fall back to the base model name carried + # by a trainer or supplied via model_metadata for a raw S3 checkpoint. + base_model_name = self._base_model_name() + return bool(base_model_name) and "nova" in base_model_name.lower() def _is_nova_model_for_telemetry(self) -> bool: """Check if the model is a Nova model for telemetry tracking.""" @@ -3888,17 +3898,24 @@ def build( self.serve_settings = self._get_serve_setting() reusable_endpoint = self._find_reusable_endpoint() if reusable_endpoint: - logger.warning( - "Reusing existing endpoint %r (matched model-source tag and " - "deployment configuration). No new resources will be created. " - "Pass reuse_resources=False to force new resources.", - reusable_endpoint, - ) - self._reused_endpoint_name = reusable_endpoint - # Point built_model at the existing Model backing the reused - # endpoint, so deploy() has a valid built_model to return. - self.built_model = self._get_model_for_endpoint(reusable_endpoint) - return self.built_model + # Only short-circuit if the backing Model can be resolved. An + # inference-component endpoint has no model on its production + # variant (the model lives in the IC), and a stale endpoint whose + # config was deleted also cannot be resolved. In those cases fall + # through and build the Model normally — an IC deploy references + # built_model.model_name and must have a real Model to deploy. + reused_model = self._get_model_for_endpoint(reusable_endpoint) + if reused_model is not None: + logger.warning( + "Reusing existing endpoint %r (matched model-source tag " + "and deployment configuration). No new resources will be " + "created. Pass reuse_resources=False to force new " + "resources.", + reusable_endpoint, + ) + self._reused_endpoint_name = reusable_endpoint + self.built_model = reused_model + return self.built_model deployables = {} @@ -4776,6 +4793,14 @@ def deploy( inference_config, ResourceRequirements ) or bool(getattr(self, "_deployables", None)) + if reuse_resources and is_inference_component_deploy: + logger.warning( + "reuse_resources has no effect for inference component " + "deployments; inference components manage their own reuse by " + "name (a new component is added to the endpoint, or an existing " + "one is updated in place)." + ) + # Resource reuse is opt-in per call. When requested, reuse an existing # endpoint built from the same model source (with matching config). If # build() already resolved one (build's reuse gate), use that; otherwise @@ -4831,6 +4856,7 @@ def deploy( endpoint_name=endpoint_name, instance_type=instance_type or self.instance_type, initial_instance_count=initial_instance_count, + inference_component_name=kwargs.pop("inference_component_name", None), wait=wait, container_timeout_in_seconds=container_timeout_in_seconds, inference_config=inference_config_param, @@ -4992,30 +5018,36 @@ def _deploy_model_customization( Returns: Endpoint: The deployed sagemaker.core.resources.Endpoint """ - from sagemaker.core.shapes import ( - InferenceComponentSpecification, - InferenceComponentContainerSpecification, - InferenceComponentRuntimeConfig, - InferenceComponentComputeResourceRequirements, - ) - from sagemaker.core.shapes import ProductionVariant from sagemaker.core.resources import InferenceComponent from sagemaker.core.resources import Tag as CoreTag - # Nova models use direct model-on-variant, no InferenceComponents - if self._is_nova_model(): + # An inference_config of ResourceRequirements requests an inference + # component deployment; otherwise the model is placed directly on the + # production variant. + is_ic_deploy = isinstance(inference_config, ResourceRequirements) + + # Nova models without IC resources use the direct model-on-variant path. + # Nova models WITH a ResourceRequirements inference_config fall through to + # the shared single-IC path below: each Nova checkpoint is hosted as one + # inference component referencing the built Model, which carries the + # image, escrow artifacts, and env. + is_nova = self._is_nova_model() + if is_nova and not is_ic_deploy: return self._deploy_nova_model( endpoint_name=endpoint_name, initial_instance_count=initial_instance_count, wait=kwargs.get("wait", True), ) - # Fetch model package + # The model package may be absent (e.g. a Nova CPTTrainer or raw-S3 + # checkpoint), restricted (e.g. a Nova MTRL Serverless job), or a normal + # package (e.g. a Nova SFTTrainer serverless job). model_package = self._fetch_model_package() - # Restricted model packages: simple endpoint deployment + # Restricted model packages deploy model-on-variant, but only when an + # inference component was not explicitly requested. from sagemaker.serve.utils.model_package_utils import is_restricted_model_package - if is_restricted_model_package(model_package): + if not is_ic_deploy and is_restricted_model_package(model_package): if not endpoint_name: endpoint_name = f"endpoint-{uuid.uuid4().hex[:8]}" EndpointConfig.create( @@ -5036,6 +5068,19 @@ def _deploy_model_customization( endpoint.wait_for_status("InService") return endpoint + if not endpoint_name: + endpoint_name = f"endpoint-{uuid.uuid4().hex[:8]}" + + # The endpoint config's network isolation must match the built Model, or + # CreateInferenceComponent rejects the mismatch. Nova models are always + # created with network isolation enabled; for other models honor the + # value on the built Model (falling back to the builder's setting). + enable_network_isolation = bool( + is_nova + or getattr(self.built_model, "enable_network_isolation", None) + or self._enable_network_isolation + ) + # Check if endpoint exists is_existing_endpoint = self._does_endpoint_exist(endpoint_name) @@ -5050,19 +5095,35 @@ def _deploy_model_customization( ) ], 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 []) + ] endpoint = Endpoint.create( - endpoint_name=endpoint_name, endpoint_config_name=endpoint_name + endpoint_name=endpoint_name, + endpoint_config_name=endpoint_name, + tags=endpoint_tags or None, ) endpoint.wait_for_status("InService") else: endpoint = Endpoint.get(endpoint_name=endpoint_name) - peft_type = self._fetch_peft() - base_model_recipe_name = model_package.inference_specification.containers[ - 0 - ].base_model.recipe_name + # Without a model package (e.g. a Nova CPTTrainer or raw-S3 checkpoint) + # there is no PEFT/recipe metadata, so the deployment follows the + # single-IC path below. + peft_type = self._fetch_peft() if model_package is not None else None + base_model_recipe_name = ( + model_package.inference_specification.containers[0].base_model.recipe_name + if model_package is not None + else None + ) if peft_type == "LORA": # LORA deployment: base IC + adapter IC @@ -5162,8 +5223,10 @@ def _deploy_model_customization( runtime_config=InferenceComponentRuntimeConfig(copy_count=1), ) - # Create lineage tracking for new endpoints - if not is_existing_endpoint: + # Create lineage tracking for new endpoints. Lineage is keyed off the + # model package, so it is only created when one is available (a Nova + # CPTTrainer / raw-S3 checkpoint has no package). + if not is_existing_endpoint and model_package is not None: try: from sagemaker.core.resources import Action, Association, Artifact from sagemaker.core.shapes import ActionSource, MetadataProperties diff --git a/sagemaker-serve/tests/unit/test_model_builder.py b/sagemaker-serve/tests/unit/test_model_builder.py index 2622b81976..bb573da468 100644 --- a/sagemaker-serve/tests/unit/test_model_builder.py +++ b/sagemaker-serve/tests/unit/test_model_builder.py @@ -750,6 +750,86 @@ def capture_ic_create(**kwargs): compute_reqs = created_ic_spec.compute_resource_requirements self.assertIs(compute_reqs, cached_reqs) + def test_deploy_nova_inference_component(self): + """Nova + ResourceRequirements deploys via the shared single-IC path. + + A Nova checkpoint with an inference_config must NOT take the + model-on-variant path (_deploy_nova_model); it is hosted as a single + inference component referencing the built Model. The endpoint config + must set network isolation to match the Nova Model, or + CreateInferenceComponent rejects the mismatch. + """ + from sagemaker.core.resources import Endpoint, EndpointConfig, InferenceComponent + from sagemaker.core.inference_config import ResourceRequirements + + mock_endpoint = Mock() + mock_endpoint.wait_for_status = Mock() + mock_ic = Mock() + mock_ic.inference_component_arn = ( + "arn:aws:sagemaker:us-east-1:123456789012:inference-component/nova-ic" + ) + + builder = ModelBuilder( + model=self.mock_training_job, + role_arn="arn:aws:iam::123456789012:role/SageMakerRole", + sagemaker_session=self.mock_session, + image_uri="test-nova-image:latest", + instance_type="ml.p5.48xlarge", + ) + builder.built_model = Mock() + builder.built_model.model_name = "nova-model" + + inference_config = ResourceRequirements( + requests={"num_accelerators": 4, "num_cpus": 20, "memory": 35000, "copies": 1} + ) + + created_config_kwargs = {} + + def capture_config_create(**kwargs): + created_config_kwargs.update(kwargs) + return Mock() + + created_ic_kwargs = {} + + def capture_ic_create(**kwargs): + created_ic_kwargs.update(kwargs) + return mock_ic + + with patch.object(builder, "_is_nova_model", return_value=True): + with patch.object(builder, "_deploy_nova_model") as mock_deploy_nova: + with patch.object(builder, "_fetch_model_package", return_value=None): + with patch.object(builder, "_does_endpoint_exist", return_value=False): + with patch.object( + EndpointConfig, "create", side_effect=capture_config_create + ): + with patch.object(Endpoint, "create", return_value=mock_endpoint): + with patch.object( + InferenceComponent, "create", side_effect=capture_ic_create + ): + result = builder._deploy_model_customization( + endpoint_name="nova-ic-endpoint", + instance_type="ml.p5.48xlarge", + initial_instance_count=1, + inference_config=inference_config, + ) + + # Routed through the IC path, not the model-on-variant Nova path. + mock_deploy_nova.assert_not_called() + self.assertEqual(result, mock_endpoint) + + # Endpoint config network isolation matches the Nova Model. + self.assertTrue(created_config_kwargs.get("enable_network_isolation")) + + # A single IC was created referencing the built Model with the requested + # compute requirements. + self.assertEqual(created_ic_kwargs["endpoint_name"], "nova-ic-endpoint") + ic_spec = created_ic_kwargs["specification"] + self.assertEqual(ic_spec.model_name, "nova-model") + compute_reqs = ic_spec.compute_resource_requirements + self.assertEqual(compute_reqs.number_of_accelerator_devices_required, 4) + self.assertEqual(compute_reqs.number_of_cpu_cores_required, 20) + self.assertEqual(compute_reqs.min_memory_required_in_mb, 35000) + def test_deploy_passes_inference_config_to_model_customization(self): """Test that deploy() passes inference_config to _deploy_model_customization for model customization deployments.""" from sagemaker.core.inference_config import ResourceRequirements @@ -1112,9 +1192,7 @@ def test_reuse_does_not_intercept_inference_component_deploy( ): # reuse_resources=True must NOT short-circuit an IC deploy; the IC path # (inference_config is a ResourceRequirements) manages its own reuse. - from sagemaker.core.compute_resource_requirements.resource_requirements import ( - ResourceRequirements, - ) + from sagemaker.core.inference_config import ResourceRequirements with ( patch.object(ModelBuilder, "_resolve_model_source_id", return_value="s3://bucket/model"),