From c0aead5f393cc39ea230a83569bf6d6a9daf9e0f Mon Sep 17 00:00:00 2001 From: Ealynn Hsu <89547630+ehsu3@users.noreply.github.com> Date: Thu, 9 Jul 2026 16:20:34 -0400 Subject: [PATCH 1/3] Feat: show_metrics() and stream_logs() helper functions (#6002) * Feat: Add show_metrics() and stream_logs() for monitoring training jobs * Feat: show_metrics() and stream_logs() --------- Co-authored-by: Ealynn Hsu --- sagemaker-train/pyproject.toml | 1 + .../src/sagemaker/train/base_trainer.py | 262 ++++++++++- .../train/common_utils/cloudwatch_metrics.py | 431 ++++++++++++++++++ .../common_utils/test_cloudwatch_metrics.py | 281 ++++++++++++ ...uning_example_notebook_pysdk_prod_v3.ipynb | 180 +++++++- 5 files changed, 1143 insertions(+), 12 deletions(-) create mode 100644 sagemaker-train/src/sagemaker/train/common_utils/cloudwatch_metrics.py create mode 100644 sagemaker-train/tests/unit/train/common_utils/test_cloudwatch_metrics.py 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..ab353c1612 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 @@ -31,9 +34,13 @@ ) 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 +275,259 @@ 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 extracted from CloudWatch logs using matplotlib. + + 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 with columns + ["global_step", ]. + + Raises: + NotImplementedError: If the training technique does not support metric + extraction (e.g., DPO). + ValueError: If no training job has been run yet, or no logs/metrics + are found. + """ + # Gate to Nova models only + model_name = getattr(self, '_model_name', None) + if model_name and not _is_nova_model(model_name): + raise NotImplementedError( + "show_metrics() is currently only supported for Nova 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." + ) + + # Resolve job ID + 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..aa880cfc85 --- /dev/null +++ b/sagemaker-train/tests/unit/train/common_utils/test_cloudwatch_metrics.py @@ -0,0 +1,281 @@ +"""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_rejects_oss_models(self): + """show_metrics() raises NotImplementedError for non-Nova models.""" + trainer = self._make_trainer(latest_job="some-job") + trainer._model_name = "test-oss-model" + + with pytest.raises(NotImplementedError, match="only supported for Nova models"): + trainer.show_metrics() 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 +} From b3152fb3676694c8bcb16970e83cfc049bfbdbde Mon Sep 17 00:00:00 2001 From: Ealynn Hsu <89547630+ehsu3@users.noreply.github.com> Date: Fri, 10 Jul 2026 14:39:12 -0400 Subject: [PATCH 2/3] show_metrics() Enhancement: Display MLFlow metrics for OSS models (#6013) * show_metrics() Enhancement: Display MLFlow metrics for OSS models * Update unit tests * Address code comments --------- Co-authored-by: Ealynn Hsu --- .../src/sagemaker/train/base_trainer.py | 80 ++++++++++++++++--- .../common_utils/test_cloudwatch_metrics.py | 28 ++++++- 2 files changed, 91 insertions(+), 17 deletions(-) diff --git a/sagemaker-train/src/sagemaker/train/base_trainer.py b/sagemaker-train/src/sagemaker/train/base_trainer.py index ab353c1612..ff3bdcd751 100644 --- a/sagemaker-train/src/sagemaker/train/base_trainer.py +++ b/sagemaker-train/src/sagemaker/train/base_trainer.py @@ -32,6 +32,7 @@ _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 @@ -283,7 +284,11 @@ def show_metrics( start_time: Optional[Any] = None, end_time: Optional[Any] = None, ) -> Any: - """Plot training metrics extracted from CloudWatch logs using matplotlib. + """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 @@ -298,22 +303,14 @@ def show_metrics( defaults to now. Returns: - pandas.DataFrame containing the extracted metrics with columns - ["global_step", ]. + 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, or no logs/metrics - are found. + ValueError: If no training job has been run yet, no logs/metrics + are found, or MLflow is not configured for OSS models. """ - # Gate to Nova models only - model_name = getattr(self, '_model_name', None) - if model_name and not _is_nova_model(model_name): - raise NotImplementedError( - "show_metrics() is currently only supported for Nova 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( @@ -321,7 +318,64 @@ def show_metrics( "to view training metrics." ) - # Resolve job ID + # 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 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 index aa880cfc85..a67fd31c1d 100644 --- a/sagemaker-train/tests/unit/train/common_utils/test_cloudwatch_metrics.py +++ b/sagemaker-train/tests/unit/train/common_utils/test_cloudwatch_metrics.py @@ -272,10 +272,30 @@ def test_smtj_stops_on_completed(self, mock_handler_cls, mock_get_job): mock_get_job.assert_called() - def test_show_metrics_rejects_oss_models(self): - """show_metrics() raises NotImplementedError for non-Nova models.""" - trainer = self._make_trainer(latest_job="some-job") + + 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(NotImplementedError, match="only supported for Nova models"): + 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"]) From 37c371eefb60b8fe1c9a35db2578d041dc16bded Mon Sep 17 00:00:00 2001 From: LN <133025223+amazeAmazing@users.noreply.github.com> Date: Fri, 10 Jul 2026 14:51:48 -0400 Subject: [PATCH 3/3] feat(serve): add opt-in model source tag-based resource reuse (#5993) MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit * feat(serve): add opt-in model source tag-based resource reuse Add reuse_resources to ModelBuilder.build/deploy and BedrockModelBuilder.deploy. On a hit, discover an existing resource by the model-source tag and return it instead of creating a duplicate (warn, do not raise). Honored per call. - New sagemaker/serve/model_reuse.py: tag helpers + service-client discovery - Consolidate Nova manifest/checkpoint reading into sagemaker/core/training/utils.py - SageMaker: build() skips Model creation on reuse (sets built_model to the existing Model); deploy() reuses the endpoint after validating env vars/image/ instance type (PrimaryContainer with Containers[0] fallback for Nova) - Reuse gates are skipped for inference-component builds/deploys so IC create/update (via _deploy_for_ic) is never silently intercepted - Bedrock: reuse custom model + active deployment; response includes modelArn - Reuse discovery uses the cached session/bedrock clients - Support raw S3 URI model input via model_metadata BASE_MODEL_NAME - Unit tests + notebook examples * feat(serve): support Nova inference-component deployment and harden reuse Route Nova model-customization deploys through the shared single-inference -component path when a ResourceRequirements inference_config is supplied, so each Nova checkpoint (full-rank or LoRA-merged) is hosted as one inference component referencing the built Model. Nova without an inference_config keeps the direct model-on-variant path. - Broaden _is_nova_model to identify Nova from a package-less source (raw S3 checkpoint or trainer) via base_model_name, in addition to the model package recipe/hub-content name. - Set EnableNetworkIsolation on the IC endpoint config to match the built Model (always True for Nova), fixing CreateInferenceComponent rejection on mismatched network isolation. - Guard model-package-dependent logic (restricted-package path, PEFT/recipe metadata, lineage tracking) so package-less Nova checkpoints deploy cleanly. - Apply accumulated tags (including the model-source reuse tag) to endpoints created on the shared IC path so they remain discoverable. - build(reuse_resources=True) only short-circuits when the backing Model can be resolved; IC endpoints and stale/deleted configs fall through and build a real Model, preventing a None built_model on later IC deploys. - deploy() warns that reuse_resources has no effect for inference-component deployments, which manage their own reuse by component name. - Surface both the manifest.json and output.tar.gz errors when Nova checkpoint URI resolution fails, instead of masking the primary failure. Add unit tests covering the Nova IC path (routing, network isolation, IC spec) and the model-on-variant fallback. * fix(core): resolve Nova checkpoint manifest across all three output layouts Nova training jobs write their checkpoint manifest to different locations depending on the training platform: HyperPod: //manifest.json Serverless: //output/output/manifest.json Serverful: //output/output.tar.gz (manifest inside) resolve_nova_checkpoint_uri previously only tried the serverless manifest path and the serverful tar.gz, so HyperPod jobs (manifest directly under the job directory) failed to resolve. Add build_nova_hyperpod_manifest_s3_uri and try all three layouts in turn, aggregating every failure into the raised error so the real cause is not masked by the last attempt's message. Add unit tests for the HyperPod builder and for resolution from the HyperPod and serverless layouts. * feat(serve): Model-tag reuse, IC-deploy guard, and instance_type fix Simplify reuse discovery by tagging SageMaker Models (not just endpoints) with the model-source identifier, so build(reuse_resources=True) can find and skip recreating an existing Model directly — no IC-state dependency. - Tag non-Nova Models at build time with the model-source tag (matching the Nova path's existing behavior). Both Nova and OSS Models are now discoverable by tag. - Add _find_reusable_model: build(reuse_resources=True) searches Models by source tag, skipping Model creation on a hit. Also discovers the endpoint for deploy() to reuse later. - Simplify _get_model_for_endpoint back to variant-only lookup (returns None for IC endpoints). No longer needs IC-spec resolution since the Model is found directly by tag. - _reused_endpoint_matches_config returns True for IC endpoints (can't read container config from variant; Model was already matched by tag). - deploy() with reuse_resources=True on an IC deploy logs a warning that the flag has no effect (ICs manage reuse by endpoint_name + IC name). - Fix deploy() to set self.instance_type from the caller's explicit value before calling _deploy_model_customization, preventing recipe-resolved defaults from overriding the user's intent. - Add model-source tag assertion to the existing OSS deploy integ test. --- .../src/sagemaker/core/training/utils.py | 181 ++++++ .../tests/unit/test_training_utils.py | 198 ++++++ .../sagemaker/serve/bedrock_model_builder.py | 246 +++++-- .../src/sagemaker/serve/model_builder.py | 603 +++++++++++++++--- .../src/sagemaker/serve/model_reuse.py | 303 +++++++++ .../test_model_customization_deployment.py | 9 + ...est_nova_model_customization_deployment.py | 16 + .../tests/unit/test_bedrock_model_builder.py | 252 +++++++- .../tests/unit/test_model_builder.py | 411 ++++++++++++ .../tests/unit/test_model_reuse.py | 297 +++++++++ .../bedrock-modelbuilder-deployment.ipynb | 32 + .../model_builder_deployment_notebook.ipynb | 234 +++++-- 12 files changed, 2587 insertions(+), 195 deletions(-) create mode 100644 sagemaker-core/tests/unit/test_training_utils.py create mode 100644 sagemaker-serve/src/sagemaker/serve/model_reuse.py create mode 100644 sagemaker-serve/tests/unit/test_model_reuse.py diff --git a/sagemaker-core/src/sagemaker/core/training/utils.py b/sagemaker-core/src/sagemaker/core/training/utils.py index 67009c9131..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/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 +}