[pipeline] Stop unstopped pipelines at interpreter exit to avoid hang (#1591)

A program that keeps a strong reference to a `Pipeline` and never calls
`stop()` (e.g. a training framework that stores the dataloader for the
whole run) can hang at interpreter exit. The pipeline's background
event-loop thread is non-daemon, so CPython joins it inside
`threading._shutdown()` during finalization; because nothing stopped the
pipeline, that join blocks forever. The existing `weakref.finalize`
cleanup runs too late: it is an `atexit`-module callback (a later
finalization phase), after non-daemon threads have already been joined.

Fix: register a per-pipeline shutdown hook via
`threading._register_atexit` (the same private mechanism
`concurrent.futures` relies on) so it runs inside
`threading._shutdown()` — before non-daemon threads and child processes
are joined — and stops the pipeline first. The hook is a thin weakref
wrapper around the existing `_stop_impl` (the single stop-with-timeout
helper, shared with the GC finalizer), registered in `Pipeline.start()`
after the pipeline's own threads/processes are up, so it lands after the
stdlib's atexit hooks and — being LIFO — runs before them. It holds only
a weak reference so it never keeps the pipeline alive, and is a safe
no-op once the pipeline is stopped or collected. A regular
`atexit.register` hook was verified insufficient: it runs after the
non-daemon thread join and still hangs.

Also make the subprocess-region bridge's blocking queue put
interruptible. A `.to(...)` region feeder parked on a full worker queue
(backpressure once the consumer stopped) cannot be released by pool
teardown, so at exit it is joined — and hangs — by `concurrent.futures`.
It now polls a teardown event and exits promptly. This also fixes a
pre-existing hang that reproduced on an explicit `stop()`.

No public API change; holding a `Pipeline` reference and relying on GC
is unchanged.

Author: mthrok

docs/source/getting_started/intro.rst

@@ -66,9 +66,13 @@ Finally call :py:meth:`Pipeline.stop` to stop the background thread.
    [19, 21, 23]
    >>> pipeline.stop()
 
-It is important to call :py:meth:`Pipeline.stop`.
-Forgetting to do so will leave the background thread running,
-which can cause the Python interpreter to hang at exit.
+Calling :py:meth:`Pipeline.stop` promptly is good practice: it releases the
+background thread and any worker processes as soon as you are done, instead of
+leaving them running until the object is garbage collected. It is no longer
+required to avoid a hang at exit, though — a ``Pipeline`` stops itself when it is
+garbage collected, and :py:meth:`Pipeline.start` additionally registers a hook
+that stops any still-running pipeline at interpreter shutdown (see
+:py:meth:`Pipeline.start` for details).
 
 In practice, there is always a possibility that the application is
 interrupted for unexpected reasons.
@@ -97,59 +101,95 @@ To make sure that the pipeline is stopped, it is recommended to use
 Unlike processes, threads cannot be killed.
 The ``Pipeline`` object uses a thread pool, which must be shut down properly.
 
-There are seemingly unharmful patterns, which can cause a deadlock
-at the end of the Python interpreter, preventing Python from exiting.
-
-.. admonition:: Keeping unnecessary references to ``Pipeline``
-   :class: danger
-
-   It is recommended to keep the resulting ``Pipeline`` object as a
-   local variable of an iterator, and NOT TO assign it to an object
-   attribute.
+The library cleans a ``Pipeline`` up automatically (see below), so these
+patterns no longer hang the interpreter at exit. They can still keep the
+background thread and worker processes alive longer than necessary, so they are
+worth avoiding for prompt resource release.
+
+.. admonition:: Holding a reference to a ``Pipeline``
+   :class: note
+
+   A ``Pipeline`` cleans itself up automatically, so holding a reference to one
+   is safe. When the object is garbage collected, a ``weakref.finalize`` drains
+   and stops its background thread — and any worker processes or subinterpreters
+   it spawned — even if you never called :py:meth:`Pipeline.stop`. And if a
+   reference survives until the program ends, :py:meth:`Pipeline.start` has
+   registered a hook that stops any still-running pipeline at the very start of
+   interpreter finalization, so the process does not hang at exit (see
+   :py:meth:`Pipeline.start` for how that ordering works). Holding a reference is
+   in fact **required** to re-iterate a ``continuous=True`` source across epochs:
+   the same ``Pipeline`` must stay alive to be iterated again.
+
+   You should still release the reference (or call :py:meth:`Pipeline.stop`) once
+   you are done, so the worker processes and memory are freed promptly rather than
+   lingering until GC — but this is resource hygiene, not a requirement to avoid a
+   hang.
 
    .. code-block::
 
       class DataLoader:
-          ...
+          def __init__(self) -> None:
+              # Safe (and, for a continuous source, required) to keep the
+              # pipeline as an attribute so it is reused across epochs.
+              self._pipeline = self.get_pipeline(...)
 
           def __iter__(self) -> Iterator[T]:
-              # 👍 Leave the `pipeline` variable as a local variable.
-              pipeline = self.get_pipeline(...)
-              # So that the `pipeline` will get garbage collected after the
-              # iterator goes out of the scope.
+              yield from self._pipeline.get_iterator(...)
 
-              with pipeline.auto_stop():
-                  yield from pipeline.get_iterator(...)
+          def close(self) -> None:
+              # Optional: drop the reference when done so the pipeline's
+              # resources are freed promptly instead of at GC / interpreter exit.
+              self._pipeline = None
 
-              # The reference count of the `pipeline` object goes to zero
-              # here, so it will be garbage collected.
+   Some frameworks stash the dataloader on a long-lived object.
+   `TorchTNT <https://pytorch.org/tnt/>`_, for example, keeps a strong reference
+   to the dataloader on its ``State`` (``PhaseState._dataloader``) until the
+   process exits. This no longer hangs the run — the shutdown hook cleans the
+   pipeline up at exit — but if you want its worker processes and memory released
+   as soon as training ends (e.g. between fit and eval) rather than at exit, you
+   can clear those references and force a collection with a callback:
 
-   .. code-block::
+   .. code-block:: python
 
-      class DataLoader:
-          ...
+      import gc
 
-          def __iter__(self) -> Iterator[T]:
-              # 🚫 Do not assign the pipeline to the object.
-              self.pipeline = self.get_pipeline(...)
-              #
-              # The pipeline won't get garbage collected until
-              # the DataLoader instance goes out of scope,
-              # which might cause dead-lock when Python tries to exit.
+      from torchtnt.framework.callback import Callback
+      from torchtnt.framework.state import State
+      from torchtnt.framework.unit import TEvalUnit, TPredictUnit, TTestUnit, TTrainUnit
+
+
+      class DetachDataloaderCallback(Callback):
+          """Optional: drop TNT's dataloader references at the end of training so
+          the SPDL ``Pipeline`` (and its workers) are released promptly, instead of
+          lingering on ``State`` until the pipeline is cleaned up at exit."""
+
+          def on_train_end(self, state: State, unit: TTrainUnit) -> None:
+              self._detach(state)
 
-              with self.pipeline.auto_stop():
-                  yield from self.pipeline.get_iterator(...)
+          def on_exception(
+              self,
+              state: State,
+              unit: TTrainUnit | TEvalUnit | TPredictUnit | TTestUnit,
+              exc: BaseException,
+          ) -> None:
+              # on_train_end does not fire on failure/preemption, so reap here too.
+              self._detach(state)
 
-              # The `pipeline` object won't get garbage collected here.
+          def _detach(self, state: State) -> None:
+              for phase_state in (state.train_state, state.eval_state):
+                  if phase_state is not None:
+                      phase_state._dataloader = None  # pyre-ignore[8]
+              gc.collect()
 
 .. admonition:: Calling ``iter`` on Pipeline
-   :class: danger
+   :class: note
 
-   We recommend to not call the :py:func:`iter` function
-   on a ``Pipeline`` object.
-   It can prevent the :py:meth:`Pipeline.stop` method from being called
-   at the right time.
-   It in turn might make the Python interpreter hang at exit.
+   Prefer not to call the :py:func:`iter` function on a ``Pipeline`` object and
+   keep the resulting iterator around. Doing so delays :py:meth:`Pipeline.stop`
+   until the iterator is collected, keeping the background thread and worker
+   processes running longer than necessary. (It will not hang the interpreter at
+   exit — the hook registered by :py:meth:`Pipeline.start` covers that — but it
+   holds resources needlessly.)
 
    Say you wrap a ``Pipeline`` to create a class that resembles conventional
    ``DataLoader``.
@@ -194,7 +234,6 @@ at the end of the Python interpreter, preventing Python from exiting.
       # the pipeline won't be shutdown until the `ite` variable
       # goes out of scope. When does that happen??
 
-   The ``Pipeline.stop`` is not called until the garbage collector deletes
-   the object.
-   It might cause a deadlock, and prevents Python interpreter from
-   exiting.
+   Until then, ``Pipeline.stop`` is deferred to whenever the garbage collector
+   deletes the object, so the background thread and workers keep holding
+   resources longer than needed.

src/spdl/pipeline/_components/_subprocess_pipe.py

@@ -39,6 +39,7 @@
 
 import asyncio
 import queue as _queue
+import threading
 import time
 from concurrent.futures import Executor, ThreadPoolExecutor
 from typing import Any
@@ -78,14 +79,28 @@
 # observe cancellation between polls instead of parking a thread on an indefinite get.
 _GET_TIMEOUT: float = 0.5
 
+# How long a blocking put may wait before re-checking the teardown flag. A put onto a *full*
+# worker queue (backpressure once the consumer stops) cannot be released by pool teardown -- an
+# mp.Queue putter blocked on the queue's semaphore is not woken by closing the queue or
+# terminating the consumer workers -- so an indefinite put would park this pool thread forever and
+# hang interpreter exit (a non-daemon executor thread joined by concurrent.futures at shutdown).
+_PUT_TIMEOUT: float = 0.5
+
 # Fixed bound (15 min) on how long the collector waits for any worker message before assuming a
 # worker died abruptly and raising, instead of hanging forever. Comfortably above any per-stage
 # latency in a data loader. Not user-configurable for now.
 _WORKER_STALL_TIMEOUT: float = 900.0
 
 
-def _put(q: Any, msg: tuple[int, Any]) -> None:
-    q.put(msg)
+def _put(q: Any, msg: tuple[int, Any], stop: threading.Event) -> None:
+    # Bounded, interruptible put: poll so this pool thread wakes to observe teardown (``stop``)
+    # and exit, rather than parking forever on a full queue (see ``_PUT_TIMEOUT``).
+    while not stop.is_set():
+        try:
+            q.put(msg, timeout=_PUT_TIMEOUT)
+            return
+        except _queue.Full:
+            continue
 
 
 def _drain_one(q: Any) -> tuple[int, Any] | None:
@@ -123,6 +138,7 @@ async def _feed(
     executor: Executor,
     abort: asyncio.Event,
     feeder_idle: asyncio.Event,
+    put_stop: threading.Event,
 ) -> None:
     """Round-robin items across the per-worker queues, then end every worker's session.
 
@@ -165,13 +181,18 @@ async def _feed(
                 feeder_idle.clear()
             if is_eof(item):
                 break
-            await loop.run_in_executor(executor, _put, in_qs[i % n], (_ITEM, item))
+            await loop.run_in_executor(
+                executor, _put, in_qs[i % n], (_ITEM, item), put_stop
+            )
             i += 1
     finally:
         abort_wait.cancel()
     # Concurrent so a full/slow worker queue does not block the markers to the others.
     await asyncio.gather(
-        *(loop.run_in_executor(executor, _put, q, (_SESSION_END, None)) for q in in_qs)
+        *(
+            loop.run_in_executor(executor, _put, q, (_SESSION_END, None), put_stop)
+            for q in in_qs
+        )
     )
 
 
@@ -234,6 +255,7 @@ async def _feed_continuous(
     executor: Executor,
     epoch_barrier: asyncio.Event,
     feeder_idle: asyncio.Event,
+    put_stop: threading.Event,
 ) -> None:
     """Round-robin items to per-worker queues; broadcast and barrier each epoch boundary.
 
@@ -251,7 +273,7 @@ async def _feed_continuous(
     async def _broadcast(msg: tuple[int, Any]) -> None:
         # Concurrent so a full/slow worker queue does not block the broadcast to the others.
         await asyncio.gather(
-            *(loop.run_in_executor(executor, _put, q, msg) for q in in_qs)
+            *(loop.run_in_executor(executor, _put, q, msg, put_stop) for q in in_qs)
         )
 
     i = 0
@@ -268,7 +290,9 @@ async def _broadcast(msg: tuple[int, Any]) -> None:
             await _broadcast((_EPOCH, None))
             await epoch_barrier.wait()
             continue
-        await loop.run_in_executor(executor, _put, in_qs[i % n], (_ITEM, item))
+        await loop.run_in_executor(
+            executor, _put, in_qs[i % n], (_ITEM, item), put_stop
+        )
         i += 1
 
 
@@ -352,12 +376,16 @@ async def _subprocess_pipeline(
     # it to tell input starvation (no work dispatched, no worker message expected) apart from an
     # unresponsive worker, so its stall guard does not fire spuriously on a slow/idle source.
     feeder_idle = asyncio.Event()
+    # Signals threads parked in a blocking ``_put`` to stop and exit on teardown. Needed because a
+    # put onto a full worker queue cannot be released by pool teardown (see ``_PUT_TIMEOUT``), so
+    # without it such a thread would outlive this stage and hang interpreter exit.
+    put_stop = threading.Event()
     async with _queue_stage_hook(output_queue):
         if handle.continuous:
             epoch_barrier = asyncio.Event()
             feeder = create_task(
                 _feed_continuous(
-                    input_queue, in_qs, executor, epoch_barrier, feeder_idle
+                    input_queue, in_qs, executor, epoch_barrier, feeder_idle, put_stop
                 )
             )
             collector = create_task(
@@ -374,7 +402,7 @@ async def _subprocess_pipeline(
             # Set by the collector on a worker error so the feeder stops forwarding new items.
             abort = asyncio.Event()
             feeder = create_task(
-                _feed(input_queue, in_qs, executor, abort, feeder_idle)
+                _feed(input_queue, in_qs, executor, abort, feeder_idle, put_stop)
             )
             collector = create_task(
                 _collect(out_q, num_workers, output_queue, executor, abort, feeder_idle)
@@ -387,6 +415,8 @@ async def _subprocess_pipeline(
             await asyncio.gather(feeder, collector, return_exceptions=True)
             raise
         finally:
-            # Don't wait on threads still parked in a blocking get/put — the pool teardown
-            # unblocks them. ``cancel_futures`` discards anything not yet started.
+            # Release any thread parked in a blocking ``_put`` so it exits instead of outliving
+            # this stage, then drop the executor. ``cancel_futures`` discards anything not yet
+            # started; a still-parked get self-releases within ``_GET_TIMEOUT``.
+            put_stop.set()
             executor.shutdown(wait=False, cancel_futures=True)