[Serve] BugFix: LiveError on controller

[cblmemo]

Fixes #4907. As #4907 mentioned, there is two rich status related bug in our SkyServe system:

1. AttributeError: 'NoneType' object has no attribute 'start'
2. LiveError('Only one live display may be active at once')

The first issue is due to:

1. a side thread logger try to stop the status through safe_logger (code block 1, L203);
2. during emitting logs, the status has been set to None (code block 2, L150);
3. the side thread logger try to start the status again (code block 1, L206), causing this error.

The second issue is due to:

1. a side thread logger stop the status (code block 1, L203);
2. main thread exit a status (__exit__ function);
3. main thread start a new status (__enter__ function), the first live status is started;
4. the side thread logger restart the status, the second live status is started, causing this error.

Both of them can be resolved by adding a lock to the __exit__ function.

skypilot/sky/utils/rich_utils.py

Lines 195 to 206 in 6b938ce

	@contextlib.contextmanager
	def safe_logger():
	with _logging_lock:
	client_status_obj = _statuses['client']
	client_status_live = (client_status_obj is not None and
	client_status_obj._live.is_started) # pylint: disable=protected-access
	if client_status_live:
	client_status_obj.stop()
	yield
	if client_status_live:
	client_status_obj.start()

skypilot/sky/utils/rich_utils.py

Lines 143 to 152 in 6b938ce

	def __exit__(self, exc_type, exc_val, exc_tb):
	global _status_nesting_level
	_status_nesting_level -= 1
	if _status_nesting_level <= 0:
	_status_nesting_level = 0
	if _statuses[self.status_type] is not None:
	_statuses[self.status_type].__exit__(exc_type, exc_val, exc_tb)
	_statuses[self.status_type] = None
	else:
	_statuses[self.status_type].update(self.previous_message)

Reproducible script:

import time
import threading
import random
import traceback
from sky import sky_logging
from sky.utils import rich_utils, annotations, ux_utils

# stat = rich_utils.safe_status
stat = rich_utils.client_status

def f():
    logger = sky_logging.init_logger("sky.b")


    def func(msg: str):
        while True:
            logger.info(f'info {msg}')
            logger.debug(f'debug {msg}')
            time.sleep(random.random() * 0.1)

    with stat(ux_utils.spinner_message("enter to debug")):
        logger.info("main info")
        logger.debug("main debug")

    with stat(ux_utils.spinner_message("debug")):
        for i in range(10):
            threading.Thread(target=func, args=(f"hello {i}",), daemon=True).start()

    i = 0
    while True:
        try:
            with stat(ux_utils.spinner_message(f"debug_loop {i}")):
                time.sleep(0.1 * random.random())
                i += 1
        except Exception as e:
            print('=1=1=1=1=1=1=1=1', i, e)
            traceback.print_exc()
            raise e

    input()

f()

Tested (run the relevant ones):

• Code formatting: install pre-commit (auto-check on commit) or bash format.sh
• Any manual or new tests for this PR (please specify below)
• 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)

[cblmemo]

TODO: Add comments in the code

[Michaelvll]

Thanks @cblmemo for fixing this quickly! LGTM! Please add some comments in the code for future reference. : )

[JGSweets]

Yes, thanks for the quick fix @cblmemo!