sagemaker-core 1.0.9__py3-none-any.whl → 1.0.10__py3-none-any.whl

sagemaker_core/main/code_injection/shape_dag.py

@@ -1380,6 +1380,11 @@ SHAPE_DAG = {
                 "shape": "AppLifecycleManagement",
                 "type": "structure",
             },
+            {
+                "name": "BuiltInLifecycleConfigArn",
+                "shape": "StudioLifecycleConfigArn",
+                "type": "string",
+            },
         ],
         "type": "structure",
     },
@@ -3793,6 +3798,11 @@ SHAPE_DAG = {
             {"name": "CreationTime", "shape": "Timestamp", "type": "timestamp"},
             {"name": "FailureReason", "shape": "FailureReason", "type": "string"},
             {"name": "ResourceSpec", "shape": "ResourceSpec", "type": "structure"},
+            {
+                "name": "BuiltInLifecycleConfigArn",
+                "shape": "StudioLifecycleConfigArn",
+                "type": "string",
+            },
         ],
         "type": "structure",
     },
@@ -7699,6 +7709,11 @@ SHAPE_DAG = {
                 "type": "structure",
             },
             {"name": "EmrSettings", "shape": "EmrSettings", "type": "structure"},
+            {
+                "name": "BuiltInLifecycleConfigArn",
+                "shape": "StudioLifecycleConfigArn",
+                "type": "string",
+            },
         ],
         "type": "structure",
     },

sagemaker_core/main/logs.py

@@ -0,0 +1,167 @@
+import boto3
+import botocore
+
+from boto3.session import Session
+import botocore.client
+from botocore.config import Config
+from typing import Generator, Tuple, List
+from sagemaker_core.main.utils import SingletonMeta
+
+
+class CloudWatchLogsClient(metaclass=SingletonMeta):
+    """
+    A singleton class for creating a CloudWatchLogs client.
+    """
+
+    client: botocore.client = None
+
+    def __init__(self):
+        if not self.client:
+            session = Session()
+            self.client = session.client(
+                "logs",
+                session.region_name,
+                config=Config(retries={"max_attempts": 10, "mode": "standard"}),
+            )
+
+
+class LogStreamHandler:
+    log_group_name: str = None
+    log_stream_name: str = None
+    stream_id: int = None
+    next_token: str = None
+    cw_client = None
+
+    def __init__(self, log_group_name: str, log_stream_name: str, stream_id: int):
+        self.log_group_name = log_group_name
+        self.log_stream_name = log_stream_name
+        self.cw_client = CloudWatchLogsClient().client
+        self.stream_id = stream_id
+
+    def get_latest_log_events(self) -> Generator[Tuple[str, dict], None, None]:
+        """
+        This method gets all the latest log events for this stream that exist at this moment in time.
+
+        cw_client.get_log_events() always returns a nextForwardToken even if the current batch of events is empty.
+        You can keep calling cw_client.get_log_events() with the same token until a new batch of log events exist.
+
+        API Reference: https://boto3.amazonaws.com/v1/documentation/api/latest/reference/services/logs/client/get_log_events.html
+
+        Returns:
+            Generator[tuple[str, dict], None, None]: Generator that yields a tuple that consists for two values
+                str: stream_name,
+                dict: event dict in format
+                    {
+                        "ingestionTime": number,
+                        "message": "string",
+                        "timestamp": number
+                    }
+        """
+        while True:
+            if not self.next_token:
+                token_args = {}
+            else:
+                token_args = {"nextToken": self.next_token}
+
+            response = self.cw_client.get_log_events(
+                logGroupName=self.log_group_name,
+                logStreamName=self.log_stream_name,
+                startFromHead=True,
+                **token_args,
+            )
+
+            self.next_token = response["nextForwardToken"]
+            if not response["events"]:
+                break
+
+            for event in response["events"]:
+                yield self.log_stream_name, event
+
+
+class MultiLogStreamHandler:
+    log_group_name: str = None
+    log_stream_name_prefix: str = None
+    expected_stream_count: int = None
+    streams: List[LogStreamHandler] = []
+    cw_client = None
+
+    def __init__(
+        self, log_group_name: str, log_stream_name_prefix: str, expected_stream_count: int
+    ):
+        self.log_group_name = log_group_name
+        self.log_stream_name_prefix = log_stream_name_prefix
+        self.expected_stream_count = expected_stream_count
+        self.cw_client = CloudWatchLogsClient().client
+
+    def get_latest_log_events(self) -> Generator[Tuple[str, dict], None, None]:
+        """
+        This method gets all the latest log events from each stream that exist at this moment.
+
+        Returns:
+             Generator[tuple[str, dict], None, None]: Generator that yields a tuple that consists for two values
+                str: stream_name,
+                dict: event dict in format -
+                    {
+                        "ingestionTime": number,
+                        "message": "string",
+                        "timestamp": number
+                    }
+        """
+        if not self.ready():
+            return []
+
+        for stream in self.streams:
+            yield from stream.get_latest_log_events()
+
+    def ready(self) -> bool:
+        """
+        Checks whether or not MultiLogStreamHandler is ready to serve new log events at this moment.
+
+        If self.streams is already set, return True.
+        Otherwise, check if the current number of log streams in the log group match the exptected stream count.
+
+        Returns:
+            bool: Whether or not MultiLogStreamHandler is ready to serve new log events.
+        """
+
+        if len(self.streams) >= self.expected_stream_count:
+            return True
+
+        try:
+            response = self.cw_client.describe_log_streams(
+                logGroupName=self.log_group_name,
+                logStreamNamePrefix=self.log_stream_name_prefix + "/",
+                orderBy="LogStreamName",
+            )
+            stream_names = [stream["logStreamName"] for stream in response["logStreams"]]
+
+            next_token = response.get("nextToken")
+            while next_token:
+                response = self.cw_client.describe_log_streams(
+                    logGroupName=self.log_group_name,
+                    logStreamNamePrefix=self.log_stream_name_prefix + "/",
+                    orderBy="LogStreamName",
+                    nextToken=next_token,
+                )
+                stream_names.extend([stream["logStreamName"] for stream in response["logStreams"]])
+                next_token = response.get("nextToken", None)
+
+            if len(stream_names) >= self.expected_stream_count:
+                self.streams = [
+                    LogStreamHandler(self.log_group_name, log_stream_name, index)
+                    for index, log_stream_name in enumerate(stream_names)
+                ]
+
+                return True
+            else:
+                # Log streams are created whenever a container starts writing to stdout/err,
+                # so if the stream count is less than the expected number, return False
+                return False
+
+        except botocore.exceptions.ClientError as e:
+            # On the very first training job run on an account, there's no log group until
+            # the container starts logging, so ignore any errors thrown about that
+            if e.response["Error"]["Code"] == "ResourceNotFoundException":
+                return False
+            else:
+                raise

