diff --git a/sagemaker-train/src/sagemaker/train/base_trainer.py b/sagemaker-train/src/sagemaker/train/base_trainer.py index 9af37e9e91..2d4e981365 100644 --- a/sagemaker-train/src/sagemaker/train/base_trainer.py +++ b/sagemaker-train/src/sagemaker/train/base_trainer.py @@ -1,5 +1,4 @@ import copy -import os import time import yaml from abc import ABC, abstractmethod @@ -15,11 +14,9 @@ import yaml import boto3 -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, 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 @@ -41,6 +38,7 @@ from sagemaker.train.common_utils.notifications import enable_notifications, delete_notification_rule, list_notification_rules 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.common_utils.log_streamer import LogStreamer, stream_log_loop from sagemaker.core.telemetry.telemetry_logging import _telemetry_emitter, TelemetryParamType from sagemaker.core.telemetry.constants import Feature from sagemaker.train.defaults import TrainDefaults @@ -703,10 +701,6 @@ def stream_logs(self, poll: int = 5, start_time: Optional[Any] = None, tail_line def _stream_logs_smtj(self, training_job, poll: int, start_time_ms=None, tail_lines: Optional[int] = None) -> None: """Stream logs for an SMTJ training job.""" - from sagemaker.train.common_utils.log_streamer import ( - LogStreamer, - stream_log_loop, - ) if hasattr(training_job, 'training_job_name'): job_name = training_job.training_job_name @@ -736,7 +730,12 @@ def _get_status() -> str: stream_log_loop(streamer, poll, _get_status, tail_lines=tail_lines) def _stream_logs_smhp(self, training_job, compute, poll: int, start_time_ms=None, tail_lines: Optional[int] = None) -> None: - """Stream logs for a HyperPod job using filter_log_events polling.""" + """Stream logs for a HyperPod job using LogStreamer with filter mode. + + Delegates to stream_log_loop for consistent behavior with SMTJ/MTRL paths. + HyperPod jobs have no simple status API, so the status function always + returns "InProgress" — the loop exits via KeyboardInterrupt or tail_lines. + """ if isinstance(training_job, str): job_id = training_job @@ -748,8 +747,6 @@ def _stream_logs_smhp(self, training_job, compute, poll: int, start_time_ms=None sagemaker_session = TrainDefaults.get_sagemaker_session( sagemaker_session=self.sagemaker_session ) - region_name = sagemaker_session.boto_session.region_name - logs_client = sagemaker_session.boto_session.client("logs", region_name=region_name) log_group = _get_smhp_log_group(compute.cluster_name, sagemaker_session.sagemaker_client) logger.info(f"Streaming logs for HyperPod job: {job_id}") @@ -758,78 +755,29 @@ def _stream_logs_smhp(self, training_job, compute, poll: int, start_time_ms=None logger.info("Press Ctrl+C to stop streaming.") # Pick start time (user-provided > training job start time > now) - if start_time_ms is not None: - last_timestamp = start_time_ms - elif hasattr(training_job, 'training_start_time') and training_job.training_start_time: - try: - last_timestamp = int(training_job.training_start_time.timestamp() * 1000) - except Exception: - last_timestamp = int(time.time() * 1000) - else: - last_timestamp = int(time.time() * 1000) - seen_event_ids = set() - lines_printed = 0 - _CW_PREFIX = "[CloudWatch] " + if start_time_ms is None: + if hasattr(training_job, 'training_start_time') and training_job.training_start_time: + try: + start_time_ms = int(training_job.training_start_time.timestamp() * 1000) + except Exception: + start_time_ms = int(time.time() * 1000) + else: + start_time_ms = int(time.time() * 1000) - empty_cycles = 0 - while True: - try: - params = { - "logGroupName": log_group, - "logStreamNamePrefix": "SagemakerHyperPodTrainingJob", - "filterPattern": f'"{job_id}"', - "startTime": last_timestamp, - } - response = logs_client.filter_log_events(**params) - events = response.get("events", []) - - if events: - empty_cycles = 0 - for event in events: - event_id = event.get("eventId", "") - if event_id not in seen_event_ids: - seen_event_ids.add(event_id) - message = event.get("message", "").rstrip() - if message: - print(f"{_CW_PREFIX}{message}") - lines_printed += 1 - if tail_lines and lines_printed >= tail_lines: - logger.info(f"Reached tail_lines limit ({tail_lines}). Stopping log stream.") - return - ts = event.get("timestamp", 0) - if ts > last_timestamp: - last_timestamp = ts - if not events: - empty_cycles += 1 - if empty_cycles == 3: - logger.info("No log events yet, still waiting...") - except ClientError as e: - error_code = e.response.get("Error", {}).get("Code", "") - if error_code == "AccessDeniedException": - raise - if error_code == "ResourceNotFoundException": - empty_cycles += 1 - if empty_cycles == 1: - logger.info("Waiting for log group to become available...") - elif empty_cycles >= 60: - logger.warning( - "Log group %s still not found after %d attempts. " - "Check IAM permissions for logs:FilterLogEvents.", - log_group, - empty_cycles, - ) - else: - logger.debug(f"Error fetching HP logs: {e}") - except Exception as e: - logger.debug(f"Error fetching HP logs: {e}") + streamer = LogStreamer( + log_group=log_group, + job_name=job_id, + sagemaker_session=sagemaker_session, + filter_pattern=f'"{job_id}"', + start_time_ms=start_time_ms, + ) - # Note: HyperPod jobs don't have a simple status API to poll for completion. - # This polls till the user interrupts with Ctrl+C. - try: - time.sleep(poll) - except KeyboardInterrupt: - logger.info("Log streaming stopped by user.") - return + # HyperPod jobs have no simple status API — always report "InProgress" + # so the loop runs until KeyboardInterrupt or tail_lines completes. + def _get_status() -> str: + return "InProgress" + + stream_log_loop(streamer, poll, _get_status, tail_lines=tail_lines) def _validate_instance_count(self, instance_count, sagemaker_session, compute): """Validate instance/node count against allowed values from SMHP recipe. 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 570a3331e3..3d2ad26c02 100644 --- a/sagemaker-train/src/sagemaker/train/common_utils/log_streamer.py +++ b/sagemaker-train/src/sagemaker/train/common_utils/log_streamer.py @@ -125,6 +125,132 @@ def poll_once(self) -> list[tuple[int, str]]: logger.debug("Transient CloudWatch error: %s", e) return [] + def poll_tail(self, n: int) -> list[tuple[int, str]]: + """Fetch the last N log events in chronological order. + + Behaves like ``tail -n`` or ``kubectl logs --tail=N``. + + For stream mode (SMTJ): uses get_log_events backward pagination to + fetch the last N events per stream, merges by timestamp, returns the + globally last N. + + For filter mode (SMHP): uses filter_log_events with startFromHead=False + to fetch recent events. Requires startTime to be set for CloudWatch to + scope the search efficiently. + + :param n: Number of most recent log events to return. + :returns: List of (timestamp_ms, message) tuples in chronological order. + """ + if self._filter_pattern is not None: + return self._tail_filter_mode(n) + return self._tail_stream_mode(n) + + def _tail_stream_mode(self, n: int) -> list[tuple[int, str]]: + """Get last N events via get_log_events backward pagination. + + For multi-stream jobs, fetches last N from each stream, merges by + timestamp, and returns the globally last N events in chronological order. + """ + if self._stream_handlers is None: + self._stream_handlers = self._discover_streams() + if not self._stream_handlers: + return [] + + all_results = [] + for handler in self._stream_handlers: + events = [] + next_token = None + + while len(events) < n: + kwargs = { + "logGroupName": self._log_group, + "logStreamName": handler["stream_name"], + "limit": n - len(events), + "startFromHead": False, + } + if next_token: + kwargs["nextToken"] = next_token + + response = self._logs_client.get_log_events(**kwargs) + if response.get("events"): + events.extend(response["events"]) + + backward_token = response.get("nextBackwardToken") + if backward_token and backward_token != next_token: + next_token = backward_token + else: + break + + for event in events: + message = event.get("message", "").rstrip() + ts = event.get("timestamp", 0) + if message: + all_results.append((ts, message)) + + # Sort by timestamp across all streams, take the last N globally + all_results.sort(key=lambda x: x[0]) + return all_results[-n:] + + def _tail_filter_mode(self, n: int) -> list[tuple[int, str]]: + """Get last N events via filter_log_events with startFromHead=False. + + Uses reverse-chronological order to get the most recent events first. + Requires startTime to be set for CloudWatch to scope the search. + Paginates without limit (faster scanning), then slices client-side. + + Note: startFromHead=False with logStreamNamePrefix may require several + pagination calls before CloudWatch locates the matching streams. + """ + # CloudWatch requires startTime on or after 2024-01-01 for + # startFromHead=False with filter_log_events. + _JAN_1_2024_MS = 1704067200000 + if self._last_timestamp_ms and self._last_timestamp_ms < _JAN_1_2024_MS: + raise ValueError( + "stream_logs does not support tail_lines when start_time is before 2024-01-01." + ) + params = { + "logGroupName": self._log_group, + "logStreamNamePrefix": _SMHP_STREAM_PREFIX, + "filterPattern": self._filter_pattern, + "startFromHead": False, + } + if self._last_timestamp_ms is not None: + params["startTime"] = self._last_timestamp_ms + else: + logger.warning( + "No start_time provided for tail_lines. Scanning without time " + "bounds may take a while to identify matching log streams." + ) + + results = [] + next_token = None + + # filter_log_events bounds pages by scan volume, not result count. + # Must follow nextToken until N matching events are collected. + while True: + if next_token: + params["nextToken"] = next_token + + response = self._logs_client.filter_log_events(**params) + for event in response.get("events", []): + message = event.get("message", "").rstrip() + ts = event.get("timestamp", 0) + if message: + results.append((ts, message)) + + # Stop once we have enough events + if len(results) >= n: + break + + next_token = response.get("nextToken") + if not next_token: + break + + # Events come in reverse chronological order; take first N and reverse + results = results[:n] + results.reverse() + return results + def _poll_filter_mode(self) -> list[tuple[int, str]]: """Poll using filter_log_events (HyperPod style).""" params = { @@ -238,22 +364,23 @@ 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 tail_lines: Optional maximum number of most recent log lines to - print. When specified, streaming stops after this many lines have - been displayed. + :param tail_lines: Optional number of most recent log events to return. + Fetches the last N events (like ``tail -n`` or ``kubectl logs --tail``), + regardless of whether the job is still running or completed. + If not provided, streams all logs until the job completes. """ _CW_PREFIX = "[CloudWatch] " - lines_printed = 0 - def _print_event(ts_ms: int, message: str) -> bool: - """Print a log event. Returns True if tail_lines limit reached.""" - nonlocal lines_printed + def _print_event(ts_ms: int, message: str): + """Print a formatted CloudWatch log event.""" print(f"{_CW_PREFIX}[{_format_timestamp(ts_ms)}] {message}") - lines_printed += 1 - if tail_lines and lines_printed >= tail_lines: - logger.info("Reached tail_lines limit (%d). Stopping log stream.", tail_lines) - return True - return False + + # When tail_lines is set, fetch the last N events and return immediately. + if tail_lines: + events = streamer.poll_tail(tail_lines) + for ts_ms, message in events: + _print_event(ts_ms, message) + return status = status_fn() if status in TERMINAL_STATUSES: @@ -264,8 +391,7 @@ def _print_event(ts_ms: int, message: str) -> bool: if not events: break for ts_ms, message in events: - if _print_event(ts_ms, message): - return + _print_event(ts_ms, message) except ClientError: pass logger.info("Job finished with status: %s", status) @@ -303,8 +429,7 @@ def _print_event(ts_ms: int, message: str) -> bool: if events: empty_cycles = 0 for ts_ms, message in events: - if _print_event(ts_ms, message): - return + _print_event(ts_ms, message) else: empty_cycles += 1 if empty_cycles == warn_cycle: @@ -316,8 +441,7 @@ def _print_event(ts_ms: int, message: str) -> bool: status = status_fn() if status in TERMINAL_STATUSES: for ts_ms, message in streamer.poll_once(): - if _print_event(ts_ms, message): - return + _print_event(ts_ms, message) logger.info("Job finished with status: %s", status) return diff --git a/sagemaker-train/tests/unit/train/test_log_streamer.py b/sagemaker-train/tests/unit/train/test_log_streamer.py index f09040e2f0..1fc95fd320 100644 --- a/sagemaker-train/tests/unit/train/test_log_streamer.py +++ b/sagemaker-train/tests/unit/train/test_log_streamer.py @@ -443,3 +443,249 @@ def test_filter_mode_access_denied_propagates(self): ) with pytest.raises(ClientError): streamer.poll_once() + + +class TestPollTailStreamMode: + """Unit tests for LogStreamer.poll_tail() in stream mode (SMTJ).""" + + def test_poll_tail_returns_last_n_events(self): + """poll_tail(n) returns the last N events in chronological order.""" + session = _make_mock_session() + mock_logs = session.boto_session.client.return_value + + # describe_log_streams returns one stream + mock_logs.get_paginator.return_value.paginate.return_value = [ + {"logStreams": [{"logStreamName": "job/algo-1-123"}]} + ] + + # get_log_events: first call (startFromHead=False) returns events via backward pagination + mock_logs.get_log_events.side_effect = [ + # First call: startFromHead=False, no token → returns empty + backward token + {"events": [], "nextBackwardToken": "btoken1"}, + # Second call: with backward token + limit → returns last N events + { + "events": [ + {"timestamp": 300, "message": "third"}, + {"timestamp": 200, "message": "second"}, + {"timestamp": 100, "message": "first"}, + ], + "nextBackwardToken": "btoken1", + }, + ] + + streamer = LogStreamer( + log_group="/aws/sagemaker/TrainingJobs", + job_name="job", + sagemaker_session=session, + ) + + result = streamer.poll_tail(3) + + assert len(result) == 3 + # Should be sorted by timestamp (chronological) + assert result[0] == (100, "first") + assert result[1] == (200, "second") + assert result[2] == (300, "third") + + def test_poll_tail_multi_stream_merges_by_timestamp(self): + """poll_tail merges events across multiple streams and returns globally last N.""" + session = _make_mock_session() + mock_logs = session.boto_session.client.return_value + + mock_logs.get_paginator.return_value.paginate.return_value = [ + {"logStreams": [ + {"logStreamName": "job/algo-1-123"}, + {"logStreamName": "job/algo-2-456"}, + ]} + ] + + # Stream 1: events at ts 100, 300 + # Stream 2: events at ts 200, 400 + mock_logs.get_log_events.side_effect = [ + # Stream 1: first call (startFromHead=False) + {"events": [], "nextBackwardToken": "bt1"}, + # Stream 1: second call with token + { + "events": [ + {"timestamp": 300, "message": "s1-late"}, + {"timestamp": 100, "message": "s1-early"}, + ], + "nextBackwardToken": "bt1", + }, + # Stream 2: first call (startFromHead=False) + {"events": [], "nextBackwardToken": "bt2"}, + # Stream 2: second call with token + { + "events": [ + {"timestamp": 400, "message": "s2-latest"}, + {"timestamp": 200, "message": "s2-mid"}, + ], + "nextBackwardToken": "bt2", + }, + ] + + streamer = LogStreamer( + log_group="/aws/sagemaker/TrainingJobs", + job_name="job", + sagemaker_session=session, + ) + + result = streamer.poll_tail(3) + + # Should take globally last 3 by timestamp: 200, 300, 400 + assert len(result) == 3 + assert result[0] == (200, "s2-mid") + assert result[1] == (300, "s1-late") + assert result[2] == (400, "s2-latest") + + +class TestPollTailFilterMode: + """Unit tests for LogStreamer.poll_tail() in filter mode (SMHP).""" + + def test_poll_tail_filter_mode_returns_events(self): + """poll_tail in filter mode paginates with startFromHead=False.""" + session = _make_mock_session() + mock_logs = session.boto_session.client.return_value + + # First page: 0 events (API scanning streams), second page: events found + mock_logs.filter_log_events.side_effect = [ + {"events": [], "nextToken": "page2"}, + { + "events": [ + {"timestamp": 300, "message": "newest"}, + {"timestamp": 200, "message": "middle"}, + {"timestamp": 100, "message": "oldest"}, + ], + }, + ] + + streamer = LogStreamer( + log_group="/aws/sagemaker/Clusters/c/id", + job_name="job", + sagemaker_session=session, + filter_pattern='"job"', + start_time_ms=1784833626000, # Must be on or after 2024-01-01 + ) + + result = streamer.poll_tail(3) + + # Events come reverse-chronological, should be reversed to chronological + assert len(result) == 3 + assert result[0] == (100, "oldest") + assert result[1] == (200, "middle") + assert result[2] == (300, "newest") + + # Verify startFromHead=False was passed + call_args = mock_logs.filter_log_events.call_args_list[0] + assert call_args[1]["startFromHead"] is False + + def test_poll_tail_filter_mode_warns_without_start_time(self): + """poll_tail logs a warning when no start_time is set.""" + session = _make_mock_session() + mock_logs = session.boto_session.client.return_value + + mock_logs.filter_log_events.return_value = { + "events": [{"timestamp": 100, "message": "msg"}], + } + + streamer = LogStreamer( + log_group="/aws/sagemaker/Clusters/c/id", + job_name="job", + sagemaker_session=session, + filter_pattern='"job"', + start_time_ms=None, # No start time + ) + + with patch("sagemaker.train.common_utils.log_streamer.logger") as mock_logger: + streamer.poll_tail(1) + mock_logger.warning.assert_called_once() + assert "start_time" in mock_logger.warning.call_args[0][0] + + +class TestStreamLogLoopTailLines: + """Unit tests for stream_log_loop with tail_lines parameter.""" + + def test_tail_lines_calls_poll_tail_and_returns(self, capsys): + """When tail_lines is set, stream_log_loop calls poll_tail and prints.""" + streamer = MagicMock() + streamer.poll_tail.return_value = [ + (1000, "line one"), + (2000, "line two"), + (3000, "line three"), + ] + status_fn = MagicMock(return_value="Completed") + + stream_log_loop(streamer, poll=5, status_fn=status_fn, tail_lines=3) + + streamer.poll_tail.assert_called_once_with(3) + # status_fn should NOT be called — tail_lines returns immediately + status_fn.assert_not_called() + + captured = capsys.readouterr() + assert "line one" in captured.out + assert "line two" in captured.out + assert "line three" in captured.out + + def test_tail_lines_none_does_not_call_poll_tail(self): + """When tail_lines is None, stream_log_loop uses normal streaming.""" + streamer = MagicMock() + streamer.poll_once.return_value = [] + status_fn = MagicMock(return_value="Completed") + + stream_log_loop(streamer, poll=5, status_fn=status_fn, tail_lines=None) + + streamer.poll_tail.assert_not_called() + status_fn.assert_called() + + + def test_poll_tail_filter_mode_raises_for_pre_2024_start_time(self): + """poll_tail raises ValueError when start_time is before 2024-01-01.""" + session = _make_mock_session() + + streamer = LogStreamer( + log_group="/aws/sagemaker/Clusters/c/id", + job_name="job", + sagemaker_session=session, + filter_pattern='"job"', + start_time_ms=1000, # Way before 2024-01-01 + ) + + with pytest.raises(ValueError, match="before 2024-01-01"): + streamer.poll_tail(5) + + def test_poll_tail_filter_mode_paginates_multiple_pages(self): + """poll_tail paginates until enough events are found.""" + session = _make_mock_session() + mock_logs = session.boto_session.client.return_value + + # Simulate: first 3 pages return empty (scanning streams), 4th has events + mock_logs.filter_log_events.side_effect = [ + {"events": [], "nextToken": "page2"}, + {"events": [], "nextToken": "page3"}, + {"events": [], "nextToken": "page4"}, + { + "events": [ + {"timestamp": 300, "message": "third"}, + {"timestamp": 200, "message": "second"}, + {"timestamp": 100, "message": "first"}, + ], + }, + ] + + streamer = LogStreamer( + log_group="/aws/sagemaker/Clusters/c/id", + job_name="job", + sagemaker_session=session, + filter_pattern='"job"', + start_time_ms=1784833626000, + ) + + result = streamer.poll_tail(3) + + assert len(result) == 3 + # Should be reversed to chronological order + assert result[0] == (100, "first") + assert result[1] == (200, "second") + assert result[2] == (300, "third") + # Should have made 4 API calls + assert mock_logs.filter_log_events.call_count == 4