Repository: skypilot-org/skypilot
Pull request #5804: [API server] Add heartbeat to streaming response
Author: login: aylei
Created: 2025-05-29T08:45:21Z | Merged: 2025-05-29T23:23:16Z

<!-- Describe the changes in this PR -->

close https://github.com/skypilot-org/skypilot/issues/5750

Explanation: https://github.com/skypilot-org/skypilot/issues/5750#issuecomment-2918711119

<!-- Describe the tests ran -->
<!-- Unit tests (tests/test_*.py) are part of GitHub CI; below are tests that launch on the cloud. -->

Tested (run the relevant ones):

- [x] Code formatting: install pre-commit (auto-check on commit) or `bash format.sh`
- [x] Manually tested on a kubernetes cluster with cloudflare loadbalancer, the long provision request no longer get connection termination issue
- [ ] All smoke tests: `/smoke-test` (CI) or `pytest tests/test_smoke.py` (local)
- [ ] Relevant individual tests: `/smoke-test -k test_name` (CI) or `pytest tests/test_smoke.py::test_name` (local)
- [ ] Backward compatibility: `/quicktest-core` (CI) or `pytest tests/smoke_tests/test_backward_compat.py` (local)

<!-- CI commands (/-prefixed) can only be triggered by repo members -->



Diff:
diff --git a/sky/server/stream_utils.py b/sky/server/stream_utils.py
index 89c45617631..cf686cc61c7 100644
--- a/sky/server/stream_utils.py
+++ b/sky/server/stream_utils.py
@@ -15,6 +15,8 @@
 
 logger = sky_logging.init_logger(__name__)
 
+_HEARTBEAT_INTERVAL = 30
+
 
 async def _yield_log_file_with_payloads_skipped(
         log_file) -> AsyncGenerator[str, None]:
