diff --git a/sagemaker-core/src/sagemaker/core/training/utils.py b/sagemaker-core/src/sagemaker/core/training/utils.py index 67009c9131..061cf4c746 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,147 @@ 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 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) + 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-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..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, @@ -151,6 +155,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 +186,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 +215,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 +227,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 +355,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,20 +1105,54 @@ 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.""" + """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: - return False - 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.""" @@ -1696,6 +1764,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 +1777,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 +1908,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 +2837,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 +2863,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 +3799,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 +3823,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 +3882,41 @@ 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: + # 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 = {} if not self.modelbuilder_list and not isinstance( @@ -4418,6 +4713,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 +4747,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 +4785,62 @@ 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)) + + 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 + # 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") @@ -4488,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, @@ -4649,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( @@ -4693,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) @@ -4707,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 @@ -4819,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 @@ -4903,8 +5309,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 +5323,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 +5346,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 +5361,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..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 @@ -889,3 +969,334 @@ 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.inference_config 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 +}