Move the implementation detail of iterate_in_subprocess to development note (#1076)

Differential Revision: D86097638

Co-authored-by: Moto Hira <[email protected]>

Author: moto-meta

docs/source/index.rst

@@ -23,15 +23,15 @@ Please use the following BibTex for citing our project if you find it useful.
 .. code-block:: text
 
    @misc{hira2025scalableperformantdataloading,
-      title={Scalable and Performant Data Loading}, 
+      title={Scalable and Performant Data Loading},
       author={Moto Hira and Christian Puhrsch and Valentin Andrei and Roman Malinovskyy and Gael Le Lan and Abhinandan Krishnan and Joseph Cummings and Miguel Martin and Gokul Gunasekaran and Yuta Inoue and Alex J Turner and Raghuraman Krishnamoorthi},
       year={2025},
       eprint={2504.20067},
       archivePrefix={arXiv},
       primaryClass={cs.DC},
-      url={https://arxiv.org/abs/2504.20067}, 
+      url={https://arxiv.org/abs/2504.20067},
    }
-  
+
 .. toctree::
    :hidden:
 

docs/source/notes/index.rst

@@ -4,3 +4,4 @@ Development Notes
 .. toctree::
 
    pipeline_impl
+   remote_iterable_protocol

docs/source/notes/remote_iterable_protocol.rst

@@ -0,0 +1,95 @@
+Remote Iterable Protocol
+========================
+
+.. py:currentmodule:: spdl.pipeline._iter_utils
+
+Manipulting an iterable object in a remote location (subprocess or subinterpreter) requires
+somewhat elaborated state control.
+The following section go over the implementation detail.
+
+Worker State
+------------
+
+The iterable object is manipulated in the worker process.
+The worker process has three states, ``Initialization``, ``Stand By`` and ``Iteration``.
+
+The Initialization state performs global initialization and create the iterable object.
+When the Initialization completes, the worker transition to ``Stand By`` mode,
+where it waits for a command from the parent process.
+
+The command can be ``START_ITERATION`` or ``ABORT``.
+When the ``START_ITERATION`` is received, the worker process transition to the Iteration mode.
+
+In the Iteration mode, the worker creates an iterator object from the iterable, then executes it.
+The resulting data are put in the queue, which the parent process is watching.
+
+The following diagram illustrates worker's state transition in simplified manner.
+Detailed diagram alongside the actual implementation is found in :py:func:`_execute_iterable`.
+
+.. mermaid::
+
+    stateDiagram-v2
+        state Parent {
+            p1: Start Iteration
+            p2: Iterate on the result
+            state pf <<fork>>
+            state pj <<join>>
+
+            [*] --> p1
+            p1 --> pf
+            pf --> pj: Wait for worker process
+            pj -->  p2
+            p2 --> [*]
+        }
+
+        state Worker {
+            state wf <<fork>>
+            w0: Initialization
+            w1: Stand By
+            w2: Iteration
+
+            [*]--> w0
+            w0 --> w1: Success
+            w0 --> [*]: Fail
+            w1 --> wf: Iteration started
+            wf --> w2
+            w2 --> w1: Iteration completed
+
+            w1 --> [*]: Abort
+            w2 --> [*]: Fail / Abort
+        }
+        pf --> w1: Issue START_ITERATION command
+        wf --> pj: Notify ITERATION_STARTED
+        w2 --> p2: Results passed via queue
+
+Helper functions and data structures
+-------------------------------------
+
+The follosing functions and data structures are used to implement
+the :py:func:`~spdl.pipeline.iterate_in_subprocess` function.
+They are not public interface, but the logic is sufficiently elaborated,
+and it is helpful to have them in the documentation, so they are listed here.
+
+.. autoclass:: _Cmd
+   :noindex:
+   :members:
+
+.. autoclass:: _Status
+   :noindex:
+   :members:
+
+.. autofunction:: _enter_iteration_mode()
+   :noindex:
+
+.. autofunction:: _execute_iterable()
+   :noindex:
+
+.. autofunction:: _drain()
+   :noindex:
+
+.. autofunction:: _iterate_results()
+   :noindex:
+
+.. autoclass:: _SubprocessIterable()
+   :noindex:
+   :members: __iter__

src/spdl/pipeline/_iter_utils/_common.py

@@ -63,7 +63,7 @@ def empty(self) -> bool:
 class _Cmd(enum.IntEnum):
     """_Cmd()
     Command issued from the parent process to the worker process in
-    :py:func:`iterate_in_subprocess`.
+    :py:func:`~spdl.pipeline.iterate_in_subprocess`.
     """
 
     ABORT = enum.auto()
@@ -94,7 +94,7 @@ class _Cmd(enum.IntEnum):
 # Final status of the iterator
 class _Status(enum.IntEnum):
     """_Status()
-    Status reported by the worker process in :py:func:`iterate_in_subprocess`.
+    Status reported by the worker process in :py:func:`~spdl.pipeline.iterate_in_subprocess`.
     """
 
     UNEXPECTED_CMD_RECIEVED = enum.auto()
@@ -155,7 +155,7 @@ class _Msg(Generic[T]):
 def _drain(q: _Queue[Any]) -> None:
     """Drain a queue by removing all items.
 
-    Works with both multiprocessing.Queue and concurrent.interpreters.Queue.
+    Works with both :py:func:`multiprocessing.Queue` and :py:func:`concurrent.interpreters.Queue`.
     """
     while True:
         try:
@@ -170,7 +170,7 @@ def _execute_iterable(
     fn: Callable[[], Iterable[T]],
     initializers: Sequence[Callable[[], None]] | None,
 ) -> None:
-    """Worker implementation for :py:func:`iterate_in_subprocess`.
+    """Worker implementation for :py:func:`~spdl.pipeline.iterate_in_subprocess`.
 
     The following diagram illustrates the state transition with more details.
 

src/spdl/pipeline/_iter_utils/_subprocess.py

@@ -131,98 +131,7 @@ def iterate_in_subprocess(
        - :py:func:`run_pipeline_in_subprocess` for runinng a :py:class:`Pipeline` in
          a subprocess
        - :ref:`parallelism-performance` for the context in which this function was created.
-
-
-    Implementation Detail
-    ---------------------
-
-    .. py:currentmodule:: spdl.pipeline._iter_utils
-
-    Manipulting an iterable object in a subprocess requires somewhat elaborated state
-    control.
-    The following section go over the implementation detail.
-
-    **Wroker State**
-
-    The iterable object is manipulated in the worker process.
-    The worker process has three states, "Initialization", "Stand By" and "Iteration".
-    The Initialization state performs global initialization and create the iterable object.
-    When the Initialization completes, the worker transition to Stand By mode, where
-    it waits for a command from the parent process. The command can be "START_ITERATION"
-    or "ABORT".
-    When the "START_ITERATION" is received, the worker process transition to the
-    Iteration mode. In the Iteration mode, the worker creates an iterator object from
-    the iterable, then executes it.
-    The resulting data are put in the queue, which the parent process is watching.
-
-    The following diagram illustrates worker's state transition in simplified manner.
-    Detailed diagram alongside the actual implementation is found in
-    :py:func:`_execute_iterable`.
-
-    .. mermaid::
-
-       stateDiagram-v2
-        state Parent {
-            p1: Start Iteration
-            p2: Iterate on the result
-            state pf <<fork>>
-            state pj <<join>>
-
-            [*] --> p1
-            p1 --> pf
-            pf --> pj: Wait for worker process
-            pj -->  p2
-            p2 --> [*]
-        }
-
-        state Worker {
-            state wf <<fork>>
-            w0: Initialization
-            w1: Stand By
-            w2: Iteration
-
-            [*]--> w0
-            w0 --> w1: Success
-            w0 --> [*]: Fail
-            w1 --> wf: Iteration started
-            wf --> w2
-            w2 --> w1: Iteration completed
-
-            w1 --> [*]: Abort
-            w2 --> [*]: Fail / Abort
-        }
-        pf --> w1: Issue START_ITERATION command
-        wf --> pj: Notify ITERATION_STARTED
-        w2 --> p2: Results passed via queue
-
-    Helper functions and data structures
-    ^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^
-
-    The follosing functions and data structures are used to implement the
-    :py:func:`iterate_in_subprocess` function.
-    They are not public interface, but the logic is sufficiently elaborated,
-    and it is helpful to have them in the documentation, so they are listed here.
-
-    .. autoclass:: _Cmd
-       :noindex:
-       :members:
-
-    .. autoclass:: _Status
-       :noindex:
-       :members:
-
-    .. autofunction:: _execute_iterable()
-       :noindex:
-
-    .. autofunction:: _enter_iteration_mode()
-       :noindex:
-
-    .. autofunction:: _iterate_results()
-       :noindex:
-
-    .. autoclass:: _SubprocessIterable()
-       :noindex:
-       :members: __iter__
+       - :doc:`../notes/remote_iterable_protocol` for implementation details
     """
     initializers = (
         None