From 793b59da838519dca297aa7de9924cbcd080ae68 Mon Sep 17 00:00:00 2001
From: Ealynn Hsu <ealynnh@amazon.com>
Date: Tue, 14 Jul 2026 13:52:18 +0000
Subject: [PATCH 1/6] [WIP] Job notifications setup

---
 .../train/common_utils/notifications.py       | 332 ++++++++++++++++++
 .../train/common_utils/test_notifications.py  | 211 +++++++++++
 2 files changed, 543 insertions(+)
 create mode 100644 sagemaker-train/src/sagemaker/train/common_utils/notifications.py
 create mode 100644 sagemaker-train/tests/unit/train/common_utils/test_notifications.py

diff --git a/sagemaker-train/src/sagemaker/train/common_utils/notifications.py b/sagemaker-train/src/sagemaker/train/common_utils/notifications.py
new file mode 100644
index 0000000000..234f8900da
--- /dev/null
+++ b/sagemaker-train/src/sagemaker/train/common_utils/notifications.py
@@ -0,0 +1,332 @@
+"""Job notification utilities for SageMaker training jobs.
+
+Manages EventBridge rules that route SageMaker Training Job status change
+events to user-provided SNS topics. Supports SMTJ (serverless and serverful)
+training jobs only.
+"""
+
+from __future__ import absolute_import
+
+import hashlib
+import json
+import logging
+from typing import Dict, List, Optional
+
+logger = logging.getLogger(__name__)
+
+# Rule name prefix used to identify SDK-created rules
+_RULE_NAME_PREFIX = "sm-pysdk-job-notif"
+
+# Default events to notify on
+_DEFAULT_EVENTS = ["Completed", "Failed", "Stopped"]
+
+# Valid event status values
+_VALID_EVENTS = {"Completed", "Failed", "Stopped", "InProgress"}
+
+
+def _get_rule_name(sns_topic_arn: str) -> str:
+    """Generate a deterministic rule name from the SNS topic ARN.
+
+    Args:
+        sns_topic_arn: The SNS topic ARN.
+
+    Returns:
+        Rule name like "sm-pysdk-job-notif-a3f8b2c1".
+    """
+    arn_hash = hashlib.sha256(sns_topic_arn.encode()).hexdigest()[:8]
+    return f"{_RULE_NAME_PREFIX}-{arn_hash}"
+
+
+def _normalize_events(events: Optional[List[str]]) -> List[str]:
+    """Normalize and validate event status values.
+
+    Args:
+        events: List of event names (e.g., ["completed", "failed"]).
+            If None, returns all default events.
+
+    Returns:
+        List of capitalized event status strings.
+
+    Raises:
+        ValueError: If an invalid event name is provided.
+    """
+    if not events:
+        return _DEFAULT_EVENTS.copy()
+
+    normalized = []
+    for event in events:
+        capitalized = event.capitalize()
+        if capitalized == "Inprogress":
+            capitalized = "InProgress"
+        if capitalized not in _VALID_EVENTS:
+            raise ValueError(
+                f"Invalid notification event: '{event}'. "
+                f"Valid events: {sorted(_VALID_EVENTS)}"
+            )
+        normalized.append(capitalized)
+
+    return normalized
+
+
+def _build_event_pattern(
+    events: List[str],
+    job_name_prefix: Optional[str] = None,
+) -> str:
+    """Build the EventBridge event pattern JSON for SMTJ job status changes.
+
+    Args:
+        events: List of TrainingJobStatus values to match.
+        job_name_prefix: Optional job name prefix filter.
+
+    Returns:
+        JSON string of the event pattern.
+    """
+    pattern: Dict = {
+        "source": ["aws.sagemaker"],
+        "detail-type": ["SageMaker Training Job State Change"],
+        "detail": {
+            "TrainingJobStatus": events,
+        },
+    }
+
+    if job_name_prefix:
+        pattern["detail"]["TrainingJobName"] = [{"prefix": job_name_prefix}]
+
+    return json.dumps(pattern)
+
+
+def _validate_notifications_permissions(events_client) -> None:
+    """Validate the caller has permissions to manage EventBridge rules.
+
+    Args:
+        events_client: boto3 EventBridge client.
+
+    Raises:
+        PermissionError: If the caller lacks required EventBridge permissions.
+    """
+    try:
+        events_client.list_rules(NamePrefix=_RULE_NAME_PREFIX, Limit=1)
+    except events_client.exceptions.ClientError as e:
+        error_code = e.response.get("Error", {}).get("Code", "")
+        if error_code in ("AccessDeniedException", "AccessDenied"):
+            raise PermissionError(
+                "Missing permissions to manage EventBridge rules. "
+                "Ensure your caller identity has: "
+                "events:PutRule, events:PutTargets, events:ListRules, "
+                "events:RemoveTargets, events:DeleteRule."
+            ) from e
+        raise
+
+
+def _validate_sns_topic(sns_client, topic_arn: str) -> None:
+    """Validate that the SNS topic exists and is accessible.
+
+    Args:
+        sns_client: boto3 SNS client.
+        topic_arn: The SNS topic ARN to validate.
+
+    Raises:
+        ValueError: If the topic doesn't exist or isn't accessible.
+    """
+    try:
+        sns_client.get_topic_attributes(TopicArn=topic_arn)
+    except Exception as e:
+        error_msg = str(e)
+        if "NotFound" in error_msg or "does not exist" in error_msg.lower():
+            raise ValueError(
+                f"SNS topic not found: {topic_arn}. "
+                "Ensure the topic exists and you have sns:GetTopicAttributes permission."
+            ) from e
+        if "AuthorizationError" in error_msg or "AccessDenied" in error_msg:
+            raise PermissionError(
+                f"Cannot access SNS topic: {topic_arn}. "
+                "Ensure you have sns:GetTopicAttributes permission."
+            ) from e
+        raise
+
+
+def enable_notifications(
+    sns_topic_arn: str,
+    sagemaker_session,
+    events: Optional[List[str]] = None,
+    event_bus_arn: Optional[str] = None,
+    job_name_prefix: Optional[str] = None,
+) -> str:
+    """Create or update an EventBridge rule for training job notifications.
+
+    Args:
+        sns_topic_arn: ARN of the SNS topic to receive notifications.
+        sagemaker_session: SageMaker session (provides boto_session).
+        events: List of job statuses to notify on. Defaults to
+            ["Completed", "Failed", "Stopped"].
+        event_bus_arn: Optional EventBridge bus ARN. Defaults to the
+            account's default bus.
+        job_name_prefix: Optional job name prefix to filter notifications.
+
+    Returns:
+        The ARN of the created/updated EventBridge rule.
+
+    Raises:
+        ValueError: If sns_topic_arn is invalid or topic doesn't exist.
+        PermissionError: If caller lacks required permissions.
+    """
+    if not sns_topic_arn or not sns_topic_arn.startswith("arn:aws:sns:"):
+        raise ValueError(
+            f"Invalid SNS topic ARN: '{sns_topic_arn}'. "
+            "Must be a valid ARN like 'arn:aws:sns:us-east-1:012345678910:my-topic'."
+        )
+
+    region_name = sagemaker_session.boto_session.region_name
+    events_client = sagemaker_session.boto_session.client("events", region_name=region_name)
+    sns_client = sagemaker_session.boto_session.client("sns", region_name=region_name)
+
+    _validate_notifications_permissions(events_client)
+    _validate_sns_topic(sns_client, sns_topic_arn)
+
+    normalized_events = _normalize_events(events)
+    rule_name = _get_rule_name(sns_topic_arn)
+    event_pattern = _build_event_pattern(normalized_events, job_name_prefix)
+
+    put_rule_kwargs = {
+        "Name": rule_name,
+        "EventPattern": event_pattern,
+        "State": "ENABLED",
+        "Description": f"SageMaker PySDK training job notifications -> {sns_topic_arn}",
+    }
+    if event_bus_arn:
+        put_rule_kwargs["EventBusName"] = event_bus_arn
+
+    response = events_client.put_rule(**put_rule_kwargs)
+    rule_arn = response["RuleArn"]
+    logger.info(f"EventBridge rule created/updated: {rule_name} ({rule_arn})")
+
+    # Add SNS topic as target with formatted message
+    target_id = f"{rule_name}-sns-target"
+    # Use a JSON object template — SNS renders JSON objects with line breaks
+    input_template = (
+        '{"Job": "<job_name>",'
+        ' "Status": "<status>",'
+        ' "Time": "<time>",'
+        ' "Region": "<region>",'
+        ' "Account": "<account>",'
+        ' "Failure Reason": "<failure_reason>",'
+        ' "Console": "https://<region>.console.aws.amazon.com/sagemaker/home?region=<region>#/jobs/<job_name>"}'
+    )
+    put_targets_kwargs = {
+        "Rule": rule_name,
+        "Targets": [{
+            "Id": target_id,
+            "Arn": sns_topic_arn,
+            "InputTransformer": {
+                "InputPathsMap": {
+                    "job_name": "$.detail.TrainingJobName",
+                    "status": "$.detail.TrainingJobStatus",
+                    "time": "$.time",
+                    "region": "$.region",
+                    "failure_reason": "$.detail.FailureReason",
+                    "account": "$.account",
+                },
+                "InputTemplate": input_template,
+            },
+        }],
+    }
+    if event_bus_arn:
+        put_targets_kwargs["EventBusName"] = event_bus_arn
+
+    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
+
+    logger.info(f"Notifications enabled: {normalized_events} -> {sns_topic_arn}")
+    return rule_arn
+
+
+def delete_notification_rules(
+    sagemaker_session,
+    rule_arn: str,
+    event_bus_arn: Optional[str] = None,
+) -> str:
+    """Delete an SDK-created EventBridge notification rule.
+
+    Args:
+        sagemaker_session: SageMaker session (provides boto_session).
+        rule_arn: The ARN of the rule to delete.
+        event_bus_arn: Optional EventBridge bus ARN. Defaults to "default".
+
+    Returns:
+        The name of the deleted rule.
+    """
+    region_name = sagemaker_session.boto_session.region_name
+    events_client = sagemaker_session.boto_session.client("events", region_name=region_name)
+
+    event_bus_name = event_bus_arn or "default"
+    rule_name = rule_arn.rsplit("/", 1)[-1] if "/" in rule_arn else rule_arn
+
+    # Remove SNS targets from the EB rule
+    try:
+        targets_response = events_client.list_targets_by_rule(
+            Rule=rule_name, EventBusName=event_bus_name
+        )
+        target_ids = [t["Id"] for t in targets_response.get("Targets", [])]
+        if target_ids:
+            events_client.remove_targets(
+                Rule=rule_name, Ids=target_ids, EventBusName=event_bus_name
+            )
+    except Exception as e:
+        logger.debug(f"Error removing targets for rule {rule_name}: {e}")
+
+    # Delete the rule
+    try:
+        events_client.delete_rule(Name=rule_name, EventBusName=event_bus_name)
+    except Exception as e:
+        if "ResourceNotFoundException" in str(type(e).__name__) or "does not exist" in str(e).lower():
+            raise ValueError(
+                f"Rule '{rule_name}' not found on event bus '{event_bus_name}'. "
+                "If the rule was created on a custom event bus, pass the same "
+                "event_bus_arn to delete_notification_rules()."
+            ) from e
+        logger.warning(f"Failed to delete rule {rule_name}: {e}")
+
+    logger.info(f"Deleted notification rule: {rule_name}")
+    return rule_name
+
+
+def list_notification_rules(
+    sagemaker_session,
+    event_bus_arn: Optional[str] = None,
+) -> List[Dict[str, str]]:
+    """List all SDK-created EventBridge notification rules (identified by the sm-pysdk-job-notif prefix).
+
+    Args:
+        sagemaker_session: SageMaker session (provides boto_session).
+        event_bus_arn: Optional EventBridge bus ARN. Defaults to "default".
+
+    Returns:
+        List of dicts with 'name', 'arn', and 'state' for each rule.
+    """
+    region_name = sagemaker_session.boto_session.region_name
+    events_client = sagemaker_session.boto_session.client("events", region_name=region_name)
+
+    event_bus_name = event_bus_arn or "default"
+    rules = []
+
+    paginator = events_client.get_paginator("list_rules")
+    for page in paginator.paginate(NamePrefix=_RULE_NAME_PREFIX, EventBusName=event_bus_name):
+        for rule in page.get("Rules", []):
+            rules.append({
+                "name": rule["Name"],
+                "arn": rule["Arn"],
+                "state": rule.get("State", "UNKNOWN"),
+            })
+
+    return rules
+
+
+
diff --git a/sagemaker-train/tests/unit/train/common_utils/test_notifications.py b/sagemaker-train/tests/unit/train/common_utils/test_notifications.py
new file mode 100644
index 0000000000..9c72fc112d
--- /dev/null
+++ b/sagemaker-train/tests/unit/train/common_utils/test_notifications.py
@@ -0,0 +1,211 @@
+"""Unit tests for job notifications module."""
+
+from __future__ import absolute_import
+
+from unittest.mock import MagicMock, patch
+
+import pytest
+
+from sagemaker.train.common_utils.notifications import (
+    _build_event_pattern,
+    _get_rule_name,
+    _normalize_events,
+    delete_notification_rules,
+    enable_notifications,
+    list_notification_rules,
+)
+
+
+class TestRuleNaming:
+
+    def test_deterministic_from_arn(self):
+        """Same ARN always produces the same rule name."""
+        arn = "arn:aws:sns:us-east-1:123456789:my-topic"
+        assert _get_rule_name(arn) == _get_rule_name(arn)
+
+    def test_different_arns_different_names(self):
+        """Different ARNs produce different rule names."""
+        arn1 = "arn:aws:sns:us-east-1:123456789:topic-a"
+        arn2 = "arn:aws:sns:us-east-1:123456789:topic-b"
+        assert _get_rule_name(arn1) != _get_rule_name(arn2)
+
+    def test_prefix_present(self):
+        """Rule name starts with the SDK prefix."""
+        name = _get_rule_name("arn:aws:sns:us-east-1:123:topic")
+        assert name.startswith("sm-pysdk-job-notif-")
+
+
+class TestNormalizeEvents:
+
+    def test_defaults_to_all_terminal(self):
+        assert _normalize_events(None) == ["Completed", "Failed", "Stopped"]
+
+    def test_capitalizes_input(self):
+        assert _normalize_events(["completed", "failed"]) == ["Completed", "Failed"]
+
+    def test_handles_inprogress(self):
+        assert _normalize_events(["inprogress"]) == ["InProgress"]
+
+    def test_invalid_event_raises(self):
+        with pytest.raises(ValueError, match="Invalid notification event"):
+            _normalize_events(["invalid_status"])
+
+
+class TestBuildEventPattern:
+
+    def test_basic_pattern(self):
+        import json
+        pattern = json.loads(_build_event_pattern(["Completed", "Failed"]))
+
+        assert pattern["source"] == ["aws.sagemaker"]
+        assert pattern["detail-type"] == ["SageMaker Training Job State Change"]
+        assert pattern["detail"]["TrainingJobStatus"] == ["Completed", "Failed"]
+        assert "TrainingJobName" not in pattern["detail"]
+
+    def test_with_job_name_prefix(self):
+        import json
+        pattern = json.loads(_build_event_pattern(["Completed"], job_name_prefix="my-team-"))
+
+        assert pattern["detail"]["TrainingJobName"] == [{"prefix": "my-team-"}]
+
+
+class TestEnableNotifications:
+
+    def _session(self):
+        s = MagicMock()
+        s.boto_session.region_name = "us-east-1"
+        return s
+
+    def test_invalid_arn_raises(self):
+        with pytest.raises(ValueError, match="Invalid SNS topic ARN"):
+            enable_notifications("not-an-arn", self._session())
+
+    def test_creates_rule_and_target(self):
+        session = self._session()
+        events_client = MagicMock()
+        sns_client = MagicMock()
+        session.boto_session.client.side_effect = lambda svc, **kw: (
+            events_client if svc == "events" else sns_client
+        )
+
+        events_client.put_rule.return_value = {"RuleArn": "arn:aws:events:us-east-1:123:rule/sm-pysdk-notif-abc"}
+        events_client.put_targets.return_value = {"FailedEntryCount": 0}
+        events_client.list_rules.return_value = {"Rules": []}
+
+        arn = enable_notifications(
+            sns_topic_arn="arn:aws:sns:us-east-1:123:my-topic",
+            sagemaker_session=session,
+        )
+
+        assert arn == "arn:aws:events:us-east-1:123:rule/sm-pysdk-notif-abc"
+        events_client.put_rule.assert_called_once()
+        events_client.put_targets.assert_called_once()
+
+    def test_with_custom_events_and_prefix(self):
+        session = self._session()
+        events_client = MagicMock()
+        sns_client = MagicMock()
+        session.boto_session.client.side_effect = lambda svc, **kw: (
+            events_client if svc == "events" else sns_client
+        )
+
+        events_client.put_rule.return_value = {"RuleArn": "arn:rule"}
+        events_client.put_targets.return_value = {"FailedEntryCount": 0}
+        events_client.list_rules.return_value = {"Rules": []}
+
+        enable_notifications(
+            sns_topic_arn="arn:aws:sns:us-east-1:123:my-topic",
+            sagemaker_session=session,
+            events=["completed"],
+            job_name_prefix="ealynnh-",
+        )
+
+        call_kwargs = events_client.put_rule.call_args[1]
+        import json
+        pattern = json.loads(call_kwargs["EventPattern"])
+        assert pattern["detail"]["TrainingJobStatus"] == ["Completed"]
+        assert pattern["detail"]["TrainingJobName"] == [{"prefix": "ealynnh-"}]
+
+
+class TestDeleteNotificationRules:
+
+    def test_deletes_specific_rule(self):
+        session = MagicMock()
+        session.boto_session.region_name = "us-east-1"
+        events_client = MagicMock()
+        session.boto_session.client.return_value = events_client
+
+        events_client.list_targets_by_rule.return_value = {
+            "Targets": [{"Id": "target-1"}]
+        }
+
+        deleted = delete_notification_rules(
+            sagemaker_session=session,
+            rule_arn="arn:aws:events:us-east-1:123:rule/sm-pysdk-job-notif-abc123",
+        )
+
+        assert deleted == "sm-pysdk-job-notif-abc123"
+        events_client.remove_targets.assert_called_once()
+        events_client.delete_rule.assert_called_once()
+
+
+class TestBaseTrainerNotifications:
+
+    def _make_trainer(self, compute=None):
+        from sagemaker.train.base_trainer import BaseTrainer
+
+        class _StubTrainer(BaseTrainer):
+            _customization_technique = "SFT"
+            def train(self, *args, **kwargs):
+                pass
+
+        trainer = _StubTrainer.__new__(_StubTrainer)
+        trainer.compute = compute
+        trainer.sagemaker_session = None
+        trainer._latest_training_job = None
+        return trainer
+
+    def test_hyperpod_raises_not_implemented(self):
+        from sagemaker.core.training.configs import HyperPodCompute
+        trainer = self._make_trainer(
+            compute=HyperPodCompute(cluster_name="c", instance_type="ml.p5.48xlarge")
+        )
+
+        with pytest.raises(NotImplementedError, match="not supported for HyperPod"):
+            trainer._setup_notifications({"sns_topic_arn": "arn:aws:sns:us-east-1:123:topic"})
+
+    def test_missing_sns_arn_raises(self):
+        trainer = self._make_trainer()
+
+        with pytest.raises(ValueError, match="requires 'sns_topic_arn'"):
+            trainer._setup_notifications({"events": ["completed"]})
+
+    def test_invalid_config_type_raises(self):
+        trainer = self._make_trainer()
+
+        with pytest.raises(ValueError, match="must be a dict"):
+            trainer._setup_notifications("not-a-dict")
+
+
+class TestListNotificationRules:
+
+    def test_lists_sdk_rules(self):
+        session = MagicMock()
+        session.boto_session.region_name = "us-east-1"
+        events_client = MagicMock()
+        session.boto_session.client.return_value = events_client
+
+        paginator = MagicMock()
+        paginator.paginate.return_value = [
+            {"Rules": [
+                {"Name": "sm-pysdk-job-notif-aaa", "Arn": "arn:aws:events:us-east-1:123:rule/sm-pysdk-job-notif-aaa", "State": "ENABLED"},
+                {"Name": "sm-pysdk-job-notif-bbb", "Arn": "arn:aws:events:us-east-1:123:rule/sm-pysdk-job-notif-bbb", "State": "ENABLED"},
+            ]}
+        ]
+        events_client.get_paginator.return_value = paginator
+
+        rules = list_notification_rules(sagemaker_session=session)
+
+        assert len(rules) == 2
+        assert rules[0]["name"] == "sm-pysdk-job-notif-aaa"
+        assert rules[0]["arn"] == "arn:aws:events:us-east-1:123:rule/sm-pysdk-job-notif-aaa"

