fix: usuability update from feedback

docs/model_customization/index.rst

@@ -23,3 +23,4 @@ Key Benefits
    open_weight_model_customization
    nova
    evaluation
+   notifications_setup

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

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,19 +815,33 @@ 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,
             training_type=self.training_type,
             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: