Fix CloudWatchLogStorage with sparse logs (#2501)

When there are time gaps in logs, CloudWatch API may return less events
than requested or even no events at all. We should poll logs until
either enough events are accumulated or we reach the end, which is
indicated by the same token:

https://docs.aws.amazon.com/AmazonCloudWatchLogs/latest/APIReference/API_GetLogEvents.html

> Partially full or empty pages don't necessarily mean that pagination
> is finished. As long as the nextBackwardToken or nextForwardToken
> returned is NOT equal to the nextToken that you passed into the API
> call, there might be more log events available.

Happens with both startFromHead=false and startFromHead=true, but way
more often when startFromHead=false.

boto/boto3#3718 (comment)

> The reason behind this behavior is that get_log_events returns
> a paginated response, with pages being split by time. The responses
> you're getting with no events are periods of time where no events
> happened, and you need to keep paginating using the nextBackwardToken

Fixes: #2500

Author: un-def

src/dstack/_internal/server/services/logs/aws.py

@@ -86,28 +86,22 @@ def poll_logs(self, project: ProjectModel, request: PollLogsRequest) -> JobSubmi
                     raise
                 logger.debug("Stream %s not found, returning dummy response", stream)
                 cw_events = []
-        cw_events_iter: Iterator[_CloudWatchLogEvent]
-        if request.descending:
-            # Regardless of the startFromHead value log events are arranged in chronological order,
-            # from earliest to latest.
-            cw_events_iter = reversed(cw_events)
-        else:
-            cw_events_iter = iter(cw_events)
         logs = [
             LogEvent(
                 timestamp=unix_time_ms_to_datetime(cw_event["timestamp"]),
                 log_source=LogEventSource.STDOUT,
                 message=cw_event["message"],
             )
-            for cw_event in cw_events_iter
+            for cw_event in cw_events
         ]
         return JobSubmissionLogs(logs=logs)
 
     def _get_log_events(self, stream: str, request: PollLogsRequest) -> List[_CloudWatchLogEvent]:
+        limit = request.limit
         parameters = {
             "logGroupName": self._group,
             "logStreamName": stream,
-            "limit": request.limit,
+            "limit": limit,
         }
         start_from_head = not request.descending
         parameters["startFromHead"] = start_from_head
@@ -119,25 +113,32 @@ def _get_log_events(self, stream: str, request: PollLogsRequest) -> List[_CloudW
             # No need to substract one millisecond in this case, though, seems that endTime is
             # exclusive, that is, time interval boundaries are [startTime, entTime)
             parameters["endTime"] = datetime_to_unix_time_ms(request.end_time)
-        response = self._client.get_log_events(**parameters)
-        events: List[_CloudWatchLogEvent] = response["events"]
-        if start_from_head or events:
-            return events
-        # Workaround for https://github.com/boto/boto3/issues/3718
-        # Required only when startFromHead = false (the default value).
-        next_token: str = response["nextBackwardToken"]
+        # "Partially full or empty pages don't necessarily mean that pagination is finished.
+        # As long as the nextBackwardToken or nextForwardToken returned is NOT equal to the
+        # nextToken that you passed into the API call, there might be more log events available."
+        events: List[_CloudWatchLogEvent] = []
+        next_token: Optional[str] = None
+        next_token_key = "nextForwardToken" if start_from_head else "nextBackwardToken"
         # Limit max tries to avoid a possible infinite loop if the API is misbehaving
         tries_left = 10
         while tries_left:
-            parameters["nextToken"] = next_token
+            if next_token is not None:
+                parameters["nextToken"] = next_token
             response = self._client.get_log_events(**parameters)
-            events = response["events"]
-            if events or response["nextBackwardToken"] == next_token:
+            if start_from_head:
+                events.extend(response["events"])
+            else:
+                # Regardless of the startFromHead value log events are arranged in
+                # chronological order, from earliest to latest.
+                events.extend(reversed(response["events"]))
+            if len(events) >= limit:
+                return events[:limit]
+            if response[next_token_key] == next_token:
                 return events
-            next_token = response["nextBackwardToken"]
+            next_token = response[next_token_key]
             tries_left -= 1
-        logger.warning("too many empty responses from stream %s, returning dummy response", stream)
-        return []
+        logger.warning("too many requests to stream %s, returning partial response", stream)
+        return events
 
     def write_logs(
         self,

src/tests/_internal/server/services/test_logs.py

@@ -72,7 +72,7 @@ def mock_client(self, monkeypatch: pytest.MonkeyPatch) -> Mock:
         mock.get_log_events.return_value = {
             "events": [],
             "nextBackwardToken": "bwd",
-            "nextFormartToken": "fwd",
+            "nextForwardToken": "fwd",
         }
         return mock
 
@@ -183,6 +183,7 @@ async def test_poll_logs_non_empty_response(
             {"timestamp": 1696586513234, "message": "SGVsbG8="},
             {"timestamp": 1696586513235, "message": "V29ybGQ="},
         ]
+        poll_logs_request.limit = 2
         job_submission_logs = log_storage.poll_logs(project, poll_logs_request)
 
         assert job_submission_logs.logs == [
@@ -199,20 +200,21 @@ async def test_poll_logs_non_empty_response(
         ]
 
     @pytest.mark.asyncio
+    @pytest.mark.parametrize("descending", [False, True])
     async def test_poll_logs_empty_response(
         self,
         project: ProjectModel,
         log_storage: CloudWatchLogStorage,
         mock_client: Mock,
         poll_logs_request: PollLogsRequest,
+        descending: bool,
     ):
-        # Check that we don't use the workaround when descending=False -> startFromHead=True
-        # https://github.com/dstackai/dstack/issues/1647
         mock_client.get_log_events.return_value["events"] = []
+        poll_logs_request.descending = descending
         job_submission_logs = log_storage.poll_logs(project, poll_logs_request)
 
         assert job_submission_logs.logs == []
-        mock_client.get_log_events.assert_called_once()
+        assert mock_client.get_log_events.call_count == 2
 
     @pytest.mark.asyncio
     async def test_poll_logs_descending_non_empty_response_on_first_call(
@@ -227,6 +229,7 @@ async def test_poll_logs_descending_non_empty_response_on_first_call(
             {"timestamp": 1696586513235, "message": "V29ybGQ="},
         ]
         poll_logs_request.descending = True
+        poll_logs_request.limit = 2
         job_submission_logs = log_storage.poll_logs(project, poll_logs_request)
 
         assert job_submission_logs.logs == [
@@ -243,16 +246,18 @@ async def test_poll_logs_descending_non_empty_response_on_first_call(
         ]
 
     @pytest.mark.asyncio
-    async def test_poll_logs_descending_two_first_calls_return_empty_response(
+    async def test_poll_logs_descending_some_responses_are_empty(
         self,
         project: ProjectModel,
         log_storage: CloudWatchLogStorage,
         mock_client: Mock,
         poll_logs_request: PollLogsRequest,
     ):
         # The first two calls return empty event lists, though the token is not the same, meaning
-        # there are more events.
-        # https://github.com/dstackai/dstack/issues/1647
+        # there are more events, see: https://github.com/dstackai/dstack/issues/1647
+        # As the third call returns less events than requested (2 < 3), we continue to poll until
+        # accumulate enough events (2 + 2) and return exactly the requested number of events (3),
+        # see: https://github.com/dstackai/dstack/issues/2500
         mock_client.get_log_events.side_effect = [
             {
                 "events": [],
@@ -272,8 +277,22 @@ async def test_poll_logs_descending_two_first_calls_return_empty_response(
                 "nextBackwardToken": "bwd3",
                 "nextForwardToken": "fwd",
             },
+            {
+                "events": [],
+                "nextBackwardToken": "bwd4",
+                "nextForwardToken": "fwd",
+            },
+            {
+                "events": [
+                    {"timestamp": 1696586513232, "message": "aW5pdCAx"},
+                    {"timestamp": 1696586513233, "message": "aW5pdCAy"},
+                ],
+                "nextBackwardToken": "bwd5",
+                "nextForwardToken": "fwd",
+            },
         ]
         poll_logs_request.descending = True
+        poll_logs_request.limit = 3
         job_submission_logs = log_storage.poll_logs(project, poll_logs_request)
 
         assert job_submission_logs.logs == [
@@ -287,8 +306,13 @@ async def test_poll_logs_descending_two_first_calls_return_empty_response(
                 log_source=LogEventSource.STDOUT,
                 message="SGVsbG8=",
             ),
+            LogEvent(
+                timestamp=datetime(2023, 10, 6, 10, 1, 53, 233000, tzinfo=timezone.utc),
+                log_source=LogEventSource.STDOUT,
+                message="aW5pdCAy",
+            ),
         ]
-        assert mock_client.get_log_events.call_count == 3
+        assert mock_client.get_log_events.call_count == 5
 
     @pytest.mark.asyncio
     async def test_poll_logs_descending_empty_response_with_same_token(
@@ -352,7 +376,7 @@ def _response_producer(*args, **kwargs):
         job_submission_logs = log_storage.poll_logs(project, poll_logs_request)
 
         assert job_submission_logs.logs == []
-        assert mock_client.get_log_events.call_count == 11  # initial call + 10 tries
+        assert mock_client.get_log_events.call_count == 10
 
     @pytest.mark.asyncio
     async def test_poll_logs_request_params_asc_no_diag_no_dates(
@@ -366,11 +390,13 @@ async def test_poll_logs_request_params_asc_no_diag_no_dates(
         poll_logs_request.limit = 5
         poll_logs_request.diagnose = False
         log_storage.poll_logs(project, poll_logs_request)
-        mock_client.get_log_events.assert_called_once_with(
+        assert mock_client.get_log_events.call_count == 2
+        mock_client.get_log_events.assert_called_with(
             logGroupName="test-group",
             logStreamName="test-proj/test-run/1b0e1b45-2f8c-4ab6-8010-a0d1a3e44e0e/job",
             limit=5,
             startFromHead=True,
+            nextToken="fwd",
         )
 
     @pytest.mark.asyncio
@@ -394,13 +420,15 @@ async def test_poll_logs_request_params_desc_diag_with_dates(
         poll_logs_request.limit = 10
         poll_logs_request.diagnose = True
         log_storage.poll_logs(project, poll_logs_request)
-        mock_client.get_log_events.assert_called_once_with(
+        assert mock_client.get_log_events.call_count == 2
+        mock_client.get_log_events.assert_called_with(
             logGroupName="test-group",
             logStreamName="test-proj/test-run/1b0e1b45-2f8c-4ab6-8010-a0d1a3e44e0e/runner",
             limit=10,
             startFromHead=False,
             startTime=1696586513235,
             endTime=1696672913234,
+            nextToken="bwd",
         )
 
     @pytest.mark.asyncio