From 91bdacd7e3fe5f4939ad3d37985048f0cbba8b88 Mon Sep 17 00:00:00 2001 From: Syed Jafri Date: Wed, 22 Jul 2026 17:31:51 -0700 Subject: [PATCH 1/4] fix: stream_logs_smhp extract training job from obj --- sagemaker-train/src/sagemaker/train/base_trainer.py | 7 ++++++- 1 file changed, 6 insertions(+), 1 deletion(-) diff --git a/sagemaker-train/src/sagemaker/train/base_trainer.py b/sagemaker-train/src/sagemaker/train/base_trainer.py index d3c5957c12..468dacd610 100644 --- a/sagemaker-train/src/sagemaker/train/base_trainer.py +++ b/sagemaker-train/src/sagemaker/train/base_trainer.py @@ -701,7 +701,12 @@ def _stream_logs_smtj(self, training_job, poll: int) -> None: def _stream_logs_smhp(self, training_job, compute, poll: int, start_time_ms=None) -> None: """Stream logs for a HyperPod job using filter_log_events polling.""" - job_id = training_job if isinstance(training_job, str) else str(training_job) + if isinstance(training_job, str): + job_id = training_job + elif hasattr(training_job, 'training_job_name'): + job_id = training_job.training_job_name + else: + job_id = str(training_job) sagemaker_session = TrainDefaults.get_sagemaker_session( sagemaker_session=self.sagemaker_session From 5b45b3c85a32e3501b1cfc01c118484ac1aef734 Mon Sep 17 00:00:00 2001 From: Syed Jafri Date: Fri, 24 Jul 2026 15:41:42 -0700 Subject: [PATCH 2/4] fix: render mlflow metrics as png to handle large number of metrics --- .../train/common_utils/metrics_visualizer.py | 52 ++++++++++++++----- 1 file changed, 40 insertions(+), 12 deletions(-) diff --git a/sagemaker-train/src/sagemaker/train/common_utils/metrics_visualizer.py b/sagemaker-train/src/sagemaker/train/common_utils/metrics_visualizer.py index fe837a91fc..cf12afb442 100644 --- a/sagemaker-train/src/sagemaker/train/common_utils/metrics_visualizer.py +++ b/sagemaker-train/src/sagemaker/train/common_utils/metrics_visualizer.py @@ -231,13 +231,24 @@ def plot_training_metrics( history = client.get_metric_history(run_id, metric_name) if history: metric_data[metric_name] = history - - # Plot + num_metrics = len(metric_data) + if num_metrics == 0: + logger.warning("No metric history found for run %s. Nothing to plot.", run_id) + return + + # Build the figure at full size (2 columns, all metrics). The inline + # backend can't rasterize this when it gets too tall, so instead of + # calling display(fig) directly we render to an in-memory PNG and embed + # it inside a scrollable HTML
. This lets the notebook display all + # metrics without choking on pixel-count limits. + import io + import base64 + rows = (num_metrics + 1) // 2 - fig, axes = plt.subplots(rows, 2, figsize=(figsize[0], figsize[1] * rows)) - axes = axes.flatten() if num_metrics > 1 else [axes] - + fig, axes = plt.subplots(rows, 2, figsize=(figsize[0], figsize[1] * rows), squeeze=False) + axes = axes.flatten() + for idx, (metric_name, history) in enumerate(metric_data.items()): steps = [h.step for h in history] values = [h.value for h in history] @@ -246,14 +257,31 @@ def plot_training_metrics( axes[idx].set_ylabel('Value') axes[idx].set_title(metric_name, fontweight='bold') axes[idx].grid(True, alpha=0.3) - - for idx in range(len(metric_data), len(axes)): + + for idx in range(num_metrics, len(axes)): axes[idx].set_visible(False) - - plt.suptitle(f'Training Metrics: {training_job.training_job_name}', fontweight='bold', fontsize=14) - plt.tight_layout(rect=[0, 0, 1, 0.98]) # Leave small space for suptitle - display(fig) - plt.close() + + fig.suptitle( + f'Training Metrics: {training_job.training_job_name}', + fontweight='bold', fontsize=14, + ) + fig.tight_layout(rect=[0, 0, 1, 0.98]) + + # Render to PNG buffer for inline display + buf = io.BytesIO() + fig.savefig(buf, format="png", dpi=80, bbox_inches="tight") + plt.close(fig) + buf.seek(0) + + # Embed as a scrollable HTML image in the notebook + b64 = base64.b64encode(buf.getvalue()).decode() + from IPython.display import HTML + display(HTML( + f'
' + f'' + f'
' + )) def get_available_metrics(training_job: TrainingJob) -> List[str]: From 0ec41c02b76b9e234929c93bc5768ac7e253bfa3 Mon Sep 17 00:00:00 2001 From: Syed Jafri Date: Mon, 27 Jul 2026 13:15:22 -0700 Subject: [PATCH 3/4] code cleanup: move io, base64 to top level imports --- .../src/sagemaker/train/common_utils/metrics_visualizer.py | 7 +++---- 1 file changed, 3 insertions(+), 4 deletions(-) diff --git a/sagemaker-train/src/sagemaker/train/common_utils/metrics_visualizer.py b/sagemaker-train/src/sagemaker/train/common_utils/metrics_visualizer.py index cf12afb442..ec5bb2dc0a 100644 --- a/sagemaker-train/src/sagemaker/train/common_utils/metrics_visualizer.py +++ b/sagemaker-train/src/sagemaker/train/common_utils/metrics_visualizer.py @@ -1,7 +1,10 @@ """MLflow metrics visualization utilities for SageMaker training jobs.""" +import io +import base64 import logging from typing import Optional, List, Dict, Any + from sagemaker.core.resources import TrainingJob logger = logging.getLogger(__name__) @@ -209,7 +212,6 @@ def plot_training_metrics( import mlflow from mlflow.tracking import MlflowClient from IPython.display import display - import logging logging.getLogger('botocore.credentials').setLevel(logging.WARNING) @@ -242,9 +244,6 @@ def plot_training_metrics( # calling display(fig) directly we render to an in-memory PNG and embed # it inside a scrollable HTML
. This lets the notebook display all # metrics without choking on pixel-count limits. - import io - import base64 - rows = (num_metrics + 1) // 2 fig, axes = plt.subplots(rows, 2, figsize=(figsize[0], figsize[1] * rows), squeeze=False) axes = axes.flatten() From 3b4be795d71f3a7b6778787935f47d0d6923a0ed Mon Sep 17 00:00:00 2001 From: Syed Jafri Date: Mon, 27 Jul 2026 16:03:20 -0700 Subject: [PATCH 4/4] fix: add helper method to create sns topic --- .../train/common_utils/notifications.py | 95 +++++++++++++++++-- 1 file changed, 87 insertions(+), 8 deletions(-) diff --git a/sagemaker-train/src/sagemaker/train/common_utils/notifications.py b/sagemaker-train/src/sagemaker/train/common_utils/notifications.py index c7e65c2792..c9990fcb73 100644 --- a/sagemaker-train/src/sagemaker/train/common_utils/notifications.py +++ b/sagemaker-train/src/sagemaker/train/common_utils/notifications.py @@ -150,6 +150,85 @@ def _validate_sns_topic(sns_client, topic_arn: str) -> None: raise +def _build_topic_policy(topic_arn: str, account_id: str) -> str: + """Build the SNS topic access policy for EventBridge notifications. + + Includes the default owner statement (standard SNS actions scoped to the + topic's own account via ``AWS:SourceAccount``) plus an explicit statement + granting the EventBridge service principal ``SNS:Publish``. + + Args: + topic_arn: The ARN of the SNS topic. + account_id: The AWS account ID that owns the topic. + + Returns: + JSON string of the topic access policy. + """ + policy = { + "Version": "2008-10-17", + "Id": "__default_policy_ID", + "Statement": [ + { + "Sid": "__default_statement_ID", + "Effect": "Allow", + "Principal": {"AWS": "*"}, + "Action": [ + "SNS:Publish", + "SNS:RemovePermission", + "SNS:SetTopicAttributes", + "SNS:DeleteTopic", + "SNS:ListSubscriptionsByTopic", + "SNS:GetTopicAttributes", + "SNS:AddPermission", + "SNS:Subscribe", + ], + "Resource": topic_arn, + "Condition": {"StringEquals": {"AWS:SourceAccount": account_id}}, + }, + { + "Sid": "AllowEventBridgePublish", + "Effect": "Allow", + "Principal": {"Service": "events.amazonaws.com"}, + "Action": "SNS:Publish", + "Resource": topic_arn, + "Condition": {"StringEquals": {"AWS:SourceAccount": account_id}}, + }, + ], + } + return json.dumps(policy) + + +def create_notification_topic(topic_name: str, sagemaker_session) -> str: + """Create an SNS topic configured for EventBridge notifications. + + Creates a new SNS topic and sets an access policy that allows EventBridge + in the same account to publish to it. The returned ARN can be passed + straight to :func:`enable_notifications`. + + Args: + topic_name: Name of the SNS topic to create. + sagemaker_session: SageMaker session (provides boto_session). + + Returns: + The ARN of the created SNS topic. + """ + region_name = sagemaker_session.boto_session.region_name + sns_client = sagemaker_session.boto_session.client("sns", region_name=region_name) + + response = sns_client.create_topic(Name=topic_name) + topic_arn = response["TopicArn"] + # topic ARN: arn:aws:sns::: + account_id = topic_arn.split(":")[4] + + policy = _build_topic_policy(topic_arn, account_id) + sns_client.set_topic_attributes( + TopicArn=topic_arn, AttributeName="Policy", AttributeValue=policy + ) + + logger.info(f"Created SNS topic with EventBridge publish policy: {topic_arn}") + return topic_arn + + def enable_notifications( sns_topic_arn: str, sagemaker_session, @@ -240,14 +319,14 @@ def enable_notifications( try: events_client.put_targets(**put_targets_kwargs) - except Exception as e: - raise ValueError( - f"Failed to attach SNS topic to EventBridge rule. " - f"Ensure your SNS topic has a resource policy allowing " - f"EventBridge to publish to it. Add this to the topic's access policy:\n" - f'{{"Effect": "Allow", "Principal": {{"Service": "events.amazonaws.com"}}, ' - f'"Action": "sns:Publish", "Resource": "{sns_topic_arn}"}}' - ) from e + except ClientError as e: + error_code = e.response.get("Error", {}).get("Code", "") + if error_code in ("AccessDeniedException", "AccessDenied"): + raise PermissionError( + "Missing permission events:PutTargets to attach the SNS topic " + "to the EventBridge rule." + ) from e + raise logger.info(f"Notifications enabled: {normalized_events} -> {sns_topic_arn}") return rule_arn