sagemaker_core/main/resources.py

@@ -41,6 +41,7 @@ from sagemaker_core.main.intelligent_defaults_helper import (
     load_default_configs_for_resource_name,
     get_config_value,
 )
+from sagemaker_core.main.logs import MultiLogStreamHandler
 from sagemaker_core.main.shapes import *
 from sagemaker_core.main.exceptions import *
 
@@ -812,7 +813,8 @@ class Algorithm(Base):
                 Group(progress, status),
                 title="Wait Log Panel",
                 border_style=Style(color=Color.BLUE.value),
-            )
+            ),
+            transient=True,
         ):
             while True:
                 self.refresh()
@@ -977,6 +979,7 @@ class App(Base):
         creation_time: The creation time of the application.  After an application has been shut down for 24 hours, SageMaker deletes all metadata for the application. To be considered an update and retain application metadata, applications must be restarted within 24 hours after the previous application has been shut down. After this time window, creation of an application is considered a new application rather than an update of the previous application.
         failure_reason: The failure reason.
         resource_spec: The instance type and the Amazon Resource Name (ARN) of the SageMaker image created on the instance.
+        built_in_lifecycle_config_arn: The lifecycle configuration that runs before the default lifecycle configuration
 
     """
 
@@ -992,6 +995,7 @@ class App(Base):
     creation_time: Optional[datetime.datetime] = Unassigned()
     failure_reason: Optional[str] = Unassigned()
     resource_spec: Optional[ResourceSpec] = Unassigned()
+    built_in_lifecycle_config_arn: Optional[str] = Unassigned()
 
     def get_name(self) -> str:
         attributes = vars(self)
@@ -1270,7 +1274,8 @@ class App(Base):
                 Group(progress, status),
                 title="Wait Log Panel",
                 border_style=Style(color=Color.BLUE.value),
-            )
+            ),
+            transient=True,
         ):
             while True:
                 self.refresh()
@@ -2652,7 +2657,11 @@ class AutoMLJob(Base):
         logger.info(f"Stopping {self.__class__.__name__} - {self.get_name()}")
 
     @Base.add_validate_call
-    def wait(self, poll: int = 5, timeout: Optional[int] = None) -> None:
+    def wait(
+        self,
+        poll: int = 5,
+        timeout: Optional[int] = None,
+    ) -> None:
         """
         Wait for a AutoMLJob resource.
 
@@ -2682,7 +2691,8 @@ class AutoMLJob(Base):
                 Group(progress, status),
                 title="Wait Log Panel",
                 border_style=Style(color=Color.BLUE.value),
-            )
+            ),
+            transient=True,
         ):
             while True:
                 self.refresh()
@@ -3130,7 +3140,11 @@ class AutoMLJobV2(Base):
         return self
 
     @Base.add_validate_call
-    def wait(self, poll: int = 5, timeout: Optional[int] = None) -> None:
+    def wait(
+        self,
+        poll: int = 5,
+        timeout: Optional[int] = None,
+    ) -> None:
         """
         Wait for a AutoMLJobV2 resource.
 
@@ -3160,7 +3174,8 @@ class AutoMLJobV2(Base):
                 Group(progress, status),
                 title="Wait Log Panel",
                 border_style=Style(color=Color.BLUE.value),
-            )
+            ),
+            transient=True,
         ):
             while True:
                 self.refresh()
@@ -3534,7 +3549,8 @@ class Cluster(Base):
                 Group(progress, status),
                 title="Wait Log Panel",
                 border_style=Style(color=Color.BLUE.value),
-            )
+            ),
+            transient=True,
         ):
             while True:
                 self.refresh()
@@ -4493,7 +4509,11 @@ class CompilationJob(Base):
         logger.info(f"Stopping {self.__class__.__name__} - {self.get_name()}")
 
     @Base.add_validate_call
-    def wait(self, poll: int = 5, timeout: Optional[int] = None) -> None:
+    def wait(
+        self,
+        poll: int = 5,
+        timeout: Optional[int] = None,
+    ) -> None:
         """
         Wait for a CompilationJob resource.
 
@@ -4523,7 +4543,8 @@ class CompilationJob(Base):
                 Group(progress, status),
                 title="Wait Log Panel",
                 border_style=Style(color=Color.BLUE.value),
-            )
+            ),
+            transient=True,
         ):
             while True:
                 self.refresh()
@@ -6581,7 +6602,8 @@ class Domain(Base):
                 Group(progress, status),
                 title="Wait Log Panel",
                 border_style=Style(color=Color.BLUE.value),
-            )
+            ),
+            transient=True,
         ):
             while True:
                 self.refresh()
@@ -7530,7 +7552,11 @@ class EdgePackagingJob(Base):
         logger.info(f"Stopping {self.__class__.__name__} - {self.get_name()}")
 
     @Base.add_validate_call
-    def wait(self, poll: int = 5, timeout: Optional[int] = None) -> None:
+    def wait(
+        self,
+        poll: int = 5,
+        timeout: Optional[int] = None,
+    ) -> None:
         """
         Wait for a EdgePackagingJob resource.
 
@@ -7560,7 +7586,8 @@ class EdgePackagingJob(Base):
                 Group(progress, status),
                 title="Wait Log Panel",
                 border_style=Style(color=Color.BLUE.value),
-            )
+            ),
+            transient=True,
         ):
             while True:
                 self.refresh()
@@ -8031,7 +8058,8 @@ class Endpoint(Base):
                 Group(progress, status),
                 title="Wait Log Panel",
                 border_style=Style(color=Color.BLUE.value),
-            )
+            ),
+            transient=True,
         ):
             while True:
                 self.refresh()
@@ -9547,7 +9575,8 @@ class FeatureGroup(Base):
                 Group(progress, status),
                 title="Wait Log Panel",
                 border_style=Style(color=Color.BLUE.value),
-            )
+            ),
+            transient=True,
         ):
             while True:
                 self.refresh()
@@ -10415,7 +10444,8 @@ class FlowDefinition(Base):
                 Group(progress, status),
                 title="Wait Log Panel",
                 border_style=Style(color=Color.BLUE.value),
-            )
+            ),
+            transient=True,
         ):
             while True:
                 self.refresh()
@@ -10910,7 +10940,8 @@ class Hub(Base):
                 Group(progress, status),
                 title="Wait Log Panel",
                 border_style=Style(color=Color.BLUE.value),
-            )
+            ),
+            transient=True,
         ):
             while True:
                 self.refresh()
@@ -11299,7 +11330,8 @@ class HubContent(Base):
                 Group(progress, status),
                 title="Wait Log Panel",
                 border_style=Style(color=Color.BLUE.value),
-            )
+            ),
+            transient=True,
         ):
             while True:
                 self.refresh()
@@ -11874,7 +11906,8 @@ class HumanTaskUi(Base):
                 Group(progress, status),
                 title="Wait Log Panel",
                 border_style=Style(color=Color.BLUE.value),
-            )
+            ),
+            transient=True,
         ):
             while True:
                 self.refresh()
@@ -12341,7 +12374,11 @@ class HyperParameterTuningJob(Base):
         logger.info(f"Stopping {self.__class__.__name__} - {self.get_name()}")
 
     @Base.add_validate_call
-    def wait(self, poll: int = 5, timeout: Optional[int] = None) -> None:
+    def wait(
+        self,
+        poll: int = 5,
+        timeout: Optional[int] = None,
+    ) -> None:
         """
         Wait for a HyperParameterTuningJob resource.
 
@@ -12371,7 +12408,8 @@ class HyperParameterTuningJob(Base):
                 Group(progress, status),
                 title="Wait Log Panel",
                 border_style=Style(color=Color.BLUE.value),
-            )
+            ),
+            transient=True,
         ):
             while True:
                 self.refresh()
@@ -12943,7 +12981,8 @@ class Image(Base):
                 Group(progress, status),
                 title="Wait Log Panel",
                 border_style=Style(color=Color.BLUE.value),
-            )
+            ),
+            transient=True,
         ):
             while True:
                 self.refresh()
@@ -13549,7 +13588,8 @@ class ImageVersion(Base):
                 Group(progress, status),
                 title="Wait Log Panel",
                 border_style=Style(color=Color.BLUE.value),
-            )
+            ),
+            transient=True,
         ):
             while True:
                 self.refresh()
@@ -13966,7 +14006,8 @@ class InferenceComponent(Base):
                 Group(progress, status),
                 title="Wait Log Panel",
                 border_style=Style(color=Color.BLUE.value),
-            )
+            ),
+            transient=True,
         ):
             while True:
                 self.refresh()
@@ -14601,7 +14642,8 @@ class InferenceExperiment(Base):
                 Group(progress, status),
                 title="Wait Log Panel",
                 border_style=Style(color=Color.BLUE.value),
-            )
+            ),
+            transient=True,
         ):
             while True:
                 self.refresh()
@@ -14974,7 +15016,11 @@ class InferenceRecommendationsJob(Base):
         logger.info(f"Stopping {self.__class__.__name__} - {self.get_name()}")
 
     @Base.add_validate_call
-    def wait(self, poll: int = 5, timeout: Optional[int] = None) -> None:
+    def wait(
+        self,
+        poll: int = 5,
+        timeout: Optional[int] = None,
+    ) -> None:
         """
         Wait for a InferenceRecommendationsJob resource.
 
@@ -15004,7 +15050,8 @@ class InferenceRecommendationsJob(Base):
                 Group(progress, status),
                 title="Wait Log Panel",
                 border_style=Style(color=Color.BLUE.value),
-            )
+            ),
+            transient=True,
         ):
             while True:
                 self.refresh()
@@ -15536,7 +15583,11 @@ class LabelingJob(Base):
         logger.info(f"Stopping {self.__class__.__name__} - {self.get_name()}")
 
     @Base.add_validate_call
-    def wait(self, poll: int = 5, timeout: Optional[int] = None) -> None:
+    def wait(
+        self,
+        poll: int = 5,
+        timeout: Optional[int] = None,
+    ) -> None:
         """
         Wait for a LabelingJob resource.
 
@@ -15566,7 +15617,8 @@ class LabelingJob(Base):
                 Group(progress, status),
                 title="Wait Log Panel",
                 border_style=Style(color=Color.BLUE.value),
-            )
+            ),
+            transient=True,
         ):
             while True:
                 self.refresh()
@@ -16314,7 +16366,8 @@ class MlflowTrackingServer(Base):
                 Group(progress, status),
                 title="Wait Log Panel",
                 border_style=Style(color=Color.BLUE.value),
-            )
+            ),
+            transient=True,
         ):
             while True:
                 self.refresh()
@@ -17572,7 +17625,8 @@ class ModelCard(Base):
                 Group(progress, status),
                 title="Wait Log Panel",
                 border_style=Style(color=Color.BLUE.value),
