[Fix][Core] Wait a while before stopping the ray_print_logs thread to…

… prevent pending logs

Closes: #48701
Signed-off-by: Chi-Sheng Liu <[email protected]>

python/ray/_private/worker.py

@@ -951,8 +951,14 @@ def print_logs(self):
             last_polling_batch_size = 0
             job_id_hex = self.current_job_id.hex()
             while True:
-                # Exit if we received a signal that we should stop.
-                if self.threads_stopped.is_set():
+                # Exit if we received a signal that we should stop or if main thread
+                # is dead. But wait until the message queue size decreases to lower
+                # than or equal to 1 to avoid missing logs.
+                if (
+                    self.threads_stopped.is_set()
+                    or not threading.main_thread().is_alive()
+                    and subscriber.queue_size <= 1
+                ):
                     return
 
                 data = subscriber.poll()
@@ -2529,7 +2535,6 @@ def connect(
             worker.logger_thread = threading.Thread(
                 target=worker.print_logs, name="ray_print_logs"
             )
-            worker.logger_thread.daemon = True
             worker.logger_thread.start()
 
     # Setup tracing here

python/ray/_raylet.pyx

@@ -2804,6 +2804,11 @@ cdef class _GcsSubscriber:
         """
         return self.inner.get().last_batch_size()
 
+    @property
+    def queue_size(self):
+        """Size of the message queue of the subscriber."""
+        return self.inner.get().queue_size()
+
     def close(self):
         """Closes the subscriber and its active subscription."""
         with nogil:

python/ray/includes/common.pxd

@@ -623,7 +623,9 @@ cdef extern from "ray/gcs/pubsub/gcs_pub_sub.h" nogil:
 
         CRayStatus Subscribe()
 
-        int64_t last_batch_size()
+        int64_t last_batch_size() const
+
+        int64_t queue_size() const
 
         CRayStatus PollError(
             c_string* key_id, int64_t timeout_ms, CErrorTableData* data)

src/ray/gcs/pubsub/gcs_pub_sub.cc

@@ -450,7 +450,12 @@ Status PythonGcsSubscriber::Close() {
   return Status::OK();
 }
 
-int64_t PythonGcsSubscriber::last_batch_size() {
+int64_t PythonGcsSubscriber::last_batch_size() const {
+  absl::MutexLock lock(&mu_);
+  return last_batch_size_;
+}
+
+int64_t PythonGcsSubscriber::queue_size() const {
   absl::MutexLock lock(&mu_);
   return last_batch_size_;
 }

src/ray/gcs/pubsub/gcs_pub_sub.h

@@ -191,7 +191,9 @@ class RAY_EXPORT PythonGcsSubscriber {
   /// Closes the subscriber and its active subscription.
   Status Close();
 
-  int64_t last_batch_size();
+  int64_t last_batch_size() const;
+
+  int64_t queue_size() const;
 
  private:
   Status DoPoll(int64_t timeout_ms, rpc::PubMessage *message);