fix(vllm-omni): avoid PDEATHSIG for thread-launched workers

[Wx727]

Summary

Linux associates PR_SET_PDEATHSIG with the thread that starts the child process. vLLM-Omni diffusion workers are started by temporary stage-initialization threads, so they were killed when those threads exited even though the owning process remained alive.

Only arm PDEATHSIG when Process.start() is called from the parent process's main thread. The existing anchor PID/PPID watchdog remains enabled for all children.

Related Issue

Fixes #275

Test Plan

Tested on one NVIDIA H20 with Python 3.12, PyTorch 2.11.0+cu129, vLLM 0.20.0+cu129, and vLLM-Omni 0.20.0.

Model: stabilityai/stable-diffusion-3.5-medium
Recipe: diffusion/sd3/sd3_vllmomni
Dataset: datasets/pickscore/train.txt

python -m unirl.train_diffusion \
  --config-name=diffusion/sd3/sd3_vllmomni \
  num_devices=1 +devices_per_node=1 \
  batch_size=1 sampling.samples_per_prompt=2 \
  sampling.height=384 sampling.width=384 \
  stack.micro_batch_size=1 +num_rollouts=1

Result: generation completed successfully and rollout 1/1 finished.

AR TP-worker termination was not tested; this validation was limited to the one-GPU diffusion reproduction from #275.

Compatibility / Risk

No API, configuration, checkpoint, or data-format changes.

Children started by helper threads now rely on the existing anchor PID/PPID watchdog. Main-thread-started workers retain PDEATHSIG.

Reviewer Notes

The regression was introduced by c1da41b in #260. No overlapping fix PR was found.

Implementation and PR drafting were AI-assisted. I reviewed the diff and ran the GPU reproduction above.

Checklist

• I reviewed the changed code and removed unrelated/generated artifacts.
• I updated tests, docs, and configs where needed, or explained why not.

[CjhHa1]

Confirmed the diagnosis independently, and the thread-scoping explanation in this PR is the correct one.

We hit the same failure from a different direction: a 4-replica layout="separate" vLLM-Omni diffusion topology running diffusion/bagel/bagel_vllmomni with the plain synchronous DiffusionTrainer. Every replica logged AsyncOmniEngine initialized in Ns immediately followed by Diffusion worker(s) died unexpectedly: ['DiffusionWorker-0'], and the first generate failed with DiffusionExecutor is closed. Toggling the prctl call took worker deaths from 5 to 0. So the bug is neither single-GPU- nor SD3-specific, and it reproduces without the async recipe.

Two notes worth putting on the record:

1. The anchor watchdog keeps a process-level reparent check (original_ppid != 1 and os.getppid() != original_ppid -> os._exit(1)). That check is effectively a coarse process-scoped analogue of PDEATHSIG, and the fact that it never fired during our runs is what independently confirms this PR's framing: the spawning process stays alive, only the stage-init thread exits. Had the parent process actually exited, that branch would have reaped the worker within one 5s poll regardless of prctl. Leaving it as-is is correct.

2. install_fate_sharing has exactly one call site, so the new required keyword-only argument is safe. The sibling _DiffrlPatchedTarget / wrap_mp_process_for_children copy in unirl/rollout/engine/sglang_diffusion/_patches/hijack.py never armed PDEATHSIG, so there is no parallel fix needed there.

One non-blocking caveat for future readers: gating on Process.start()'s calling thread is right for spawn and fork because the creating thread is the one calling start(). It would not hold under forkserver, where the child is forked by the forkserver process. vLLM-Omni forces spawn, so this is not a concern today.