From e281495c7a2d99eef1b9815af737b5f2ea7279e0 Mon Sep 17 00:00:00 2001
From: Ealynn Hsu <ealynnh@amazon.com>
Date: Tue, 14 Jul 2026 13:53:10 +0000
Subject: [PATCH 2/6] [WIP] Job notif update

---
 .../src/sagemaker/train/base_trainer.py       | 109 ++++++++++++++++++
 1 file changed, 109 insertions(+)

diff --git a/sagemaker-train/src/sagemaker/train/base_trainer.py b/sagemaker-train/src/sagemaker/train/base_trainer.py
index ff3bdcd751..dd9096d3f6 100644
--- a/sagemaker-train/src/sagemaker/train/base_trainer.py
+++ b/sagemaker-train/src/sagemaker/train/base_trainer.py
@@ -34,6 +34,7 @@
 )
 from sagemaker.train.common_utils.metrics_visualizer import plot_training_metrics
 from sagemaker.train.common_utils.mlflow_config_utils import resolve_mlflow_tracking_fields
+from sagemaker.train.common_utils.notifications import enable_notifications, delete_notification_rules, 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.defaults import TrainDefaults
@@ -102,6 +103,7 @@ def __init__(
         training_image: Optional[str] = None,
         base_model_name: Optional[str] = None,
         disable_output_compression: Optional[bool] = False,
+        notifications: Optional[Dict[str, Any]] = None,
     ):
         self.sagemaker_session = sagemaker_session
         self.role = role
