fix: treat Updating as in progress in core live-logging wait

Session.wait_for_endpoint(live_logging=True) stopped waiting as soon
as the endpoint left Creating, so an endpoint that was Updating was
reported as a failed deployment. Use the same in-progress statuses
as _deploy_done (Creating, Updating), which the sagemaker-serve
live-logging check already does.

Also state that the ResourceIterator repeated-token guard only
catches a NextToken that comes back unchanged, the same check
botocore paginators do.

Author: lucasjia-aws

sagemaker-core/src/sagemaker/core/helper/session_helper.py

@@ -3013,10 +3013,11 @@ def reset(self):
 def _live_logging_deploy_done(
     sagemaker_client, endpoint_name, paginator, paginator_config, poll, not_found_budget=None
 ):
-    """Return the ``DescribeEndpoint`` response once the endpoint leaves ``Creating``.
+    """Return the ``DescribeEndpoint`` response once the endpoint is no longer in progress.
 
     Streams the endpoint's CloudWatch logs on every poll and returns ``None`` while
-    the endpoint is still being created. The log group only exists once a container
+    the endpoint is still ``Creating`` or ``Updating``, the same in-progress statuses
+    as ``_deploy_done``. The log group only exists once a container
     has started, so an endpoint that fails before any instance is provisioned (for
     example on InsufficientInstanceCapacity) never gets one: a missing log group
     must not keep a finished deployment waiting.
@@ -3026,6 +3027,7 @@ def _live_logging_deploy_done(
             consecutive polls a missing endpoint is tolerated. Without it a
             missing endpoint is waited on indefinitely.
     """
+    in_progress_statuses = ["Creating", "Updating"]
     stop = False
     endpoint_status = None
     try:
@@ -3042,7 +3044,7 @@ def _live_logging_deploy_done(
         not_found_budget.reset()
 
     # if endpoint is in an invalid state -> set stop to true, sleep, and flush the logs
-    if endpoint_status != "Creating":
+    if endpoint_status not in in_progress_statuses:
         stop = True
         if endpoint_status == "InService":
             LOGGER.info(

sagemaker-core/src/sagemaker/core/utils/utils.py

@@ -524,7 +524,8 @@ def __next__(self) -> T:
                 # A page can be empty and still carry a NextToken (the service
                 # filters results after paginating), so an empty page only ends
                 # the iteration when there is no further page. A NextToken that
-                # repeats would page forever, so it also ends the iteration.
+                # comes back unchanged would page forever, so it also ends the
+                # iteration (the same consecutive-repeat check botocore paginators do).
                 if not self.next_token or self.next_token == previous_token:
                     raise StopIteration
 

sagemaker-core/tests/unit/helper/test_session_helper.py

@@ -429,14 +429,19 @@ def test_finished_endpoint_without_log_group_returns_desc(self, mock_sleep):
 
             assert _live_logging_deploy_done(client, "my-endpoint", paginator, {}, 5) == desc
 
-    def test_creating_endpoint_without_log_group_keeps_waiting(self):
-        """Creating with no log group yet is still in progress."""
-        client = Mock()
-        client.describe_endpoint.return_value = {"EndpointStatus": "Creating"}
-        paginator = Mock()
-        paginator.paginate.side_effect = self.RESOURCE_NOT_FOUND
-
-        assert _live_logging_deploy_done(client, "my-endpoint", paginator, {}, 5) is None
+    @patch("time.sleep")
+    def test_in_progress_endpoint_keeps_waiting(self, mock_sleep):
+        """Creating or Updating is still in progress, with or without a log group."""
+        for status in ("Creating", "Updating"):
+            for log_error in (self.RESOURCE_NOT_FOUND, None):
+                client = Mock()
+                client.describe_endpoint.return_value = {"EndpointStatus": status}
+                paginator = Mock()
+                paginator.paginate.side_effect = log_error
+                paginator.paginate.return_value = []
+
+                assert _live_logging_deploy_done(client, "my-endpoint", paginator, {}, 5) is None
+        mock_sleep.assert_not_called()
 
     def test_other_log_errors_are_raised(self):
         """Log-fetch errors other than a missing log group still propagate."""