Repository: skypilot-org/skypilot
Pull request: [Core] Handle mid-sequence chunking in log streaming
Closes #4905.

Our chunking may chunk in the middle of a UTF-8 sequence, which would cause decoding to break. We now store the undecoded bytes in a buffer and wait to receive the next chunk before trying to decode again.

Problem statement:
decode_rich_status can chunk in the middle of utf-8 sequences
On some occasions my logs are not retrievable and they end with:
```
File "/home/user/.keybotic/venvs/skypilot/bin/sky", line 8, in <module>
    sys.exit(cli())
  File "/home/user/.keybotic/venvs/skypilot/lib/python3.8/site-packages/click/core.py", line 1161, in __call__
    return self.main(*args, **kwargs)
  File "/home/user/.keybotic/venvs/skypilot/lib/python3.8/site-packages/click/core.py", line 1082, in main
    rv = self.invoke(ctx)
  File "/home/user/.keybotic/venvs/skypilot/lib/python3.8/site-packages/sky/utils/common_utils.py", line 434, in _record
    return f(*args, **kwargs)
  File "/home/user/.keybotic/venvs/skypilot/lib/python3.8/site-packages/sky/cli.py", line 875, in invoke
    return super().invoke(ctx)
  File "/home/user/.keybotic/venvs/skypilot/lib/python3.8/site-packages/click/core.py", line 1697, in invoke
    return _process_result(sub_ctx.command.invoke(sub_ctx))
  File "/home/user/.keybotic/venvs/skypilot/lib/python3.8/site-packages/sky/utils/common_utils.py", line 434, in _record
    return f(*args, **kwargs)
  File "/home/user/.keybotic/venvs/skypilot/lib/python3.8/site-packages/sky/cli.py", line 875, in invoke
    return super().invoke(ctx)
  File "/home/user/.keybotic/venvs/skypilot/lib/python3.8/site-packages/click/core.py", line 1697, in invoke
    return _process_result(sub_ctx.command.invoke(sub_ctx))
  File "/home/user/.keybotic/venvs/skypilot/lib/python3.8/site-packages/click/core.py", line 1443, in invoke
    return ctx.invoke(self.callback, **ctx.params)
  File "/home/user/.keybotic/venvs/skypilot/lib/python3.8/site-packages/click/core.py", line 788, in invoke
    return __callback(*args, **kwargs)
  File "/home/user/.keybotic/venvs/skypilot/lib/python3.8/site-packages/sky/utils/common_utils.py", line 454, in _record
    return f(*args, **kwargs)
  File "/home/user/.keybotic/venvs/skypilot/lib/python3.8/site-packages/sky/cli.py", line 4092, in jobs_logs
    managed_jobs.tail_logs(name=name,
  File "/home/user/.keybotic/venvs/skypilot/lib/python3.8/site-packages/sky/utils/common_utils.py", line 454, in _record
    return f(*args, **kwargs)
  File "/home/user/.keybotic/venvs/skypilot/lib/python3.8/site-packages/sky/server/common.py", line 285, in wrapper
    return func(*args, **kwargs)
  File "/home/user/.keybotic/venvs/skypilot/lib/python3.8/site-packages/sky/jobs/client/sdk.py", line 214, in tail_logs
    sdk.stream_response(request_id, response, output_stream)
  File "/home/user/.keybotic/venvs/skypilot/lib/python3.8/site-packages/sky/client/sdk.py", line 72, in stream_response
    for line in rich_utils.decode_rich_status(response):
  File "/home/user/.keybotic/venvs/skypilot/lib/python3.8/site-packages/sky/utils/rich_utils.py", line 239, in decode_rich_status
    encoded_msg = encoded_msg.decode('utf-8')
UnicodeDecodeError: 'utf-8' codec can't decode bytes in position 3962-3963: unexpected end of data
```