-            )
+            ),
+            transient=True,
         ):
             while True:
                 self.refresh()
@@ -17946,7 +18000,11 @@ class ModelCardExportJob(Base):
         return self
 
     @Base.add_validate_call
-    def wait(self, poll: int = 5, timeout: Optional[int] = None) -> None:
+    def wait(
+        self,
+        poll: int = 5,
+        timeout: Optional[int] = None,
+    ) -> None:
         """
         Wait for a ModelCardExportJob resource.
 
@@ -17976,7 +18034,8 @@ class ModelCardExportJob(Base):
                 Group(progress, status),
                 title="Wait Log Panel",
                 border_style=Style(color=Color.BLUE.value),
-            )
+            ),
+            transient=True,
         ):
             while True:
                 self.refresh()
@@ -18933,7 +18992,8 @@ class ModelPackage(Base):
                 Group(progress, status),
                 title="Wait Log Panel",
                 border_style=Style(color=Color.BLUE.value),
-            )
+            ),
+            transient=True,
         ):
             while True:
                 self.refresh()
@@ -19400,7 +19460,8 @@ class ModelPackageGroup(Base):
                 Group(progress, status),
                 title="Wait Log Panel",
                 border_style=Style(color=Color.BLUE.value),
-            )
+            ),
+            transient=True,
         ):
             while True:
                 self.refresh()
@@ -20789,7 +20850,8 @@ class MonitoringSchedule(Base):
                 Group(progress, status),
                 title="Wait Log Panel",
                 border_style=Style(color=Color.BLUE.value),
-            )
+            ),
+            transient=True,
         ):
             while True:
                 self.refresh()
@@ -21358,7 +21420,8 @@ class NotebookInstance(Base):
                 Group(progress, status),
                 title="Wait Log Panel",
                 border_style=Style(color=Color.BLUE.value),
-            )
+            ),
+            transient=True,
         ):
             while True:
                 self.refresh()
@@ -22197,7 +22260,11 @@ class OptimizationJob(Base):
         logger.info(f"Stopping {self.__class__.__name__} - {self.get_name()}")
 
     @Base.add_validate_call
-    def wait(self, poll: int = 5, timeout: Optional[int] = None) -> None:
+    def wait(
+        self,
+        poll: int = 5,
+        timeout: Optional[int] = None,
+    ) -> None:
         """
         Wait for a OptimizationJob resource.
 
@@ -22227,7 +22294,8 @@ class OptimizationJob(Base):
                 Group(progress, status),
                 title="Wait Log Panel",
                 border_style=Style(color=Color.BLUE.value),
-            )
+            ),
+            transient=True,
         ):
             while True:
                 self.refresh()
@@ -22698,7 +22766,8 @@ class Pipeline(Base):
                 Group(progress, status),
                 title="Wait Log Panel",
                 border_style=Style(color=Color.BLUE.value),
-            )
+            ),
+            transient=True,
         ):
             while True:
                 self.refresh()
@@ -23096,7 +23165,8 @@ class PipelineExecution(Base):
                 Group(progress, status),
                 title="Wait Log Panel",
                 border_style=Style(color=Color.BLUE.value),
-            )
+            ),
+            transient=True,
         ):
             while True:
                 self.refresh()
@@ -24073,13 +24143,19 @@ class ProcessingJob(Base):
         logger.info(f"Stopping {self.__class__.__name__} - {self.get_name()}")
 
     @Base.add_validate_call
-    def wait(self, poll: int = 5, timeout: Optional[int] = None) -> None:
+    def wait(
+        self,
+        poll: int = 5,
+        timeout: Optional[int] = None,
+        logs: Optional[bool] = False,
+    ) -> None:
         """
         Wait for a ProcessingJob resource.
 
         Parameters:
             poll: The number of seconds to wait between each poll.
             timeout: The maximum number of seconds to wait before timing out.
+            logs: Whether to print logs while waiting.
 
         Raises:
             TimeoutExceededError:  If the resource does not reach a terminal state before the timeout.
@@ -24098,18 +24174,32 @@ class ProcessingJob(Base):
         progress.add_task("Waiting for ProcessingJob...")
         status = Status("Current status:")
 
+        instance_count = self.processing_resources.cluster_config.instance_count
+        if logs:
+            multi_stream_logger = MultiLogStreamHandler(
+                log_group_name=f"/aws/sagemaker/ProcessingJobs",
+                log_stream_name_prefix=self.get_name(),
+                expected_stream_count=instance_count,
+            )
+
         with Live(
             Panel(
                 Group(progress, status),
                 title="Wait Log Panel",
                 border_style=Style(color=Color.BLUE.value),
-            )
+            ),
+            transient=True,
         ):
             while True:
                 self.refresh()
                 current_status = self.processing_job_status
                 status.update(f"Current status: [bold]{current_status}")
 
+                if logs and multi_stream_logger.ready():
+                    stream_log_events = multi_stream_logger.get_latest_log_events()
+                    for stream_id, event in stream_log_events:
+                        logger.info(f"{stream_id}:\n{event['message']}")
+
                 if current_status in terminal_states:
                     logger.info(f"Final Resource Status: [bold]{current_status}")
 
@@ -24537,7 +24627,8 @@ class Project(Base):
                 Group(progress, status),
                 title="Wait Log Panel",
                 border_style=Style(color=Color.BLUE.value),
-            )
+            ),
+            transient=True,
         ):
             while True:
                 self.refresh()
@@ -25182,7 +25273,8 @@ class Space(Base):
                 Group(progress, status),
                 title="Wait Log Panel",
                 border_style=Style(color=Color.BLUE.value),
-            )
+            ),
+            transient=True,
         ):
             while True:
                 self.refresh()
@@ -26433,13 +26525,19 @@ class TrainingJob(Base):
         logger.info(f"Stopping {self.__class__.__name__} - {self.get_name()}")
 
     @Base.add_validate_call
-    def wait(self, poll: int = 5, timeout: Optional[int] = None) -> None:
+    def wait(
+        self,
+        poll: int = 5,
+        timeout: Optional[int] = None,
+        logs: Optional[bool] = False,
+    ) -> None:
         """
         Wait for a TrainingJob resource.
 
         Parameters:
             poll: The number of seconds to wait between each poll.
             timeout: The maximum number of seconds to wait before timing out.