@@ -90,6 +92,8 @@ async def log_streamer(request_id: Optional[str],
             for line_str in lines:
                 yield line_str
 
+        last_heartbeat_time = asyncio.get_event_loop().time()
+
         while True:
             # Sleep 0 to yield control to allow other coroutines to run,
             # while keeps the loop tight to make log stream responsive.
@@ -106,15 +110,32 @@ async def log_streamer(request_id: Optional[str],
                         break
                 if not follow:
                     break
+
+                current_time = asyncio.get_event_loop().time()
+                if current_time - last_heartbeat_time >= _HEARTBEAT_INTERVAL:
+                    # Currently just used to keep the connection busy, refer to
+                    # https://github.com/skypilot-org/skypilot/issues/5750 for
+                    # more details.
+                    yield message_utils.encode_payload(
+                        rich_utils.Control.HEARTBEAT.encode(''))
+                    last_heartbeat_time = current_time
+
                 # Sleep shortly to avoid storming the DB and CPU, this has
                 # little impact on the responsivness here since we are waiting
                 # for a new line to come in.
                 await asyncio.sleep(0.1)
                 continue
+
+            # Refresh the heartbeat time, this is a trivial optimization for
+            # performance but it helps avoid unnecessary heartbeat strings
+            # being printed when the client runs in an old version.
+            last_heartbeat_time = asyncio.get_event_loop().time()
             line_str = line.decode('utf-8')
             if plain_logs:
                 is_payload, line_str = message_utils.decode_payload(
                     line_str, raise_for_mismatch=False)
+                # TODO(aylei): implement heartbeat mechanism for plain logs,
+                # sending invisible characters might be okay.
                 if is_payload:
                     continue
             yield line_str
diff --git a/sky/utils/rich_utils.py b/sky/utils/rich_utils.py
index 53c759e1778..e804228ace6 100644
--- a/sky/utils/rich_utils.py
+++ b/sky/utils/rich_utils.py
@@ -57,6 +57,7 @@ class Control(enum.Enum):
     STOP = 'rich_stop'
     EXIT = 'rich_exit'
     UPDATE = 'rich_update'
+    HEARTBEAT = 'heartbeat'
 
     def encode(self, msg: str) -> str:
         return f'<{self.value}>{msg}</{self.value}>'
@@ -385,6 +386,10 @@ def decode_rich_status(
                         decoding_status.__exit__(None, None, None)
                     elif control == Control.START:
                         decoding_status.start()
+                    elif control == Control.HEARTBEAT:
+                        # Heartbeat is not displayed to the user, so we do not
+                        # need to update the status.
+                        pass
     finally:
         if decoding_status is not None:
             decoding_status.__exit__(None, None, None)
diff --git a/tests/test_api.py b/tests/test_api.py
index 26e1bf0b584..d7d57f7665a 100644
--- a/tests/test_api.py
+++ b/tests/test_api.py
@@ -1,5 +1,15 @@
+import asyncio
+import pathlib
+import tempfile
+import time
+from unittest import mock
+
 import sky
 from sky.clouds.cloud import Cloud
+from sky.server import stream_utils
+from sky.server.requests import requests as requests_lib
+from sky.utils import message_utils
+from sky.utils import rich_utils
 
 
 def test_sky_launch(enable_all_clouds):
@@ -20,3 +30,196 @@ def dryrun_task_with_cloud(cloud: Cloud):
     dryrun_task_with_cloud(sky.K8s())
 
     dryrun_task_with_cloud(sky.Kubernetes())
+
+
+def test_api_stream_heartbeat(monkeypatch):
+    """Test that stream_and_get sends heartbeats when logs are not emitted."""
+
+    # Patch the heartbeat interval to 0.1 seconds for faster testing
+    monkeypatch.setattr(stream_utils, '_HEARTBEAT_INTERVAL', 0.1)
+
+    # Track heartbeats generated
+    heartbeats_generated = []
+
+    # Test the actual log_streamer function with mocked file I/O
+    async def test_heartbeat_generation():
+        """Test that log_streamer generates heartbeats when no logs are written."""
+
+        with tempfile.NamedTemporaryFile(mode='w', delete=False) as f:
+            f.write("Initial log line\n")
+            f.flush()
+            temp_log_path = f.name
+
+        class MockRequest:
+
+            def __init__(self):
+                self.request_id = "test-request"
+                self.status = requests_lib.RequestStatus.RUNNING
+                self.name = "test_heartbeat"
+                self.schedule_type = requests_lib.ScheduleType.LONG
+                self.status_msg = None
+
+        def mock_get_request(request_id):
+            return MockRequest()
+
+        monkeypatch.setattr('sky.server.requests.requests.get_request',
+                            mock_get_request)
+
+        log_path = pathlib.Path(temp_log_path)
+
+        try:
+            streamed_items = []
+            start_time = asyncio.get_event_loop().time()
+
+            # Create the async generator explicitly so we can close it properly
+            log_stream = stream_utils.log_streamer(request_id="test-request",
+                                                   log_path=log_path,
+                                                   follow=True)
+
+            try:
+                async for item in log_stream:
+                    streamed_items.append(item)
+                    current_time = asyncio.get_event_loop().time()
+
+                    try:
+                        is_payload, decoded = message_utils.decode_payload(
+                            item, raise_for_mismatch=False)
+                        if is_payload:
+                            control, _ = rich_utils.Control.decode(decoded)
+                            if control == rich_utils.Control.HEARTBEAT:
+                                heartbeats_generated.append(current_time)
+                    except Exception:
+                        pass
+
+                    if current_time - start_time > 0.55:
+                        break
+            finally:
+                # Properly close the async generator to avoid pending task errors
+                await log_stream.aclose()
+
+        finally:
+            import os
+            os.unlink(temp_log_path)
+
+        return streamed_items
+
+    # Create and run the async test
+    loop = asyncio.new_event_loop()
+    asyncio.set_event_loop(loop)
+
+    try:
+        streamed_items = loop.run_until_complete(test_heartbeat_generation())
+    finally:
+        loop.close()
+
+    assert len(heartbeats_generated) >= 5, (
+        f"Expected at least 5 heartbeats, got {len(heartbeats_generated)}. "
+        f"Total streamed items: {len(streamed_items)}")
+
+
+def test_heartbeat_not_displayed_to_users(monkeypatch):
+    """Test that heartbeat messages are filtered out from user display."""
+
+    # Patch the heartbeat interval to 0.1 seconds for faster testing
+    monkeypatch.setattr(stream_utils, '_HEARTBEAT_INTERVAL', 0.1)
+
+    # Track what gets displayed to users vs heartbeats
+    user_visible_messages = []
+    heartbeat_messages = []
+
+    async def test_heartbeat_filtering():
+        """Test that heartbeats are filtered from user-visible output."""
+
+        with tempfile.NamedTemporaryFile(mode='w', delete=False) as f:
+            f.write("User log message 1\n")
+            f.flush()
+            temp_log_path = f.name
+
+        class MockRequest:
+
+            def __init__(self):
+                self.request_id = "test-request"
+                self.status = requests_lib.RequestStatus.RUNNING
+                self.name = "test_heartbeat_display"
+                self.schedule_type = requests_lib.ScheduleType.LONG
+                self.status_msg = None
+
+        def mock_get_request(request_id):
+            return MockRequest()
+
+        monkeypatch.setattr('sky.server.requests.requests.get_request',
+                            mock_get_request)
+
+        log_path = pathlib.Path(temp_log_path)
+
+        try:
+            start_time = asyncio.get_event_loop().time()
+
+            # Create the async generator
+            log_stream = stream_utils.log_streamer(request_id="test-request",
+                                                   log_path=log_path,
+                                                   follow=True)
+
+            try:
+                async for item in log_stream:
+                    current_time = asyncio.get_event_loop().time()
+
+                    try:
+                        is_payload, decoded = message_utils.decode_payload(
+                            item, raise_for_mismatch=False)
+                        if is_payload:
+                            control, _ = rich_utils.Control.decode(decoded)
+                            if control == rich_utils.Control.HEARTBEAT:
+                                # Heartbeat should NOT be displayed to users
+                                heartbeat_messages.append(item)
+                            else:
+                                user_visible_messages.append(item)
+                        else:
+                            user_visible_messages.append(item)
+                    except Exception:
+                        # Non-control messages should be displayed
+                        user_visible_messages.append(item)
+
+                    if current_time - start_time > 0.55:
+                        break
+            finally:
+                await log_stream.aclose()
+
+        finally:
+            import os
+            os.unlink(temp_log_path)
+
+        return user_visible_messages, heartbeat_messages
+
+    loop = asyncio.new_event_loop()
+    asyncio.set_event_loop(loop)
+
+    try:
+        user_messages, heartbeats = loop.run_until_complete(
+            test_heartbeat_filtering())
+    finally:
+        loop.close()
+
+    assert len(heartbeats) >= 3, (
+        f"Expected at least 3 heartbeats, got {len(heartbeats)}")
+
+    # Verify user messages contain actual log content, not heartbeats
+    user_content = ''.join(
+        [msg for msg in user_messages if isinstance(msg, str)])
+    assert "User log message 1" in user_content, (
+        "Expected user log content to be visible")
+
+    # Verify no heartbeat control messages leaked into user display
+    for msg in user_messages:
+        if isinstance(msg, str):
+            continue
+        try:
+            is_payload, decoded = message_utils.decode_payload(
+                msg, raise_for_mismatch=False)
+            if is_payload:
+                control, _ = rich_utils.Control.decode(decoded)
+                assert control != rich_utils.Control.HEARTBEAT, (
+                    "Heartbeat message should not be in user-visible output")
+        except Exception:
+            # Non-control messages are fine in user output
+            pass


Reviews:
user: login: aylei
state: COMMENTED
author_association: COLLABORATOR
submitted_at: 2025-05-29T08:46:10Z
user: login: aylei
state: COMMENTED
author_association: COLLABORATOR
submitted_at: 2025-05-29T10:17:36Z
user: login: concretevitamin
state: COMMENTED
author_association: MEMBER
submitted_at: 2025-05-29T13:25:09Z
user: login: aylei
state: COMMENTED
author_association: COLLABORATOR
submitted_at: 2025-05-29T14:21:59Z
user: login: Michaelvll
body: Looks awesome! Thanks for figuring out the solution here @aylei! LGTM
state: APPROVED
author_association: COLLABORATOR
submitted_at: 2025-05-29T18:20:26Z
user: login: Michaelvll
state: COMMENTED
author_association: COLLABORATOR
submitted_at: 2025-05-29T18:21:24Z
user: login: aylei
state: COMMENTED
author_association: COLLABORATOR
submitted_at: 2025-05-29T23:16:18Z

Review comments:
diff_hunk: @@ -106,15 +110,28 @@ async def log_streamer(request_id: Optional[str],
                         break
                 if not follow:
                     break
+
+                current_time = asyncio.get_event_loop().time()
+                if current_time - last_heartbeat_time >= _HEARTBEAT_INTERVAL:
+                    # Currently just used to keep the connection busy, refer to
+                    # https://github.com/skypilot-org/skypilot/issues/5750 for
+                    # more details.
+                    yield message_utils.encode_payload(
+                        rich_utils.Control.HEARTBEAT.encode(''))
+                    last_heartbeat_time = current_time
+
                 # Sleep shortly to avoid storming the DB and CPU, this has
                 # little impact on the responsivness here since we are waiting
                 # for a new line to come in.
                 await asyncio.sleep(0.1)
                 continue
+
             line_str = line.decode('utf-8')
             if plain_logs:
                 is_payload, line_str = message_utils.decode_payload(
                     line_str, raise_for_mismatch=False)
+                # TODO(aylei): implement heartbeat mechanism for plain logs,
+                # sending invisible characters might be okay.
                 if is_payload:
path: sky/server/stream_utils.py
user: login: aylei
body: Not look into the plain_logs part yet, logged as a followup
created_at: 2025-05-29T08:46:06Z
author_association: COLLABORATOR
diff_hunk: @@ -106,15 +110,28 @@ async def log_streamer(request_id: Optional[str],
                         break
                 if not follow:
                     break
+
+                current_time = asyncio.get_event_loop().time()
+                if current_time - last_heartbeat_time >= _HEARTBEAT_INTERVAL:
+                    # Currently just used to keep the connection busy, refer to
+                    # https://github.com/skypilot-org/skypilot/issues/5750 for
+                    # more details.
+                    yield message_utils.encode_payload(
+                        rich_utils.Control.HEARTBEAT.encode(''))
+                    last_heartbeat_time = current_time
path: sky/server/stream_utils.py
user: login: aylei
body: If the server send HEARTBEAT but client runs in an old version that does not handle heartbeat, the following line will be printed **if there is no log for every 30 seconds**:

```
<heartbeat></heartbeat>
```
created_at: 2025-05-29T10:14:41Z
author_association: COLLABORATOR
diff_hunk: @@ -106,15 +110,28 @@ async def log_streamer(request_id: Optional[str],
                         break
                 if not follow:
                     break
+
+                current_time = asyncio.get_event_loop().time()
+                if current_time - last_heartbeat_time >= _HEARTBEAT_INTERVAL:
+                    # Currently just used to keep the connection busy, refer to
+                    # https://github.com/skypilot-org/skypilot/issues/5750 for
+                    # more details.
+                    yield message_utils.encode_payload(
+                        rich_utils.Control.HEARTBEAT.encode(''))
+                    last_heartbeat_time = current_time
path: sky/server/stream_utils.py
user: login: concretevitamin
body: Is this a user visible log? If so, is there a way to not leak this impl detail, or perhaps can change it to a more readable message?
created_at: 2025-05-29T13:25:09Z
author_association: MEMBER
diff_hunk: @@ -106,15 +110,28 @@ async def log_streamer(request_id: Optional[str],
                         break
                 if not follow:
                     break
+
+                current_time = asyncio.get_event_loop().time()
+                if current_time - last_heartbeat_time >= _HEARTBEAT_INTERVAL:
+                    # Currently just used to keep the connection busy, refer to
+                    # https://github.com/skypilot-org/skypilot/issues/5750 for
+                    # more details.
+                    yield message_utils.encode_payload(
+                        rich_utils.Control.HEARTBEAT.encode(''))
+                    last_heartbeat_time = current_time
path: sky/server/stream_utils.py
user: login: aylei
body: I've thought several options when implementing this PR:

1. Make it verbose: something like `<heartbeat>Keep alive message, this is displayed in old versions of client, upgrade your client to 0.9.4 or later to hide this message</heartbeat>`;
2. Bump API version, so old client cannot talk to new server anymore;
3. Add a new request parameter and set this parameter in client to indicates the client can handle heartbeat, only send heartbeat in this case;

The current trade-off assumes:

- heartbeat will not a be sent in common case;
- there would be a version bump introduced by other API change eventually, which eliminate this problem entirely;

Maybe we can add option 1 in this PR, which is trivial work and can make things less confusing by adding a small amount of bandwidth usage, which looks fine.
created_at: 2025-05-29T14:21:59Z
author_association: COLLABORATOR
diff_hunk: @@ -106,15 +110,28 @@ async def log_streamer(request_id: Optional[str],
                         break
                 if not follow:
                     break
+
+                current_time = asyncio.get_event_loop().time()
+                if current_time - last_heartbeat_time >= _HEARTBEAT_INTERVAL:
+                    # Currently just used to keep the connection busy, refer to
+                    # https://github.com/skypilot-org/skypilot/issues/5750 for
+                    # more details.
+                    yield message_utils.encode_payload(
+                        rich_utils.Control.HEARTBEAT.encode(''))
+                    last_heartbeat_time = current_time
path: sky/server/stream_utils.py
user: login: Michaelvll
body: Since this only happens for a client in an old version, I think it should be fine to not have a special handling, and we have a recent version bump anyway.
created_at: 2025-05-29T16:52:27Z
author_association: COLLABORATOR
diff_hunk: @@ -1,5 +1,15 @@
+import asyncio
path: tests/test_api.py
user: login: Michaelvll
body: We may need add a bit test for heartbeat message is not displayed to users
created_at: 2025-05-29T18:21:24Z
author_association: COLLABORATOR
diff_hunk: @@ -1,5 +1,15 @@
+import asyncio
path: tests/test_api.py
user: login: aylei
body: Good catch, added!
created_at: 2025-05-29T23:16:18Z
author_association: COLLABORATOR

PR comments:
user: login: aylei
created_at: 2025-05-29T10:17:47Z
body:  /quicktest-core

```
E 05-29 12:11:38 sdk.py:1601] sky.exceptions.ClusterStatusFetchingError: Failed to query SSH cluster 'sky-jobs-controller-bba7fae4' status: [sky.exceptions.ClusterStatusFetchingError] Failed to query cluster 'sky-jobs-controller-bba7fae4-bba7fae4' status. Cannot connect to SSH Node Pool test-ssh-node-pools. Please check if the SSH Node Pool is up and accessible. To debug, run `sky check ssh` to check the status of the SSH Node Pool.
```

The error looks not relevant to this change, is the SSH cluster down in CI env? @zpoint 
author_association: COLLABORATOR
user: login: zpoint
created_at: 2025-05-30T01:39:25Z
body: > The error looks not relevant to this change, is the SSH cluster down in CI env? @zpoint

Yes, I will fix that today.
author_association: COLLABORATOR