diff --git a/sagemaker-core/src/sagemaker/core/training/utils.py b/sagemaker-core/src/sagemaker/core/training/utils.py index 67009c9131..0d03d2fcfb 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,180 @@ 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_hyperpod_manifest_s3_uri(s3_output_path: str, training_job_name: str) -> str: + """Build the HyperPod manifest.json S3 URI for a Nova training job. + + HyperPod jobs write the manifest directly under the job directory: + ``//manifest.json``. + + 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}/manifest.json" + + +def build_nova_manifest_s3_uri(s3_output_path: str, training_job_name: str) -> str: + """Build the serverless manifest.json S3 URI for a Nova training job. + + Serverless jobs write the manifest under a nested output directory: + ``//output/output/manifest.json``. + + 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 any known + output layout. + """ + # Nova jobs write their manifest to different locations depending on the + # training platform: + # HyperPod: //manifest.json + # Serverless: //output/output/manifest.json + # Serverful: //output/output.tar.gz (manifest is inside) + # Try each in turn and surface every failure if none resolve, so the real + # cause is not masked by a misleading message from the last attempt. + hyperpod_manifest_uri = build_nova_hyperpod_manifest_s3_uri( + s3_output_path, training_job_name + ) + serverless_manifest_uri = build_nova_manifest_s3_uri(s3_output_path, training_job_name) + tar_gz_uri = build_nova_output_tar_gz_s3_uri(s3_output_path, training_job_name) + + attempts = [ + ("HyperPod manifest.json", hyperpod_manifest_uri, read_nova_checkpoint_uri_from_manifest), + ("serverless manifest.json", serverless_manifest_uri, read_nova_checkpoint_uri_from_manifest), + ("serverful output.tar.gz", tar_gz_uri, _read_checkpoint_uri_from_tar_gz), + ] + + errors = [] + for label, uri, reader in attempts: + try: + return reader(s3_client, uri) + except Exception as error: # noqa: PERF203 - each attempt may fail independently + errors.append(f"{label} at {uri} failed: {error}") + + raise ValueError( + "Could not resolve the Nova checkpoint URI from any known output layout. " + + " ".join(errors) + ) 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..fafca66d8c --- /dev/null +++ b/sagemaker-core/tests/unit/test_training_utils.py @@ -0,0 +1,198 @@ +# Copyright Amazon.com, Inc. or its affiliates. All Rights Reserved. +# +# Licensed under the Apache License, Version 2.0 (the "License"). You +# may not use this file except in compliance with the License. A copy of +# the License is located at +# +# http://aws.amazon.com/apache2.0/ +# +# or in the "license" file accompanying this file. This file is +# distributed on an "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF +# ANY KIND, either express or implied. See the License for the specific +# language governing permissions and limitations under the License. +"""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_hyperpod_manifest_s3_uri, + 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_hyperpod_manifest_s3_uri(): + result = build_nova_hyperpod_manifest_s3_uri("s3://bucket/output/", "my-job") + assert result == "s3://bucket/output/my-job/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_from_hyperpod_layout(): + """HyperPod writes the manifest at //manifest.json.""" + client = Mock() + client.exceptions = Mock() + no_such_key = type("NoSuchKey", (Exception,), {}) + client.exceptions.NoSuchKey = no_such_key + hyperpod_key = "output/my-job/manifest.json" + + def get_object(Bucket, Key): + if Key != hyperpod_key: + raise no_such_key() + body = Mock() + body.read.return_value = json.dumps( + {"checkpoint_s3_bucket": CHECKPOINT_URI} + ).encode("utf-8") + 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_from_serverless_layout(): + """Serverless writes the manifest at //output/output/manifest.json.""" + client = Mock() + client.exceptions = Mock() + no_such_key = type("NoSuchKey", (Exception,), {}) + client.exceptions.NoSuchKey = no_such_key + serverless_key = "output/my-job/output/output/manifest.json" + + def get_object(Bucket, Key): + if Key != serverless_key: + raise no_such_key() + body = Mock() + body.read.return_value = json.dumps( + {"checkpoint_s3_bucket": CHECKPOINT_URI} + ).encode("utf-8") + 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_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..5595277e75 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,215 @@ 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 a model-on-variant endpoint. + + Reads the endpoint's config production variant to find the model name, + then fetches that Model. Returns None for IC-based endpoints (no + ModelName on variant) or if the config cannot be read. + """ + 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].get("ModelName") + if not model_name: + return None + 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_model(self) -> Optional["Model"]: + """Find an existing SageMaker Model tagged with the same model source. + + Scans Models by creation time and checks for a matching model-source tag. + Returns the Model resource if found (and it still exists), None otherwise. + """ + source_id = self._resolve_model_source_id() + if not source_id: + return None + + from sagemaker.serve.model_reuse import normalize_tag_value + tag_value = normalize_tag_value(source_id) + sagemaker_client = self.sagemaker_session.sagemaker_client + + try: + next_token = None + while True: + kwargs = {"SortBy": "CreationTime", "SortOrder": "Descending"} + if next_token: + kwargs["NextToken"] = next_token + response = sagemaker_client.list_models(**kwargs) + for model_summary in response.get("Models", []): + model_name = model_summary.get("ModelName") + model_arn = model_summary.get("ModelArn") + if not model_arn: + continue + tags = sagemaker_client.list_tags(ResourceArn=model_arn).get("Tags", []) + if any( + t.get("Key") == MODEL_SOURCE_TAG_KEY and t.get("Value") == tag_value + for t in tags + ): + return Model.get(model_name=model_name, region=self.region) + next_token = response.get("NextToken") + if not next_token: + break + except Exception as e: + logger.warning("Could not search Models for reuse: %s", 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_name = variant.get("ModelName") + if not model_name: + # IC-based endpoint — can't validate container config from the + # variant. Proceed with reuse (the Model was already matched by + # tag in _find_reusable_model). + return True + model_desc = sagemaker_client.describe_model(ModelName=model_name) + # 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 +2886,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 +2912,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 @@ -2706,9 +3009,16 @@ def _build_single_modelbuilder( ) model_name = self.model_name or f"model-{uuid.uuid4().hex[:10]}" - # Create model + source_id = self._resolve_model_source_id() + model_tags = None + if source_id: + source_tag = build_source_tag(source_id) + model_tags = [{"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] + execution_role_arn=self.role_arn, + model_name=model_name, + containers=[container_def], + tags=model_tags, ) return self.built_model @@ -3545,6 +3855,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 +3879,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 +3938,33 @@ 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 Model built from the same source is + # found (by model-source tag), skip creating a new one. Also discover the + # endpoint for deploy() to reuse later. + if reuse_resources and not is_inference_component_build: + self.serve_settings = self._get_serve_setting() + reused_model = self._find_reusable_model() + if reused_model is not None: + reusable_endpoint = self._find_reusable_endpoint() + if reusable_endpoint: + self._reused_endpoint_name = reusable_endpoint + logger.warning( + "Reusing existing Model %r (matched model-source tag). " + "No new Model will be created. Pass reuse_resources=False " + "to force a new Model.", + reused_model.model_name, + ) + self.built_model = reused_model + return self.built_model + deployables = {} if not self.modelbuilder_list and not isinstance( @@ -4418,6 +4761,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 +4795,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,12 +4833,73 @@ 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 " + "endpoint_name (infrastructure reuse) and " + "inference_component_name (IC update). The flag is ignored." + ) + + # 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") if not self.instance_type and not instance_type: self.instance_type = self._fetch_default_instance_type_for_custom_model() + # Ensure self.instance_type reflects the caller's intent so the + # endpoint config creation in _deploy_model_customization picks it up. + if instance_type: + self.instance_type = instance_type + # Pass inference_config if it's ResourceRequirements inference_config_param = None if isinstance(inference_config, ResourceRequirements): @@ -4486,8 +4907,8 @@ def deploy( return self._deploy_model_customization( 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 +5070,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 +5120,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 +5147,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 +5275,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 +5361,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 +5375,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 +5398,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 +5413,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_model_customization_deployment.py b/sagemaker-serve/tests/integ/test_model_customization_deployment.py index 565b76ad39..3af07c32c8 100644 --- a/sagemaker-serve/tests/integ/test_model_customization_deployment.py +++ b/sagemaker-serve/tests/integ/test_model_customization_deployment.py @@ -147,6 +147,14 @@ 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" + # Verify model-source tag is present on the endpoint for reuse discovery. + from sagemaker.serve.model_reuse import MODEL_SOURCE_TAG_KEY + 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" + ) + if peft_type == "LORA": # Verify base IC was created base_ic_name = f"{endpoint_name}-inference-component" @@ -697,3 +705,4 @@ def test_model_customization_workflow(training_job_name): raise + 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/sagemaker-train/pyproject.toml b/sagemaker-train/pyproject.toml index 460a01a5cf..3537d3b4fb 100644 --- a/sagemaker-train/pyproject.toml +++ b/sagemaker-train/pyproject.toml @@ -68,6 +68,7 @@ notebook = [ "ipywidgets>=8.0.0", "rich>=13.0.0", "matplotlib>=3.5.0", + "pandas", ] [tool.setuptools.packages.find] diff --git a/sagemaker-train/src/sagemaker/train/base_trainer.py b/sagemaker-train/src/sagemaker/train/base_trainer.py index a453720f53..ff3bdcd751 100644 --- a/sagemaker-train/src/sagemaker/train/base_trainer.py +++ b/sagemaker-train/src/sagemaker/train/base_trainer.py @@ -1,7 +1,9 @@ import copy import os +import time import yaml from abc import ABC, abstractmethod +from datetime import datetime as _datetime from typing import Optional, Dict, Any, List, Union import json import logging @@ -15,7 +17,8 @@ import boto3 from sagemaker.core.helper.session_helper import Session -from sagemaker.core.training.configs import Tag, Networking, InputData, Channel, OutputDataConfig +from sagemaker.core.training.configs import Tag, Networking, InputData, Channel, OutputDataConfig, HyperPodCompute +from sagemaker.core.utils.logs import MultiLogStreamHandler from sagemaker.core.shapes import shapes from sagemaker.core.resources import TrainingJob from sagemaker.train.common_utils.recipe_utils import _is_nova_model, resolve_recipe, get_resolved_recipe_from_context, NoRecipeError @@ -29,11 +32,16 @@ _validate_hyperparameter_values, _get_smhp_replicas_enum, ) +from sagemaker.train.common_utils.metrics_visualizer import plot_training_metrics from sagemaker.train.common_utils.mlflow_config_utils import resolve_mlflow_tracking_fields from sagemaker.train.common_utils.validator import validate_hyperpod_compute +from sagemaker.train.common_utils.cloudwatch_metrics import fetch_and_plot_metrics, _get_smhp_log_group from sagemaker.train.defaults import TrainDefaults from sagemaker.train.utils import _get_unique_name +logger = logging.getLogger(__name__) + + class BaseTrainer(ABC): """Abstract base class for all SageMaker training workflows. @@ -268,6 +276,312 @@ def _apply_recipe_to_hyperparameters(self, final_hyperparameters: Dict[str, Any] return final_hyperparameters + def show_metrics( + self, + metrics: Optional[List[str]] = None, + starting_step: Optional[int] = None, + ending_step: Optional[int] = None, + start_time: Optional[Any] = None, + end_time: Optional[Any] = None, + ) -> Any: + """Plot training metrics from CloudWatch logs (Nova) or MLflow (OSS). + + For Nova models, parses CloudWatch logs for training_loss, lr, and reward_score. + For non-Nova (OSS) models, pulls metrics from MLflow (requires mlflow_resource_arn + to be configured on the trainer or auto-resolved). + + Args: + metrics: Optional list of metric names to plot. If None, plots all + available metrics for the training technique. + starting_step: Only plot metrics from this global step onwards. + ending_step: Only plot metrics up to this global step. + start_time: Optional start time for log retrieval. Accepts a + datetime object or epoch milliseconds (int). When not provided, + auto-resolved from the training job's start time. + end_time: Optional end time for log retrieval. Accepts a + datetime object or epoch milliseconds (int). When not provided, + defaults to now. + + Returns: + pandas.DataFrame containing the extracted metrics. + + Raises: + NotImplementedError: If the training technique does not support metric + extraction (e.g., DPO). + ValueError: If no training job has been run yet, no logs/metrics + are found, or MLflow is not configured for OSS models. + """ + # Validate that we have a training job to get metrics from + if not hasattr(self, '_latest_training_job') or self._latest_training_job is None: + raise ValueError( + "No training job found. Call .train() first, then call .show_metrics() " + "to view training metrics." + ) + + # Route based on model type + model_name = getattr(self, '_model_name', None) + is_nova = _is_nova_model(model_name) if model_name else False + + if is_nova: + return self._show_metrics_cloudwatch(metrics, starting_step, ending_step, start_time, end_time) + else: + return self._show_metrics_mlflow(metrics, starting_step, ending_step) + + def _show_metrics_mlflow( + self, + metrics: Optional[List[str]] = None, + starting_step: Optional[int] = None, + ending_step: Optional[int] = None, + ) -> None: + """Pull and plot training metrics from MLflow for non-Nova models.""" + training_job = self._latest_training_job + + # Resolve the TrainingJob object if it's a string + if isinstance(training_job, str): + logger.info(f"Resolving training job: {training_job}") + training_job = TrainingJob.get(training_job_name=training_job) + + # Validate MLflow is configured + mlflow_config = getattr(training_job, 'mlflow_config', None) + if not mlflow_config or not getattr(mlflow_config, 'mlflow_resource_arn', None): + raise ValueError( + "show_metrics() for non-Nova models requires MLflow to be configured. " + "Either pass mlflow_resource_arn when creating the trainer, or ensure " + "your account has an MLflow app set up." + ) + + mlflow_details = getattr(training_job, 'mlflow_details', None) + if not mlflow_details or not getattr(mlflow_details, 'mlflow_run_id', None): + raise ValueError( + "No MLflow run ID found on the training job. " + "MLflow metrics are only available after the job completes. " + "If the job is still running, wait for it to finish and try again. " + f"MLflow app ARN: {mlflow_config.mlflow_resource_arn}" + ) + + logger.info( + f"Fetching metrics from MLflow app: {mlflow_config.mlflow_resource_arn}, " + f"run: {mlflow_details.mlflow_run_id}" + ) + + plot_training_metrics(training_job, metrics=metrics) + + def _show_metrics_cloudwatch( + self, + metrics: Optional[List[str]] = None, + starting_step: Optional[int] = None, + ending_step: Optional[int] = None, + start_time: Optional[Any] = None, + end_time: Optional[Any] = None, + ) -> Any: + """Parse and plot training metrics from CloudWatch logs (Nova models).""" + + training_job = self._latest_training_job + if hasattr(training_job, 'training_job_name'): + job_id = training_job.training_job_name + elif isinstance(training_job, str): + job_id = training_job + else: + job_id = str(training_job) + + # Determine platform from compute config + compute = getattr(self, 'compute', None) + + # Get customization technique + customization_technique = getattr(self, '_customization_technique', None) + if not customization_technique: + raise ValueError( + "Could not determine training technique. " + "show_metrics() requires a trainer with a known customization technique." + ) + + # Resolve session + sagemaker_session = TrainDefaults.get_sagemaker_session( + sagemaker_session=self.sagemaker_session + ) + + # Resolve start_time: user-provided > training job metadata > None + start_time_ms = None + if start_time is not None: + if isinstance(start_time, _datetime): + start_time_ms = int(start_time.timestamp() * 1000) + else: + start_time_ms = int(start_time) + elif hasattr(training_job, 'training_start_time') and training_job.training_start_time: + try: + start_time_ms = int(training_job.training_start_time.timestamp() * 1000) + except Exception: + pass + + # Resolve end_time: user-provided > None (defaults to now in fetch layer) + end_time_ms = None + if end_time is not None: + if isinstance(end_time, _datetime): + end_time_ms = int(end_time.timestamp() * 1000) + else: + end_time_ms = int(end_time) + + return fetch_and_plot_metrics( + job_id=job_id, + compute=compute, + customization_technique=customization_technique, + sagemaker_session=sagemaker_session, + metrics=metrics, + starting_step=starting_step, + ending_step=ending_step, + start_time=start_time_ms, + end_time=end_time_ms, + ) + + def stream_logs(self, poll: int = 5, start_time: Optional[Any] = None) -> None: + """Stream CloudWatch logs in real-time (like ``kubectl logs -f``). + + Continuously polls for new log events and prints them as they arrive. + Blocks until the training job reaches a terminal state (SMTJ) or + the user interrupts with Ctrl+C (HyperPod). + + Args: + poll: Polling interval in seconds between log fetches. Defaults to 5. + start_time: Optional start time to stream logs from. Accepts a + datetime object or epoch milliseconds (int). Useful when + attaching to a job that's already running. If not provided, + auto-resolved from the training job's start time (SMTJ) or + defaults to now (HyperPod). + + Raises: + ValueError: If no training job has been run yet. + """ + if not hasattr(self, '_latest_training_job') or self._latest_training_job is None: + raise ValueError( + "No training job found. Call .train(wait=False) first, " + "then call .stream_logs() to stream logs in real-time." + ) + + # Resolve start_time for SMHP jobs + start_time_ms = None + if start_time is not None: + if isinstance(start_time, _datetime): + start_time_ms = int(start_time.timestamp() * 1000) + else: + start_time_ms = int(start_time) + + training_job = self._latest_training_job + compute = getattr(self, 'compute', None) + + if isinstance(compute, HyperPodCompute): + self._stream_logs_smhp(training_job, compute, poll, start_time_ms) + else: + self._stream_logs_smtj(training_job, poll) + + def _stream_logs_smtj(self, training_job, poll: int) -> None: + """Stream logs for an SMTJ training job using MultiLogStreamHandler.""" + + # Resolve job name + if hasattr(training_job, 'training_job_name'): + job_name = training_job.training_job_name + else: + job_name = str(training_job) + + log_group = "/aws/sagemaker/TrainingJobs" + instance_count = 1 + if hasattr(self, 'compute') and self.compute and hasattr(self.compute, 'instance_count'): + instance_count = self.compute.instance_count or 1 + + handler = MultiLogStreamHandler( + log_group_name=log_group, + log_stream_name_prefix=job_name, + expected_stream_count=instance_count, + ) + + logger.info(f"Streaming logs for job: {job_name}") + logger.info(f"Log group: {log_group}") + + terminal_statuses = {"Completed", "Failed", "Stopped"} + + while True: + for stream_name, event in handler.get_latest_log_events(): + message = event.get("message", "").rstrip() + if message: + logger.info(message) + + # Check job status + try: + job = TrainingJob.get(training_job_name=job_name) + status = job.training_job_status + if status in terminal_statuses: + # Final flush + for stream_name, event in handler.get_latest_log_events(): + message = event.get("message", "").rstrip() + if message: + logger.info(message) + logger.info(f"Job {job_name} finished with status: {status}") + return + except Exception: + pass + + time.sleep(poll) + + def _stream_logs_smhp(self, training_job, compute, poll: int, start_time_ms=None) -> None: + """Stream logs for a HyperPod job using filter_log_events polling.""" + + job_id = training_job if isinstance(training_job, str) else str(training_job) + + sagemaker_session = TrainDefaults.get_sagemaker_session( + sagemaker_session=self.sagemaker_session + ) + region_name = sagemaker_session.boto_session.region_name + logs_client = sagemaker_session.boto_session.client("logs", region_name=region_name) + log_group = _get_smhp_log_group(compute.cluster_name, sagemaker_session.sagemaker_client) + + logger.info(f"Streaming logs for HyperPod job: {job_id}") + logger.info(f"Cluster: {compute.cluster_name}") + logger.info(f"Log group: {log_group}") + logger.info("Press Ctrl+C to stop streaming.") + + # Pick start time (user-provided > training job start time > now) + if start_time_ms is not None: + last_timestamp = start_time_ms + elif hasattr(training_job, 'training_start_time') and training_job.training_start_time: + try: + last_timestamp = int(training_job.training_start_time.timestamp() * 1000) + except Exception: + last_timestamp = int(time.time() * 1000) + else: + last_timestamp = int(time.time() * 1000) + seen_event_ids = set() + + while True: + try: + params = { + "logGroupName": log_group, + "logStreamNamePrefix": "SagemakerHyperPodTrainingJob", + "filterPattern": f'"{job_id}"', + "startTime": last_timestamp, + } + response = logs_client.filter_log_events(**params) + events = response.get("events", []) + + for event in events: + event_id = event.get("eventId", "") + if event_id not in seen_event_ids: + seen_event_ids.add(event_id) + message = event.get("message", "").rstrip() + if message: + logger.info(message) + ts = event.get("timestamp", 0) + if ts > last_timestamp: + last_timestamp = ts + except Exception as e: + logger.debug(f"Error fetching HP logs: {e}") + + # Note: HyperPod jobs don't have a simple status API to poll for completion. + # This polls till the user interrupts with Ctrl+C. + try: + time.sleep(poll) + except KeyboardInterrupt: + logger.info("Log streaming stopped by user.") + return + def _validate_instance_count(self, instance_count, sagemaker_session): """Validate instance/node count against allowed values from SMHP recipe.""" smhp_replicas_enum = _get_smhp_replicas_enum( diff --git a/sagemaker-train/src/sagemaker/train/common_utils/cloudwatch_metrics.py b/sagemaker-train/src/sagemaker/train/common_utils/cloudwatch_metrics.py new file mode 100644 index 0000000000..ed1906675f --- /dev/null +++ b/sagemaker-train/src/sagemaker/train/common_utils/cloudwatch_metrics.py @@ -0,0 +1,431 @@ +"""CloudWatch log-based training metrics extraction and plotting. + +Parses CloudWatch log events from SageMaker Training Jobs and HyperPod clusters +to extract step-level training metrics (loss, reward scores) and plot them. + +Supports: +- SFT/CPT on SMTJ and SMHP: reduced_train_loss +- RLVR on SMTJ: critic/rewards/mean +- RLVR on SMHP: train_rm_score +""" + +from __future__ import absolute_import + +import logging +import re +from datetime import datetime +from typing import Any, Dict, List, Optional + +from sagemaker.core.training.configs import HyperPodCompute + + +logger = logging.getLogger(__name__) + +GLOBAL_STEP_REGEX = r"global_step[=:]\s*([\d.]+)" +TRAINING_LOSS_REGEX = r"reduced_train_loss[=:]\s*(-?[\d.]+(?:[eE][+-]?\d+)?)" +LEARNING_RATE_REGEX = r"(? str: + """Return the CW log group for SageMaker Training Jobs.""" + return "/aws/sagemaker/TrainingJobs" + + +def _get_smhp_log_group(cluster_name: str, sagemaker_client) -> str: + """Return the CW log group for a HyperPod cluster. + + The log group follows the pattern: + /aws/sagemaker/Clusters/{cluster_name}/{cluster_id} + """ + response = sagemaker_client.describe_cluster(ClusterName=cluster_name) + cluster_arn = response["ClusterArn"] + cluster_id = cluster_arn.split("/")[-1] + return f"/aws/sagemaker/Clusters/{cluster_name}/{cluster_id}" + + +def _fetch_smtj_logs( + job_name: str, + logs_client, + log_group: str, + start_time: Optional[int] = None, + end_time: Optional[int] = None, +) -> List[Dict[str, Any]]: + """Fetch CloudWatch log events for an SMTJ training job.""" + # Find the job's dedicated log stream + try: + response = logs_client.describe_log_streams( + logGroupName=log_group, + logStreamNamePrefix=job_name, + ) + except Exception as e: + logger.warning(f"Could not describe log streams for job '{job_name}': {e}") + return [] + + streams = response.get("logStreams", []) + if not streams: + logger.warning(f"No log stream found for job '{job_name}' in {log_group}") + return [] + + log_stream_name = streams[0]["logStreamName"] + + # Read events with get_log_events + all_events: List[Dict[str, Any]] = [] + next_token = None + end_time_ms = end_time or int(datetime.now().timestamp() * 1000) + + while True: + params: Dict[str, Any] = { + "logGroupName": log_group, + "logStreamName": log_stream_name, + "startFromHead": False, + "endTime": end_time_ms, + } + if start_time: + params["startTime"] = start_time + if next_token: + params["nextToken"] = next_token + + response = logs_client.get_log_events(**params) + events = response.get("events", []) + all_events.extend(events) + + current_token = next_token + next_token = response.get("nextBackwardToken") + if next_token == current_token: + break + + return all_events + + +def _fetch_smhp_logs( + job_id: str, + logs_client, + log_group: str, + start_time: Optional[int] = None, + end_time: Optional[int] = None, +) -> List[Dict[str, Any]]: + """Fetch CloudWatch log events for a HyperPod training job. + + HyperPod doesn't separate log streams by job — uses filter_log_events + with the job ID as the filter pattern. + """ + all_events: List[Dict[str, Any]] = [] + next_token = None + end_time_ms = end_time or int(datetime.now().timestamp() * 1000) + + while True: + params: Dict[str, Any] = { + "logGroupName": log_group, + "logStreamNamePrefix": "SagemakerHyperPodTrainingJob", + "filterPattern": f'"{job_id}"', + "endTime": end_time_ms, + } + if start_time: + params["startTime"] = start_time + if next_token: + params["nextToken"] = next_token + + try: + response = logs_client.filter_log_events(**params) + except Exception as e: + logger.warning(f"Could not filter log events for HP job '{job_id}': {e}") + return all_events + + events = response.get("events", []) + all_events.extend(events) + + next_token = response.get("nextToken") + if not next_token: + break + + return all_events + + +def parse_metrics_from_logs( + logs: List[Dict[str, Any]], + platform: str, + customization_technique: str, + metrics: Optional[List[str]] = None, +) -> "pandas.DataFrame": + """Parse training metrics from CloudWatch log events. + + Scans each log line for global_step, then extracts the relevant metric + value from the same line based on the platform and training technique. + + Args: + logs: List of CW log event dicts (each with a "message" key). + platform: "smtj" or "smhp". + customization_technique: "SFT", "CPT", or "RLVR". + metrics: Optional list of metric names to extract. If None, extracts all + available metrics for the given platform/technique combination. + + Returns: + pandas DataFrame with columns ["global_step", , ...]. + + Raises: + ImportError: If pandas is not installed. + NotImplementedError: If the technique is not supported. + ValueError: If a requested metric is not available. + """ + try: + import pandas + except ImportError: + raise ImportError( + "pandas is required for metric extraction. " + "Install it with: pip install pandas\n" + ) + + technique = customization_technique.upper() + + if technique in _UNSUPPORTED_TECHNIQUES: + raise NotImplementedError( + f"Training metrics extraction is not supported for {technique} jobs. " + f"Supported techniques: SFT, CPT, RLVR." + ) + + available = AVAILABLE_METRICS.get(platform, {}).get(technique) + if not available: + raise NotImplementedError( + f"No metric patterns defined for technique '{technique}' on platform '{platform}'. " + f"Supported: {list(AVAILABLE_METRICS.get(platform, {}).keys())}" + ) + + if not metrics: + metrics = list(available.keys()) + + patterns = [] + for metric_name in metrics: + if metric_name not in available: + raise ValueError( + f"Metric '{metric_name}' is not available for {technique} on {platform}. " + f"Available metrics: {list(available.keys())}" + ) + patterns.append(available[metric_name]) + + # Parse log lines — extract each metric independently per line. + # A line must have global_step plus at least one requested metric to be included. + all_rows: List[List] = [] + log_lines = [line for log in logs for line in log.get("message", "").splitlines()] + + for line in log_lines: + step_match = re.search(GLOBAL_STEP_REGEX, line) + if not step_match: + continue + + step_value = int(float(step_match.group(1))) + row_values: List = [step_value] + found_any = False + + for pattern in patterns: + match = re.search(pattern, line) + if match: + row_values.append(float(match.group(1))) + found_any = True + else: + row_values.append(None) + + if found_any: + all_rows.append(row_values) + + return pandas.DataFrame(all_rows, columns=["global_step"] + metrics) + + +def plot_metrics( + metrics_df: "pandas.DataFrame", + title: str = "Training Metrics", + starting_step: Optional[int] = None, + ending_step: Optional[int] = None, +) -> None: + """Plot training metrics using matplotlib. + + Args: + metrics_df: DataFrame with "global_step" column and one or more metric columns. + title: Plot title. + starting_step: Filter to steps >= this value. + ending_step: Filter to steps <= this value. + + Raises: + ImportError: If matplotlib is not installed. + ValueError: If no metrics found in the specified range. + """ + try: + import matplotlib.pyplot as plt + except ImportError: + raise ImportError( + "matplotlib is required for plotting metrics. " + "Install it with: pip install matplotlib\n" + ) + + if metrics_df.empty: + raise ValueError("No metrics data available to plot.") + + # Deduplicate and filter by step range + df = metrics_df.drop_duplicates(subset=["global_step"], keep="last").copy() + + if starting_step is not None: + df = df[df["global_step"] >= starting_step] + if ending_step is not None: + df = df[df["global_step"] <= ending_step] + + if df.empty: + range_desc = f"[{starting_step or 'start'} - {ending_step or 'end'}]" + raise ValueError(f"No metrics found in the specified step range {range_desc}") + + df = df.sort_values("global_step").reset_index(drop=True) + + # Plot each metric in its own subplot (stacked vertically) + metric_columns = [col for col in df.columns if col != "global_step"] + # Only plot columns that have at least one non-null value + metric_columns = [col for col in metric_columns if df[col].notna().any()] + + num_metrics = len(metric_columns) + if num_metrics == 0: + raise ValueError("No plottable metric data found.") + + fig, axes = plt.subplots(num_metrics, 1, figsize=(10, 4 * num_metrics), squeeze=False) + + for idx, col in enumerate(metric_columns): + ax = axes[idx, 0] + col_data = df[["global_step", col]].dropna(subset=[col]) + ax.plot(col_data["global_step"], col_data[col], linewidth=1.5) + ax.set_xlabel("Global Step") + ax.set_ylabel(col) + ax.set_title(col) + ax.grid(True, alpha=0.3) + + fig.suptitle(title, fontweight="bold", fontsize=13) + fig.tight_layout() + plt.show() + + +def fetch_and_plot_metrics( + job_id: str, + compute, + customization_technique: str, + sagemaker_session, + metrics: Optional[List[str]] = None, + starting_step: Optional[int] = None, + ending_step: Optional[int] = None, + start_time: Optional[int] = None, + end_time: Optional[int] = None, +) -> "pandas.DataFrame": + """Fetch CW logs, parse metrics, and plot them. + + This is the main entry point used by BaseTrainer.show_metrics(). + + Args: + job_id: Training job name (SMTJ) or HyperPod job name. + compute: Determines whether to use SMHP or SMTJ strategy. + customization_technique: "SFT", "CPT", or "RLVR". + sagemaker_session: SageMaker session (provides boto_session). + metrics: Optional list of metric names to extract. + starting_step: Filter to steps >= this value. + ending_step: Filter to steps <= this value. + start_time: Optional epoch ms to filter logs from (speeds up retrieval). + end_time: Optional epoch ms to filter logs until. + + Returns: + pandas DataFrame with the extracted metrics. + + Raises: + NotImplementedError: If the technique is not supported (e.g., DPO). + ValueError: If no logs or metrics are found. + """ + technique = customization_technique.upper() + + # Determine platform from compute type + platform = "smhp" if isinstance(compute, HyperPodCompute) else "smtj" + + if technique in _UNSUPPORTED_TECHNIQUES: + raise NotImplementedError( + f"show_metrics() is not supported for {technique} jobs. " + f"Supported training techniques: SFT, CPT, RLVR." + ) + + # Validate technique is recognized before doing any expensive log fetching + available = AVAILABLE_METRICS.get(platform, {}).get(technique) + if not available: + supported = list(AVAILABLE_METRICS.get(platform, {}).keys()) + raise ValueError( + f"'{customization_technique}' is not a supported training technique. " + f"Supported techniques for platform '{platform}': {supported}" + ) + + region_name = sagemaker_session.boto_session.region_name + logs_client = sagemaker_session.boto_session.client("logs", region_name=region_name) + + # Fetch logs based on compute type + if isinstance(compute, HyperPodCompute): + if not compute.cluster_name: + raise ValueError( + "cluster_name is required for HyperPod metrics. " + "This should be available from the compute configuration." + ) + log_group = _get_smhp_log_group(compute.cluster_name, sagemaker_session.sagemaker_client) + log_events = _fetch_smhp_logs( + job_id=job_id, + logs_client=logs_client, + log_group=log_group, + start_time=start_time, + end_time=end_time, + ) + else: + log_group = _get_smtj_log_group() + log_events = _fetch_smtj_logs( + job_name=job_id, + logs_client=logs_client, + log_group=log_group, + start_time=start_time, + end_time=end_time, + ) + + if not log_events: + raise ValueError( + f"No CloudWatch logs found for job '{job_id}' in log group '{log_group}'. " + f"The job may still be starting, or logs may not be available yet." + ) + + # Parse metrics from logs + metrics_df = parse_metrics_from_logs( + logs=log_events, + platform=platform, + customization_technique=technique, + metrics=metrics, + ) + + if metrics_df.empty: + raise ValueError( + f"No training metrics could be extracted from logs for job '{job_id}'. " + f"The job may not have started training steps yet." + ) + + # Sort by global_step before plotting and returning + metrics_df = metrics_df.sort_values("global_step").reset_index(drop=True) + + # Plot + title = f"Training Metrics: {job_id}" + plot_metrics( + metrics_df=metrics_df, + title=title, + starting_step=starting_step, + ending_step=ending_step, + ) + + return metrics_df diff --git a/sagemaker-train/tests/unit/train/common_utils/test_cloudwatch_metrics.py b/sagemaker-train/tests/unit/train/common_utils/test_cloudwatch_metrics.py new file mode 100644 index 0000000000..a67fd31c1d --- /dev/null +++ b/sagemaker-train/tests/unit/train/common_utils/test_cloudwatch_metrics.py @@ -0,0 +1,301 @@ +"""Unit tests for cloudwatch_metrics module.""" + +from __future__ import absolute_import + +from datetime import datetime, timezone +from unittest.mock import MagicMock, patch + +import pandas +import pytest + +from sagemaker.core.training.configs import Compute, HyperPodCompute +from sagemaker.train.base_trainer import BaseTrainer +from sagemaker.train.common_utils.cloudwatch_metrics import ( + _fetch_smhp_logs, + _fetch_smtj_logs, + fetch_and_plot_metrics, + parse_metrics_from_logs, +) + + +FAKE_SFT_LOGS = [ + {"message": "Training epoch 0, iteration 0/9 | lr: 6.25e-07 | global_batch_size: 32 | global_step: 1 | reduced_train_loss: 9.240 | ..."}, + {"message": "Training epoch 0, iteration 1/9 | lr: 1.25e-06 | global_batch_size: 32 | global_step: 2 | reduced_train_loss: 7.750 | ..."}, + {"message": "Training epoch 0, iteration 2/9 | lr: 1.87e-06 | global_batch_size: 32 | global_step: 3 | reduced_train_loss: 6.615 | ..."}, + {"message": "Some other log line without any metrics"}, +] + +FAKE_RLVR_SMTJ_LOGS = [ + {"message": "global_step=1 critic/rewards/mean=0.123"}, + {"message": "global_step=2 critic/rewards/mean=0.456"}, + {"message": "PPO iteration complete, buffers flushed"}, +] + +FAKE_RLVR_SMHP_LOGS = [ + {"message": "global_step: 1 train_rm_score: 0.55"}, + {"message": "global_step: 2 train_rm_score: 0.72"}, +] + + +class TestParseMetrics: + + def test_sft_extracts_loss_and_lr(self): + df = parse_metrics_from_logs(FAKE_SFT_LOGS, "smtj", "SFT") + + assert len(df) == 3 + assert list(df.columns) == ["global_step", "training_loss", "lr"] + assert df["training_loss"].iloc[0] == pytest.approx(9.240) + assert df["lr"].iloc[0] == pytest.approx(6.25e-07) + + def test_rlvr_smtj_extracts_reward_score(self): + df = parse_metrics_from_logs(FAKE_RLVR_SMTJ_LOGS, "smtj", "RLVR") + + assert len(df) == 2 + assert list(df.columns) == ["global_step", "reward_score"] + assert df["reward_score"].iloc[0] == pytest.approx(0.123) + + def test_rlvr_smhp_extracts_reward_score(self): + df = parse_metrics_from_logs(FAKE_RLVR_SMHP_LOGS, "smhp", "RLVR") + + assert len(df) == 2 + assert df["reward_score"].iloc[0] == pytest.approx(0.55) + + def test_subset_metrics_lr_only(self): + df = parse_metrics_from_logs(FAKE_SFT_LOGS, "smtj", "SFT", metrics=["lr"]) + + assert list(df.columns) == ["global_step", "lr"] + assert "training_loss" not in df.columns + + def test_partial_metrics_fills_nan(self): + """Lines missing a metric get NaN for that column.""" + logs = [ + {"message": "global_step=1 reduced_train_loss=5.0"}, + {"message": "global_step=2 reduced_train_loss=4.0 lr=1e-5"}, + ] + df = parse_metrics_from_logs(logs, "smtj", "SFT") + + assert len(df) == 2 + assert pandas.isna(df["lr"].iloc[0]) + assert df["lr"].iloc[1] == pytest.approx(1e-5) + + def test_dpo_raises_not_implemented(self): + with pytest.raises(NotImplementedError, match="not supported for DPO"): + parse_metrics_from_logs([], "smtj", "DPO") + + def test_empty_logs_returns_empty_dataframe(self): + df = parse_metrics_from_logs([], "smtj", "SFT") + assert df.empty + + +class TestFetchLogs: + + def test_smtj_fetches_from_dedicated_stream(self): + mock_client = MagicMock() + mock_client.describe_log_streams.return_value = { + "logStreams": [{"logStreamName": "my-job/algo-1"}] + } + mock_client.get_log_events.side_effect = [ + {"events": [{"message": "global_step=1 reduced_train_loss=5.0"}], "nextBackwardToken": "t1"}, + {"events": [], "nextBackwardToken": "t1"}, + ] + + events = _fetch_smtj_logs("my-job", mock_client, "/aws/sagemaker/TrainingJobs") + assert len(events) == 1 + + def test_smtj_no_stream_returns_empty(self): + mock_client = MagicMock() + mock_client.describe_log_streams.return_value = {"logStreams": []} + + events = _fetch_smtj_logs("nonexistent-job", mock_client, "/aws/sagemaker/TrainingJobs") + assert events == [] + + def test_smhp_uses_filter_with_job_id(self): + mock_client = MagicMock() + mock_client.filter_log_events.return_value = { + "events": [{"message": "global_step: 1 train_rm_score: 0.5"}], + } + + events = _fetch_smhp_logs("hp-job-123", mock_client, "/aws/sagemaker/Clusters/c/id") + assert len(events) == 1 + call_kwargs = mock_client.filter_log_events.call_args[1] + assert '"hp-job-123"' in call_kwargs["filterPattern"] + + +class TestFetchAndPlotMetrics: + + def _session(self): + s = MagicMock() + s.boto_session.region_name = "us-east-1" + return s + + @patch("sagemaker.train.common_utils.cloudwatch_metrics.plot_metrics") + @patch("sagemaker.train.common_utils.cloudwatch_metrics._fetch_smtj_logs") + def test_smtj_sft_end_to_end(self, mock_fetch, mock_plot): + mock_fetch.return_value = FAKE_SFT_LOGS + + df = fetch_and_plot_metrics( + "my-job", Compute(instance_type="ml.p5.48xlarge", instance_count=1), + "SFT", self._session(), + ) + + assert len(df) == 3 + assert "training_loss" in df.columns + mock_plot.assert_called_once() + + @patch("sagemaker.train.common_utils.cloudwatch_metrics.plot_metrics") + @patch("sagemaker.train.common_utils.cloudwatch_metrics._fetch_smhp_logs") + @patch("sagemaker.train.common_utils.cloudwatch_metrics._get_smhp_log_group") + def test_smhp_rlvr_end_to_end(self, mock_lg, mock_fetch, mock_plot): + mock_lg.return_value = "/aws/sagemaker/Clusters/c/id" + mock_fetch.return_value = FAKE_RLVR_SMHP_LOGS + + df = fetch_and_plot_metrics( + "hp-job", HyperPodCompute(cluster_name="c", instance_type="ml.p5.48xlarge", node_count=1), + "RLVR", self._session(), + ) + + assert len(df) == 2 + assert "reward_score" in df.columns + + def test_invalid_technique_raises_before_fetching(self): + with pytest.raises(ValueError, match="not a supported training technique"): + fetch_and_plot_metrics( + "job", Compute(instance_type="ml.p5.48xlarge", instance_count=1), + "RFT", self._session(), + ) + + @patch("sagemaker.train.common_utils.cloudwatch_metrics._fetch_smtj_logs") + def test_no_logs_found_raises(self, mock_fetch): + mock_fetch.return_value = [] + + with pytest.raises(ValueError, match="No CloudWatch logs found"): + fetch_and_plot_metrics( + "missing-job", Compute(instance_type="ml.p5.48xlarge", instance_count=1), + "SFT", self._session(), + ) + + @patch("sagemaker.train.common_utils.cloudwatch_metrics.plot_metrics") + @patch("sagemaker.train.common_utils.cloudwatch_metrics._fetch_smtj_logs") + def test_result_sorted_by_step(self, mock_fetch, mock_plot): + """Out-of-order events (startFromHead=False) are sorted in returned DataFrame.""" + mock_fetch.return_value = [ + {"message": "global_step=5 reduced_train_loss=1.0"}, + {"message": "global_step=1 reduced_train_loss=5.0"}, + ] + + df = fetch_and_plot_metrics( + "job", Compute(instance_type="ml.p5.48xlarge", instance_count=1), + "SFT", self._session(), metrics=["training_loss"], + ) + + assert df["global_step"].tolist() == [1, 5] + +class TestStreamLogs: + """Tests for BaseTrainer.stream_logs() dispatch and behavior.""" + + def _make_trainer(self, compute=None, latest_job=None): + """Create a minimal trainer stub for stream_logs testing.""" + class _StubTrainer(BaseTrainer): + _customization_technique = "SFT" + + def train(self, *args, **kwargs): + pass + + trainer = _StubTrainer.__new__(_StubTrainer) + trainer.compute = compute + trainer.sagemaker_session = None + trainer._latest_training_job = latest_job + return trainer + + def test_no_job_raises_valueerror(self): + """stream_logs() raises if no training job exists.""" + trainer = self._make_trainer(latest_job=None) + + with pytest.raises(ValueError, match="No training job found"): + trainer.stream_logs() + + @patch("sagemaker.train.base_trainer.TrainDefaults.get_sagemaker_session") + @patch("sagemaker.train.common_utils.cloudwatch_metrics._get_smhp_log_group") + def test_smhp_dispatches_with_start_time(self, mock_log_group, mock_session): + """SMHP stream_logs passes start_time to the polling loop.""" + mock_log_group.return_value = "/aws/sagemaker/Clusters/c/id" + mock_sess = MagicMock() + mock_sess.boto_session.region_name = "us-east-1" + mock_logs_client = MagicMock() + mock_sess.boto_session.client.return_value = mock_logs_client + mock_session.return_value = mock_sess + + mock_logs_client.filter_log_events.return_value = { + "events": [{"eventId": "e1", "message": "hello", "timestamp": 1000}] + } + + trainer = self._make_trainer( + compute=HyperPodCompute(cluster_name="c", instance_type="ml.p5.48xlarge", node_count=1), + latest_job="my-hp-job", + ) + + # Simulate KeyboardInterrupt on first sleep to stop the loop + with patch("time.sleep", side_effect=KeyboardInterrupt): + trainer.stream_logs( + start_time=datetime(2026, 7, 8, 14, 0, 0, tzinfo=timezone.utc) + ) + + # Verify filter_log_events was called with the user-provided startTime + call_kwargs = mock_logs_client.filter_log_events.call_args[1] + expected_ts = int(datetime(2026, 7, 8, 14, 0, 0, tzinfo=timezone.utc).timestamp() * 1000) + assert call_kwargs["startTime"] == expected_ts + + @patch("sagemaker.core.resources.TrainingJob.get") + @patch("sagemaker.train.base_trainer.MultiLogStreamHandler") + def test_smtj_stops_on_completed(self, mock_handler_cls, mock_get_job): + """SMTJ stream_logs exits when job status is Completed.""" + # Mock the handler to return one event then empty + mock_handler = MagicMock() + mock_handler.get_latest_log_events.side_effect = [ + iter([("stream", {"message": "Training complete"})]), + iter([]), # final flush + ] + mock_handler_cls.return_value = mock_handler + + # Mock job status as Completed + mock_job = MagicMock() + mock_job.training_job_status = "Completed" + mock_get_job.return_value = mock_job + + trainer = self._make_trainer( + compute=Compute(instance_type="ml.p5.48xlarge", instance_count=1), + latest_job=MagicMock(training_job_name="my-smtj-job"), + ) + + # Should return without hanging (job is already Completed) + trainer.stream_logs() + + mock_get_job.assert_called() + + + def test_show_metrics_oss_without_mlflow_raises(self): + """show_metrics() raises ValueError for non-Nova models without MLflow configured.""" + trainer = self._make_trainer(latest_job=MagicMock( + training_job_name="some-job", + mlflow_config=None, + mlflow_details=None, + )) + trainer._model_name = "test-oss-model" + + with pytest.raises(ValueError, match="requires MLflow to be configured"): + trainer.show_metrics() + + @patch("sagemaker.train.base_trainer.plot_training_metrics") + def test_show_metrics_oss_with_mlflow_delegates(self, mock_plot): + """show_metrics() for OSS models with MLflow configured calls plot_training_metrics.""" + mock_job = MagicMock() + mock_job.training_job_name = "oss-sft-job" + mock_job.mlflow_config.mlflow_resource_arn = "arn:aws:sagemaker:us-east-1:012345678910:mlflow-app/app-123" + mock_job.mlflow_details.mlflow_run_id = "run-abc123" + + trainer = self._make_trainer(latest_job=mock_job) + trainer._model_name = "test-oss-model" + + trainer.show_metrics(metrics=["loss"]) + + mock_plot.assert_called_once_with(mock_job, metrics=["loss"]) 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 +} diff --git a/v3-examples/model-customization-examples/sft_finetuning_example_notebook_pysdk_prod_v3.ipynb b/v3-examples/model-customization-examples/sft_finetuning_example_notebook_pysdk_prod_v3.ipynb index 342edd34b5..cc45e38c99 100644 --- a/v3-examples/model-customization-examples/sft_finetuning_example_notebook_pysdk_prod_v3.ipynb +++ b/v3-examples/model-customization-examples/sft_finetuning_example_notebook_pysdk_prod_v3.ipynb @@ -335,16 +335,35 @@ { "cell_type": "markdown", "id": "1f34ece2", - "source": "#### Fetch the Model Package ARN from a previous Training Job\n\nAfter a serverless training job completes, it produces an `OutputModelPackageArn` that can be used as the base model for iterative (continued) fine-tuning. Here's how to retrieve it:", - "metadata": {} + "metadata": {}, + "source": [ + "#### Fetch the Model Package ARN from a previous Training Job\n", + "\n", + "After a serverless training job completes, it produces an `OutputModelPackageArn` that can be used as the base model for iterative (continued) fine-tuning. Here's how to retrieve it:" + ] }, { "cell_type": "code", + "execution_count": null, "id": "aa54acf7", - "source": "from sagemaker.core.resources import TrainingJob\n\n# Get the completed training job\nprevious_job = TrainingJob.get(training_job_name=\"\")\n\n# The output model package ARN is available on completed serverless training jobs\noutput_model_package_arn = previous_job.output_model_package_arn\nprint(f\"Output Model Package ARN: {output_model_package_arn}\")\n\n# Use this ARN to get the ModelPackage for iterative training\nfrom sagemaker.core.resources import ModelPackage\n\nmodel_package = ModelPackage.get(model_package_name=output_model_package_arn)\npretty_print(model_package)", "metadata": {}, - "execution_count": null, - "outputs": [] + "outputs": [], + "source": [ + "from sagemaker.core.resources import TrainingJob\n", + "\n", + "# Get the completed training job\n", + "previous_job = TrainingJob.get(training_job_name=\"\")\n", + "\n", + "# The output model package ARN is available on completed serverless training jobs\n", + "output_model_package_arn = previous_job.output_model_package_arn\n", + "print(f\"Output Model Package ARN: {output_model_package_arn}\")\n", + "\n", + "# Use this ARN to get the ModelPackage for iterative training\n", + "from sagemaker.core.resources import ModelPackage\n", + "\n", + "model_package = ModelPackage.get(model_package_name=output_model_package_arn)\n", + "pretty_print(model_package)" + ] }, { "cell_type": "markdown", @@ -436,7 +455,14 @@ "cell_type": "markdown", "id": "1c8263d4", "metadata": {}, - "source": "## Finetuning with Serverful Compute (TrainingJobCompute)\n\nUse `TrainingJobCompute` to run training on dedicated instances. This enables:\n- Recipe overrides and validation\n- Iterative training from S3 checkpoints via `base_model_name`\n- Uncompressed output with `disable_output_compression=True`" + "source": [ + "## Finetuning with Serverful Compute (TrainingJobCompute)\n", + "\n", + "Use `TrainingJobCompute` to run training on dedicated instances. This enables:\n", + "- Recipe overrides and validation\n", + "- Iterative training from S3 checkpoints via `base_model_name`\n", + "- Uncompressed output with `disable_output_compression=True`" + ] }, { "cell_type": "code", @@ -444,15 +470,47 @@ "id": "5a9b3598", "metadata": {}, "outputs": [], - "source": "from sagemaker.core.training.configs import TrainingJobCompute\n\n# Base training on serverful compute\nsft_trainer_serverful = SFTTrainer(\n model=\"meta-textgeneration-llama-3-2-1b-instruct\",\n training_type=TrainingType.LORA,\n compute=TrainingJobCompute(instance_type=\"ml.p5.48xlarge\", instance_count=4),\n training_dataset=dataset.arn,\n s3_output_path=\"s3://my-bucket/output/\",\n disable_output_compression=True,\n accept_eula=True,\n)\n\ntraining_job = sft_trainer_serverful.train(wait=False)\nprint(f\"Serverful SFT job submitted: {training_job.training_job_name}\")" + "source": [ + "from sagemaker.core.training.configs import TrainingJobCompute\n", + "\n", + "# Base training on serverful compute\n", + "sft_trainer_serverful = SFTTrainer(\n", + " model=\"meta-textgeneration-llama-3-2-1b-instruct\",\n", + " training_type=TrainingType.LORA,\n", + " compute=TrainingJobCompute(instance_type=\"ml.p5.48xlarge\", instance_count=4),\n", + " training_dataset=dataset.arn,\n", + " s3_output_path=\"s3://my-bucket/output/\",\n", + " disable_output_compression=True,\n", + " accept_eula=True,\n", + ")\n", + "\n", + "training_job = sft_trainer_serverful.train(wait=False)\n", + "print(f\"Serverful SFT job submitted: {training_job.training_job_name}\")" + ] }, { "cell_type": "code", + "execution_count": null, "id": "59155cc1", - "source": "# Iterative training from S3 checkpoint (serverful)\n# Requires base_model_name to identify model for recipe lookup\nsft_trainer_iterative = SFTTrainer(\n model=\"s3://my-bucket/output/my-sft-job/output/model/\",\n base_model_name=\"meta-textgeneration-llama-3-2-1b-instruct\",\n training_type=TrainingType.LORA,\n compute=TrainingJobCompute(instance_type=\"ml.p5.48xlarge\", instance_count=4),\n training_dataset=dataset.arn,\n s3_output_path=\"s3://my-bucket/iterative-output/\",\n disable_output_compression=True,\n accept_eula=True,\n)\n\ntraining_job = sft_trainer_iterative.train(wait=False)\nprint(f\"Iterative SFT job submitted: {training_job.training_job_name}\")", "metadata": {}, - "execution_count": null, - "outputs": [] + "outputs": [], + "source": [ + "# Iterative training from S3 checkpoint (serverful)\n", + "# Requires base_model_name to identify model for recipe lookup\n", + "sft_trainer_iterative = SFTTrainer(\n", + " model=\"s3://my-bucket/output/my-sft-job/output/model/\",\n", + " base_model_name=\"meta-textgeneration-llama-3-2-1b-instruct\",\n", + " training_type=TrainingType.LORA,\n", + " compute=TrainingJobCompute(instance_type=\"ml.p5.48xlarge\", instance_count=4),\n", + " training_dataset=dataset.arn,\n", + " s3_output_path=\"s3://my-bucket/iterative-output/\",\n", + " disable_output_compression=True,\n", + " accept_eula=True,\n", + ")\n", + "\n", + "training_job = sft_trainer_iterative.train(wait=False)\n", + "print(f\"Iterative SFT job submitted: {training_job.training_job_name}\")" + ] }, { "cell_type": "markdown", @@ -506,6 +564,106 @@ " wait=True,\n", ")" ] + }, + { + "cell_type": "markdown", + "metadata": {}, + "source": [ + "## Training Monitoring\n", + "\n", + "After submitting a training job, you can monitor its progress using `show_metrics()` and `stream_logs()`." + ] + }, + { + "cell_type": "markdown", + "metadata": {}, + "source": [ + "### Stream Logs\n", + "\n", + "Stream CloudWatch logs in real-time. For SMTJ jobs, streaming automatically stops when the job completes. For HyperPod jobs, press Ctrl+C to stop." + ] + }, + { + "cell_type": "code", + "execution_count": null, + "metadata": {}, + "outputs": [], + "source": [ + "# Define the job you want to run. \n", + "sft_trainer_nova = SFTTrainer(\n", + " #model=\"test-nova-lite-v2\", \n", + " #model=\"nova-textgeneration-micro\",\n", + " model=\"nova-textgeneration-lite-v2\",\n", + " training_type=TrainingType.LORA, \n", + " model_package_group=\"sdk-test-finetuned-models\", \n", + " mlflow_experiment_name=\"test-nova-finetuned-models-exp\", \n", + " mlflow_run_name=\"test-nova-finetuned-models-run\", \n", + " training_dataset=\"arn:aws:sagemaker:us-east-1:<>:hub-content/sdktest/DataSet/sft-nova-test-dataset/0.0.1\",\n", + " s3_output_path=\"s3:///output/\" # TODO: replace with your S3 output path\n", + ")\n", + "\n", + "# Start a job without waiting, then stream logs live\n", + "training_job = sft_trainer_nova.train(wait=False)\n", + "sft_trainer_nova.stream_logs()" + ] + }, + { + "cell_type": "markdown", + "metadata": {}, + "source": [ + "### Show Training Metrics\n", + "\n", + "Plot training metrics (loss, learning rate) from CloudWatch logs. This can be called during or after training." + ] + }, + { + "cell_type": "code", + "execution_count": null, + "metadata": {}, + "outputs": [], + "source": [ + "# Plot all available metrics (training_loss and lr for SFT/CPT, reward for RLVR)\n", + "df = sft_trainer_nova.show_metrics()\n", + "\n", + "# Plot only training_loss\n", + "df = sft_trainer_nova.show_metrics(metrics=[\"training_loss\"])\n", + "\n", + "# Provide a start/end time range to speed up log retrieval for SMHP jobs.\n", + "df = trainer.show_metrics(\n", + " start_time=datetime(2026, 1, 1, 0, 0, 0, tzinfo=timezone.utc),\n", + " end_time=datetime(2026, 1, 2, 0, 0, 0, tzinfo=timezone.utc)\n", + ")\n", + "\n", + "# Prints out the dataframe in a table-format.\n", + "print(df)" + ] + }, + { + "cell_type": "markdown", + "metadata": {}, + "source": [ + "### Reconnecting to a Past Job\n", + "\n", + "If your kernel restarted, you can re-attach to an existing job and view its metrics.\n", + "For HyperPod jobs, provide `start_time` to speed up log retrieval." + ] + }, + { + "cell_type": "code", + "execution_count": null, + "metadata": {}, + "outputs": [], + "source": [ + "from datetime import datetime, timezone\n", + "\n", + "# Re-create trainer with same config, attach existing job\n", + "sft_trainer_nova._latest_training_job = \"\"\n", + "\n", + "# For HyperPod, provide start_time to bound the log search\n", + "df = sft_trainer_nova.show_metrics(\n", + " start_time=datetime(2026, 7, 8, 14, 0, 0, tzinfo=timezone.utc)\n", + ")" + ] } ], "metadata": { @@ -529,4 +687,4 @@ }, "nbformat": 4, "nbformat_minor": 5 -} \ No newline at end of file +}