+            logs: Whether to print logs while waiting.
 
         Raises:
             TimeoutExceededError:  If the resource does not reach a terminal state before the timeout.
@@ -26458,18 +26556,41 @@ class TrainingJob(Base):
         progress.add_task("Waiting for TrainingJob...")
         status = Status("Current status:")
 
+        instance_count = (
+            sum(
+                instance_group.instance_count
+                for instance_group in self.resource_config.instance_groups
+            )
+            if self.resource_config.instance_groups
+            and not isinstance(self.resource_config.instance_groups, Unassigned)
+            else self.resource_config.instance_count
+        )
+
+        if logs:
+            multi_stream_logger = MultiLogStreamHandler(
+                log_group_name=f"/aws/sagemaker/TrainingJobs",
+                log_stream_name_prefix=self.get_name(),
+                expected_stream_count=instance_count,
+            )
+
         with Live(
             Panel(
                 Group(progress, status),
                 title="Wait Log Panel",
                 border_style=Style(color=Color.BLUE.value),
-            )
+            ),
+            transient=True,
         ):
             while True:
                 self.refresh()
                 current_status = self.training_job_status
                 status.update(f"Current status: [bold]{current_status}")
 
+                if logs and multi_stream_logger.ready():
+                    stream_log_events = multi_stream_logger.get_latest_log_events()
+                    for stream_id, event in stream_log_events:
+                        logger.info(f"{stream_id}:\n{event['message']}")
+
                 if current_status in terminal_states:
                     logger.info(f"Final Resource Status: [bold]{current_status}")
 
@@ -26884,13 +27005,19 @@ class TransformJob(Base):
         logger.info(f"Stopping {self.__class__.__name__} - {self.get_name()}")
 
     @Base.add_validate_call
-    def wait(self, poll: int = 5, timeout: Optional[int] = None) -> None:
+    def wait(
+        self,
+        poll: int = 5,
+        timeout: Optional[int] = None,
+        logs: Optional[bool] = False,
+    ) -> None:
         """
         Wait for a TransformJob resource.
 
         Parameters:
             poll: The number of seconds to wait between each poll.
             timeout: The maximum number of seconds to wait before timing out.
+            logs: Whether to print logs while waiting.
 
         Raises:
             TimeoutExceededError:  If the resource does not reach a terminal state before the timeout.
@@ -26909,18 +27036,32 @@ class TransformJob(Base):
         progress.add_task("Waiting for TransformJob...")
         status = Status("Current status:")
 
+        instance_count = self.transform_resources.instance_count
+        if logs:
+            multi_stream_logger = MultiLogStreamHandler(
+                log_group_name=f"/aws/sagemaker/TransformJobs",
+                log_stream_name_prefix=self.get_name(),
+                expected_stream_count=instance_count,
+            )
+
         with Live(
             Panel(
                 Group(progress, status),
                 title="Wait Log Panel",
                 border_style=Style(color=Color.BLUE.value),
-            )
+            ),
+            transient=True,
         ):
             while True:
                 self.refresh()
                 current_status = self.transform_job_status
                 status.update(f"Current status: [bold]{current_status}")
 
+                if logs and multi_stream_logger.ready():
+                    stream_log_events = multi_stream_logger.get_latest_log_events()
+                    for stream_id, event in stream_log_events:
+                        logger.info(f"{stream_id}:\n{event['message']}")
+
                 if current_status in terminal_states:
                     logger.info(f"Final Resource Status: [bold]{current_status}")
 
@@ -27736,7 +27877,8 @@ class TrialComponent(Base):
                 Group(progress, status),
                 title="Wait Log Panel",
                 border_style=Style(color=Color.BLUE.value),
-            )
+            ),
+            transient=True,
         ):
             while True:
                 self.refresh()
@@ -28397,7 +28539,8 @@ class UserProfile(Base):
                 Group(progress, status),
                 title="Wait Log Panel",
                 border_style=Style(color=Color.BLUE.value),
-            )
+            ),
+            transient=True,
         ):
             while True:
                 self.refresh()
@@ -28882,7 +29025,8 @@ class Workforce(Base):
                 Group(progress, status),
                 title="Wait Log Panel",
                 border_style=Style(color=Color.BLUE.value),
-            )
+            ),
+            transient=True,
         ):
             while True:
                 self.refresh()

sagemaker_core/main/shapes.py

@@ -3306,12 +3306,14 @@ class CodeEditorAppSettings(Base):
     custom_images: A list of custom SageMaker images that are configured to run as a Code Editor app.
     lifecycle_config_arns: The Amazon Resource Name (ARN) of the Code Editor application lifecycle configuration.
     app_lifecycle_management: Settings that are used to configure and manage the lifecycle of CodeEditor applications.
+    built_in_lifecycle_config_arn: The lifecycle configuration that runs before the default lifecycle configuration. It can override changes made in the default lifecycle configuration.
     """
 
     default_resource_spec: Optional[ResourceSpec] = Unassigned()
     custom_images: Optional[List[CustomImage]] = Unassigned()
     lifecycle_config_arns: Optional[List[str]] = Unassigned()
     app_lifecycle_management: Optional[AppLifecycleManagement] = Unassigned()
+    built_in_lifecycle_config_arn: Optional[str] = Unassigned()
 
 
 class CodeRepository(Base):
