[diffusion] fix: don't self-kill diffusion worker when PID 1 is the real parent (#31361)
Co-authored-by: Mick <mickjagger19@icloud.com>
This commit is contained in:
@@ -0,0 +1,55 @@
|
|||||||
|
import signal
|
||||||
|
import unittest
|
||||||
|
from unittest.mock import Mock, patch
|
||||||
|
|
||||||
|
from sglang.multimodal_gen import utils
|
||||||
|
|
||||||
|
|
||||||
|
class TestKillItselfWhenParentDied(unittest.TestCase):
|
||||||
|
def _run_linux_case(self, parent_pids):
|
||||||
|
libc = Mock()
|
||||||
|
libc.prctl.return_value = 0
|
||||||
|
|
||||||
|
with (
|
||||||
|
patch.object(utils.sys, "platform", "linux"),
|
||||||
|
patch.object(utils.ctypes, "CDLL", return_value=libc) as cdll,
|
||||||
|
patch.object(utils.os, "getppid", side_effect=parent_pids) as getppid,
|
||||||
|
patch.object(utils.os, "getpid", return_value=12345),
|
||||||
|
patch.object(utils.os, "kill") as kill,
|
||||||
|
):
|
||||||
|
utils.kill_itself_when_parent_died()
|
||||||
|
|
||||||
|
return libc, cdll, getppid, kill
|
||||||
|
|
||||||
|
def test_stable_pid_1_parent_does_not_self_kill(self):
|
||||||
|
libc, cdll, getppid, kill = self._run_linux_case([1, 1])
|
||||||
|
|
||||||
|
cdll.assert_called_once_with("libc.so.6", use_errno=True)
|
||||||
|
libc.prctl.assert_called_once_with(1, signal.SIGKILL)
|
||||||
|
self.assertEqual(getppid.call_count, 2)
|
||||||
|
kill.assert_not_called()
|
||||||
|
|
||||||
|
def test_reparent_to_init_self_kills(self):
|
||||||
|
_, _, _, kill = self._run_linux_case([1000, 1])
|
||||||
|
|
||||||
|
kill.assert_called_once_with(12345, signal.SIGKILL)
|
||||||
|
|
||||||
|
def test_reparent_to_subreaper_self_kills(self):
|
||||||
|
_, _, _, kill = self._run_linux_case([1000, 42])
|
||||||
|
|
||||||
|
kill.assert_called_once_with(12345, signal.SIGKILL)
|
||||||
|
|
||||||
|
def test_non_linux_noop(self):
|
||||||
|
with (
|
||||||
|
patch.object(utils.sys, "platform", "darwin"),
|
||||||
|
patch.object(utils.ctypes, "CDLL") as cdll,
|
||||||
|
patch.object(utils.os, "kill") as kill,
|
||||||
|
):
|
||||||
|
utils.kill_itself_when_parent_died()
|
||||||
|
|
||||||
|
cdll.assert_not_called()
|
||||||
|
kill.assert_not_called()
|
||||||
|
|
||||||
|
|
||||||
|
if __name__ == "__main__":
|
||||||
|
unittest.main()
|
||||||
@@ -533,11 +533,17 @@ def kill_itself_when_parent_died() -> None:
|
|||||||
|
|
||||||
# keep GPU workers tied to the CLI process even if the parent is SIGKILLed
|
# keep GPU workers tied to the CLI process even if the parent is SIGKILLed
|
||||||
PR_SET_PDEATHSIG = 1
|
PR_SET_PDEATHSIG = 1
|
||||||
|
# Capture parent before arming PDEATHSIG: if the parent already died in the
|
||||||
|
# fork->prctl window, PDEATHSIG won't fire, so detect the reparent explicitly.
|
||||||
|
parent_pid = os.getppid()
|
||||||
libc = ctypes.CDLL("libc.so.6", use_errno=True)
|
libc = ctypes.CDLL("libc.so.6", use_errno=True)
|
||||||
if libc.prctl(PR_SET_PDEATHSIG, signal.SIGKILL) != 0:
|
if libc.prctl(PR_SET_PDEATHSIG, signal.SIGKILL) != 0:
|
||||||
err = ctypes.get_errno()
|
err = ctypes.get_errno()
|
||||||
raise OSError(err, os.strerror(err))
|
raise OSError(err, os.strerror(err))
|
||||||
if os.getppid() == 1:
|
# getppid() changing means we were reparented (parent gone). Comparing to the
|
||||||
|
# captured pid instead of "== 1" avoids self-killing when PID 1 is the real
|
||||||
|
# parent (e.g. running as a container's init process).
|
||||||
|
if os.getppid() != parent_pid:
|
||||||
os.kill(os.getpid(), signal.SIGKILL)
|
os.kill(os.getpid(), signal.SIGKILL)
|
||||||
|
|
||||||
|
|
||||||
|
|||||||
Reference in New Issue
Block a user