@@ -114,6 +116,10 @@ def __init__(
         self.training_image = training_image
         self.base_model_name = base_model_name
         self.disable_output_compression = disable_output_compression
+
+        # Set up notifications if configured
+        if notifications:
+            self._setup_notifications(notifications)
         self._checkpoint_s3_uri = None
 
     def _is_nova_model_for_telemetry(self) -> bool:
@@ -433,6 +439,109 @@ def _show_metrics_cloudwatch(
             end_time=end_time_ms,
         )
 
+    def _setup_notifications(self, notifications: Optional[Dict[str, Any]]) -> Optional[str]:
+        """Set up EventBridge notifications for the training job.
+
+        Called internally by trainer.train() after job submission when a
+        notifications config is provided.
+
+        Args:
+            notifications: Notification configuration dict with keys:
+                - sns_topic_arn (str, required): ARN of the SNS topic.
+                - events (list[str], optional): Job statuses to notify on.
+                    Defaults to ["Completed", "Failed", "Stopped"].
+                - event_bus_arn (str, optional): EventBridge bus ARN.
+                    Defaults to the account's default bus.
+                - job_name_prefix (str, optional): Only notify for jobs
+                    with names matching this prefix.
+
+        Returns:
+            The EventBridge rule ARN if notifications were set up, None otherwise.
+
+        Raises:
+            NotImplementedError: If compute is HyperPodCompute.
+            ValueError: If the config is invalid.
+            PermissionError: If the caller lacks required permissions.
+        """
+        if not notifications:
+            return None
+
+        # Validate compute type
+        if isinstance(getattr(self, 'compute', None), HyperPodCompute):
+            raise NotImplementedError(
+                "Job notifications are not supported for HyperPod compute."
+            )
+
+        # Validate config
+        if not isinstance(notifications, dict):
+            raise ValueError(
+                "notifications must be a dict with at least 'sns_topic_arn'. "
+                "Example: {'sns_topic_arn': 'arn:aws:sns:us-east-1:123:my-topic'}"
+            )
+
+        sns_topic_arn = notifications.get("sns_topic_arn")
+        if not sns_topic_arn:
+            raise ValueError(
+                "notifications config requires 'sns_topic_arn'. "
+                "Example: {'sns_topic_arn': 'arn:aws:sns:us-east-1:123:my-topic'}"
+            )
+
+        sagemaker_session = TrainDefaults.get_sagemaker_session(
+            sagemaker_session=self.sagemaker_session
+        )
+
+        rule_arn = enable_notifications(
+            sns_topic_arn=sns_topic_arn,
+            sagemaker_session=sagemaker_session,
+            events=notifications.get("events"),
+            event_bus_arn=notifications.get("event_bus_arn"),
+            job_name_prefix=notifications.get("job_name_prefix"),
+        )
+
+        return rule_arn
+
+    def delete_notification_rules(
+        self,
+        rule_arn: str,
+        event_bus_arn: Optional[str] = None,
+    ) -> str:
+        """Delete an SDK-created EventBridge notification rule.
+
+        Args:
+            rule_arn: The ARN of the rule to delete.
+            event_bus_arn: Optional EventBridge bus ARN. Defaults to "default".
+
+        Returns:
+            The name of the deleted rule.
+        """
+        sagemaker_session = TrainDefaults.get_sagemaker_session(
+            sagemaker_session=self.sagemaker_session
+        )
+
+        return delete_notification_rules(
+            sagemaker_session=sagemaker_session,
+            rule_arn=rule_arn,
+            event_bus_arn=event_bus_arn,
+        )
+
+    def list_notification_rules(
+        self,
+        event_bus_arn: Optional[str] = None,
+    ) -> List[Dict[str, str]]:
+        """List all SDK-created EventBridge notification rules.
+
+        Returns:
+            List of dicts with 'name', 'arn', and 'state' for each rule.
+        """
+        sagemaker_session = TrainDefaults.get_sagemaker_session(
+            sagemaker_session=self.sagemaker_session
+        )
+
+        return list_notification_rules(
+            sagemaker_session=sagemaker_session,
+            event_bus_arn=event_bus_arn,
+        )
+
     def stream_logs(self, poll: int = 5, start_time: Optional[Any] = None) -> None:
         """Stream CloudWatch logs in real-time (like ``kubectl logs -f``).
 

From 44c8930e8bd2d006cf8d880a0049412fd263be91 Mon Sep 17 00:00:00 2001
From: Ealynn Hsu <ealynnh@amazon.com>
Date: Thu, 16 Jul 2026 20:41:44 +0000
Subject: [PATCH 3/6] Adding tests, dedupe logic, and example

---
 .../train/common_utils/notifications.py       | 23 ++---
 .../train/common_utils/test_notifications.py  | 40 +++++++--
 .../sft_finetuning_serverful_smtj.ipynb       | 83 ++++++++++---------
 3 files changed, 92 insertions(+), 54 deletions(-)

diff --git a/sagemaker-train/src/sagemaker/train/common_utils/notifications.py b/sagemaker-train/src/sagemaker/train/common_utils/notifications.py
index 234f8900da..d08f0575a8 100644
--- a/sagemaker-train/src/sagemaker/train/common_utils/notifications.py
+++ b/sagemaker-train/src/sagemaker/train/common_utils/notifications.py
@@ -16,25 +16,28 @@
 
 # Rule name prefix used to identify SDK-created rules
 _RULE_NAME_PREFIX = "sm-pysdk-job-notif"
-
-# Default events to notify on
 _DEFAULT_EVENTS = ["Completed", "Failed", "Stopped"]
-
-# Valid event status values
 _VALID_EVENTS = {"Completed", "Failed", "Stopped", "InProgress"}
 
 
-def _get_rule_name(sns_topic_arn: str) -> str:
-    """Generate a deterministic rule name from the SNS topic ARN.
+def _get_rule_name(sns_topic_arn: str, events: List[str], job_name_prefix: Optional[str] = None) -> str:
+    """Generate a deterministic rule name from the full notification config.
+
+    Hashes the topic ARN + events + prefix so that:
+    - Identical configs -> same rule name
+    - Any config difference -> different rule name
 
     Args:
         sns_topic_arn: The SNS topic ARN.
+        events: Normalized list of event statuses.
+        job_name_prefix: Optional job name prefix filter.
 
     Returns:
         Rule name like "sm-pysdk-job-notif-a3f8b2c1".
     """
-    arn_hash = hashlib.sha256(sns_topic_arn.encode()).hexdigest()[:8]
-    return f"{_RULE_NAME_PREFIX}-{arn_hash}"
+    config_str = f"{sns_topic_arn}|{','.join(sorted(events))}|{job_name_prefix or ''}"
+    config_hash = hashlib.sha256(config_str.encode()).hexdigest()[:8]
+    return f"{_RULE_NAME_PREFIX}-{config_hash}"
 
 
 def _normalize_events(events: Optional[List[str]]) -> List[str]:
@@ -184,7 +187,7 @@ def enable_notifications(
     _validate_sns_topic(sns_client, sns_topic_arn)
 
     normalized_events = _normalize_events(events)
-    rule_name = _get_rule_name(sns_topic_arn)
+    rule_name = _get_rule_name(sns_topic_arn, normalized_events, job_name_prefix)
     event_pattern = _build_event_pattern(normalized_events, job_name_prefix)
 
     put_rule_kwargs = {
@@ -202,7 +205,7 @@ def enable_notifications(
 
     # Add SNS topic as target with formatted message
     target_id = f"{rule_name}-sns-target"
-    # Use a JSON object template — SNS renders JSON objects with line breaks
+    # Use a JSON object template
     input_template = (
         '{"Job": "<job_name>",'
         ' "Status": "<status>",'
diff --git a/sagemaker-train/tests/unit/train/common_utils/test_notifications.py b/sagemaker-train/tests/unit/train/common_utils/test_notifications.py
index 9c72fc112d..f3c9c0e02e 100644
--- a/sagemaker-train/tests/unit/train/common_utils/test_notifications.py
+++ b/sagemaker-train/tests/unit/train/common_utils/test_notifications.py
@@ -18,22 +18,52 @@
 
 class TestRuleNaming:
 
-    def test_deterministic_from_arn(self):
-        """Same ARN always produces the same rule name."""
+    def test_deterministic_from_same_config(self):
+        """Same config always produces the same rule name."""
         arn = "arn:aws:sns:us-east-1:123456789:my-topic"
-        assert _get_rule_name(arn) == _get_rule_name(arn)
+        events = ["Completed", "Failed"]
+        assert _get_rule_name(arn, events) == _get_rule_name(arn, events)
 
     def test_different_arns_different_names(self):
         """Different ARNs produce different rule names."""
+        events = ["Completed", "Failed"]
         arn1 = "arn:aws:sns:us-east-1:123456789:topic-a"
         arn2 = "arn:aws:sns:us-east-1:123456789:topic-b"
-        assert _get_rule_name(arn1) != _get_rule_name(arn2)
+        assert _get_rule_name(arn1, events) != _get_rule_name(arn2, events)
+
+    def test_same_topic_different_events_different_names(self):
+        """Same topic with different events produces different rule names."""
+        arn = "arn:aws:sns:us-east-1:123456789:my-topic"
+        assert _get_rule_name(arn, ["Completed"]) != _get_rule_name(arn, ["Completed", "Failed"])
+
+    def test_same_topic_different_prefix_different_names(self):
+        """Same topic with different job_name_prefix produces different rule names."""
+        arn = "arn:aws:sns:us-east-1:123456789:my-topic"
+        events = ["Completed", "Failed"]
+        assert _get_rule_name(arn, events, "team-a-") != _get_rule_name(arn, events, "team-b-")
+
+    def test_same_topic_same_prefix_same_name(self):
+        """Same topic + same events + same prefix = same rule (idempotent)."""
+        arn = "arn:aws:sns:us-east-1:123456789:my-topic"
+        events = ["Completed", "Failed"]
+        assert _get_rule_name(arn, events, "my-prefix-") == _get_rule_name(arn, events, "my-prefix-")
+
+    def test_event_order_does_not_matter(self):
+        """Events are sorted internally, so order doesn't affect the hash."""
+        arn = "arn:aws:sns:us-east-1:123456789:my-topic"
+        assert _get_rule_name(arn, ["Failed", "Completed"]) == _get_rule_name(arn, ["Completed", "Failed"])
 
     def test_prefix_present(self):
         """Rule name starts with the SDK prefix."""
-        name = _get_rule_name("arn:aws:sns:us-east-1:123:topic")
+        name = _get_rule_name("arn:aws:sns:us-east-1:123:topic", ["Completed"])
         assert name.startswith("sm-pysdk-job-notif-")
 
+    def test_no_prefix_vs_none_prefix_same(self):
+        """No prefix and None prefix produce the same rule."""
+        arn = "arn:aws:sns:us-east-1:123:topic"
+        events = ["Completed"]
+        assert _get_rule_name(arn, events) == _get_rule_name(arn, events, None)
+
 
 class TestNormalizeEvents:
 
diff --git a/v3-examples/model-customization-examples/sft_finetuning_serverful_smtj.ipynb b/v3-examples/model-customization-examples/sft_finetuning_serverful_smtj.ipynb
index f130931e54..92a830844b 100644
--- a/v3-examples/model-customization-examples/sft_finetuning_serverful_smtj.ipynb
+++ b/v3-examples/model-customization-examples/sft_finetuning_serverful_smtj.ipynb
@@ -1,17 +1,4 @@
 {
- "nbformat": 4,
- "nbformat_minor": 5,
- "metadata": {
-  "kernelspec": {
-   "display_name": "Python 3",
-   "language": "python",
-   "name": "python3"
-  },
-  "language_info": {
-   "name": "python",
-   "version": "3.12.0"
-  }
- },
  "cells": [
   {
    "cell_type": "markdown",
@@ -40,17 +27,19 @@
   },
   {
    "cell_type": "code",
+   "execution_count": null,
    "metadata": {},
+   "outputs": [],
    "source": [
     "import json\n",
     "import boto3"
-   ],
-   "outputs": [],
-   "execution_count": null
+   ]
   },
   {
    "cell_type": "code",
+   "execution_count": null,
    "metadata": {},
+   "outputs": [],
    "source": [
     "# === Fill in your AWS resources ===\n",
     "REGION = \"<your-region>\"  # e.g. \"us-east-1\"\n",
@@ -59,9 +48,7 @@
     "\n",
     "S3_OUTPUT_PATH = f\"s3://{S3_BUCKET}/sft-smtj/output\"\n",
     "TRAINING_DATASET = f\"s3://{S3_BUCKET}/datasets/sft_training_data.jsonl\""
-   ],
-   "outputs": [],
-   "execution_count": null
+   ]
   },
   {
    "cell_type": "markdown",
@@ -75,7 +62,9 @@
   },
   {
    "cell_type": "code",
+   "execution_count": null,
    "metadata": {},
+   "outputs": [],
    "source": [
     "from sagemaker.train import SFTTrainer\n",
     "from sagemaker.train.common import TrainingType\n",
@@ -93,9 +82,7 @@
     "    role=ROLE_ARN,\n",
     "    base_job_name=\"sft-smtj\",\n",
     ")"
-   ],
-   "outputs": [],
-   "execution_count": null
+   ]
   },
   {
    "cell_type": "markdown",
@@ -115,7 +102,9 @@
   },
   {
    "cell_type": "code",
+   "execution_count": null,
    "metadata": {},
+   "outputs": [],
    "source": [
     "import yaml\n",
     "\n",
@@ -134,13 +123,13 @@
     "\n",
     "print(\"Recipe file contents:\")\n",
     "print(yaml.dump(recipe_config, default_flow_style=False))"
-   ],
-   "outputs": [],
-   "execution_count": null
+   ]
   },
   {
    "cell_type": "code",
+   "execution_count": null,
    "metadata": {},
+   "outputs": [],
    "source": [
     "# Create trainer with recipe + overrides\n",
     "# Here we override learning_rate (1e-5 -> 5e-6) and num_epochs (3 -> 5)\n",
@@ -158,15 +147,18 @@
     "    recipe=\"my_sft_recipe.yaml\",\n",
     "    overrides={\"training_config\": {\"learning_rate\": 5e-6, \"num_epochs\": 5}},\n",
     "    base_job_name=\"sft-recipe-override-smtj\",\n",
+    "    # Add SMTJ job notifications sent on job completion to an SNS topic (user creates this manually)\n",
+    "    # notifications={\n",
+    "    #     \"sns_topic_arn\": \"arn:aws:sns:us-east-1:012345678910:user-sns-topic\", \n",
+    "    #     \"job_name_prefix\": \"username-\"\n",
+    "    # },\n",
     ")\n",
     "\n",
     "# Inspect the resolved recipe to confirm overrides were applied\n",
     "resolved = sft_trainer_with_recipe.get_resolved_recipe()\n",
     "print(\"Resolved training_config:\")\n",
     "print(json.dumps(resolved.get(\"training_config\", resolved), indent=2))"
-   ],
-   "outputs": [],
-   "execution_count": null
+   ]
   },
   {
    "cell_type": "markdown",
@@ -177,7 +169,9 @@
   },
   {
    "cell_type": "code",
+   "execution_count": null,
    "metadata": {},
+   "outputs": [],
    "source": [
     "sft_trainer.hyperparameters.max_steps = 50\n",
     "sft_trainer.hyperparameters.learning_rate = 5e-6\n",
@@ -186,9 +180,7 @@
     "# Submit (non-blocking)\n",
     "training_job = sft_trainer.train(wait=False)\n",
     "print(f\"Training job submitted: {training_job}\")"
-   ],
-   "outputs": [],
-   "execution_count": null
+   ]
   },
   {
    "cell_type": "markdown",
@@ -199,16 +191,16 @@
   },
   {
    "cell_type": "code",
+   "execution_count": null,
    "metadata": {},
+   "outputs": [],
    "source": [
     "from sagemaker.core.resources import TrainingJob\n",
     "\n",
     "job = TrainingJob.get(training_job_name=training_job.training_job_name)\n",
     "print(f\"Status: {job.training_job_status}\")\n",
     "print(f\"Secondary Status: {job.secondary_status}\")"
-   ],
-   "outputs": [],
-   "execution_count": null
+   ]
   },
   {
    "cell_type": "markdown",
@@ -222,7 +214,9 @@
   },
   {
    "cell_type": "code",
+   "execution_count": null,
    "metadata": {},
+   "outputs": [],
    "source": [
     "import yaml\n",
     "from sagemaker.train import ModelTrainer\n",
@@ -250,9 +244,20 @@
     "\n",
     "resolved = model_trainer.get_resolved_recipe()\n",
     "print(json.dumps(resolved.get(\"training_config\", resolved), indent=2))"
-   ],
-   "outputs": [],
-   "execution_count": null
+   ]
   }
