Improve logging, error messages and minor bug fixes

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
 
 

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')

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)
 

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 {}

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

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