@@ -4235,6 +4237,7 @@ class JupyterLabAppSettings(Base):
     code_repositories: A list of Git repositories that SageMaker automatically displays to users for cloning in the JupyterLab application.
     app_lifecycle_management: Indicates whether idle shutdown is activated for JupyterLab applications.
     emr_settings: The configuration parameters that specify the IAM roles assumed by the execution role of SageMaker (assumable roles) and the cluster instances or job execution environments (execution roles or runtime roles) to manage and access resources required for running Amazon EMR clusters or Amazon EMR Serverless applications.
+    built_in_lifecycle_config_arn: The lifecycle configuration that runs before the default lifecycle configuration. It can override changes made in the default lifecycle configuration.
     """
 
     default_resource_spec: Optional[ResourceSpec] = Unassigned()
@@ -4243,6 +4246,7 @@ class JupyterLabAppSettings(Base):
     code_repositories: Optional[List[CodeRepository]] = Unassigned()
     app_lifecycle_management: Optional[AppLifecycleManagement] = Unassigned()
     emr_settings: Optional[EmrSettings] = Unassigned()
+    built_in_lifecycle_config_arn: Optional[str] = Unassigned()
 
 
 class DefaultEbsStorageSettings(Base):

sagemaker_core/main/utils.py

@@ -160,6 +160,12 @@ def enable_textual_rich_console_and_traceback():
         textual_rich_console_and_traceback_enabled = True
 
 
+def get_rich_handler():
+    handler = RichHandler(markup=True)
+    handler.setFormatter(logging.Formatter("%(message)s"))
+    return handler
+
+
 def get_textual_rich_logger(name: str, log_level: str = "INFO") -> logging.Logger:
     """
     Get a logger with textual rich handler.
@@ -175,7 +181,7 @@ def get_textual_rich_logger(name: str, log_level: str = "INFO") -> logging.Logge
 
     """
     enable_textual_rich_console_and_traceback()
-    handler = RichHandler(markup=True)
+    handler = get_rich_handler()
     logging.basicConfig(level=getattr(logging, log_level), handlers=[handler])
     logger = logging.getLogger(name)
 
@@ -217,8 +223,8 @@ def configure_logging(log_level=None):
     # reset any currently associated handlers with log level
     for handler in _logger.handlers:
         _logger.removeHandler(handler)
-    console_handler = RichHandler(markup=True)
-    _logger.addHandler(console_handler)
+    rich_handler = get_rich_handler()
+    _logger.addHandler(rich_handler)
 
 
 def is_snake_case(s: str):

sagemaker_core/tools/constants.py

@@ -20,6 +20,8 @@ OBJECT_METHODS = set(
 
 TERMINAL_STATES = set(["Completed", "Stopped", "Deleted", "Failed", "Succeeded", "Cancelled"])
 
+RESOURCE_WITH_LOGS = set(["TrainingJob", "ProcessingJob", "TransformJob"])
+
 CONFIGURABLE_ATTRIBUTE_SUBSTRINGS = [
     "kms",
     "s3",

sagemaker_core/tools/resources_codegen.py

@@ -29,6 +29,7 @@ from sagemaker_core.tools.constants import (
     CONFIG_SCHEMA_FILE_NAME,
     PYTHON_TYPES_TO_BASIC_JSON_TYPES,
     CONFIGURABLE_ATTRIBUTE_SUBSTRINGS,
+    RESOURCE_WITH_LOGS,
 )
 from sagemaker_core.tools.method import Method, MethodType
 from sagemaker_core.main.utils import (
@@ -71,6 +72,8 @@ from sagemaker_core.tools.templates import (
     GET_ALL_METHOD_WITH_ARGS_TEMPLATE,
     UPDATE_METHOD_TEMPLATE_WITHOUT_DECORATOR,
     RESOURCE_METHOD_EXCEPTION_DOCSTRING,
+    INIT_WAIT_LOGS_TEMPLATE,
+    PRINT_WAIT_LOGS,
 )
 from sagemaker_core.tools.data_extractor import (
     load_combined_shapes_data,
@@ -188,6 +191,7 @@ class ResourcesCodeGen:
             "from sagemaker_core.main.utils import SageMakerClient, ResourceIterator, Unassigned, get_textual_rich_logger, "
             "snake_to_pascal, pascal_to_snake, is_not_primitive, is_not_str_dict, is_primitive_list, serialize",
             "from sagemaker_core.main.intelligent_defaults_helper import load_default_configs_for_resource_name, get_config_value",
+            "from sagemaker_core.main.logs import MultiLogStreamHandler",
             "from sagemaker_core.main.shapes import *",
             "from sagemaker_core.main.exceptions import *",
         ]
@@ -1541,6 +1545,28 @@ class ResourcesCodeGen:
 
         return "'(Unknown)'"
 
+    def _get_instance_count_ref(self, resource_name: str) -> str:
+        """Get the instance count reference for a resource object.
+        Args:
+            resource_name (str): The resource name.
+        Returns:
+            str: The instance count reference for resource object
+        """
+
+        if resource_name == "TrainingJob":
+            return """(
+                sum(instance_group.instance_count for instance_group in self.resource_config.instance_groups)
+                if self.resource_config.instance_groups and not isinstance(self.resource_config.instance_groups, Unassigned)
+                else self.resource_config.instance_count
+            )
+            """
+        elif resource_name == "TransformJob":
+            return "self.transform_resources.instance_count"
+        elif resource_name == "ProcessingJob":
+            return "self.processing_resources.cluster_config.instance_count"
+
+        raise ValueError(f"Instance count reference not found for resource {resource_name}")
+
     def generate_wait_method(self, resource_name: str) -> str:
         """Auto-Generate WAIT Method for a waitable resource.
 
@@ -1573,11 +1599,32 @@ class ResourcesCodeGen:
         )
         formatted_failed_block = add_indent(formatted_failed_block, 16)
 