- ]
-}
\ No newline at end of file
+ ],
+ "metadata": {
+  "kernelspec": {
+   "display_name": "Python 3",
+   "language": "python",
+   "name": "python3"
+  },
+  "language_info": {
+   "name": "python",
+   "version": "3.12.0"
+  }
+ },
+ "nbformat": 4,
+ "nbformat_minor": 5
+}

From 8a8276305fcf809877cb8991cf78b5b4c518171e Mon Sep 17 00:00:00 2001
From: Ealynn Hsu <ealynnh@amazon.com>
Date: Fri, 17 Jul 2026 19:30:31 +0000
Subject: [PATCH 4/6] Update ARN example values, clean up SM session
 definitions, rename func

---
 .../src/sagemaker/train/base_trainer.py       | 28 ++++---------
 .../train/common_utils/notifications.py       |  4 +-
 .../train/common_utils/test_notifications.py  | 40 +++++++++----------
 3 files changed, 30 insertions(+), 42 deletions(-)

diff --git a/sagemaker-train/src/sagemaker/train/base_trainer.py b/sagemaker-train/src/sagemaker/train/base_trainer.py
index dd9096d3f6..4233937c1c 100644
--- a/sagemaker-train/src/sagemaker/train/base_trainer.py
+++ b/sagemaker-train/src/sagemaker/train/base_trainer.py
@@ -34,7 +34,7 @@
 )
 from sagemaker.train.common_utils.metrics_visualizer import plot_training_metrics
 from sagemaker.train.common_utils.mlflow_config_utils import resolve_mlflow_tracking_fields
