Python runtime: keep waiting log lines in order without spinning

Follow-up to the change that made the bindings' log function stop
running Python code. Lines that cannot go to Python at once wait for
Python's main thread (a pending call) or for a thread of the bindings. A
few cases were wrong:

- The main thread could spin. When some of its lines waited for the
  bindings' thread and it logged again with sys.stderr replaced, its
  pending call put the lines back and asked to run again at every
  bytecode check, until the other thread finished. A Python loop then
  ran hundreds of times slower. Now such lines go behind the waiting
  ones to the bindings' thread, even if sys.stderr is replaced for a
  moment, since keeping the order matters more.
- On the main thread, a binding call that ran in an except block right
  after a call that raised wrote its lines before the earlier lines
  still waiting for the pending call. Now a call on the main thread
  writes those lines first, then its own.
- With fd 2 non-blocking and full, the bindings' thread dropped its
  lines when the write returned EAGAIN. Now it waits for room.
- The bindings' thread wrote all its lines in one write, so on a pipe
  another thread's line could land inside one of them. Now it writes a
  line at a time, which keeps a line of up to PIPE_BUF bytes whole.
- If that thread was still writing to fd 2 when Python finished
  exiting, it went back to a thread state Python may have freed. Now it
  stops without touching Python once exit has begun, and exit waits
  while it takes the GIL back.

The README no longer says that every thread that logs outside a call
hands its line over: only a thread that holds the GIL does, and the list
of hangs that remain now includes the other case.

New tests check the order of the main thread's lines after a call that
raised, that the main thread does not slow down while its lines wait,
and that no line is lost when fd 2 is non-blocking and full.

Author: shoumikhin

extension/pybindings/README.md

@@ -241,28 +241,33 @@ look wrong there. If Python has replaced `sys.stderr`, for example with pytest's
 `capsys`, `contextlib.redirect_stderr` or in a Jupyter notebook, the lines go to
 `sys.stderr` instead, a little later. The lines a binding call logs wait until
 the call returns. A line logged outside any binding call, for example by the LLM
-runner, goes to `sys.stderr` before the next Python code runs if Python's main
-thread logged it. A line from another thread goes there a moment later, from a
-thread that the bindings start for this. Outside a call, only a thread that
-holds the GIL can tell where `sys.stderr` goes, so a line from a thread that
-does not hold it, such as one the LLM runner logs while it generates, goes to
-the process's stderr. The lines wait rather than go to Python at once because a
-backend can log while it holds a lock, and writing to Python at that moment
-could deadlock with another thread.
+runner, goes to `sys.stderr` soon after. If Python's main thread logged it, that
+thread writes it itself, before the lines of its next binding call. But if
+earlier lines of the main thread still wait for the bindings' thread, this line
+follows them there. A line from another thread also goes to the bindings'
+thread, which the bindings start for this. It writes to whatever `sys.stderr` is
+at that moment. Outside a call, only a thread that holds the GIL can tell where
+`sys.stderr` goes, so a line from a thread that does not hold it, such as one
+the LLM runner logs while it generates, goes to the process's stderr. The lines
+wait rather than go to Python at once because a backend can log while it holds a
+lock, and writing to Python at that moment could deadlock with another thread.
 
 On Linux and macOS, a line also waits once the process's stderr is a pipe or
 socket that stays full for 100 ms, so a thread that holds the GIL does not hang
 on it while a Python thread in the same process empties it, as under
 `wurlitzer.pipes()`. The rest of a binding call's lines then wait for the call
-to return. A thread that logs outside a call does not wait for room itself. It
-hands the line to the bindings' thread, which writes it without the GIL. While
-it writes, a binding call's lines wait for the call to return. A terminal is
-written at once, since its reader is normally another process. The process can
-still hang if another writer fills the process's stderr right after the check
-for room, or if a Python thread in the same process reads a terminal that is the
-process's stderr. On macOS, a pipe that has not grown past 16 KiB takes a line
-only once the line fits in 512 bytes beside what the pipe holds, because macOS
-reports room in such a pipe that it may not have.
+to return. A thread that holds the GIL and logs outside a call does not wait for
+room itself. It hands the line to the bindings' thread, which writes it a line
+at a time without the GIL. While it writes, a binding call's lines wait for the
+call to return. A terminal is written at once, since its reader is normally
+another process. The process can still hang if another writer fills the
+process's stderr right after the check for room, if a thread that does not hold
+the GIL logs outside a call while a Python thread in the same process empties a
+full stderr and a binding call holds the GIL waiting for that thread, or if a
+Python thread in the same process reads a terminal that is the process's stderr.
+On macOS, a pipe that has not grown past 16 KiB takes a line only once the line
+fits in 512 bytes beside what the pipe holds, because macOS reports room in such
+a pipe that it may not have.
 
 Up to 1 MiB of lines wait at a time; past that, about the newest 512 KiB to
 1 MiB are kept, with a note of how many earlier bytes were left out. A fatal

