Add logs for subset of jobs

Author: benieric

example_notebooks/sagemaker_core_overview.ipynb

@@ -426,7 +426,7 @@
     "    stopping_condition=StoppingCondition(\n",
     "        max_runtime_in_seconds=600\n",
     "    )\n",
-    ")\n"
+    ")"
    ]
   },
   {
@@ -445,7 +445,7 @@
    "source": [
     "# Wait for TrainingJob with SageMakerCore\n",
     "\n",
-    "training_job.wait()"
+    "training_job.wait(logs=True)"
    ]
   },
   {

src/sagemaker_core/main/logs.py

@@ -2,49 +2,53 @@
 import botocore
 
 from boto3.session import Session
+import botocore.client
 from botocore.config import Config
-from typing import Generator
-from sagemaker_core.generated.utils import SingletonMeta, logger
+from typing import Generator, Tuple, List
+from sagemaker_core.main.utils import SingletonMeta
 
 
 class CloudWatchLogsClient(metaclass=SingletonMeta):
     """
     A singleton class for creating a SageMaker client.
     """
+
     def __init__(self):
         session = Session()
         self.client = session.client(
-            "logs", 
-            session.region_name, 
-            config=Config(retries={"max_attempts": 10, "mode": "standard"})
+            "logs",
+            session.region_name,
+            config=Config(retries={"max_attempts": 10, "mode": "standard"}),
         )
 
 
 class LogStreamHandler:
     log_groupt_name: str = None
     log_stream_name: str = None
+    stream_id: int = None
     next_token: str = None
-    cw_client: CloudWatchLogsClient = None
-        
-    def __init__(self, log_group_name: str, log_stream_name: str):
+    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()
+        self.cw_client = CloudWatchLogsClient().client
+        self.stream_id = stream_id
 
-    def get_latest_log_events(self) -> Generator[tuple[str, dict], None, None]:
+    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. 
+
+        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
@@ -54,109 +58,107 @@ def get_latest_log_events(self) -> Generator[tuple[str, dict], None, None]:
             if not self.next_token:
                 token_args = {}
             else:
-                token_args = {
-                    "nextToken": self.next_token
-                }
-            
+                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
+                **token_args,
             )
 
             self.next_token = response["nextForwardToken"]
             if not response["events"]:
                 break
-            
+
             for event in response["events"]:
-                yield self.log_stream_name, event
+                yield self.stream_id, event
 
 
 class MultiLogStreamHandler:
     log_group_name: str = None
     log_stream_name_prefix: str = None
     expected_stream_count: int = None
-    streams: list[LogStreamHandler] = []
-    cw_client: CloudWatchLogsClient = None
-    
-    def __init__(self, log_group_name: str, log_stream_name_prefix: str, expected_stream_count: int):
+    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()
+        self.cw_client = CloudWatchLogsClient().client
 
-    def get_latest_log_events(self) -> Generator[tuple[str, dict], None, None]:
+    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 - 
-                    { 
+                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()            
-            
 
+        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.        
+        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')
+            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, 
+                    logGroupName=self.log_group_name,
                     logStreamNamePrefix=self.log_stream_name_prefix + "/",
                     orderBy="LogStreamName",
-                    nextToken=next_token
+                    nextToken=next_token,
                 )
-                stream_names.extend([stream['logStreamName'] for stream in response['logStreams']])
-                next_token = response.get('nextToken', None)
-                
+                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) 
-                    for log_stream_name in stream_names
+                    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   
-            
+                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':
+            if e.response["Error"]["Code"] == "ResourceNotFoundException":
                 return False
             else:
-                raise
+                raise