-from sagemaker.train.common_utils.notifications import enable_notifications, delete_notification_rules, list_notification_rules
+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.defaults import TrainDefaults
@@ -476,23 +476,19 @@ def _setup_notifications(self, notifications: Optional[Dict[str, Any]]) -> Optio
         if not isinstance(notifications, dict):
             raise ValueError(
                 "notifications must be a dict with at least 'sns_topic_arn'. "
-                "Example: {'sns_topic_arn': 'arn:aws:sns:us-east-1:123:my-topic'}"
+                "Example: {'sns_topic_arn': 'arn:aws:sns:us-east-1:123456789012:my-topic'}"
             )
 
         sns_topic_arn = notifications.get("sns_topic_arn")
         if not sns_topic_arn:
             raise ValueError(
                 "notifications config requires 'sns_topic_arn'. "
-                "Example: {'sns_topic_arn': 'arn:aws:sns:us-east-1:123:my-topic'}"
+                "Example: {'sns_topic_arn': 'arn:aws:sns:us-east-1:123456789012:my-topic'}"
             )
 
-        sagemaker_session = TrainDefaults.get_sagemaker_session(
-            sagemaker_session=self.sagemaker_session
-        )
-
         rule_arn = enable_notifications(
             sns_topic_arn=sns_topic_arn,
-            sagemaker_session=sagemaker_session,
+            sagemaker_session=TrainDefaults.get_sagemaker_session(sagemaker_session=self.sagemaker_session),
             events=notifications.get("events"),
             event_bus_arn=notifications.get("event_bus_arn"),
             job_name_prefix=notifications.get("job_name_prefix"),
@@ -500,7 +496,7 @@ def _setup_notifications(self, notifications: Optional[Dict[str, Any]]) -> Optio
 
         return rule_arn
 
-    def delete_notification_rules(
+    def delete_notification_rule(
         self,
         rule_arn: str,
         event_bus_arn: Optional[str] = None,
@@ -514,12 +510,8 @@ def delete_notification_rules(
         Returns:
             The name of the deleted rule.
         """
-        sagemaker_session = TrainDefaults.get_sagemaker_session(
-            sagemaker_session=self.sagemaker_session
-        )
-
-        return delete_notification_rules(
-            sagemaker_session=sagemaker_session,
+        return delete_notification_rule(
+            sagemaker_session=TrainDefaults.get_sagemaker_session(sagemaker_session=self.sagemaker_session),
             rule_arn=rule_arn,
             event_bus_arn=event_bus_arn,
         )
@@ -533,12 +525,8 @@ def list_notification_rules(
         Returns:
             List of dicts with 'name', 'arn', and 'state' for each rule.
         """
-        sagemaker_session = TrainDefaults.get_sagemaker_session(
-            sagemaker_session=self.sagemaker_session
-        )
-
         return list_notification_rules(
-            sagemaker_session=sagemaker_session,
+            sagemaker_session=TrainDefaults.get_sagemaker_session(sagemaker_session=self.sagemaker_session),
             event_bus_arn=event_bus_arn,
         )
 
diff --git a/sagemaker-train/src/sagemaker/train/common_utils/notifications.py b/sagemaker-train/src/sagemaker/train/common_utils/notifications.py
index d08f0575a8..3566f89223 100644
--- a/sagemaker-train/src/sagemaker/train/common_utils/notifications.py
+++ b/sagemaker-train/src/sagemaker/train/common_utils/notifications.py
@@ -251,7 +251,7 @@ def enable_notifications(
     return rule_arn
 
 
-def delete_notification_rules(
+def delete_notification_rule(
     sagemaker_session,
     rule_arn: str,
     event_bus_arn: Optional[str] = None,
@@ -293,7 +293,7 @@ def delete_notification_rules(
             raise ValueError(
                 f"Rule '{rule_name}' not found on event bus '{event_bus_name}'. "
                 "If the rule was created on a custom event bus, pass the same "
-                "event_bus_arn to delete_notification_rules()."
+                "event_bus_arn to delete_notification_rule()."
             ) from e
         logger.warning(f"Failed to delete rule {rule_name}: {e}")
 
diff --git a/sagemaker-train/tests/unit/train/common_utils/test_notifications.py b/sagemaker-train/tests/unit/train/common_utils/test_notifications.py
index f3c9c0e02e..483c6529e5 100644
--- a/sagemaker-train/tests/unit/train/common_utils/test_notifications.py
+++ b/sagemaker-train/tests/unit/train/common_utils/test_notifications.py
@@ -10,7 +10,7 @@
     _build_event_pattern,
     _get_rule_name,
     _normalize_events,
-    delete_notification_rules,
+    delete_notification_rule,
     enable_notifications,
     list_notification_rules,
 )
@@ -20,47 +20,47 @@ class TestRuleNaming:
 
     def test_deterministic_from_same_config(self):
         """Same config always produces the same rule name."""
-        arn = "arn:aws:sns:us-east-1:123456789:my-topic"
+        arn = "arn:aws:sns:us-east-1:123456789012:my-topic"
         events = ["Completed", "Failed"]
         assert _get_rule_name(arn, events) == _get_rule_name(arn, events)
 
     def test_different_arns_different_names(self):
         """Different ARNs produce different rule names."""
         events = ["Completed", "Failed"]
-        arn1 = "arn:aws:sns:us-east-1:123456789:topic-a"
-        arn2 = "arn:aws:sns:us-east-1:123456789:topic-b"
+        arn1 = "arn:aws:sns:us-east-1:123456789012:topic-a"
+        arn2 = "arn:aws:sns:us-east-1:123456789012:topic-b"
         assert _get_rule_name(arn1, events) != _get_rule_name(arn2, events)
 
     def test_same_topic_different_events_different_names(self):
         """Same topic with different events produces different rule names."""
-        arn = "arn:aws:sns:us-east-1:123456789:my-topic"
+        arn = "arn:aws:sns:us-east-1:123456789012:my-topic"
         assert _get_rule_name(arn, ["Completed"]) != _get_rule_name(arn, ["Completed", "Failed"])
 
     def test_same_topic_different_prefix_different_names(self):
         """Same topic with different job_name_prefix produces different rule names."""
-        arn = "arn:aws:sns:us-east-1:123456789:my-topic"
+        arn = "arn:aws:sns:us-east-1:123456789012:my-topic"
         events = ["Completed", "Failed"]
         assert _get_rule_name(arn, events, "team-a-") != _get_rule_name(arn, events, "team-b-")
 
     def test_same_topic_same_prefix_same_name(self):
         """Same topic + same events + same prefix = same rule (idempotent)."""
-        arn = "arn:aws:sns:us-east-1:123456789:my-topic"
+        arn = "arn:aws:sns:us-east-1:123456789012:my-topic"
         events = ["Completed", "Failed"]
         assert _get_rule_name(arn, events, "my-prefix-") == _get_rule_name(arn, events, "my-prefix-")
 
     def test_event_order_does_not_matter(self):
         """Events are sorted internally, so order doesn't affect the hash."""
-        arn = "arn:aws:sns:us-east-1:123456789:my-topic"
+        arn = "arn:aws:sns:us-east-1:123456789012:my-topic"
         assert _get_rule_name(arn, ["Failed", "Completed"]) == _get_rule_name(arn, ["Completed", "Failed"])
 
     def test_prefix_present(self):
         """Rule name starts with the SDK prefix."""
-        name = _get_rule_name("arn:aws:sns:us-east-1:123:topic", ["Completed"])
+        name = _get_rule_name("arn:aws:sns:us-east-1:123456789012:topic", ["Completed"])
         assert name.startswith("sm-pysdk-job-notif-")
 
     def test_no_prefix_vs_none_prefix_same(self):
         """No prefix and None prefix produce the same rule."""
-        arn = "arn:aws:sns:us-east-1:123:topic"
+        arn = "arn:aws:sns:us-east-1:123456789012:topic"
         events = ["Completed"]
         assert _get_rule_name(arn, events) == _get_rule_name(arn, events, None)
 
@@ -118,16 +118,16 @@ def test_creates_rule_and_target(self):
             events_client if svc == "events" else sns_client
         )
 
-        events_client.put_rule.return_value = {"RuleArn": "arn:aws:events:us-east-1:123:rule/sm-pysdk-notif-abc"}
+        events_client.put_rule.return_value = {"RuleArn": "arn:aws:events:us-east-1:123456789012:rule/sm-pysdk-notif-abc"}
         events_client.put_targets.return_value = {"FailedEntryCount": 0}
         events_client.list_rules.return_value = {"Rules": []}
 
         arn = enable_notifications(
-            sns_topic_arn="arn:aws:sns:us-east-1:123:my-topic",
+            sns_topic_arn="arn:aws:sns:us-east-1:123456789012:my-topic",
             sagemaker_session=session,
         )
 
-        assert arn == "arn:aws:events:us-east-1:123:rule/sm-pysdk-notif-abc"
+        assert arn == "arn:aws:events:us-east-1:123456789012:rule/sm-pysdk-notif-abc"
         events_client.put_rule.assert_called_once()
         events_client.put_targets.assert_called_once()
 
@@ -144,7 +144,7 @@ def test_with_custom_events_and_prefix(self):
         events_client.list_rules.return_value = {"Rules": []}
 
         enable_notifications(
-            sns_topic_arn="arn:aws:sns:us-east-1:123:my-topic",
+            sns_topic_arn="arn:aws:sns:us-east-1:123456789012:my-topic",
             sagemaker_session=session,
             events=["completed"],
             job_name_prefix="ealynnh-",
@@ -169,9 +169,9 @@ def test_deletes_specific_rule(self):
             "Targets": [{"Id": "target-1"}]
         }
 
-        deleted = delete_notification_rules(
+        deleted = delete_notification_rule(
             sagemaker_session=session,
-            rule_arn="arn:aws:events:us-east-1:123:rule/sm-pysdk-job-notif-abc123",
+            rule_arn="arn:aws:events:us-east-1:123456789012:rule/sm-pysdk-job-notif-abc123",
         )
 
         assert deleted == "sm-pysdk-job-notif-abc123"
@@ -202,7 +202,7 @@ def test_hyperpod_raises_not_implemented(self):
         )
 
         with pytest.raises(NotImplementedError, match="not supported for HyperPod"):
-            trainer._setup_notifications({"sns_topic_arn": "arn:aws:sns:us-east-1:123:topic"})
+            trainer._setup_notifications({"sns_topic_arn": "arn:aws:sns:us-east-1:123456789012:topic"})
 
     def test_missing_sns_arn_raises(self):
         trainer = self._make_trainer()
@@ -228,8 +228,8 @@ def test_lists_sdk_rules(self):
         paginator = MagicMock()
         paginator.paginate.return_value = [
             {"Rules": [
-                {"Name": "sm-pysdk-job-notif-aaa", "Arn": "arn:aws:events:us-east-1:123:rule/sm-pysdk-job-notif-aaa", "State": "ENABLED"},
-                {"Name": "sm-pysdk-job-notif-bbb", "Arn": "arn:aws:events:us-east-1:123:rule/sm-pysdk-job-notif-bbb", "State": "ENABLED"},
+                {"Name": "sm-pysdk-job-notif-aaa", "Arn": "arn:aws:events:us-east-1:123456789012:rule/sm-pysdk-job-notif-aaa", "State": "ENABLED"},
+                {"Name": "sm-pysdk-job-notif-bbb", "Arn": "arn:aws:events:us-east-1:123456789012:rule/sm-pysdk-job-notif-bbb", "State": "ENABLED"},
             ]}
         ]
         events_client.get_paginator.return_value = paginator
@@ -238,4 +238,4 @@ def test_lists_sdk_rules(self):
 
         assert len(rules) == 2
         assert rules[0]["name"] == "sm-pysdk-job-notif-aaa"
-        assert rules[0]["arn"] == "arn:aws:events:us-east-1:123:rule/sm-pysdk-job-notif-aaa"
+        assert rules[0]["arn"] == "arn:aws:events:us-east-1:123456789012:rule/sm-pysdk-job-notif-aaa"

From 3c69274afbef426bdfe5d255991d8b48432abbe9 Mon Sep 17 00:00:00 2001
From: Ealynn Hsu <ealynnh@amazon.com>
Date: Fri, 17 Jul 2026 20:49:35 +0000
Subject: [PATCH 5/6] Update trainers to include 'notifications' param and add
 arn regex check

---
 .../src/sagemaker/train/common_utils/notifications.py    | 3 ++-
 sagemaker-train/src/sagemaker/train/cpt_trainer.py       | 9 +++++++--
 sagemaker-train/src/sagemaker/train/dpo_trainer.py       | 9 +++++++--
 .../src/sagemaker/train/multi_turn_rl_trainer.py         | 9 +++++++--
 sagemaker-train/src/sagemaker/train/rlaif_trainer.py     | 9 +++++++--
 sagemaker-train/src/sagemaker/train/rlvr_trainer.py      | 7 ++++++-
 sagemaker-train/src/sagemaker/train/sft_trainer.py       | 9 +++++++--
 7 files changed, 43 insertions(+), 12 deletions(-)

diff --git a/sagemaker-train/src/sagemaker/train/common_utils/notifications.py b/sagemaker-train/src/sagemaker/train/common_utils/notifications.py
index 3566f89223..ba579f1788 100644
--- a/sagemaker-train/src/sagemaker/train/common_utils/notifications.py
+++ b/sagemaker-train/src/sagemaker/train/common_utils/notifications.py
@@ -10,6 +10,7 @@
 import hashlib
 import json
 import logging
+import re
 from typing import Dict, List, Optional
 
 logger = logging.getLogger(__name__)
@@ -173,7 +174,7 @@ def enable_notifications(
         ValueError: If sns_topic_arn is invalid or topic doesn't exist.
         PermissionError: If caller lacks required permissions.
     """
-    if not sns_topic_arn or not sns_topic_arn.startswith("arn:aws:sns:"):
+    if not sns_topic_arn or not re.match(r"^arn:aws[a-z\-]*:sns:[a-z0-9\-]+:\d{12}:.+$", sns_topic_arn):
         raise ValueError(
             f"Invalid SNS topic ARN: '{sns_topic_arn}'. "
             "Must be a valid ARN like 'arn:aws:sns:us-east-1:012345678910:my-topic'."
diff --git a/sagemaker-train/src/sagemaker/train/cpt_trainer.py b/sagemaker-train/src/sagemaker/train/cpt_trainer.py
index b7ec2ff7fb..d3e3cf2489 100644
--- a/sagemaker-train/src/sagemaker/train/cpt_trainer.py
+++ b/sagemaker-train/src/sagemaker/train/cpt_trainer.py
@@ -12,7 +12,7 @@
 # language governing permissions and limitations under the License.
 """CPTTrainer — Continued Pre-Training on foundation models using SageMaker HyperPod."""
 
-from typing import Optional, Union
+from typing import Any, Dict, Optional, Union
 import logging
 
 from sagemaker.train.base_trainer import BaseTrainer
@@ -106,6 +106,10 @@ class CPTTrainer(BaseTrainer):
             Programmatic overrides dict.
         training_image (Optional[str]):
             Custom training container image URI. If not provided, auto-resolved from Hub.
+        notifications (Optional[Dict[str, Any]]):
+            Configuration for SNS notifications on job status changes. Requires 'sns_topic_arn'.
+            Optional keys: 'events' ["Completed", "Failed", "Stopped"], 'event_bus_arn',
+            and 'job_name_prefix'. If not specified, no notifications are sent.
     """
 
     _customization_technique = CustomizationTechnique.CPT.value
@@ -131,9 +135,10 @@ def __init__(
         data_mixing_config: Optional[DataMixingConfig] = None,
         base_model_name: Optional[str] = None,
         disable_output_compression: Optional[bool] = False,
+        notifications: Optional[Dict[str, Any]] = None,
         **kwargs,
     ):
-        super().__init__(training_image=training_image, base_model_name=base_model_name, disable_output_compression=disable_output_compression, **kwargs)
+        super().__init__(training_image=training_image, base_model_name=base_model_name, disable_output_compression=disable_output_compression, notifications=notifications, **kwargs)
 
         self.model, self._model_name, self.model_source = _resolve_model_with_checkpoint(
             model, self.base_model_name, compute, self.sagemaker_session,
diff --git a/sagemaker-train/src/sagemaker/train/dpo_trainer.py b/sagemaker-train/src/sagemaker/train/dpo_trainer.py
index 387e7b226a..db7a288632 100644
--- a/sagemaker-train/src/sagemaker/train/dpo_trainer.py
+++ b/sagemaker-train/src/sagemaker/train/dpo_trainer.py
@@ -1,4 +1,4 @@
-from typing import Optional, Union
+from typing import Any, Dict, Optional, Union
 import logging
 from sagemaker.ai_registry.dataset import DataSet
 from sagemaker.train.base_trainer import BaseTrainer
@@ -107,6 +107,10 @@ class DPOTrainer(BaseTrainer):
         is_multimodal (Optional[bool]):
             Whether the training dataset contains multimodal data. If None (default),
             auto-detected from the training dataset at train time.
+        notifications (Optional[Dict[str, Any]]):
+            Configuration for SNS notifications on job status changes. Requires 'sns_topic_arn'.
+            Optional keys: 'events' ["Completed", "Failed", "Stopped"], 'event_bus_arn',
+            and 'job_name_prefix'. If not specified, no notifications are sent.
     """
 
     _customization_technique = CustomizationTechnique.DPO.value
@@ -132,9 +136,10 @@ def __init__(
             is_multimodal: Optional[bool] = None,
             base_model_name: Optional[str] = None,
         disable_output_compression: Optional[bool] = False,
+        notifications: Optional[Dict[str, Any]] = None,
         **kwargs,
     ):
-        super().__init__(base_model_name=base_model_name, disable_output_compression=disable_output_compression, **kwargs)
+        super().__init__(base_model_name=base_model_name, disable_output_compression=disable_output_compression, notifications=notifications, **kwargs)
 
         self.model, self._model_name, self.model_source = _resolve_model_with_checkpoint(
             model, self.base_model_name, compute, self.sagemaker_session,
diff --git a/sagemaker-train/src/sagemaker/train/multi_turn_rl_trainer.py b/sagemaker-train/src/sagemaker/train/multi_turn_rl_trainer.py
index 5c85e2f6ae..13b9d53a0d 100644
--- a/sagemaker-train/src/sagemaker/train/multi_turn_rl_trainer.py
+++ b/sagemaker-train/src/sagemaker/train/multi_turn_rl_trainer.py
@@ -17,7 +17,7 @@
 import json
 import logging
 import re
-from typing import Optional, Union
+from typing import Any, Dict, Optional, Union
 
 import boto3
 
@@ -167,6 +167,10 @@ class MultiTurnRLTrainer(BaseTrainer):
         kms_key_arn: KMS key ID for output encryption (optional).
         accept_eula: Boolean for EULA acceptance (optional).
         **kwargs: Passed to BaseTrainer (sagemaker_session, role, base_job_name, tags).
+        notifications (Optional[Dict[str, Any]]):
+            Configuration for SNS notifications on job status changes. Requires 'sns_topic_arn'.
+            Optional keys: 'events' ["Completed", "Failed", "Stopped"], 'event_bus_arn',
+            and 'job_name_prefix'. If not specified, no notifications are sent.
     """
 
     def __init__(
@@ -187,9 +191,10 @@ def __init__(
         accept_eula: bool = False,
         recipe: Optional[str] = None,
         overrides: Optional[dict] = None,
+        notifications: Optional[Dict[str, Any]] = None,
         **kwargs,
     ):
-        super().__init__(**kwargs)
+        super().__init__(notifications=notifications, **kwargs)
         self._recipe_path = recipe
         self._overrides = overrides
         self._resolved_recipe_cache = None
diff --git a/sagemaker-train/src/sagemaker/train/rlaif_trainer.py b/sagemaker-train/src/sagemaker/train/rlaif_trainer.py
index 50722f6d80..0949ee0ea3 100644
--- a/sagemaker-train/src/sagemaker/train/rlaif_trainer.py
+++ b/sagemaker-train/src/sagemaker/train/rlaif_trainer.py
@@ -1,4 +1,4 @@
-from typing import Optional, Union
+from typing import Any, Dict, Optional, Union
 import logging
 from sagemaker.train.base_trainer import BaseTrainer
 from sagemaker.train.common import TrainingType, CustomizationTechnique, JOB_TYPE
@@ -118,6 +118,10 @@ class RLAIFTrainer(BaseTrainer):
         is_multimodal (Optional[bool]):
             Whether the training dataset contains multimodal data. If None (default),
             auto-detected from the training dataset at train time.
+        notifications (Optional[Dict[str, Any]]):
+            Configuration for SNS notifications on job status changes. Requires 'sns_topic_arn'.
+            Optional keys: 'events' ["Completed", "Failed", "Stopped"], 'event_bus_arn',
+            and 'job_name_prefix'. If not specified, no notifications are sent.
     """
 
     def __init__(
@@ -140,9 +144,10 @@ def __init__(
         accept_eula: bool = False,
         stopping_condition: Optional[StoppingCondition] = None,
         is_multimodal: Optional[bool] = None,
+        notifications: Optional[Dict[str, Any]] = None,
         **kwargs,
     ):
-        super().__init__(**kwargs)
+        super().__init__(notifications=notifications, **kwargs)
 
         # Resolve model and model name
         self.model, self._model_name = _resolve_model_and_name(model, self.sagemaker_session)
diff --git a/sagemaker-train/src/sagemaker/train/rlvr_trainer.py b/sagemaker-train/src/sagemaker/train/rlvr_trainer.py
index 0c0670425f..1130912b95 100644
--- a/sagemaker-train/src/sagemaker/train/rlvr_trainer.py
+++ b/sagemaker-train/src/sagemaker/train/rlvr_trainer.py
@@ -132,6 +132,10 @@ class RLVRTrainer(BaseTrainer):
         is_multimodal (Optional[bool]):
             Whether the training dataset contains multimodal data. If None (default),
             auto-detected from the training dataset at train time.
+        notifications (Optional[Dict[str, Any]]):
+            Configuration for SNS notifications on job status changes. Requires 'sns_topic_arn'.
+            Optional keys: 'events' ["Completed", "Failed", "Stopped"], 'event_bus_arn',
+            and 'job_name_prefix'. If not specified, no notifications are sent.
     """
 
     _customization_technique = CustomizationTechnique.RLVR.value
@@ -158,9 +162,10 @@ def __init__(
         is_multimodal: Optional[bool] = None,
         base_model_name: Optional[str] = None,
         disable_output_compression: Optional[bool] = False,
+        notifications: Optional[Dict[str, Any]] = None,
         **kwargs,
     ):
-        super().__init__(base_model_name=base_model_name, disable_output_compression=disable_output_compression, **kwargs)
+        super().__init__(base_model_name=base_model_name, disable_output_compression=disable_output_compression, notifications=notifications, **kwargs)
 
         self.model, self._model_name, self.model_source = _resolve_model_with_checkpoint(
             model, self.base_model_name, compute, self.sagemaker_session,
diff --git a/sagemaker-train/src/sagemaker/train/sft_trainer.py b/sagemaker-train/src/sagemaker/train/sft_trainer.py
index f16e097d8e..858b8c294e 100644
--- a/sagemaker-train/src/sagemaker/train/sft_trainer.py
+++ b/sagemaker-train/src/sagemaker/train/sft_trainer.py
@@ -1,4 +1,4 @@
-from typing import Optional, Union
+from typing import Any, Dict, Optional, Union
 import logging
 from sagemaker.train.base_trainer import BaseTrainer
 from sagemaker.train.common import TrainingType, CustomizationTechnique, JOB_TYPE
@@ -136,6 +136,10 @@ class SFTTrainer(BaseTrainer):
             Whether to disable compression of model output artifacts. When True,
             model artifacts are stored uncompressed in S3 (compression_type="NONE").
             Recommended for large model outputs. Defaults to False (gzip compression).
+        notifications (Optional[Dict[str, Any]]):
+            Configuration for SNS notifications on job status changes. Requires 'sns_topic_arn'.
+            Optional keys: 'events' ["Completed", "Failed", "Stopped"], 'event_bus_arn',
+            and 'job_name_prefix'. If not specified, no notifications are sent.
     """
 
     _customization_technique = CustomizationTechnique.SFT.value
@@ -162,9 +166,10 @@ def __init__(
         data_mixing_config: Optional[DataMixingConfig] = None,
         base_model_name: Optional[str] = None,
         disable_output_compression: Optional[bool] = False,
+        notifications: Optional[Dict[str, Any]] = None,
         **kwargs,
     ):
-        super().__init__(base_model_name=base_model_name, disable_output_compression=disable_output_compression, **kwargs)
+        super().__init__(base_model_name=base_model_name, disable_output_compression=disable_output_compression, notifications=notifications, **kwargs)
 
         self.model, self._model_name, self.model_source = _resolve_model_with_checkpoint(
             model, self.base_model_name, compute, self.sagemaker_session,

From e19096d13085838bbaef26ee8202db22e28115e6 Mon Sep 17 00:00:00 2001
From: Ealynn Hsu <ealynnh@amazon.com>
Date: Fri, 17 Jul 2026 21:30:26 +0000
Subject: [PATCH 6/6] Address PR comments, use botocore errors, return
 notification rule arn

---
 sagemaker-train/src/sagemaker/train/base_trainer.py    | 10 +++++++++-
 .../src/sagemaker/train/common_utils/notifications.py  |  3 ++-
 2 files changed, 11 insertions(+), 2 deletions(-)

diff --git a/sagemaker-train/src/sagemaker/train/base_trainer.py b/sagemaker-train/src/sagemaker/train/base_trainer.py
index 4233937c1c..009d1c2cb6 100644
--- a/sagemaker-train/src/sagemaker/train/base_trainer.py
+++ b/sagemaker-train/src/sagemaker/train/base_trainer.py
@@ -76,6 +76,12 @@ class BaseTrainer(ABC):
         training_image (Optional[str]):
             Custom training container image URI. If not provided, the image is
             auto-resolved from the model's recipe metadata in SageMaker Hub.
+        notifications (Optional[Dict[str, Any]]):
+            Configuration for SNS notifications on job status changes. Requires 'sns_topic_arn'.
+            Optional keys: 'events' ["Completed", "Failed", "Stopped"], 'event_bus_arn',
+            and 'job_name_prefix'. If not specified, no notifications are sent.
+        notification_rule_arn (str):
+            String of the EventBridge rule that is set up when enabling job notifications.
     """
     
     # Class-level attributes with default values
@@ -116,10 +122,11 @@ def __init__(
         self.training_image = training_image
         self.base_model_name = base_model_name
         self.disable_output_compression = disable_output_compression
+        self.notification_rule_arn = None
 
         # Set up notifications if configured
         if notifications:
-            self._setup_notifications(notifications)
+            self.notification_rule_arn = self._setup_notifications(notifications)
         self._checkpoint_s3_uri = None
 
     def _is_nova_model_for_telemetry(self) -> bool:
@@ -494,6 +501,7 @@ def _setup_notifications(self, notifications: Optional[Dict[str, Any]]) -> Optio
             job_name_prefix=notifications.get("job_name_prefix"),
         )
 
+        logger.debug("Notification rule ARN: %s", rule_arn)
         return rule_arn
 
     def delete_notification_rule(
diff --git a/sagemaker-train/src/sagemaker/train/common_utils/notifications.py b/sagemaker-train/src/sagemaker/train/common_utils/notifications.py
index ba579f1788..c7e65c2792 100644
--- a/sagemaker-train/src/sagemaker/train/common_utils/notifications.py
+++ b/sagemaker-train/src/sagemaker/train/common_utils/notifications.py
@@ -12,6 +12,7 @@
 import logging
 import re
 from typing import Dict, List, Optional
+from botocore.exceptions import ClientError
 
 logger = logging.getLogger(__name__)
 
@@ -110,7 +111,7 @@ def _validate_notifications_permissions(events_client) -> None:
     """
     try:
         events_client.list_rules(NamePrefix=_RULE_NAME_PREFIX, Limit=1)
-    except events_client.exceptions.ClientError as e:
+    except ClientError as e:
         error_code = e.response.get("Error", {}).get("Code", "")
         if error_code in ("AccessDeniedException", "AccessDenied"):
             raise PermissionError(