extension/pybindings/pybindings.cpp

@@ -181,8 +181,9 @@ struct DeferredLogs final {
   uint64_t batch = 1;
   uint64_t batches_written = 0;
   bool writer_running = false;
-  // A new thread state crashes a Python that has finished exiting, so exit
-  // waits while the writer takes the GIL with one.
+  // Exit waits while the writer takes the GIL, since a Python that has
+  // finished exiting crashes on a new thread state and may have freed the
+  // writer's saved one.
   bool writer_entering = false;
   bool exiting = false;
 };
@@ -222,6 +223,12 @@ void write_to_fd_2(const char* data, size_t size) {
     if (written < 0 && errno == EINTR) {
       continue;
     }
+    if (written < 0 && (errno == EAGAIN || errno == EWOULDBLOCK)) {
+      // fd 2 is non-blocking and full.
+      pollfd stderr_fd{2, POLLOUT, 0};
+      poll(&stderr_fd, 1, -1);
+      continue;
+    }
     if (written <= 0) {
       return;
     }
@@ -389,7 +396,20 @@ void write_logs_to_python(PendingLogs& logs) noexcept {
   if (logs.lines.empty()) {
     return;
   }
-  // A call's lines go after the lines this thread handed to the writer thread.
+  if (PyThread_get_thread_ident() == main_thread_id) {
+    // The lines that the main thread's pending call has yet to write come
+    // first, since a call that raised runs no pending call before the next.
+    DeferredLogs& deferred = deferred_logs();
+    std::lock_guard<std::mutex> lock(deferred.mutex);
+    note_left_out_bytes(deferred.main_thread_logs);
+    try {
+      logs.lines.insert(0, deferred.main_thread_logs.lines);
+      deferred.main_thread_logs = PendingLogs();
+      deferred.main_thread_write_scheduled = false;
+    } catch (const std::bad_alloc&) {
+    }
+  }
+  // A call's lines go after the lines this thread handed over.
   if (logs.daemon || own_deferred_lines_wait()) {
     defer_lines(logs.lines);
     return;
@@ -431,10 +451,38 @@ int write_main_thread_logs(void*) noexcept {
   return 0;
 }
 
+// Writes the lines to fd 2 one at a time with the GIL released, so that a line
+// no longer than PIPE_BUF stays whole beside other threads' lines. Returns
+// false, without the GIL, if Python began to exit meanwhile: Python may then
+// free this thread's state, so the thread must not touch Python again.
+bool write_lines_to_fd_2_without_gil(const std::string& lines) {
+  PyThreadState* const state = PyEval_SaveThread();
+  size_t begin = 0;
+  while (begin < lines.size()) {
+    const size_t end = lines.find('\n', begin) + 1;
+    write_to_fd_2(lines.data() + begin, end - begin);
+    begin = end;
+  }
+  DeferredLogs& deferred = deferred_logs();
+  {
+    std::lock_guard<std::mutex> lock(deferred.mutex);
+    if (deferred.exiting) {
+      return false;
+    }
+    deferred.writer_entering = true;
+  }
+  PyEval_RestoreThread(state);
+  std::lock_guard<std::mutex> lock(deferred.mutex);
+  deferred.writer_entering = false;
+  deferred.writer_changed.notify_all();
+  return true;
+}
+
 // Writes through Python's C API only, and is not noexcept, since Python may end
 // this thread while it waits for the GIL during exit, and the unwinding must
-// neither touch a Python object nor stop at a noexcept frame.
-void write_lines_through_python(const PendingLogs& logs) {
+// neither touch a Python object nor stop at a noexcept frame. Returns false as
+// write_lines_to_fd_2_without_gil does.
+bool write_lines_through_python(const PendingLogs& logs) {
   PyObject* const text = PyUnicode_DecodeUTF8(
       logs.lines.data(), logs.lines.size(), "backslashreplace");
   PyObject* const stream = PySys_GetObject("stderr");
@@ -444,12 +492,11 @@ void write_lines_through_python(const PendingLogs& logs) {
   const bool written = result != nullptr;
   Py_XDECREF(result);
   Py_XDECREF(text);
-  if (!written) {
-    PyErr_Clear();
-    PyThreadState* const state = PyEval_SaveThread();
-    write_to_fd_2(logs.lines.data(), logs.lines.size());
-    PyEval_RestoreThread(state);
+  if (written) {
+    return true;
   }
+  PyErr_Clear();
+  return write_lines_to_fd_2_without_gil(logs.lines);
 }
 
 // Runs on a thread of its own until no line waits.
@@ -476,14 +523,16 @@ void write_deferred_logs() {
       deferred.writer_changed.notify_all();
     }
     note_left_out_bytes(logs);
-    if (stderr_is_process_stderr()) {
-      // Waits for room without the GIL or any lock of Python's, so that a full
-      // fd 2 never holds up Python's exit.
-      PyThreadState* const state = PyEval_SaveThread();
-      write_to_fd_2(logs.lines.data(), logs.lines.size());
-      PyEval_RestoreThread(state);
-    } else {
-      write_lines_through_python(logs);
+    // Lines for the process's stderr wait for room in fd 2 without the GIL or
+    // any lock of Python's, so that a full fd 2 never holds up Python's exit.
+    const bool python_still_running = stderr_is_process_stderr()
+        ? write_lines_to_fd_2_without_gil(logs.lines)
+        : write_lines_through_python(logs);
+    if (!python_still_running) {
+      std::lock_guard<std::mutex> lock(deferred.mutex);
+      deferred.writer_running = false;
+      deferred.writer_changed.notify_all();
+      return;
     }
     PyGILState_Release(gil);
     std::lock_guard<std::mutex> lock(deferred.mutex);
@@ -508,8 +557,10 @@ void defer_lines(const std::string& lines) noexcept {
     // The main thread logged these lines from native code, which is inside a
     // sys.stderr.write only if that write runs such code. So it can write them
     // before its next bytecode, while a short redirect of sys.stderr still
-    // lasts. Its earlier lines stay ahead.
+    // lasts. While the writer thread holds earlier lines of it, these follow
+    // them there instead, since keeping a thread's order beats the redirect.
     if (PyThread_get_thread_ident() == main_thread_id &&
+        last_deferred_batch <= deferred.batches_written &&
         (deferred.main_thread_write_scheduled ||
          !sys_stderr_is_checked_fd_2())) {
       keep_line(deferred.main_thread_logs, lines);
@@ -539,6 +590,8 @@ bool deferred_lines_wait() {
       !deferred.exiting;
 }
 
+// Whether lines this thread handed over still wait, which its later lines must
+// not pass.
 bool own_deferred_lines_wait() {
   DeferredLogs& deferred = deferred_logs();
   std::lock_guard<std::mutex> lock(deferred.mutex);

extension/pybindings/test/test_pybindings.py

@@ -707,6 +707,28 @@ def log_twice():
     print("in order:", 0 <= first < text.find("Program identifier"))
     sys.exit()
 
+if sys.argv[1] == "replaced, a call after a failed load":
+    in_order = True
+    # The first round comes out in order even without the wait, so it repeats.
+    for _ in range(5):
+        captured = io.StringIO()
+        sys.stderr = captured
+        try:
+            TextLLMRunner(missing_program, sys.argv[2])
+        except RuntimeError:
+            # A call that raised runs no pending call before the next call, yet
+            # this call's line must come after the runner's lines.
+            try:
+                portable_lib._load_program_from_buffer(b"x" * 128)
+            except RuntimeError:
+                pass
+        sys.stderr = sys.__stderr__
+        text = captured.getvalue()
+        first = text.find("[llm_runner_helper.cpp:")
+        in_order = in_order and 0 <= first < text.find("Program identifier")
+    print("in order:", in_order)
+    sys.exit()
+
 with (
     contextlib.redirect_stderr(captured)
     if sys.argv[1] == "replaced after import"
@@ -729,6 +751,7 @@ def log_twice():
                 "replaced before import",
                 "replaced after import",
                 "replaced, another thread",
+                "replaced, a call after a failed load",
                 "process stderr",
             ):
                 completed = subprocess.run(
@@ -744,8 +767,8 @@ def log_twice():
                     self.skipTest("the LLM runner bindings are not built")
                 if mode == "process stderr":
                     self.assertIn("[llm_runner_helper.cpp:", completed.stderr)
-                elif mode == "replaced, another thread":
-                    self.assertIn("in order: True", completed.stdout)
+                elif mode.startswith("replaced, a"):
+                    self.assertIn("in order: True", completed.stdout, mode)
                 else:
                     self.assertIn("lines through Python:", completed.stdout, mode)
                     self.assertNotIn("lines through Python: 0", completed.stdout, mode)
@@ -1076,6 +1099,111 @@ def build_runners():
                     if "no LLM runner" not in completed.stdout:
                         self.assertIn("done", completed.stdout, f"{mode}, {kind}")
 
+    def test_lines_waiting_for_a_full_stderr(self):
+        if self.kernel_mode != "portable":
+            self.skipTest("the script below uses the portable bindings")
+        if sys.platform == "win32":
+            self.skipTest("uses a POSIX pipe as fd 2")
+
+        script = """
+import io
+import os
+import sys
+import threading
+import time
+
+from executorch.extension.pybindings import portable_lib
+
+try:
+    from executorch.extension.llm.runner import TextLLMRunner
+except (ImportError, RuntimeError):
+    # The package raises RuntimeError when its extension is not built.
+    print("no LLM runner")
+    sys.exit()
+missing_program = os.path.join(os.path.dirname(sys.argv[2]), "missing.pte")
+
+
+def spin(seconds):
+    rounds = 0
+    end = time.perf_counter() + seconds
+    while time.perf_counter() < end:
+        for _ in range(100):
+            rounds += 1
+    return rounds
+
+
+normal_speed = spin(0.5)
+read_end, write_end = os.pipe()
+os.set_blocking(write_end, False)
+for size in (4096, 512, 64, 1):
+    try:
+        while True:
+            os.write(write_end, b"x" * size)
+    except BlockingIOError:
+        pass
+saved_stderr = os.dup(2)
+os.dup2(write_end, 2)
+if sys.argv[1] == "blocking":
+    os.set_blocking(2, True)
+drained = bytearray()
+
+
+def drain():
+    # Starts once the lines below wait, and needs the GIL after every read.
+    time.sleep(2)
+    os.set_blocking(read_end, False)
+    end = time.monotonic() + 3
+    while time.monotonic() < end:
+        try:
+            drained.extend(os.read(read_end, 65536))
+        except BlockingIOError:
+            time.sleep(0.01)
+
+
+drainer = threading.Thread(target=drain)
+drainer.start()
+try:
+    # fd 2 is full, so these lines go to a thread of the bindings, which waits
+    # to write them.
+    try:
+        TextLLMRunner(missing_program, sys.argv[2])
+    except RuntimeError:
+        pass
+    if sys.argv[1] == "blocking":
+        # The main thread's later lines must wait behind those without making
+        # it spin.
+        time.sleep(0.3)
+        sys.stderr = io.StringIO()
+        try:
+            portable_lib._load_program_from_buffer(b"x" * 128)
+        except RuntimeError:
+            pass
+        print("speed:", round(spin(0.5) / normal_speed, 3))
+        sys.stderr = sys.__stderr__
+    drainer.join()
+finally:
+    os.dup2(saved_stderr, 2)
+print("runner lines:", drained.count(b"[llm_runner_helper.cpp:"))
+"""
+        with tempfile.TemporaryDirectory() as directory:
+            tokenizer = write_llm_tokenizer(directory)
+            for mode in ("blocking", "non-blocking"):
+                completed = subprocess.run(
+                    [sys.executable, "-c", script, mode, tokenizer],
+                    capture_output=True,
+                    text=True,
+                    timeout=120,
+                )
+                self.assertEqual(
+                    completed.returncode, 0, f"{mode}: {completed.stderr[-2000:]}"
+                )
+                if "no LLM runner" in completed.stdout:
+                    self.skipTest("the LLM runner bindings are not built")
+                self.assertNotIn("runner lines: 0", completed.stdout, mode)
+                if mode == "blocking":
+                    speed = float(completed.stdout.split("speed:")[1].split()[0])
+                    self.assertGreater(speed, 0.05, completed.stdout)
+
     def test_deferred_runtime_logs_do_not_nest_in_a_stderr_write(self):
         if self.kernel_mode != "portable":
             self.skipTest("the script below uses the portable bindings")