mp_prefetch worker processes keep running after the main process is killed (SIGKILL/SIGTERM)

Start with start_prefetch in grain/_src/python/dataset/transformations/process_prefetch.py, where the worker is created with daemon=True, and with the worker loop near line 240 and add_element_to_queue in ipc/queue.py, which only exit when should_stop is set. Run repro.py to see the workers survive a SIGKILL of the parent. Done means both rounds report no survivors without a manual cleanup step. The fix direction is contested (the three options in the issue), so get a maintainer decision before writing code.

Versions

grain 0.2.18 (PyPI), numpy 2.5.3, CPython 3.12.6, Linux 6.8 x86_64. On main at 2792d67 the code paths cited below are the same.

What happens

When the process iterating an mp_prefetch pipeline is killed with SIGKILL or SIGTERM, its worker processes do not exit. The kernel reparents them to the nearest subreaper (here systemd --user), and they stay alive holding their memory. In our training jobs we have seen 4 workers still alive 13 minutes after the parent died, at about 700 MB each, until we killed them by hand.

Why daemon=True does not cover it

ProcessPrefetchDatasetIterator.start_prefetch starts the worker with daemon=True (grain/_src/python/dataset/transformations/process_prefetch.py:454-478 on main). multiprocessing only terminates daemonic children from its atexit handler during a normal interpreter exit. SIGKILL never runs that handler. Neither does SIGTERM under Python's default disposition, or the OOM killer.

The worker then has no way to notice. Its loop is while not should_stop.is_set() (process_prefetch.py:240). When the output queue is full it retries in ipc/queue.py:36 add_element_to_queue, which loops while not should_stop(). Only the parent sets should_stop, in _stop_prefetch (process_prefetch.py:563). Nothing in grain/ sets a parent-death signal or checks parent liveness: git grep -E 'pdeathsig|PDEATHSIG|prctl|getppid|parent_process\(' on main finds nothing.

Reproduction

repro.py (Grain and NumPy only). It runs the pipeline in a child process, kills that child, waits 5 s, reports the child's former children, and then kills the survivors so nothing is left behind:

"""mp_prefetch workers outlive a parent that is killed (SIGKILL / SIGTERM).

Usage: python repro.py

The script starts a child Python process that iterates an `mp_prefetch`
pipeline, kills that child with SIGKILL (then, in a second round, SIGTERM),
waits a few seconds, and reports which of the child's worker processes are
still alive and who their parent is now. It then kills the survivors itself,
so it leaves no process behind. Linux only (reads /proc).
"""

import os
import signal
import subprocess
import sys
import time

NUM_WORKERS = 2
WAIT_AFTER_KILL_S = 5.0


def child() -> None:
  """Runs an mp_prefetch pipeline, then idles holding the iterator open."""
  import grain
  import numpy as np

  ds = (
      grain.MapDataset.range(1_000_000)
      .map(lambda i: np.full((1024,), i, dtype=np.int64))
      .to_iter_dataset()
      .mp_prefetch(grain.MultiprocessingOptions(num_workers=NUM_WORKERS))
  )
  it = iter(ds)
  for _ in range(10):
    next(it)
  print("READY", flush=True)
  while True:  # The training loop is "still running" when the kill arrives.
    next(it)
    time.sleep(0.01)


def proc_status(pid: int) -> dict[str, str] | None:
  try:
    with open(f"/proc/{pid}/status") as f:
      lines = f.read().splitlines()
  except FileNotFoundError:
    return None
  fields = dict(line.split(":\t", 1) for line in lines if ":\t" in line)
  if fields.get("State", "").startswith("Z"):
    return None  # A zombie has exited; it only awaits reaping.
  return fields


def children_of(pid: int) -> list[int]:
  found = []
  for entry in os.listdir("/proc"):
    if entry.isdigit():
      status = proc_status(int(entry))
      if status is not None and status.get("PPid") == str(pid):
        found.append(int(entry))
  return sorted(found)


def cmdline(pid: int) -> str:
  try:
    with open(f"/proc/{pid}/cmdline", "rb") as f:
      return f.read().replace(b"\0", b" ").decode()
  except FileNotFoundError:
    return "<gone>"


def run_round(sig: signal.Signals) -> int:
  print(f"\n=== round: kill the parent with {sig.name} ===")
  proc = subprocess.Popen(
      [sys.executable, __file__, "--child"], stdout=subprocess.PIPE, text=True
  )
  assert proc.stdout.readline().strip() == "READY"
  workers = children_of(proc.pid)
  print(f"parent pid={proc.pid}; its child processes: {workers}")
  for pid in workers:
    print(f"  {pid}: {cmdline(pid)[:90]}")

  os.kill(proc.pid, sig)
  proc.wait()
  print(f"parent exited with returncode={proc.returncode}")
  time.sleep(WAIT_AFTER_KILL_S)

  survivors = []
  for pid in workers:
    status = proc_status(pid)
    if status is None:
      print(f"  {pid}: gone")
    else:
      survivors.append(pid)
      print(
          f"  {pid}: ALIVE {WAIT_AFTER_KILL_S:.0f}s after the parent died,"
          f" PPid={status['PPid']} ({cmdline(int(status['PPid'])).strip()}),"
          f" VmRSS={status.get('VmRSS', '?').strip()}"
      )

  # Clean up what the repro started. Kill the prefetch workers only: the
  # multiprocessing resource tracker exits by itself once they are gone, and
  # unlinks the queues' named semaphores on the way out. Killing it as well
  # would leave those entries in /dev/shm.
  trackers = [pid for pid in survivors if "resource_tracker" in cmdline(pid)]
  for pid in survivors:
    if pid not in trackers:
      os.kill(pid, signal.SIGKILL)
  deadline = time.monotonic() + 30
  while time.monotonic() < deadline and any(
      proc_status(pid) is not None for pid in survivors
  ):
    time.sleep(0.2)
  left = [pid for pid in survivors if proc_status(pid) is not None]
  for pid in left:
    os.kill(pid, signal.SIGKILL)
  print(
      f"cleanup: killed {len(survivors) - len(trackers)} worker(s); resource"
      f" tracker exited by itself: {not any(p in left for p in trackers)};"
      f" force-killed afterwards: {left}"
  )
  return len(survivors)


