[Feat]: Job Notifications for SMTJ

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_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
@@ -75,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
@@ -102,6 +109,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 +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.notification_rule_arn = self._setup_notifications(notifications)
  self._checkpoint_s3_uri = None
 
  def _is_nova_model_for_telemetry(self) -> bool:
@@ -433,6 +446,98 @@ 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: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:123456789012:my-topic'}"
+ )
+
+ rule_arn = enable_notifications(
+ sns_topic_arn=sns_topic_arn,
+ 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"),
+ )
+
+ logger.debug("Notification rule ARN: %s", rule_arn)
+ return rule_arn
+
+ def delete_notification_rule(
+ 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.
+ """
+ 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,
+ )
+
+ 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.
+ """
+ return list_notification_rules(
+ sagemaker_session=TrainDefaults.get_sagemaker_session(sagemaker_session=self.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``).
 

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_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
@@ -75,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
@@ -102,6 +109,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 +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.notification_rule_arn = self._setup_notifications(notifications)
  self._checkpoint_s3_uri = None
 
  def _is_nova_model_for_telemetry(self) -> bool:
@@ -433,6 +446,98 @@ 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: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:123456789012:my-topic'}"
+ )
+
+ rule_arn = enable_notifications(
+ sns_topic_arn=sns_topic_arn,
+ 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"),
+ )
+
+ logger.debug("Notification rule ARN: %s", rule_arn)
+ return rule_arn
+
+ def delete_notification_rule(
+ 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.
+ """
+ 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,
+ )
+
+ 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.
+ """
+ return list_notification_rules(
+ sagemaker_session=TrainDefaults.get_sagemaker_session(sagemaker_session=self.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``).
 

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_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
@@ -75,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
@@ -102,6 +109,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 +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.notification_rule_arn = self._setup_notifications(notifications)
  self._checkpoint_s3_uri = None
 
  def _is_nova_model_for_telemetry(self) -> bool:
@@ -433,6 +446,98 @@ 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: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:123456789012:my-topic'}"
+ )
+
+ rule_arn = enable_notifications(
+ sns_topic_arn=sns_topic_arn,
+ 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"),
+ )
+
+ logger.debug("Notification rule ARN: %s", rule_arn)
+ return rule_arn
+
+ def delete_notification_rule(
+ 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.
+ """
+ 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,
+ )
+
+ 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.
+ """
+ return list_notification_rules(
+ sagemaker_session=TrainDefaults.get_sagemaker_session(sagemaker_session=self.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``).
 

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_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
@@ -75,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
@@ -102,6 +109,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 +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.notification_rule_arn = self._setup_notifications(notifications)
  self._checkpoint_s3_uri = None
 
  def _is_nova_model_for_telemetry(self) -> bool:
@@ -433,6 +446,98 @@ 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: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:123456789012:my-topic'}"
+ )
+
+ rule_arn = enable_notifications(
+ sns_topic_arn=sns_topic_arn,
+ 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"),
+ )
+
+ logger.debug("Notification rule ARN: %s", rule_arn)
+ return rule_arn
+
+ def delete_notification_rule(
+ 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.
+ """
+ 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,
+ )
+
+ 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.
+ """
+ return list_notification_rules(
+ sagemaker_session=TrainDefaults.get_sagemaker_session(sagemaker_session=self.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``).
 

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_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
@@ -75,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
@@ -102,6 +109,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 +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.notification_rule_arn = self._setup_notifications(notifications)
  self._checkpoint_s3_uri = None
 
  def _is_nova_model_for_telemetry(self) -> bool:
@@ -433,6 +446,98 @@ 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: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:123456789012:my-topic'}"
+ )
+
+ rule_arn = enable_notifications(
+ sns_topic_arn=sns_topic_arn,
+ 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"),
+ )
+
+ logger.debug("Notification rule ARN: %s", rule_arn)
+ return rule_arn
+
+ def delete_notification_rule(
+ 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.
+ """
+ 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,
+ )
+
+ 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.
+ """
+ return list_notification_rules(
+ sagemaker_session=TrainDefaults.get_sagemaker_session(sagemaker_session=self.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``).
 

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_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
@@ -75,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
@@ -102,6 +109,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 +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.notification_rule_arn = self._setup_notifications(notifications)
  self._checkpoint_s3_uri = None
 
  def _is_nova_model_for_telemetry(self) -> bool:
@@ -433,6 +446,98 @@ 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: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:123456789012:my-topic'}"
+ )
+
+ rule_arn = enable_notifications(
+ sns_topic_arn=sns_topic_arn,
+ 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"),
+ )
+
+ logger.debug("Notification rule ARN: %s", rule_arn)
+ return rule_arn
+
+ def delete_notification_rule(
+ 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.
+ """
+ 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,
+ )
+
+ 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.
+ """
+ return list_notification_rules(
+ sagemaker_session=TrainDefaults.get_sagemaker_session(sagemaker_session=self.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``).
 

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_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
@@ -75,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
@@ -102,6 +109,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 +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.notification_rule_arn = self._setup_notifications(notifications)
  self._checkpoint_s3_uri = None
 
  def _is_nova_model_for_telemetry(self) -> bool:
@@ -433,6 +446,98 @@ 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: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:123456789012:my-topic'}"
+ )
+
+ rule_arn = enable_notifications(
+ sns_topic_arn=sns_topic_arn,
+ 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"),
+ )
+
+ logger.debug("Notification rule ARN: %s", rule_arn)
+ return rule_arn
+
+ def delete_notification_rule(
+ 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.
+ """
+ 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,
+ )
+
+ 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.
+ """
+ return list_notification_rules(
+ sagemaker_session=TrainDefaults.get_sagemaker_session(sagemaker_session=self.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``).
 

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_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
@@ -75,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
@@ -102,6 +109,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 +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.notification_rule_arn = self._setup_notifications(notifications)
  self._checkpoint_s3_uri = None
 
  def _is_nova_model_for_telemetry(self) -> bool:
@@ -433,6 +446,98 @@ 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: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:123456789012:my-topic'}"
+ )
+
+ rule_arn = enable_notifications(
+ sns_topic_arn=sns_topic_arn,
+ 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"),
+ )
+
+ logger.debug("Notification rule ARN: %s", rule_arn)
+ return rule_arn
+
+ def delete_notification_rule(
+ 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.
+ """
+ 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,
+ )
+
+ 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.
+ """
+ return list_notification_rules(
+ sagemaker_session=TrainDefaults.get_sagemaker_session(sagemaker_session=self.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``).