Printing the offending log buffer:
```
b'  0.9255      1.133      3.474      1.101         22        640:  20%|\xe2\x96\x88\xe2\x96\x89        | 115/581 [00:12<00:41, 11.20it/s]\n\x1b[36m(TestSky, pid=2952)\x1b[0m        1/11      2.75G     0.9236      1.129      3.474        1.1         17        640:  20%|\xe2\x96\x88\xe2\x96\x89        | 115/581 [00:12<00:41, 11.20it/s]\n       1/11      2.75G     0.9209      1.126       3.47      1.098         23        640:  20%|\xe2\x96\x88\xe2\x96\x89        | 115/581 [00:13<00:41, 11.20it/s]\n       1/11      2.75G     0.9209      1.126       3.47      1.098         23        640:  20%|\xe2\x96\x88\xe2\x96\x88        | 117/581 [00:13<00:52,  8.81it/s]\n\x1b[36m(TestSky, pid=2952)\x1b[0m        1/11      2.75G     0.9185       1.12      3.469      1.097         16        640:  20%|\xe2\x96\x88\xe2\x96\x88        | 117/581 [00:13<00:52,  8.81it/s]\n       1/11      2.75G     0.9161      1.118      3.466      1.096         27        640:  20%|\xe2\x96\x88\xe2\x96\x88        | 117/581 [00:13<00:52,  8.81it/s]\n\x1b[36m(TestSky, pid=2952)\x1b[0m        1/11      2.75G     0.9161      1.118      3.466      1.096         27        640:  20%|\xe2\x96\x88\xe2\x96\x88        | 119/581 [00:13<00:47,  9.83it/s]\n\x1b[36m(TestSky, pid=2952)\x1b[0m        1/11      2.75G     0.9141      1.115      3.462      1.094         27        640:  20%|\xe2\x96\x88\xe2\x96\x88        | 119/581 [00:13<00:47,  9.83it/s]\n       1/11      2.75G     0.9114       1.11      3.457      1.092         28        640:  20%|\xe2\x96\x88\xe2\x96\x88        | 119/581 [00:13<00:47,  9.83it/s]\n\x1b[36m(TestSky, pid=2952)\x1b[0m        1/11      2.75G     0.9114       1.11      3.457      1.092         28        640:  21%|\xe2\x96\x88\xe2\x96\x88        | 121/581 [00:13<00:42, 10.70it/s]\n\x1b[36m(TestSky, pid=2952)\x1b[0m        1/11      2.75G     0.9075      1.106      3.453       1.09         20        640:  21%|\xe2\x96\x88\xe2\x96\x88        | 121/581 [00:13<00:42, 10.70it/s]\n       1/11      2.75G     0.9053      1.103      3.448      1.089         19        640:  21%|\xe2\x96\x88\xe2\x96\x88        | 121/581 [00:13<00:42, 10.70it/s]\n\x1b[36m(TestSky, pid=2952)\x1b[0m        1/11      2.75G     0.9053      1.103      3.448      1.089         19        640:  21%|\xe2\x96\x88\xe2\x96\x88        | 123/581 [00:13<00:41, 11.09it/s]\n\x1b[36m(TestSky, pid=2952)\x1b[0m        1/11      2.75G     0.9025      1.098      3.444      1.088         18        640:  21%|\xe2\x96\x88\xe2\x96\x88        | 123/581 [00:13<00:41, 11.09it/s]\n       1/11      2.75G     0.9001      1.102       3.44      1.086         22        640:  21%|\xe2\x96\x88\xe2\x96\x88        | 123/581 [00:14<00:41, 11.09it/s]\n\x1b[36m(TestSky, pid=2952)\x1b[0m        1/11      2.75G     0.9001      1.102       3.44      1.086         22        640:  22%|\xe2\x96\x88\xe2\x96\x88\xe2\x96\x8f       | 125/581 [00:14<00:57,  7.99it/s]\n\x1b[36m(TestSky, pid=2952)\x1b[0m        1/11      2.75G     0.8983      1.097      3.438      1.084         22        640:  22%|\xe2\x96\x88\xe2\x96\x88\xe2\x96\x8f       | 125/581 [00:14<00:57,  7.99it/s]\n       1/11      2.75G     0.8974      1.095      3.433      1.083         21        640:  22%|\xe2\x96\x88\xe2\x96\x88\xe2\x96\x8f       | 125/581 [00:14<00:57,  7.99it/s]\n       1/11      2.75G     0.8974      1.095      3.433      1.083         21        640:  22%|\xe2\x96\x88\xe2\x96\x88\xe2\x96\x8f       | 127/581 [00:14<00:49,  9.12it/s]\n\x1b[36m(TestSky, pid=2952)\x1b[0m        1/11      2.75G     0.8968      1.095       3.43      1.081         25        640:  22%|\xe2\x96\x88\xe2\x96\x88\xe2\x96\x8f       | 127/581 [00:14<00:49,  9.12it/s]\n       1/11      2.75G     0.8946      1.092      3.426       1.08         23        640:  22%|\xe2\x96\x88\xe2\x96\x88\xe2\x96\x8f       | 127/581 [00:14<00:49,  9.12it/s]\n\x1b[36m(TestSky, pid=2952)\x1b[0m        1/11      2.75G     0.8946      1.092      3.426       1.08         23        640:  22%|\xe2\x96\x88\xe2\x96\x88\xe2\x96\x8f       | 129/581 [00:14<00:45, 10.04it/s]\n\x1b[36m(TestSky, pid=2952)\x1b[0m        1/11      2.75G     0.8923      1.088      3.422      1.078         25        640:  22%|\xe2\x96\x88\xe2\x96\x88\xe2\x96\x8f       | 129/581 [00:14<00:45, 10.04it/s]\n       1/11      2.75G     0.8907      1.086      3.417      1.077         33        640:  22%|\xe2\x96\x88\xe2\x96\x88\xe2\x96\x8f       | 129/581 [00:14<00:45, 10.04it/s]\n\x1b[36m(TestSky, pid=2952)\x1b[0m        1/11      2.75G     0.8907      1.086      3.417      1.077         33        640:  23%|\xe2\x96'
```