def main() -> None:
  import grain
  import numpy as np

  print(f"python {sys.version.split()[0]}, grain {grain.__version__},"
        f" numpy {np.__version__}, platform {sys.platform}")
  shm_before = set(os.listdir("/dev/shm"))
  total = sum(run_round(sig) for sig in (signal.SIGKILL, signal.SIGTERM))
  print(f"\nsurvivors across both rounds: {total}")
  time.sleep(1.0)
  print(f"/dev/shm entries left behind: {len(set(os.listdir('/dev/shm')) - shm_before)}")


if __name__ == "__main__":
  if sys.argv[1:] == ["--child"]:
    child()
  else:
    main()

Command, in an empty directory:

env -i HOME=$HOME PATH=/usr/bin:/bin:$HOME/.local/bin uv run --isolated --no-project --python 3.12 --with grain==0.2.18 --with numpy python repro.py

Output (interpreter paths shortened; seven absl "Failed to load jax profiler" warnings and the resource tracker's leaked-semaphore warnings from the cleanup step omitted):

python 3.12.6, grain 0.2.18, numpy 2.5.3, platform linux

=== round: kill the parent with SIGKILL ===
parent pid=902437; its child processes: [902486, 902489, 902490]
  902486: .../bin/python -c from multiprocessing.resource_tra
  902489: .../bin/python -c from multiprocessing.spawn import
  902490: .../bin/python -c from multiprocessing.spawn import
parent exited with returncode=-9
  902486: ALIVE 5s after the parent died, PPid=2323 (/usr/lib/systemd/systemd --user), VmRSS=16940 kB
  902489: ALIVE 5s after the parent died, PPid=2323 (/usr/lib/systemd/systemd --user), VmRSS=59384 kB
  902490: ALIVE 5s after the parent died, PPid=2323 (/usr/lib/systemd/systemd --user), VmRSS=59352 kB
cleanup: killed 2 worker(s); resource tracker exited by itself: True; force-killed afterwards: []

=== round: kill the parent with SIGTERM ===
parent pid=903320; its child processes: [903380, 903383, 903384]
  ...
parent exited with returncode=-15
  903380: ALIVE 5s after the parent died, PPid=2323 (/usr/lib/systemd/systemd --user), VmRSS=17044 kB
  903383: ALIVE 5s after the parent died, PPid=2323 (/usr/lib/systemd/systemd --user), VmRSS=59388 kB
  903384: ALIVE 5s after the parent died, PPid=2323 (/usr/lib/systemd/systemd --user), VmRSS=59320 kB
cleanup: killed 2 worker(s); resource tracker exited by itself: True; force-killed afterwards: []

survivors across both rounds: 6
/dev/shm entries left behind: 0

Both workers survive each kill. So does the multiprocessing resource tracker, which holds on until the workers exit. Without a cleanup step like the one in the repro, the workers run until someone finds them.

Impact

• Memory held after the job is gone. A job killed by the OOM killer, by a scheduler's cancellation, or by timeout leaves its workers holding whatever they had decoded. That memory is not returned until someone notices and kills them. After an OOM kill this is the worst case, because the memory the kill was meant to free stays taken.
• Stale writes. An orphaned worker keeps running its pipeline. If the pipeline writes (a cache, a decoded copy, np.lib.format.open_memmap(mode="w+")), a restarted job that reuses the same path can have rows overwritten by the previous job's orphans.

Possible directions

I am not sure which of these fits Grain best:

1. A parent-death signal set in the worker, prctl(PR_SET_PDEATHSIG, SIGKILL) on Linux. One caveat: the kernel sends it when the thread that created the child exits, not the process. Grain spawns each worker from the shard's grain-thread-prefetch-* thread (interleave.py:483/:493, prefetch.py:579), and that thread is joined on close, set_state, _set_next_index and on an element error (prefetch.py:641). In a quick test, setting it from worker_init_fn made set_state on a live iterator kill the workers and hang the next next(). It also needs a getppid() re-check after prctl, for a parent that dies before the worker gets that far. And it is Linux only.
2. A parent liveness check in the worker loop. multiprocessing.parent_process().sentinel becomes readable when the parent process exits, whichever thread started the child, and it works on macOS too. A worker could wait on it in a daemon thread, or poll it alongside should_stop. In the same test, a daemon thread blocking on that sentinel ended the workers within about 0.1 s of a SIGKILL. It changed nothing on normal close, set_state or exhaustion.
3. A process group / session for the workers, so that a supervisor can kill the group. This only helps when something outside the process does the killing.

Happy to test a patch.