Skip to content

fix: tail_lines in stream_logs returns last N events (true tail seman… #6164

New issue

Have a question about this project? Sign up for a free GitHub account to open an issue and contact its maintainers and the community.

Sign up for GitHub

By clicking “Sign up for GitHub”, you agree to our terms of service and privacy statement. We’ll occasionally send you account related emails.

Already on GitHub? Sign in to your account

Merged
zhaoqizqwang merged 1 commit into from
Aug 6, 2026
+416 −98
Merged

fix: tail_lines in stream_logs returns last N events (true tail seman… #6164

Show file tree
Hide file tree
Changes from all commits
Commits
Show all changes
1 commit
Select commit
File filter

Filter by extension

Filter by extension
Conversations
Failed to load comments.
Loading
Jump to
Jump to file
Failed to load files.
Loading
Diff view
Diff view
108 changes: 28 additions & 80 deletions sagemaker-train/src/sagemaker/train/base_trainer.py
View file
Open in desktop
Original file line number Diff line number Diff line change
@@ -1,5 +1,4 @@
import copy
import os
import time
import yaml
from abc import ABC, abstractmethod
Expand All @@ -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
Expand All @@ -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
Expand Down Expand Up @@ -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
Expand Down Expand Up @@ -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
Expand All @@ -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}")
Expand All @@ -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.
Expand Down
160 changes: 142 additions & 18 deletions sagemaker-train/src/sagemaker/train/common_utils/log_streamer.py
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -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 = {
Expand Down Expand Up @@ -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:
Expand All @@ -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)
Expand Down Expand Up @@ -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:
Expand All @@ -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

Expand Down
Loading