[Serve] BugFix: LiveError on controller

sky/utils/rich_utils.py

@@ -141,15 +141,28 @@ def __enter__(self):
         return _statuses[self.status_type]
 
     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)
+        # We use the same lock with the `safe_logger` to avoid the following 2
+        # raice conditions. We refer loggers in another thread as "thread
+        # logger" hereafter.
+        # 1. When a thread logger stopped the status in `safe_logger`, and
+        #    here we exit the status and set it to None. Then the thread logger
+        #    will raise an error when it tries to restart the status.
+        # 2. When a thread logger stopped the status in `safe_logger`, and
+        #    here we exit the status and entered a new one. Then the thread
+        #    logger will raise an error when it tries to restart the old status,
+        #    since only one LiveStatus can be started at the same time.
+        # Please refer to #4995 for more information.
+        with _logging_lock:
+            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)
 
     def update(self, *args, **kwargs):
         _statuses[self.status_type].update(*args, **kwargs)