Repository: skypilot-org/skypilot
Pull request #5804 (2025-05-29 08:45:21)

[API server] Add heartbeat to streaming response
<!-- 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 -->


aylei:  /quicktest-core

Patch:
diff --git a/sky/server/stream_utils.py b/sky/server/stream_utils.py
index 89c456176..cf686cc61 100644
--- a/sky/server/stream_utils.py
+++ b/sky/server/stream_utils.py
@@ -15,6 +15,8 @@ from sky.utils import rich_utils
 
 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/common_utils.py b/sky/utils/common_utils.py
index 00d9db4c7..636d86456 100644
--- a/sky/utils/common_utils.py
+++ b/sky/utils/common_utils.py
@@ -324,9 +324,75 @@ def get_pretty_entrypoint_cmd() -> str:
         # Turn '/.../anaconda/envs/py36/bin/sky' into 'sky', but keep other
         # things like 'examples/app.py'.
         argv[0] = basename
+
+    # Redact sensitive environment variable values
+    argv = _redact_env_values(argv)
+
     return ' '.join(argv)
 
 
+def _redact_env_values(argv: List[str]) -> List[str]:
+    """Redact sensitive values from --env arguments.
+
+    Args:
+        argv: Command line arguments
+
+    Returns:
+        Modified argv with redacted --env values, or original argv if any error
+
+    Examples:
+        ['sky', 'launch', '--env', 'HF_TOKEN=secret'] ->
+        ['sky', 'launch', '--env', 'HF_TOKEN=<redacted>']
+
+        ['sky', 'launch', '--env=HF_TOKEN=secret'] ->
+        ['sky', 'launch', '--env=HF_TOKEN=<redacted>']
+
+        ['sky', 'launch', '--env', 'HF_TOKEN'] ->
+        ['sky', 'launch', '--env', 'HF_TOKEN'] (no change)
+    """
+    try:
+        if not argv:
+            return argv or []
+
+        result = []
+        i = 0
+
+        while i < len(argv):
+            arg = argv[i]
+
+            # Ensure arg is a string
+            if not isinstance(arg, str):
+                result.append(arg)
+                i += 1
+                continue
+
+            if arg == '--env' and i + 1 < len(argv):
+                result.append(arg)
+                next_arg = argv[i + 1]
+                # Ensure next_arg is a string and handle redaction safely
+                if isinstance(next_arg, str):
+                    redacted = re.sub(r'^([^=]+)=.*', r'\1=<redacted>',
+                                      next_arg)
+                    result.append(redacted)
+                else:
+                    result.append(next_arg)
+                i += 2
+            elif arg.startswith('--env='):
+                # Redact only if there's a value after the key
+                redacted = re.sub(r'^(--env=[^=]+)=.*', r'\1=<redacted>', arg)
+                result.append(redacted)
+                i += 1
+            else:
+                result.append(arg)
+                i += 1
+
+        return result
+    except Exception:  # pylint: disable=broad-except
+        # If anything goes wrong with redaction, return original argv
+        # This ensures the command can still execute
+        return argv or []
+
+
 def user_and_hostname_hash() -> str:
     """Returns a string containing <user>-<hostname hash last 4 chars>.
 
diff --git a/sky/utils/rich_utils.py b/sky/utils/rich_utils.py
index 53c759e17..e804228ac 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)


Test patch:
diff --git a/tests/test_api.py b/tests/test_api.py
index 26e1bf0b5..d7d57f766 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 test_k8s_alias(monkeypatch, enable_all_clouds):
     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
diff --git a/tests/unit_tests/test_sky/utils/test_common_utils.py b/tests/unit_tests/test_sky/utils/test_common_utils.py
index d42d908c5..c294ea419 100644
--- a/tests/unit_tests/test_sky/utils/test_common_utils.py
+++ b/tests/unit_tests/test_sky/utils/test_common_utils.py
@@ -274,3 +274,158 @@ class TestCgroupFunctions:
         with mock.patch.dict('os.environ',
                              {'SKYPILOT_POD_MEMORY_GB_LIMIT': '2'}):
             assert common_utils.get_mem_size_gb() == 2.0
+
+
+class TestRedactEnvValues:
+    """Test environment variable value redaction in command lines."""
+
+    def test_env_separate_with_value(self):
+        """Test --env KEY=value format."""
+        argv = ['sky', 'launch', '--env', 'HF_TOKEN=secret123', 'app.yaml']
+        result = common_utils._redact_env_values(argv)
+        expected = ['sky', 'launch', '--env', 'HF_TOKEN=<redacted>', 'app.yaml']
+        assert result == expected
+
+    def test_env_separate_without_value(self):
+        """Test --env KEY format (no value to redact)."""
+        argv = ['sky', 'launch', '--env', 'HF_TOKEN', 'app.yaml']
+        result = common_utils._redact_env_values(argv)
+        expected = ['sky', 'launch', '--env', 'HF_TOKEN', 'app.yaml']
+        assert result == expected
+
+    def test_env_combined_with_value(self):
+        """Test --env=KEY=value format."""
+        argv = ['sky', 'launch', '--env=HF_TOKEN=secret123', 'app.yaml']
+        result = common_utils._redact_env_values(argv)
+        expected = ['sky', 'launch', '--env=HF_TOKEN=<redacted>', 'app.yaml']
+        assert result == expected
+
+    def test_env_combined_without_value(self):
+        """Test --env=KEY format (no value to redact)."""
+        argv = ['sky', 'launch', '--env=HF_TOKEN', 'app.yaml']
+        result = common_utils._redact_env_values(argv)
+        expected = ['sky', 'launch', '--env=HF_TOKEN', 'app.yaml']
+        assert result == expected
+
+    def test_multiple_env_args(self):
+        """Test multiple --env arguments with different formats."""
+        argv = [
+            'sky', 'launch', '--env', 'KEY1=secret1', '--env', 'KEY2=secret2',
+            'app.yaml'
+        ]
+        result = common_utils._redact_env_values(argv)
+        expected = [
+            'sky', 'launch', '--env', 'KEY1=<redacted>', '--env',
+            'KEY2=<redacted>', 'app.yaml'
+        ]
+        assert result == expected
+
+    def test_mixed_env_formats(self):
+        """Test mixed --env formats in one command."""
+        argv = [
+            'sky', 'launch', '--env', 'KEY1=secret1', '--env=KEY2=secret2',
+            '--env', 'KEY3', 'app.yaml'
+        ]
+        result = common_utils._redact_env_values(argv)
+        expected = [
+            'sky', 'launch', '--env', 'KEY1=<redacted>',
+            '--env=KEY2=<redacted>', '--env', 'KEY3', 'app.yaml'
+        ]
+        assert result == expected
+
+    def test_no_env_args(self):
+        """Test command without --env arguments."""
+        argv = ['sky', 'launch', 'app.yaml']
+        result = common_utils._redact_env_values(argv)
+        expected = ['sky', 'launch', 'app.yaml']
+        assert result == expected
+
+    def test_value_with_equals(self):
+        """Test env value that contains equals signs."""
+        argv = ['sky', 'launch', '--env', 'KEY1=value=with=equals', 'app.yaml']
+        result = common_utils._redact_env_values(argv)
+        expected = ['sky', 'launch', '--env', 'KEY1=<redacted>', 'app.yaml']
+        assert result == expected
+
+    def test_value_with_equals_combined(self):
+        """Test env value with equals signs in combined format."""
+        argv = ['sky', 'launch', '--env=KEY2=value=with=equals', 'app.yaml']
+        result = common_utils._redact_env_values(argv)
+        expected = ['sky', 'launch', '--env=KEY2=<redacted>', 'app.yaml']
+        assert result == expected
+
+    def test_empty_value(self):
+        """Test env argument with empty value."""
+        argv = ['sky', 'launch', '--env', 'KEY3=', 'app.yaml']
+        result = common_utils._redact_env_values(argv)
+        expected = ['sky', 'launch', '--env', 'KEY3=<redacted>', 'app.yaml']
+        assert result == expected
+
+    def test_empty_value_combined(self):
+        """Test env argument with empty value in combined format."""
+        argv = ['sky', 'launch', '--env=KEY4=', 'app.yaml']
+        result = common_utils._redact_env_values(argv)
+        expected = ['sky', 'launch', '--env=KEY4=<redacted>', 'app.yaml']
+        assert result == expected
+
+    def test_env_at_end_without_value(self):
+        """Test --env at end of command without following argument."""
+        argv = ['sky', 'launch', 'app.yaml', '--env']
+        result = common_utils._redact_env_values(argv)
+        expected = ['sky', 'launch', 'app.yaml', '--env']
+        assert result == expected
+
+    def test_edge_cases(self):
+        """Test edge cases that should not cause failures."""
+        # Empty list
+        assert common_utils._redact_env_values([]) == []
+
+        # None input
+        assert common_utils._redact_env_values(None) == []
+
+        # Single element
+        assert common_utils._redact_env_values(['sky']) == ['sky']
+
+        # Non-string elements (should be preserved)
+        argv = ['sky', 'launch', 123, '--env', 'KEY=value']
+        result = common_utils._redact_env_values(argv)
+        expected = ['sky', 'launch', 123, '--env', 'KEY=<redacted>']
+        assert result == expected
+
+        # Non-string after --env
+        argv = ['sky', 'launch', '--env', 123]
+        result = common_utils._redact_env_values(argv)
+        expected = ['sky', 'launch', '--env', 123]
+        assert result == expected
+
+        # Very long strings (should not cause memory issues)
+        long_value = 'x' * 10000
+        argv = ['sky', 'launch', '--env', f'KEY={long_value}']
+        result = common_utils._redact_env_values(argv)
+        expected = ['sky', 'launch', '--env', 'KEY=<redacted>']
+        assert result == expected
+
+        # Unicode characters
+        argv = ['sky', 'launch', '--env', 'KEY=密码123']
+        result = common_utils._redact_env_values(argv)
+        expected = ['sky', 'launch', '--env', 'KEY=<redacted>']
+        assert result == expected
+
+        # Special regex characters in key/value
+        argv = ['sky', 'launch', '--env', 'KEY[0]=value.*']
+        result = common_utils._redact_env_values(argv)
+        expected = ['sky', 'launch', '--env', 'KEY[0]=<redacted>']
+        assert result == expected
+
+    @mock.patch('sky.utils.common_utils.re.sub')
+    def test_error_handling_fallback(self, mock_re_sub):
+        """Test that function returns original argv if redaction fails."""
+        # Make re.sub raise an exception
+        mock_re_sub.side_effect = Exception("Simulated regex error")
+
+        argv = ['sky', 'launch', '--env', 'KEY=value']
+        result = common_utils._redact_env_values(argv)
+
+        # Should return original argv when error occurs
+        expected = ['sky', 'launch', '--env', 'KEY=value']
+        assert result == expected