diff --git a/sagemaker-train/src/sagemaker/train/base_trainer.py b/sagemaker-train/src/sagemaker/train/base_trainer.py index ff3bdcd751..009d1c2cb6 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_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``). 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..c7e65c2792 --- /dev/null +++ b/sagemaker-train/src/sagemaker/train/common_utils/notifications.py @@ -0,0 +1,337 @@ +"""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 +import re +from typing import Dict, List, Optional +from botocore.exceptions import ClientError + +logger = logging.getLogger(__name__) + +# Rule name prefix used to identify SDK-created rules +_RULE_NAME_PREFIX = "sm-pysdk-job-notif" +_DEFAULT_EVENTS = ["Completed", "Failed", "Stopped"] +_VALID_EVENTS = {"Completed", "Failed", "Stopped", "InProgress"} + + +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". + """ + 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]: + """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 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 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'." + ) + + 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, normalized_events, job_name_prefix) + 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 + input_template = ( + '{"Job": "",' + ' "Status": "",' + ' "Time": "