+        logs_arg = ""
+        logs_arg_doc = ""
+        init_wait_logs = ""
+        print_wait_logs = ""
+        if resource_name in RESOURCE_WITH_LOGS:
+            logs_arg = "logs: Optional[bool] = False,"
+            logs_arg_doc = "logs: Whether to print logs while waiting.\n"
+
+            instance_count = self._get_instance_count_ref(resource_name)
+            init_wait_logs = add_indent(
+                INIT_WAIT_LOGS_TEMPLATE.format(
+                    get_instance_count=instance_count,
+                    job_type=resource_name,
+                )
+            )
+            print_wait_logs = add_indent(PRINT_WAIT_LOGS, 12)
+
         formatted_method = WAIT_METHOD_TEMPLATE.format(
             terminal_resource_states=terminal_resource_states,
             status_key_path=status_key_path,
             failed_error_block=formatted_failed_block,
             resource_name=resource_name,
+            logs_arg=logs_arg,
+            logs_arg_doc=logs_arg_doc,
+            init_wait_logs=init_wait_logs,
+            print_wait_logs=print_wait_logs,
         )
         return formatted_method
 

sagemaker_core/tools/templates.py

@@ -262,12 +262,31 @@ if "failed" in current_status.lower():
     raise FailedStatusError(resource_type="{resource_name}", status=current_status, reason={reason})
 """
 
+INIT_WAIT_LOGS_TEMPLATE = """
+instance_count = {get_instance_count}
+if logs:
+    multi_stream_logger = MultiLogStreamHandler(
+        log_group_name=f"/aws/sagemaker/{job_type}s",
+        log_stream_name_prefix=self.get_name(),
+        expected_stream_count=instance_count
+    )
+"""
+
+PRINT_WAIT_LOGS = """
+if logs and multi_stream_logger.ready():
+    stream_log_events = multi_stream_logger.get_latest_log_events()
+    for stream_id, event in stream_log_events:
+        logger.info(f"{stream_id}:\\n{event['message']}")
+"""
+
+
 WAIT_METHOD_TEMPLATE = '''
 @Base.add_validate_call
 def wait(
     self,
     poll: int = 5,
-    timeout: Optional[int] = None
+    timeout: Optional[int] = None,
+    {logs_arg}
 ) -> None:
     """
     Wait for a {resource_name} resource.
