diff --git a/docs/model_customization/index.rst b/docs/model_customization/index.rst index 80363939c8..018586d224 100644 --- a/docs/model_customization/index.rst +++ b/docs/model_customization/index.rst @@ -23,3 +23,4 @@ Key Benefits open_weight_model_customization nova evaluation + notifications_setup diff --git a/docs/model_customization/notifications_setup.rst b/docs/model_customization/notifications_setup.rst new file mode 100644 index 0000000000..687552b5b4 --- /dev/null +++ b/docs/model_customization/notifications_setup.rst @@ -0,0 +1,250 @@ +Setting Up Training Job Notifications +======================================== + +Get notified when your training jobs complete, fail, or stop via SNS email/SMS +alerts. This guide walks through creating the prerequisite SNS topic and +configuring your trainer to send notifications. + +Architecture +------------- + +.. code-block:: text + + SageMaker Training Job ──► EventBridge Rule ──► SNS Topic ──► Email/SMS/Slack + +The SDK creates an EventBridge rule that listens for training job status changes +and routes them to your SNS topic. You provide the topic; the SDK handles the +wiring. + +Prerequisites +-------------- + +You need: + +1. An SNS topic with a policy allowing EventBridge to publish to it +2. A subscription on that topic (email, SMS, Slack, etc.) +3. IAM permissions for EventBridge rule management (see below) + +Step 1: Create an SNS Topic +~~~~~~~~~~~~~~~~~~~~~~~~~~~~~ + +**AWS Console** + +1. Go to **Amazon SNS** → **Topics** → **Create topic** +2. Choose **Standard** type +3. Name it (e.g., ``my-training-alerts``) +4. Click **Create topic** +5. Note the **Topic ARN** (e.g., ``arn:aws:sns:us-east-1:123456789012:my-training-alerts``) + +**AWS CLI** + +.. code-block:: bash + + aws sns create-topic --name my-training-alerts + +Step 2: Allow EventBridge to Publish +~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~ + +The topic needs a resource policy granting EventBridge publish access. + +**AWS Console** + +1. Go to **Amazon SNS** → **Topics** → open your topic → **Access policy** tab → **Edit** +2. Add this statement to the policy's ``Statement`` array (replace the ARN and + account ID with your own SNS topic ARN and AWS account ID): + +.. code-block:: json + + { + "Sid": "AllowEventBridgePublish", + "Effect": "Allow", + "Principal": {"Service": "events.amazonaws.com"}, + "Action": "SNS:Publish", + "Resource": "arn:aws:sns:us-east-1:123456789012:my-training-alerts", + "Condition": { + "StringEquals": {"AWS:SourceAccount": "123456789012"} + } + } + +.. note:: + + Replace ``arn:aws:sns:us-east-1:123456789012:my-training-alerts`` with your + actual SNS topic ARN, and ``123456789012`` with your AWS account ID. These + must match so that only EventBridge in your account can publish to your topic. + +**AWS CLI** + +.. code-block:: bash + + TOPIC_ARN="arn:aws:sns:us-east-1:123456789012:my-training-alerts" + ACCOUNT_ID="123456789012" + + aws sns set-topic-attributes \ + --topic-arn $TOPIC_ARN \ + --attribute-name Policy \ + --attribute-value '{ + "Version": "2008-10-17", + "Statement": [{ + "Sid": "AllowEventBridgePublish", + "Effect": "Allow", + "Principal": {"Service": "events.amazonaws.com"}, + "Action": "SNS:Publish", + "Resource": "'$TOPIC_ARN'", + "Condition": {"StringEquals": {"AWS:SourceAccount": "'$ACCOUNT_ID'"}} + }] + }' + +Step 3: Subscribe to the Topic +~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~ + +**Email** + +.. code-block:: bash + + aws sns subscribe \ + --topic-arn arn:aws:sns:us-east-1:123456789012:my-training-alerts \ + --protocol email \ + --notification-endpoint you@example.com + +Check your inbox and confirm the subscription. + +**SMS** + +.. code-block:: bash + + aws sns subscribe \ + --topic-arn arn:aws:sns:us-east-1:123456789012:my-training-alerts \ + --protocol sms \ + --notification-endpoint +15551234567 + +Step 4: Use with the SDK +~~~~~~~~~~~~~~~~~~~~~~~~~~ + +Pass the topic ARN in the ``notifications`` config when constructing your trainer: + +.. code-block:: python + + from sagemaker.train import SFTTrainer + from sagemaker.train.common import TrainingType + from sagemaker.core.training.configs import TrainingJobCompute + + trainer = SFTTrainer( + model="amazon.nova-2-lite-v1", + training_type=TrainingType.LORA, + training_dataset="s3://my-bucket/data/train.jsonl", + compute=TrainingJobCompute(instance_type="ml.p4d.24xlarge"), + notifications={ + "sns_topic_arn": "arn:aws:sns:us-east-1:123456789012:my-training-alerts", + }, + ) + + trainer.train() + +Configuration Options +~~~~~~~~~~~~~~~~~~~~~~ + +The ``notifications`` dict supports: + +.. list-table:: + :header-rows: 1 + :widths: 25 10 65 + + * - Key + - Required + - Description + * - ``sns_topic_arn`` + - Yes + - ARN of your SNS topic + * - ``events`` + - No + - List of statuses to notify on. Default: ``["Completed", "Failed", "Stopped"]``. + Valid values: ``Completed``, ``Failed``, ``Stopped``, ``InProgress``. + * - ``event_bus_arn`` + - No + - Custom EventBridge bus ARN. Defaults to the account's default event bus. + * - ``job_name_prefix`` + - No + - Only notify for jobs whose name starts with this prefix. + +Example: notify only on failures for jobs matching a prefix: + +.. code-block:: python + + notifications={ + "sns_topic_arn": "arn:aws:sns:us-east-1:123456789012:my-training-alerts", + "events": ["Failed"], + "job_name_prefix": "prod-sft-", + } + +Managing Notification Rules +----------------------------- + +List active rules: + +.. code-block:: python + + rules = trainer.list_notification_rules() + for rule in rules: + print(f"{rule['name']} ({rule['state']})") + +Delete a rule: + +.. code-block:: python + + trainer.delete_notification_rule(rule_arn="arn:aws:events:us-east-1:123456789012:rule/sm-pysdk-job-notif-abc123") + +Required IAM Permissions +-------------------------- + +The caller (your IAM role or user) needs these permissions. Replace the SNS +resource ARN with your actual topic ARN: + +.. code-block:: json + + { + "Version": "2012-10-17", + "Statement": [ + { + "Effect": "Allow", + "Action": [ + "events:PutRule", + "events:PutTargets", + "events:ListRules", + "events:ListTargetsByRule", + "events:RemoveTargets", + "events:DeleteRule" + ], + "Resource": "arn:aws:events:*:*:rule/sm-pysdk-job-notif-*" + }, + { + "Effect": "Allow", + "Action": "sns:GetTopicAttributes", + "Resource": "arn:aws:sns:*:*:my-training-alerts" + } + ] + } + +.. note:: + + The ``sns:GetTopicAttributes`` resource must match the SNS topic you created + in Step 1. You can use a wildcard (``arn:aws:sns:*:*:*``) for broader access + or scope it to your specific topic ARN for least privilege. + +Troubleshooting +----------------- + +**PermissionError: Missing permissions to manage EventBridge rules** + +Your IAM identity needs ``events:PutRule`` and ``events:PutTargets``. Ask your +admin to attach the policy above. + +**ValueError: SNS topic not found** + +Verify the topic ARN is correct and exists in the same region as your +SageMaker session. Ensure you have ``sns:GetTopicAttributes`` permission. + +**Not receiving notifications** + +1. Confirm the SNS subscription is in ``Confirmed`` state (check in the Console) +2. Verify the topic policy allows EventBridge to publish (Step 2 above) +3. Check that the ``events`` list includes the status you're waiting for diff --git a/sagemaker-train/src/sagemaker/train/base_trainer.py b/sagemaker-train/src/sagemaker/train/base_trainer.py index 2d36016256..10ec7c6a27 100644 --- a/sagemaker-train/src/sagemaker/train/base_trainer.py +++ b/sagemaker-train/src/sagemaker/train/base_trainer.py @@ -627,7 +627,7 @@ def list_notification_rules( event_bus_arn=event_bus_arn, ) - def stream_logs(self, poll: int = 5, start_time: Optional[Any] = None, num_lines: Optional[int] = None) -> None: + def stream_logs(self, poll: int = 5, start_time: Optional[Any] = None, tail_lines: Optional[int] = None) -> None: """Stream CloudWatch logs in real-time (like ``kubectl logs -f``). Continuously polls for new log events and prints them as they arrive. @@ -641,10 +641,13 @@ def stream_logs(self, poll: int = 5, start_time: Optional[Any] = None, num_lines attaching to a job that's already running. If not provided, auto-resolved from the training job's start time (SMTJ) or defaults to now (HyperPod). - num_lines: Optional maximum number of log lines to print. When - specified, streaming stops after this many lines have been - printed. Useful for long jobs where the full log is too verbose. - If not provided, streams all logs until the job completes. + tail_lines: Optional maximum number of most recent log lines to + print. Logs are returned in chronological order; when specified, + only the last ``tail_lines`` entries are shown (similar to + ``kubectl logs --tail``). Useful for quickly checking the latest + output of long-running jobs without scrolling through the full + history. If not provided, streams all logs until the job + completes. Raises: ValueError: If no training job has been run yet. @@ -678,11 +681,11 @@ def stream_logs(self, poll: int = 5, start_time: Optional[Any] = None, num_lines compute = getattr(self, 'compute', None) if isinstance(compute, HyperPodCompute): - self._stream_logs_smhp(training_job, compute, poll, start_time_ms, num_lines=num_lines) + self._stream_logs_smhp(training_job, compute, poll, start_time_ms, tail_lines=tail_lines) else: - self._stream_logs_smtj(training_job, poll, start_time_ms, num_lines=num_lines) + self._stream_logs_smtj(training_job, poll, start_time_ms, tail_lines=tail_lines) - def _stream_logs_smtj(self, training_job, poll: int, start_time_ms=None, num_lines: Optional[int] = None) -> None: + def _stream_logs_smtj(self, training_job, poll: int, start_time_ms=None, tail_lines: Optional[int] = None) -> None: """Stream logs for an SMTJ training job.""" from sagemaker.train.common_utils.log_streamer import ( LogStreamer, @@ -714,9 +717,9 @@ def _get_status() -> str: job = TrainingJob.get(training_job_name=job_name) return job.training_job_status - stream_log_loop(streamer, poll, _get_status, num_lines=num_lines) + stream_log_loop(streamer, poll, _get_status, tail_lines=tail_lines) - def _stream_logs_smhp(self, training_job, compute, poll: int, start_time_ms=None, num_lines: Optional[int] = None) -> None: + def _stream_logs_smhp(self, training_job, compute, poll: int, start_time_ms=None, tail_lines: Optional[int] = None) -> None: """Stream logs for a HyperPod job using filter_log_events polling.""" if isinstance(training_job, str): @@ -774,8 +777,8 @@ def _stream_logs_smhp(self, training_job, compute, poll: int, start_time_ms=None if message: print(f"{_CW_PREFIX}{message}") lines_printed += 1 - if num_lines and lines_printed >= num_lines: - logger.info(f"Reached num_lines limit ({num_lines}). Stopping log stream.") + if tail_lines and lines_printed >= tail_lines: + logger.info(f"Reached tail_lines limit ({tail_lines}). Stopping log stream.") return ts = event.get("timestamp", 0) if ts > last_timestamp: @@ -812,8 +815,14 @@ def _stream_logs_smhp(self, training_job, compute, poll: int, start_time_ms=None logger.info("Log streaming stopped by user.") return - def _validate_instance_count(self, instance_count, sagemaker_session): - """Validate instance/node count against allowed values from SMHP recipe.""" + def _validate_instance_count(self, instance_count, sagemaker_session, compute): + """Validate instance/node count against allowed values from SMHP recipe. + + For HyperPod compute, raises ValueError on mismatch since the recipe's + node count constraints are hard requirements for distributed training. + For SMTJ compute, logs a warning instead since SMTJ may support counts + not listed in the SMHP recipe depending on the model. + """ smhp_replicas_enum = _get_smhp_replicas_enum( model_name=self._model_name, customization_technique=self._customization_technique, @@ -821,10 +830,18 @@ def _validate_instance_count(self, instance_count, sagemaker_session): sagemaker_session=sagemaker_session, ) if smhp_replicas_enum and instance_count not in smhp_replicas_enum: - raise ValueError( - f"Node/Instance count '{instance_count}' is not supported. " - f"Allowed values: {sorted(smhp_replicas_enum)}." - ) + if isinstance(compute, HyperPodCompute): + raise ValueError( + f"Node/Instance count '{instance_count}' is not supported. " + f"Allowed values: {sorted(smhp_replicas_enum)}." + ) + else: + logger.warning( + f"Instance count '{instance_count}' is not in the recommended values " + f"{sorted(smhp_replicas_enum)} from the model recipe. " + f"This may or may not work depending on the model. " + f"Proceeding anyway for SMTJ compute." + ) return smhp_replicas_enum def _validate_instance_type(self, instance_type, sagemaker_session): @@ -954,7 +971,7 @@ def _channel_mount_path(dataset_uri, channel_name): ) # Validate instance count against allowed values from SMHP recipe. - smhp_replicas_enum = self._validate_instance_count(compute.instance_count, sagemaker_session) + smhp_replicas_enum = self._validate_instance_count(compute.instance_count, sagemaker_session, compute) if smhp_replicas_enum: override_spec.setdefault("replicas", {})["enum"] = smhp_replicas_enum @@ -1436,7 +1453,7 @@ def _train_hyperpod(self, training_dataset=None, validation_dataset=None, job_base_name = self.base_job_name or f"{self._model_name}-{self._customization_technique}" # Validate node_count against allowed values from SMHP recipe - self._validate_instance_count(compute.node_count, sagemaker_session) + self._validate_instance_count(compute.node_count, sagemaker_session, compute) # Resolve and validate the recipe (3-level merge: base → user recipe → overrides) try: diff --git a/sagemaker-train/src/sagemaker/train/common_utils/log_streamer.py b/sagemaker-train/src/sagemaker/train/common_utils/log_streamer.py index 3538c95200..570a3331e3 100644 --- a/sagemaker-train/src/sagemaker/train/common_utils/log_streamer.py +++ b/sagemaker-train/src/sagemaker/train/common_utils/log_streamer.py @@ -227,7 +227,7 @@ def stream_log_loop( streamer: LogStreamer, poll: int, status_fn: Callable[[], str], - num_lines: Optional[int] = None, + tail_lines: Optional[int] = None, ) -> None: """Run the standard log streaming loop. @@ -238,19 +238,20 @@ def stream_log_loop( :param streamer: A configured LogStreamer instance. :param poll: Seconds between polls. :param status_fn: Callable that returns the current job status string. - :param num_lines: Optional maximum number of log lines to print. - When specified, streaming stops after this many lines. + :param tail_lines: Optional maximum number of most recent log lines to + print. When specified, streaming stops after this many lines have + been displayed. """ _CW_PREFIX = "[CloudWatch] " lines_printed = 0 def _print_event(ts_ms: int, message: str) -> bool: - """Print a log event. Returns True if num_lines limit reached.""" + """Print a log event. Returns True if tail_lines limit reached.""" nonlocal lines_printed print(f"{_CW_PREFIX}[{_format_timestamp(ts_ms)}] {message}") lines_printed += 1 - if num_lines and lines_printed >= num_lines: - logger.info("Reached num_lines limit (%d). Stopping log stream.", num_lines) + if tail_lines and lines_printed >= tail_lines: + logger.info("Reached tail_lines limit (%d). Stopping log stream.", tail_lines) return True return False diff --git a/sagemaker-train/src/sagemaker/train/common_utils/notifications.py b/sagemaker-train/src/sagemaker/train/common_utils/notifications.py index eced509685..1f3f712559 100644 --- a/sagemaker-train/src/sagemaker/train/common_utils/notifications.py +++ b/sagemaker-train/src/sagemaker/train/common_utils/notifications.py @@ -3,6 +3,45 @@ Manages EventBridge rules that route SageMaker Training Job status change events to user-provided SNS topics. Supports SMTJ (serverless and serverful) training jobs only. + +Prerequisites: + Before using notifications, you need an SNS topic with a policy that allows + EventBridge to publish to it. You can create one via the AWS Console, CLI, + or CloudFormation/CDK: + + **AWS CLI example**:: + + # Create the topic + aws sns create-topic --name my-training-alerts + + # Allow EventBridge to publish to it + aws sns set-topic-attributes \\ + --topic-arn arn:aws:sns:::my-training-alerts \\ + --attribute-name Policy \\ + --attribute-value '{ + "Version": "2008-10-17", + "Statement": [{ + "Sid": "AllowEventBridgePublish", + "Effect": "Allow", + "Principal": {"Service": "events.amazonaws.com"}, + "Action": "SNS:Publish", + "Resource": "arn:aws:sns:::my-training-alerts", + "Condition": {"StringEquals": {"AWS:SourceAccount": ""}} + }] + }' + + # Subscribe your email + aws sns subscribe \\ + --topic-arn arn:aws:sns:::my-training-alerts \\ + --protocol email \\ + --notification-endpoint you@example.com + + Then pass the topic ARN when constructing a trainer:: + + trainer = SFTTrainer( + ..., + notifications={"sns_topic_arn": "arn:aws:sns:::my-training-alerts"}, + ) """ from __future__ import absolute_import @@ -150,85 +189,6 @@ def _validate_sns_topic(sns_client, topic_arn: str) -> None: raise -def _build_topic_policy(topic_arn: str, account_id: str) -> str: - """Build the SNS topic access policy for EventBridge notifications. - - Includes the default owner statement (standard SNS actions scoped to the - topic's own account via ``AWS:SourceAccount``) plus an explicit statement - granting the EventBridge service principal ``SNS:Publish``. - - Args: - topic_arn: The ARN of the SNS topic. - account_id: The AWS account ID that owns the topic. - - Returns: - JSON string of the topic access policy. - """ - policy = { - "Version": "2008-10-17", - "Id": "SageMakerNotificationsTopicPolicy", - "Statement": [ - { - "Sid": "SNSTopicAdministration", - "Effect": "Allow", - "Principal": {"AWS": "*"}, - "Action": [ - "SNS:Publish", - "SNS:RemovePermission", - "SNS:SetTopicAttributes", - "SNS:DeleteTopic", - "SNS:ListSubscriptionsByTopic", - "SNS:GetTopicAttributes", - "SNS:AddPermission", - "SNS:Subscribe", - ], - "Resource": topic_arn, - "Condition": {"StringEquals": {"AWS:SourceAccount": account_id}}, - }, - { - "Sid": "AllowEventBridgePublish", - "Effect": "Allow", - "Principal": {"Service": "events.amazonaws.com"}, - "Action": "SNS:Publish", - "Resource": topic_arn, - "Condition": {"StringEquals": {"AWS:SourceAccount": account_id}}, - }, - ], - } - return json.dumps(policy) - - -def create_notification_topic(topic_name: str, sagemaker_session) -> str: - """Create an SNS topic configured for EventBridge notifications. - - Creates a new SNS topic and sets an access policy that allows EventBridge - in the same account to publish to it. The returned ARN can be passed - straight to :func:`enable_notifications`. - - Args: - topic_name: Name of the SNS topic to create. - sagemaker_session: SageMaker session (provides boto_session). - - Returns: - The ARN of the created SNS topic. - """ - region_name = sagemaker_session.boto_session.region_name - sns_client = sagemaker_session.boto_session.client("sns", region_name=region_name) - - response = sns_client.create_topic(Name=topic_name) - topic_arn = response["TopicArn"] - # topic ARN: arn:aws:sns::: - account_id = topic_arn.split(":")[4] - - policy = _build_topic_policy(topic_arn, account_id) - sns_client.set_topic_attributes( - TopicArn=topic_arn, AttributeName="Policy", AttributeValue=policy - ) - - logger.info(f"Created SNS topic with EventBridge publish policy: {topic_arn}") - return topic_arn - - def enable_notifications( sns_topic_arn: str, sagemaker_session,