diff --git a/sagemaker-core/src/sagemaker/core/helper/iam_role_resolver.py b/sagemaker-core/src/sagemaker/core/helper/iam_role_resolver.py index db4023b573..3baa1205da 100644 --- a/sagemaker-core/src/sagemaker/core/helper/iam_role_resolver.py +++ b/sagemaker-core/src/sagemaker/core/helper/iam_role_resolver.py @@ -601,7 +601,7 @@ def resolve_and_validate_role( role_type, ) else: - logger.info("Role '%s' validated for %s. Using it.", role_arn, role_type) + logger.debug("Role '%s' validated for %s. Using it.", role_arn, role_type) return role_arn diff --git a/sagemaker-core/src/sagemaker/core/utils/utils.py b/sagemaker-core/src/sagemaker/core/utils/utils.py index 865ae0f96e..8874d5fb58 100644 --- a/sagemaker-core/src/sagemaker/core/utils/utils.py +++ b/sagemaker-core/src/sagemaker/core/utils/utils.py @@ -369,10 +369,8 @@ def __init__( self.session = session self.region_name = region_name # Read region from environment variable, default to us-west-2 - import os env_region = os.environ.get('SAGEMAKER_REGION', region_name) env_stage = os.environ.get('SAGEMAKER_STAGE', 'prod') # default to gamma - logger.info(f"Runs on sagemaker {env_stage}, region:{env_region}") endpoint_url = os.environ.get('SAGEMAKER_ENDPOINT') runtime_endpoint_url = os.environ.get('SAGEMAKER_RUNTIME_ENDPOINT') diff --git a/sagemaker-serve/src/sagemaker/serve/model_builder.py b/sagemaker-serve/src/sagemaker/serve/model_builder.py index c125a749f4..5e9903cd7a 100644 --- a/sagemaker-serve/src/sagemaker/serve/model_builder.py +++ b/sagemaker-serve/src/sagemaker/serve/model_builder.py @@ -1811,6 +1811,11 @@ def _is_model_customization(self) -> bool: # ModelTrainer with model customization if isinstance(self.model, ModelTrainer) and hasattr(self.model, "_latest_training_job"): + # Resolve string job name to TrainingJob object if needed + if isinstance(self.model._latest_training_job, str): + self.model._latest_training_job = TrainingJob.get( + training_job_name=self.model._latest_training_job + ) # Check model_package_config first (new location) if ( hasattr(self.model._latest_training_job, "model_package_config") @@ -1845,6 +1850,11 @@ 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"): + # Resolve string job name to TrainingJob object if needed + if isinstance(self.model._latest_training_job, str): + self.model._latest_training_job = TrainingJob.get( + training_job_name=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. @@ -1926,6 +1936,11 @@ def _fetch_model_package_arn(self) -> Optional[str]: return arn if hasattr(self.model, "_latest_training_job"): + # Resolve string job name to TrainingJob object if needed + if isinstance(self.model._latest_training_job, str): + self.model._latest_training_job = TrainingJob.get( + training_job_name=self.model._latest_training_job + ) # Try output_model_package_arn first (preferred) if hasattr(self.model._latest_training_job, "output_model_package_arn"): arn = self.model._latest_training_job.output_model_package_arn @@ -2049,38 +2064,24 @@ def _get_model_for_endpoint(self, endpoint_name: str) -> Optional[Model]: 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. + Uses the Resource Groups Tagging API for efficient server-side tag + filtering when available, falling back to paginated list+list_tags scan. 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 + from sagemaker.serve.model_reuse import normalize_tag_value, find_sagemaker_model_arn_by_tag 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 + model_arn = find_sagemaker_model_arn_by_tag(sagemaker_client, tag_value) + if model_arn: + # Extract model name from ARN: arn:aws:sagemaker:region:account:model/name + model_name = model_arn.rsplit("/", 1)[-1] + return Model.get(model_name=model_name, region=self.region) except Exception as e: logger.warning("Could not search Models for reuse: %s", e) diff --git a/sagemaker-serve/src/sagemaker/serve/model_reuse.py b/sagemaker-serve/src/sagemaker/serve/model_reuse.py index b534173d1f..f08dbce3a0 100644 --- a/sagemaker-serve/src/sagemaker/serve/model_reuse.py +++ b/sagemaker-serve/src/sagemaker/serve/model_reuse.py @@ -20,6 +20,8 @@ from botocore.exceptions import ClientError +import boto3 + logger = logging.getLogger(__name__) MODEL_SOURCE_TAG_KEY = "sagemaker.amazonaws.com/model-source" @@ -218,8 +220,119 @@ def _bedrock_resource_has_tag(bedrock_client, resource_arn: str, tag_value: str) ) +def _find_resource_arn_by_tagging_api( + sagemaker_client, tag_value: str, resource_type: str +) -> Optional[str]: + """Use Resource Groups Tagging API to find a resource by tag (fast path). + + Makes a single server-side filtered query instead of iterating through all + resources and calling list_tags on each one. + + Args: + sagemaker_client: A boto3 SageMaker client (used for region detection + and as fallback for creating the tagging client). + tag_value: The normalized tag value to search for. + resource_type: The resource type filter (e.g. "sagemaker:model", + "sagemaker:endpoint"). + + Returns: + The resource ARN if found, empty string "" if no match (signals fast path + completed successfully with no results), or None if the tagging API is + unavailable (signals caller should fall back to the slow scan). + """ + try: + region = sagemaker_client.meta.region_name + if not region: + return None + + tagging_client = boto3.client("resourcegroupstaggingapi", region_name=region) + + pagination_token = "" + while True: + kwargs = { + "TagFilters": [ + {"Key": MODEL_SOURCE_TAG_KEY, "Values": [tag_value]} + ], + "ResourceTypeFilters": [resource_type], + } + if pagination_token: + kwargs["PaginationToken"] = pagination_token + + response = tagging_client.get_resources(**kwargs) + for mapping in response.get("ResourceTagMappingList", []): + arn = mapping.get("ResourceARN") + if arn: + return arn + + pagination_token = response.get("PaginationToken", "") + if not pagination_token: + return "" # Fast path succeeded, no matching resource found + + except ClientError as e: + error_code = e.response.get("Error", {}).get("Code", "") + if error_code == _ACCESS_DENIED_CODE: + logger.debug( + "Resource Groups Tagging API access denied (tag:GetResources). " + "Falling back to paginated list+list_tags scan." + ) + return None # Signal caller to use fallback + logger.debug("Resource Groups Tagging API call failed: %s. Using fallback.", e) + return None + except Exception as e: + logger.debug("Resource Groups Tagging API unavailable: %s. Using fallback.", e) + return None + + +def find_sagemaker_model_arn_by_tag(sagemaker_client, tag_value: str) -> Optional[str]: + """Return the ARN of the first SageMaker Model carrying the source tag. + + Uses the Resource Groups Tagging API for efficient server-side filtering + when available, falling back to paginated list+list_tags if denied. + + Args: + sagemaker_client: A boto3 SageMaker client. + tag_value: The normalized tag value to match. + + Returns: + Model ARN if found, None otherwise. + """ + # Try the fast path first: Resource Groups Tagging API + arn = _find_resource_arn_by_tagging_api( + sagemaker_client, tag_value, resource_type="sagemaker:model" + ) + if arn is not None: + return arn if arn != "" else None + + # Fallback: paginate through all models and check tags individually + 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_arn = model_summary.get("ModelArn") + if model_arn and _sagemaker_resource_has_tag(sagemaker_client, model_arn, tag_value): + return model_arn + next_token = response.get("NextToken") + if not next_token: + return None + + 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.""" + """Return the ARN of the first SageMaker endpoint carrying the source tag. + + Uses the Resource Groups Tagging API for efficient server-side filtering + when available, falling back to paginated list+list_tags if denied. + """ + # Try the fast path first: Resource Groups Tagging API + arn = _find_resource_arn_by_tagging_api( + sagemaker_client, tag_value, resource_type="sagemaker:endpoint" + ) + if arn is not None: + return arn if arn != "" else None + + # Fallback: paginate through all endpoints and check tags individually next_token = None while True: kwargs = {"NextToken": next_token} if next_token else {} diff --git a/sagemaker-train/src/sagemaker/train/base_trainer.py b/sagemaker-train/src/sagemaker/train/base_trainer.py index 106b6b287b..2d36016256 100644 --- a/sagemaker-train/src/sagemaker/train/base_trainer.py +++ b/sagemaker-train/src/sagemaker/train/base_trainer.py @@ -18,9 +18,10 @@ from botocore.exceptions import ClientError from sagemaker.core.helper.session_helper import Session -from sagemaker.core.training.configs import Tag, Networking, InputData, Channel, OutputDataConfig, HyperPodCompute +from sagemaker.core.training.configs import Tag, Networking, InputData, Channel, OutputDataConfig, HyperPodCompute, TrainingJobCompute from sagemaker.core.utils.logs import MultiLogStreamHandler from sagemaker.core.shapes import shapes +from sagemaker.core.shapes import S3DataSource 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 from sagemaker.core.s3.utils import resolve_s3_uri_placeholders @@ -41,6 +42,7 @@ 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.model_trainer import ModelTrainer from sagemaker.train.utils import _get_unique_name logger = logging.getLogger(__name__) @@ -114,7 +116,7 @@ def __init__( disable_output_compression: Optional[bool] = False, notifications: Optional[Dict[str, Any]] = None, ): - self.sagemaker_session = sagemaker_session + self.sagemaker_session = sagemaker_session or TrainDefaults.get_sagemaker_session() self.role = role self.base_job_name = base_job_name self.tags = tags @@ -625,7 +627,7 @@ def list_notification_rules( event_bus_arn=event_bus_arn, ) - def stream_logs(self, poll: int = 5, start_time: Optional[Any] = None) -> None: + def stream_logs(self, poll: int = 5, start_time: Optional[Any] = None, num_lines: Optional[int] = None) -> None: """Stream CloudWatch logs in real-time (like ``kubectl logs -f``). Continuously polls for new log events and prints them as they arrive. @@ -639,6 +641,10 @@ def stream_logs(self, poll: int = 5, start_time: Optional[Any] = None) -> None: 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). + num_lines: Optional maximum number of log lines to print. When + specified, streaming stops after this many lines have been + printed. Useful for long jobs where the full log is too verbose. + If not provided, streams all logs until the job completes. Raises: ValueError: If no training job has been run yet. @@ -672,11 +678,11 @@ def stream_logs(self, poll: int = 5, start_time: Optional[Any] = None) -> None: compute = getattr(self, 'compute', None) if isinstance(compute, HyperPodCompute): - self._stream_logs_smhp(training_job, compute, poll, start_time_ms) + self._stream_logs_smhp(training_job, compute, poll, start_time_ms, num_lines=num_lines) else: - self._stream_logs_smtj(training_job, poll, start_time_ms) + self._stream_logs_smtj(training_job, poll, start_time_ms, num_lines=num_lines) - def _stream_logs_smtj(self, training_job, poll: int, start_time_ms=None) -> None: + def _stream_logs_smtj(self, training_job, poll: int, start_time_ms=None, num_lines: Optional[int] = None) -> None: """Stream logs for an SMTJ training job.""" from sagemaker.train.common_utils.log_streamer import ( LogStreamer, @@ -708,9 +714,9 @@ def _get_status() -> str: job = TrainingJob.get(training_job_name=job_name) return job.training_job_status - stream_log_loop(streamer, poll, _get_status) + stream_log_loop(streamer, poll, _get_status, num_lines=num_lines) - def _stream_logs_smhp(self, training_job, compute, poll: int, start_time_ms=None) -> None: + def _stream_logs_smhp(self, training_job, compute, poll: int, start_time_ms=None, num_lines: Optional[int] = None) -> None: """Stream logs for a HyperPod job using filter_log_events polling.""" if isinstance(training_job, str): @@ -743,6 +749,8 @@ def _stream_logs_smhp(self, training_job, compute, poll: int, start_time_ms=None else: last_timestamp = int(time.time() * 1000) seen_event_ids = set() + lines_printed = 0 + _CW_PREFIX = "[CloudWatch] " empty_cycles = 0 while True: @@ -764,7 +772,11 @@ def _stream_logs_smhp(self, training_job, compute, poll: int, start_time_ms=None seen_event_ids.add(event_id) message = event.get("message", "").rstrip() if message: - logger.info(message) + print(f"{_CW_PREFIX}{message}") + lines_printed += 1 + if num_lines and lines_printed >= num_lines: + logger.info(f"Reached num_lines limit ({num_lines}). Stopping log stream.") + return ts = event.get("timestamp", 0) if ts > last_timestamp: last_timestamp = ts @@ -859,22 +871,10 @@ def _train_serverful_smtj(self, training_dataset=None, validation_dataset=None, from ``self._customization_technique``) and any extra hyperparameters from ``_get_extra_smtj_hyperparameters()``. """ - import logging - import tempfile - from sagemaker.train.model_trainer import ModelTrainer - from sagemaker.core.training.configs import TrainingJobCompute, InputData, Networking - from sagemaker.core.shapes import S3DataSource - from sagemaker.train.common_utils.finetune_utils import ( - get_recipe_s3_uri, - get_training_image, - _validate_hyperparameter_values, - ) - from sagemaker.train.defaults import TrainDefaults - sagemaker_session = TrainDefaults.get_sagemaker_session( sagemaker_session=self.sagemaker_session ) - role = TrainDefaults.get_role(role=self.role, sagemaker_session=sagemaker_session) + role = self.role compute = self.compute customization_technique = self._customization_technique @@ -887,7 +887,7 @@ def _train_serverful_smtj(self, training_dataset=None, validation_dataset=None, sagemaker_session=sagemaker_session, ) - logger.info(f"SMTJ recipe S3 URI: {recipe_s3_uri}") + logger.debug(f"SMTJ recipe S3 URI: {recipe_s3_uri}") # Download recipe from S3 to a local temp file recipe_s3_uri = resolve_s3_uri_placeholders(recipe_s3_uri, sagemaker_session) @@ -1106,7 +1106,7 @@ def _yaml_safe_default(value): with open(recipe_local_path, "w") as f: f.write(recipe_content) - logger.info(f"Recipe downloaded and rendered to: {recipe_local_path}") + logger.debug(f"Recipe downloaded and rendered to: {recipe_local_path}") # Resolve training image training_image = self.training_image diff --git a/sagemaker-train/src/sagemaker/train/common_utils/cloudwatch_metrics.py b/sagemaker-train/src/sagemaker/train/common_utils/cloudwatch_metrics.py index 4bec69559a..fa7fc30a85 100644 --- a/sagemaker-train/src/sagemaker/train/common_utils/cloudwatch_metrics.py +++ b/sagemaker-train/src/sagemaker/train/common_utils/cloudwatch_metrics.py @@ -445,9 +445,17 @@ def fetch_and_plot_metrics( ) if not log_events: + time_hint = "" + if start_time or end_time: + time_hint = ( + f" No logs were found in the specified time range " + f"(start_time={start_time}, end_time={end_time}). " + f"Try adjusting the time range or omitting start_time/end_time " + f"to search the full job duration." + ) 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." + f"The job may still be starting, or logs may not be available yet.{time_hint}" ) # Parse metrics from logs diff --git a/sagemaker-train/src/sagemaker/train/common_utils/data_utils.py b/sagemaker-train/src/sagemaker/train/common_utils/data_utils.py index 3c417b56e9..7a2b7dda8b 100644 --- a/sagemaker-train/src/sagemaker/train/common_utils/data_utils.py +++ b/sagemaker-train/src/sagemaker/train/common_utils/data_utils.py @@ -140,12 +140,10 @@ def validate_data_path_exists( except ClientError as e: code = e.response["Error"]["Code"] if code == "403" or "AccessDenied" in str(e): - # Caller may not have access but the execution role might — - # log a warning and allow the job to proceed. - logger.warning( - "Cannot verify S3 %s path %s from caller identity " - "(AccessDenied). The execution role may still have access.", - label, data_path, + raise ValueError( + f"Cannot verify S3 {label} path '{data_path}': access denied. " + f"Ensure the path exists and the caller has s3:ListBucket and " + f"s3:GetObject permissions on the bucket." ) else: raise ValueError(f"Error accessing S3 {label} path {data_path}: {e}") diff --git a/sagemaker-train/src/sagemaker/train/common_utils/finetune_utils.py b/sagemaker-train/src/sagemaker/train/common_utils/finetune_utils.py index 94f5f7af3e..a3bd1fba97 100644 --- a/sagemaker-train/src/sagemaker/train/common_utils/finetune_utils.py +++ b/sagemaker-train/src/sagemaker/train/common_utils/finetune_utils.py @@ -1123,8 +1123,12 @@ def _validate_and_resolve_model_package_group(model, model_package_group_name): if isinstance(model, ModelPackage): return model.model_package_group_name - raise ValueError("model_package_group_name must be provided when model given is " - "not a ModelPackage artifact/not continued finetuning") + raise ValueError( + "model_package_group is required for serverless training (when compute is not set). " + "Either provide model_package_group to store the fine-tuned model, or set " + "compute=TrainingJobCompute(...) / HyperPodCompute(...) to use managed compute " + "where model_package_group is optional." + ) def _validate_eula_for_gated_model(model, accept_eula, is_gated_model): diff --git a/sagemaker-train/src/sagemaker/train/common_utils/log_streamer.py b/sagemaker-train/src/sagemaker/train/common_utils/log_streamer.py index c120bd9131..3538c95200 100644 --- a/sagemaker-train/src/sagemaker/train/common_utils/log_streamer.py +++ b/sagemaker-train/src/sagemaker/train/common_utils/log_streamer.py @@ -227,6 +227,7 @@ def stream_log_loop( streamer: LogStreamer, poll: int, status_fn: Callable[[], str], + num_lines: Optional[int] = None, ) -> None: """Run the standard log streaming loop. @@ -237,7 +238,22 @@ def stream_log_loop( :param streamer: A configured LogStreamer instance. :param poll: Seconds between polls. :param status_fn: Callable that returns the current job status string. + :param num_lines: Optional maximum number of log lines to print. + When specified, streaming stops after this many lines. """ + _CW_PREFIX = "[CloudWatch] " + lines_printed = 0 + + def _print_event(ts_ms: int, message: str) -> bool: + """Print a log event. Returns True if num_lines limit reached.""" + nonlocal lines_printed + print(f"{_CW_PREFIX}[{_format_timestamp(ts_ms)}] {message}") + lines_printed += 1 + if num_lines and lines_printed >= num_lines: + logger.info("Reached num_lines limit (%d). Stopping log stream.", num_lines) + return True + return False + status = status_fn() if status in TERMINAL_STATUSES: logger.info("Job already in terminal state: %s", status) @@ -247,7 +263,8 @@ def stream_log_loop( if not events: break for ts_ms, message in events: - logger.info("[%s] %s", _format_timestamp(ts_ms), message) + if _print_event(ts_ms, message): + return except ClientError: pass logger.info("Job finished with status: %s", status) @@ -285,7 +302,8 @@ def stream_log_loop( if events: empty_cycles = 0 for ts_ms, message in events: - logger.info("[%s] %s", _format_timestamp(ts_ms), message) + if _print_event(ts_ms, message): + return else: empty_cycles += 1 if empty_cycles == warn_cycle: @@ -297,7 +315,8 @@ def stream_log_loop( status = status_fn() if status in TERMINAL_STATUSES: for ts_ms, message in streamer.poll_once(): - logger.info("[%s] %s", _format_timestamp(ts_ms), message) + if _print_event(ts_ms, message): + return logger.info("Job finished with status: %s", status) return diff --git a/sagemaker-train/src/sagemaker/train/defaults.py b/sagemaker-train/src/sagemaker/train/defaults.py index c358c11a04..dbd5cdfb0d 100644 --- a/sagemaker-train/src/sagemaker/train/defaults.py +++ b/sagemaker-train/src/sagemaker/train/defaults.py @@ -164,10 +164,10 @@ def get_stopping_condition( max_pending_time_in_seconds=None, max_wait_time_in_seconds=None, ) - logger.info(f"StoppingCondition not provided. Using default:\n{stopping_condition}") + logger.debug(f"StoppingCondition not provided. Using default:\n{stopping_condition}") if stopping_condition.max_runtime_in_seconds is None: stopping_condition.max_runtime_in_seconds = DEFAULT_MAX_RUNTIME_IN_SECONDS - logger.info( + logger.debug( "Max runtime not provided. Using default:\n" f"{stopping_condition.max_runtime_in_seconds}" ) @@ -201,7 +201,7 @@ def get_output_data_config( ) if output_data_config.compression_type is None: output_data_config.compression_type = "GZIP" - logger.info( + logger.debug( f"OutputDataConfig compression type not provided. Using default:\n" f"{output_data_config.compression_type}" ) diff --git a/sagemaker-train/src/sagemaker/train/model_trainer.py b/sagemaker-train/src/sagemaker/train/model_trainer.py index 3554a74ea7..15aa4f0d6b 100644 --- a/sagemaker-train/src/sagemaker/train/model_trainer.py +++ b/sagemaker-train/src/sagemaker/train/model_trainer.py @@ -1283,7 +1283,6 @@ def from_recipe( raise ValueError("training_image must be provided when using training_image_config.") sagemaker_session = TrainDefaults.get_sagemaker_session(sagemaker_session) - role = TrainDefaults.get_role(role=role, sagemaker_session=sagemaker_session) # The training recipe is used to prepare the following args: # - source_code