@@ -275,7 +294,7 @@ def wait(
     Parameters:
         poll: The number of seconds to wait between each poll.
         timeout: The maximum number of seconds to wait before timing out.
-    
+        {logs_arg_doc}
     Raises:
         TimeoutExceededError:  If the resource does not reach a terminal state before the timeout.
         FailedStatusError:   If the resource reaches a failed state.
@@ -291,13 +310,22 @@ def wait(
     )
     progress.add_task("Waiting for {resource_name}...")
     status = Status("Current status:")
-
-    with Live(Panel(Group(progress, status), title="Wait Log Panel", border_style=Style(color=Color.BLUE.value))):
+    {init_wait_logs}
+
+    with Live(
+        Panel(
+            Group(progress, status),
+            title="Wait Log Panel",
+            border_style=Style(color=Color.BLUE.value
+            )
+        ),
+        transient=True
+    ):
         while True:
             self.refresh()
             current_status = self{status_key_path}
             status.update(f"Current status: [bold]{{current_status}}")
-
+            {print_wait_logs}
             if current_status in terminal_states:
                 logger.info(f"Final Resource Status: [bold]{{current_status}}")
 {failed_error_block}
@@ -338,7 +366,15 @@ def wait_for_status(
     progress.add_task(f"Waiting for {resource_name} to reach [bold]{{target_status}} status...")
     status = Status("Current status:")
 
-    with Live(Panel(Group(progress, status), title="Wait Log Panel", border_style=Style(color=Color.BLUE.value))):
+    with Live(
+        Panel(
+            Group(progress, status),
+            title="Wait Log Panel",
+            border_style=Style(color=Color.BLUE.value
+            )
+        ),
+        transient=True
+    ):
         while True:
             self.refresh()
             current_status = self{status_key_path}

{sagemaker_core-1.0.9.dist-info → sagemaker_core-1.0.10.dist-info}/METADATA

@@ -1,6 +1,6 @@
 Metadata-Version: 2.1
 Name: sagemaker-core
-Version: 1.0.9
+Version: 1.0.10
 Summary: An python package for sagemaker core functionalities
 Author-email: AWS <sagemaker-interests@amazon.com>
 Project-URL: Repository, https://github.com/aws/sagemaker-core.git

{sagemaker_core-1.0.9.dist-info → sagemaker_core-1.0.10.dist-info}/RECORD

@@ -6,29 +6,30 @@ sagemaker_core/main/__init__.py,sha256=47DEQpj8HBSa-_TImW-5JCeuQeRkm5NMpJWZG3hSu
 sagemaker_core/main/config_schema.py,sha256=TeGoTobT4yotEVyfguLF0IdKYlOymsDZ45ySxXiCDuw,56998
 sagemaker_core/main/exceptions.py,sha256=87DUlrmHxaWoiYNlpNY9ixxFMPRk_dIGPsA2e_xdVwQ,5602
 sagemaker_core/main/intelligent_defaults_helper.py,sha256=5SDM6UavZtp-k5LhqRL7GRIDgzFB5UsC_p7YuiSPK9A,8334
-sagemaker_core/main/resources.py,sha256=seW0HAoyQlKGwXaRN9YrvK2SGJhAsZGh7-b5IpAJ2q4,1317128
-sagemaker_core/main/shapes.py,sha256=Groaj8psouJwQpRljmZIDeRM6fyU6p2J3wgjzcvQ_5k,696205
+sagemaker_core/main/logs.py,sha256=yfEH7uP91nbE1lefymOlBr81ziBzsDSIOF2Qyd54FJE,6241
+sagemaker_core/main/resources.py,sha256=uJs7vsqZcu6-xjLcMCX-Ys8wDJQ7LAHtfaf3riMygv8,1321298
+sagemaker_core/main/shapes.py,sha256=5loClZSSYkEmdshd8UaT0d4VzMtahZSMlLpIO600g7o,696705
 sagemaker_core/main/user_agent.py,sha256=4sZybDXkzRoZnOnVDQ8p8zFTfiRJdsH7amDWInVQ4xU,2708
-sagemaker_core/main/utils.py,sha256=lXkJyiCow5uf32l0EmkimB0RKVk2BS7OM2fYoLsOfD4,18346
+sagemaker_core/main/utils.py,sha256=LCFDM6oxf6_e1i-_Dgtkm3ehl7YfoEpJ2kTTFTL6iOU,18471
 sagemaker_core/main/code_injection/__init__.py,sha256=47DEQpj8HBSa-_TImW-5JCeuQeRkm5NMpJWZG3hSuFU,0
 sagemaker_core/main/code_injection/base.py,sha256=11_Jif0nOzfbLGlXaacKf-wcizzfS64U0OSZGoVffFU,1733
 sagemaker_core/main/code_injection/codec.py,sha256=2DjmeD2uND307UqDefvVEpE0rZ8yfFU3Bi3TvQCQveI,7658
 sagemaker_core/main/code_injection/constants.py,sha256=2ICExGge8vAWx7lSTW0JGh-bH1korkvpOpDu5M63eI4,980
-sagemaker_core/main/code_injection/shape_dag.py,sha256=YHFV0V37x_DN4oSMsFiEebMBkoixFUREmD2Bp2t4uDU,658150
+sagemaker_core/main/code_injection/shape_dag.py,sha256=FyU_a0Jh3AVxTNVN2kjt0oHqqyHNjAgIy3G7kzqksm4,658657
 sagemaker_core/resources/__init__.py,sha256=EAYTFMN-nPjnPjjBbhIUeaL67FLKNPd7qbcbl9VIrws,31
 sagemaker_core/shapes/__init__.py,sha256=RnbIu9eTxKt-DNsOFJabrWIgrrtS9_SdAozP9JBl_ic,28
 sagemaker_core/tools/__init__.py,sha256=xX79JImxCVzrWMnjgntLCve2G5I-R4pRar5s20kT9Rs,56
 sagemaker_core/tools/codegen.py,sha256=mKWVi2pWnPxyIoWUEPYjEc9Gw7D9bCOrHqa00yzIZ1o,2005
-sagemaker_core/tools/constants.py,sha256=8oM0nHuzXWnOJTIdXsu50f9RsEz4rgmmZKWhLkrTP-s,3309
+sagemaker_core/tools/constants.py,sha256=a2WjUDK7gzxgilZs99vp30qh4kQ-y6JKhrwwqVAA12o,3385
 sagemaker_core/tools/data_extractor.py,sha256=pNfmTA0NUA96IgfLrla7a36Qjc1NljbwgZYaOhouKqQ,2113
 sagemaker_core/tools/method.py,sha256=4Hmu4UWpiBgUTZljYdW1KIKDduDxf_nfhCyuWgLVMWI,717
-sagemaker_core/tools/resources_codegen.py,sha256=ShsWjABBkSpPlp3gByzpe36M6hfYfBofODlUynqVXTY,82405
+sagemaker_core/tools/resources_codegen.py,sha256=ASirF9UMkGAOYZrrZxmFqK3gYQ4YbVYdAiUcWt6qdII,84360
 sagemaker_core/tools/resources_extractor.py,sha256=hN61ehZbPnhFW-2FIVDi7NsEz4rLvGr-WoglHQGfrug,14523
 sagemaker_core/tools/shapes_codegen.py,sha256=_ve959bwH8usZ6dPlpXxi2on9t0hLpcmhRWnaWHCWMQ,11745
 sagemaker_core/tools/shapes_extractor.py,sha256=4KjgDmhlPM4G1f1NeYbORKlXs1s7Q_sm_NK31S_ROQ0,11950
-sagemaker_core/tools/templates.py,sha256=nze_A01EpegYUwoR_gRv2qBNKNFruBY8L3RiIX5lz3M,22458
-sagemaker_core-1.0.9.dist-info/LICENSE,sha256=xx0jnfkXJvxRnG63LTGOxlggYnIysveWIZ6H3PNdCrQ,11357
-sagemaker_core-1.0.9.dist-info/METADATA,sha256=uLiMbw5h-h1HaXDvoQ_VR0h6AOyxddhLyNZKar6FlLY,4877
-sagemaker_core-1.0.9.dist-info/WHEEL,sha256=GV9aMThwP_4oNCtvEC2ec3qUYutgWeAzklro_0m4WJQ,91
-sagemaker_core-1.0.9.dist-info/top_level.txt,sha256=R3GAZZ1zC5JxqdE_0x2Lu_WYi2Xfke7VsiP3L5zngfA,15
-sagemaker_core-1.0.9.dist-info/RECORD,,
+sagemaker_core/tools/templates.py,sha256=yX2RQKeClgYwKS5Qu_mDpnWJIBCuj0yELrdm95aiTpk,23262
+sagemaker_core-1.0.10.dist-info/LICENSE,sha256=xx0jnfkXJvxRnG63LTGOxlggYnIysveWIZ6H3PNdCrQ,11357
+sagemaker_core-1.0.10.dist-info/METADATA,sha256=LJ5rjbEJgXzd3afuvD5kIzDQ2Yxx1d2wCobLb5Jy7OI,4878
+sagemaker_core-1.0.10.dist-info/WHEEL,sha256=GV9aMThwP_4oNCtvEC2ec3qUYutgWeAzklro_0m4WJQ,91
+sagemaker_core-1.0.10.dist-info/top_level.txt,sha256=R3GAZZ1zC5JxqdE_0x2Lu_WYi2Xfke7VsiP3L5zngfA,15
+sagemaker_core-1.0.10.dist-info/RECORD,,