I believe the chunking is splitting a multibyte utf-8 character or something.

Hints from issue discussion:

Commit message: Handle mid-sequence chunking

Patch to review:
diff --git a/sky/utils/rich_utils.py b/sky/utils/rich_utils.py
index 54e9672f0ac..2a82d5ce679 100644
--- a/sky/utils/rich_utils.py
+++ b/sky/utils/rich_utils.py
@@ -230,15 +230,37 @@ def decode_rich_status(
     decoding_status = None
     try:
         last_line = ''
+        # Buffer to store incomplete UTF-8 bytes between chunks
+        undecoded_buffer = b''
+        
         # Iterate over the response content in chunks. We do not use iter_lines
         # because it will strip the trailing newline characters, causing the
         # progress bar ending with `\r` becomes a pyramid.
-        for encoded_msg in response.iter_content(chunk_size=None):
-            if encoded_msg is None:
+        for chunk in response.iter_content(chunk_size=None):
+            if chunk is None:
                 return
-            encoded_msg = encoded_msg.decode('utf-8')
+            
+            # Append the new chunk to any leftover bytes from previous iteration
+            current_bytes = undecoded_buffer + chunk
+            undecoded_buffer = b''
+            
+            # Try to decode the combined bytes
+            try:
+                encoded_msg = current_bytes.decode('utf-8')
+            except UnicodeDecodeError as e:
+                # If decoding fails due to incomplete byte sequence at the end,
+                # save those bytes for the next chunk and decode the valid part
+                valid_byte_count = e.start
+                encoded_msg = current_bytes[:valid_byte_count].decode('utf-8')
+                undecoded_buffer = current_bytes[valid_byte_count:]
+                
             lines = encoded_msg.splitlines(keepends=True)
-
+            
+            # Skip processing if lines is empty to avoid IndexError
+            if not lines:
+                continue
+                
+            # Append any leftover text from previous chunk to first line
             lines[0] = last_line + lines[0]
             last_line = lines[-1]
             # If the last line is not ended with `\r` or `\n` (with ending
diff --git a/tests/unit_tests/sky/utils/test_rich_utils.py b/tests/unit_tests/sky/utils/test_rich_utils.py
new file mode 100644
index 00000000000..ae6c1162af0
--- /dev/null
+++ b/tests/unit_tests/sky/utils/test_rich_utils.py
@@ -0,0 +1,85 @@
+import unittest
+from unittest import mock
+
+import requests
+
+from sky.utils import rich_utils
+
+
+class RichUtilsTest(unittest.TestCase):
+
+    def test_decode_rich_status_with_split_utf8(self):
+        """Test that decode_rich_status handles UTF-8 characters split across chunks."""
+        # This is a progress bar with UTF-8 box drawing characters
+        # The character █ is a multibyte UTF-8 character (0xE2 0x96 0x88)
+        progress_bar = "20%|████        | 115/581 [00:12<00:41, 11.20it/s]\n"
+        
+        # Create a response object with a custom content iterator
+        mock_response = mock.MagicMock(spec=requests.Response)
+        
+        # Convert to bytes and deliberately split in the middle of a multibyte character
+        # The character █ is encoded as three bytes: \xe2\x96\x88
+        progress_bytes = progress_bar.encode('utf-8')
+        
+        # Find position of the first █ character and split between its bytes
+        pos = progress_bytes.find(b'\xe2\x96\x88')
+        
+        # Split the bytes so that \xe2 is in one chunk and \x96\x88 is in the next
+        chunk1 = progress_bytes[:pos+1]  # Include the first byte of █ (\xe2)
+        chunk2 = progress_bytes[pos+1:]  # Rest of the bytes including \x96\x88
+        
+        # Configure the mock to return these chunks when iterated
+        mock_response.iter_content.return_value = [chunk1, chunk2]
+        
+        # The function should handle the split character correctly
+        result = list(rich_utils.decode_rich_status(mock_response))
+        
+        # Verify the output contains the correctly decoded progress bar
+        self.assertEqual(len(result), 1)
+        self.assertEqual(result[0], progress_bar)
+        
+    def test_decode_rich_status_with_multiple_split_characters(self):
+        """Test handling multiple UTF-8 characters split across several chunks."""
+        # Progress bar with multiple multibyte characters
+        progress_bar = "50%|█████████████████████████████░░░░░░░░░░░░░| 250/500 [01:15<01:15, 3.33it/s]\n"
+        progress_bytes = progress_bar.encode('utf-8')
+        
+        # Create chunks that split at several different points
+        # Find positions of several █ characters and split between their bytes
+        multibyte_char = '█'.encode('utf-8')  # \xe2\x96\x88
+        positions = []
+        start = 0
+        while True:
+            pos = progress_bytes.find(multibyte_char, start)
+            if pos == -1 or len(positions) >= 3:
+                break
+            positions.append(pos)
+            start = pos + len(multibyte_char)
+        
+        # Create chunks that split at different multibyte character boundaries
+        chunks = []
+        last_pos = 0
+        for i, pos in enumerate(positions):
+            if i % 2 == 0:  # For even indices, split after first byte of multibyte char
+                chunks.append(progress_bytes[last_pos:pos+1])
+                last_pos = pos+1
+            else:  # For odd indices, split normally
+                chunks.append(progress_bytes[last_pos:pos])
+                last_pos = pos
+        
+        # Add the final chunk
+        chunks.append(progress_bytes[last_pos:])
+        
+        # Configure mock response
+        mock_response = mock.MagicMock(spec=requests.Response)
+        mock_response.iter_content.return_value = chunks
+        
+        # Should handle multiple split characters correctly
+        result = list(rich_utils.decode_rich_status(mock_response))
+        
+        # Verify the output
+        joined_result = ''.join(r for r in result if r is not None)
+        self.assertEqual(joined_result, progress_bar)
+
+if __name__ == '__main__':
+    unittest.main()
\ No newline at end of file


Code review comments:
diff_hunk: @@ -0,0 +1,85 @@
+import unittest
+from unittest import mock
+
+import requests
+
+from sky.utils import rich_utils
+
+
+class RichUtilsTest(unittest.TestCase):
+
+    def test_decode_rich_status_with_split_utf8(self):
+        """Test that decode_rich_status handles UTF-8 characters split across chunks."""
+        # This is a progress bar with UTF-8 box drawing characters
+        # The character █ is a multibyte UTF-8 character (0xE2 0x96 0x88)
+        progress_bar = "20%|████        | 115/581 [00:12<00:41, 11.20it/s]\n"
+        
+        # Create a response object with a custom content iterator
+        mock_response = mock.MagicMock(spec=requests.Response)
+        
+        # Convert to bytes and deliberately split in the middle of a multibyte character
+        # The character █ is encoded as three bytes: \xe2\x96\x88
+        progress_bytes = progress_bar.encode('utf-8')
+        
+        # Find position of the first █ character and split between its bytes
+        pos = progress_bytes.find(b'\xe2\x96\x88')
+        
+        # Split the bytes so that \xe2 is in one chunk and \x96\x88 is in the next
+        chunk1 = progress_bytes[:pos+1]  # Include the first byte of █ (\xe2)
+        chunk2 = progress_bytes[pos+1:]  # Rest of the bytes including \x96\x88
+        
+        # Configure the mock to return these chunks when iterated
+        mock_response.iter_content.return_value = [chunk1, chunk2]
+        
+        # The function should handle the split character correctly
+        result = list(rich_utils.decode_rich_status(mock_response))
+        
+        # Verify the output contains the correctly decoded progress bar
+        self.assertEqual(len(result), 1)
+        self.assertEqual(result[0], progress_bar)
+        
+    def test_decode_rich_status_with_multiple_split_characters(self):
+        """Test handling multiple UTF-8 characters split across several chunks."""
+        # Progress bar with multiple multibyte characters
+        progress_bar = "50%|█████████████████████████████░░░░░░░░░░░░░| 250/500 [01:15<01:15, 3.33it/s]\n"
+        progress_bytes = progress_bar.encode('utf-8')
+        
+        # Create chunks that split at several different points
+        # Find positions of several █ characters and split between their bytes
+        multibyte_char = '█'.encode('utf-8')  # \xe2\x96\x88
+        positions = []
+        start = 0
+        while True:
+            pos = progress_bytes.find(multibyte_char, start)
+            if pos == -1 or len(positions) >= 3:
+                break
+            positions.append(pos)
+            start = pos + len(multibyte_char)
+        
+        # Create chunks that split at different multibyte character boundaries
+        chunks = []
+        last_pos = 0
+        for i, pos in enumerate(positions):
+            if i % 2 == 0:  # For even indices, split after first byte of multibyte char
+                chunks.append(progress_bytes[last_pos:pos+1])
+                last_pos = pos+1
+            else:  # For odd indices, split normally
+                chunks.append(progress_bytes[last_pos:pos])
+                last_pos = pos
+        
+        # Add the final chunk
+        chunks.append(progress_bytes[last_pos:])
+        
+        # Configure mock response
+        mock_response = mock.MagicMock(spec=requests.Response)
+        mock_response.iter_content.return_value = chunks
+        
+        # Should handle multiple split characters correctly
+        result = list(rich_utils.decode_rich_status(mock_response))
+        
+        # Verify the output
+        joined_result = ''.join(r for r in result if r is not None)
+        self.assertEqual(joined_result, progress_bar)
+
+if __name__ == '__main__':
+    unittest.main()
path: tests/unit_tests/sky/utils/test_rich_utils.py
text: @user1:
This is not necessary? Also, might be good to check if this goes well with pytest

@author:
ah yeah, this is not required, removed