We had a duplicate fix in flight (#277) that deleted the prctl call outright. Closing that one in favour of this PR, since it was filed first and its stated rationale is the accurate one.

[CjhHa1]

Settling the mechanism, because the now-closed #277 claimed the opposite in its description: that a standalone probe on this kernel showed PDEATHSIG firing on parent process death but "not on creating-thread exit". If that were true, this PR's main-thread gating would be treating the wrong variable. It is not true — the probe behind that claim hit a race.

Measured on the same kernel as the bug report, 6.6.98-40.6.tl4.x86_64:

armed before creator thread exited: True
creator thread alive: False  parent process pid: 2041693
child alive 6s after creator-thread exit: False
child exitcode: -9  (-9 == SIGKILL from PDEATHSIG)
child armed record: ('armed', rc=0, readback=9, pid=2041697, ppid=2041693)
VERDICT: THREAD-scoped -- creating thread exit killed the child

The parent process stayed alive the whole time; only the creating thread returned. The child took SIGKILL. That is the documented behaviour (prctl(2): "the parent in this case is considered to be the thread that created this process"), and forget_original_parent() in kernel/exit.c sends pdeath_signal on any real-parent task exit, thread or not.

The reason a probe can conclude "process-scoped" is a race that is easy to miss: the child arms PDEATHSIG itself, asynchronously, some time after Process.start() returns. If the creating thread exits first, the reparent has already happened and the signal can never be delivered. Flipping only that one line — creator thread returns immediately instead of waiting for the child to confirm it armed — reproduces the wrong verdict exactly:

armed before creator thread exited: False
child alive 6s after creator-thread exit: True
child exitcode: None
child armed record: ('armed', rc=0, readback=9, ...)   <- armed fine, just too late
ticks observed after arming: 23
VERDICT: PROCESS-scoped -- child survived creating-thread exit

Same binary, same kernel, opposite conclusion, purely from whether the creator thread outlives the child's prctl call.

So the framing in this PR is the correct one, and the Process.start() calling-thread gate is the right place to make the decision. It also explains the reported timing in #275 — model load and dummy warmup succeed, the worker dies as the first real request arrives — since the stage-init thread finishes after warmup, which is precisely when the armed signal fires.

Withdrawing the caveat in my earlier review about needing to keep this narrow for safety reasons unrelated to the mechanism: the mechanism now checks out end to end.

[CjhHa1]

Follow-up data on why gating is the right shape, rather than dropping PDEATHSIG outright.

I had one open question about this PR: if nothing in the current topology is actually created from a parent's main thread, the gate would be a no-op and the simpler "just delete it" change would be equivalent. Measured, that is not the case.

Which thread creates what (vllm_omni 0.20.0, the pinned version)

AsyncOmniEngine._initialize_stage_replicas runs stage init in a ThreadPoolExecutor(thread_name_prefix="stage-init") used as a context manager (async_omni_engine.py:890), so those threads are joined and terminated the moment initialization finishes. That is the trigger, and it fires exactly at the "AsyncOmniEngine initialized" boundary reported in #275.

_initialize_diffusion_replica then branches on use_inline = num_stages == 1 and num_replicas == 1:

• inline (the 1-GPU SD3 repro in [Bug] vLLM-Omni diffusion worker exits on the first rollout request #275): DiffusionEngine.make_engine() runs directly on the stage-init thread, so _launch_workers() creates DiffusionWorker-N from the stage-init thread.
• non-inline: spawn_diffusion_proc() creates StageDiffusionProc from the stage-init thread; inside it run_diffusion_proc executes on that process's main thread, and initialize() -> _launch_workers() creates DiffusionWorker-N from a main thread.

Gate probe

That topology replicated with the three patch variants, real processes, real prctl:

process	created by	pre-fix (c1da41b)	this PR	delete-it
StageDiffusionProc	stage-init thread	armed -> killed (-9)	not armed -> alive	not armed -> alive
DiffusionWorker (inline)	stage-init thread	armed -> killed (-9)	not armed -> alive	not armed -> alive
DiffusionWorker-0	StageDiffusionProc main thread	armed -> killed (cascade)	armed -> alive	not armed -> alive

The gate does fire, on exactly one class, and that class survives. Also worth recording: before the fix the blast radius is wider than #275 reported — StageDiffusionProc itself is SIGKILLed too, not just the workers.

What the retained signal buys

The anchor poll is a Python thread inside the worker, so it needs the GIL. Simulating a worker wedged in a C call with ctypes.PyDLL (which, unlike CDLL, does not release the GIL around the call):

worker state when the parent dies	outcome
PDEATHSIG armed	reaped in 0.25 s
poll only, GIL wedged	still alive after 14 s -> leaked
poll only, GIL free (control)	reaped in 4.00 s

The control shows the watchdog logic itself is sound — it fails specifically when the worker cannot run Python. That is the scenario install_fate_sharing was written for in the first place (Worker_TP* reparenting to init holding ~7.3 GiB each until reaped by hand). Dropping the signal everywhere gives that case up; gating it keeps it exactly where it is safe, because a parent's main thread lives as long as the parent process does.

Caveats

The thread attribution is from reading the 0.20.0 source. The kill/survival numbers are measured on the same kernel as the bug report (6.6.98-40.6.tl4) but against a structural replication of the topology, not a GPU run. I did not verify the AR-side Worker_TP* directly — no vllm in that environment — but they are created inside StageEngineCoreProc the same way DiffusionWorker-0 is created inside StageDiffusionProc, so they should land in the same armed-and-safe bucket.