wehub-resource-sync
303 lines
9.5 KiB
Python
303 lines
9.5 KiB
Python
import logging
|
||
import os
|
||
import subprocess
|
||
import sys
|
||
import time
|
||
|
||
import pytest
|
||
|
||
import ray
|
||
from ray._common.test_utils import wait_for_condition
|
||
|
||
import psutil
|
||
|
||
logger = logging.getLogger(__name__)
|
||
|
||
|
||
@pytest.fixture
|
||
def enable_subreaper():
|
||
os.environ["RAY_kill_child_processes_on_worker_exit_with_raylet_subreaper"] = "true"
|
||
yield
|
||
del os.environ["RAY_kill_child_processes_on_worker_exit_with_raylet_subreaper"]
|
||
|
||
|
||
@pytest.fixture
|
||
def enable_pg_cleanup():
|
||
os.environ["RAY_process_group_cleanup_enabled"] = "true"
|
||
yield
|
||
del os.environ["RAY_process_group_cleanup_enabled"]
|
||
|
||
|
||
def sleep_forever():
|
||
while True:
|
||
time.sleep(10000)
|
||
|
||
|
||
def get_process_info(pid):
|
||
# may raise psutil.NoSuchProcess
|
||
process = psutil.Process(pid)
|
||
return {
|
||
"PID": process.pid,
|
||
"Name": process.name(),
|
||
"Status": process.status(),
|
||
"CPU Times": process.cpu_times(),
|
||
"Memory Info": process.memory_info(),
|
||
}
|
||
|
||
|
||
@ray.remote
|
||
class BedMaker:
|
||
def make_sleeper(self):
|
||
p = subprocess.Popen(["sleep", "1000"]) # inherits PGID
|
||
return p.pid
|
||
|
||
def spawn_daemon(self):
|
||
# Spawns a bash script (shell=True) that starts a daemon process
|
||
# which sleeps 1000s. The bash exits immediately, leaving the daemon
|
||
# running in the background. We don't want to kill the daemon when
|
||
# the actor is alive; we want to kill it after the actor is killed.
|
||
command = "nohup sleep 1000 >/dev/null 2>&1 & echo $!"
|
||
output = subprocess.check_output(command, shell=True, text=True)
|
||
return int(output.strip())
|
||
|
||
def my_pid(self):
|
||
return os.getpid()
|
||
|
||
|
||
@pytest.mark.skipif(
|
||
sys.platform != "linux",
|
||
reason="Orphan process killing only works on Linux.",
|
||
)
|
||
def test_ray_kill_can_kill_subprocess(shutdown_only):
|
||
"""
|
||
This works becuase of kill_child_processes_on_worker_exit.
|
||
Even if kill_child_processes_on_worker_exit_with_raylet_subreaper
|
||
is not set, the worker will still kill its subprocesses.
|
||
"""
|
||
ray.init()
|
||
b = BedMaker.remote()
|
||
pid = ray.get(b.make_sleeper.remote())
|
||
|
||
# ray.kill can kill subprocesses.
|
||
logger.info(get_process_info(pid)) # shows the process
|
||
ray.kill(b)
|
||
time.sleep(1)
|
||
with pytest.raises(psutil.NoSuchProcess):
|
||
logger.info(get_process_info(pid)) # subprocess killed
|
||
|
||
|
||
@pytest.mark.skipif(
|
||
sys.platform != "linux",
|
||
reason="Orphan process killing only works on Linux.",
|
||
)
|
||
def test_sigkilled_worker_can_kill_subprocess(enable_subreaper, shutdown_only):
|
||
ray.init()
|
||
# sigkill'd actor can't kill subprocesses
|
||
b = BedMaker.remote()
|
||
pid = ray.get(b.make_sleeper.remote())
|
||
actor_pid = ray.get(b.my_pid.remote())
|
||
|
||
logger.info(get_process_info(pid)) # shows the process
|
||
psutil.Process(actor_pid).kill() # sigkill
|
||
time.sleep(11) # unowned processes are killed every 10s.
|
||
with pytest.raises(psutil.NoSuchProcess):
|
||
logger.info(get_process_info(pid)) # subprocess killed
|
||
|
||
|
||
@pytest.mark.skipif(
|
||
sys.platform != "linux",
|
||
reason="Orphan process killing only works on Linux.",
|
||
)
|
||
def test_daemon_processes_not_killed_until_actor_dead(enable_subreaper, shutdown_only):
|
||
ray.init()
|
||
# sigkill'd actor can't kill subprocesses
|
||
b = BedMaker.remote()
|
||
daemon_pid = ray.get(b.spawn_daemon.remote())
|
||
actor_pid = ray.get(b.my_pid.remote())
|
||
|
||
# The pid refers to a daemon process that should not be killed, although
|
||
# it's already reparented to the core worker.
|
||
time.sleep(11) # even after a cycle of killing...
|
||
assert psutil.Process(daemon_pid).ppid() == actor_pid
|
||
|
||
psutil.Process(actor_pid).kill() # sigkill
|
||
time.sleep(11) # unowned processes are killed every 10s.
|
||
with pytest.raises(psutil.NoSuchProcess):
|
||
logger.info(get_process_info(daemon_pid)) # subprocess killed
|
||
|
||
|
||
@pytest.mark.skipif(
|
||
sys.platform != "linux",
|
||
reason="Orphan process killing only works on Linux.",
|
||
)
|
||
def test_default_sigchld_handler(enable_subreaper, shutdown_only):
|
||
"""
|
||
Core worker auto-reaps zombies via SIG_IGN. If the user wants to wait for subprocess
|
||
they can add it back.
|
||
"""
|
||
ray.init()
|
||
|
||
@ray.remote
|
||
class A:
|
||
def auto_reap(self):
|
||
"""
|
||
Auto subprocess management. Since the signal handler is set to SIG_IGN
|
||
by the flag, zombies are reaped automatically.
|
||
"""
|
||
process = subprocess.Popen(["true"])
|
||
pid = process.pid
|
||
wait_for_condition(
|
||
lambda: not psutil.pid_exists(pid), retry_interval_ms=100
|
||
)
|
||
|
||
def manual_reap(self):
|
||
"""
|
||
Manual subprocess management. Since the signal handler is set back to
|
||
default, user needs to call `process.wait()` on their own, or the zombie
|
||
process would persist.
|
||
"""
|
||
|
||
import signal
|
||
|
||
signal.signal(signal.SIGCHLD, signal.SIG_DFL)
|
||
|
||
process = subprocess.Popen(["true"])
|
||
pid = process.pid
|
||
time.sleep(1) # wait for the process to exit.
|
||
|
||
assert psutil.Process(pid).status() == psutil.STATUS_ZOMBIE
|
||
|
||
process.wait()
|
||
# after reaping, it's gone.
|
||
with pytest.raises(psutil.NoSuchProcess):
|
||
psutil.Process(pid)
|
||
|
||
a = A.remote()
|
||
# order matters, since `manual_reap` sets the signal handler.
|
||
ray.get(a.auto_reap.remote())
|
||
ray.get(a.manual_reap.remote())
|
||
|
||
|
||
@pytest.mark.skipif(
|
||
sys.platform != "linux",
|
||
reason="Orphan process killing only works on Linux.",
|
||
)
|
||
def test_sigkilled_worker_child_process_cleaned_up(enable_pg_cleanup, shutdown_only):
|
||
ray.init()
|
||
# SIGKILL the actor; PG cleanup should terminate the background child.
|
||
b = BedMaker.remote()
|
||
child_pid = ray.get(b.make_sleeper.remote())
|
||
actor_pid = ray.get(b.my_pid.remote())
|
||
|
||
logger.info(get_process_info(child_pid)) # shows the process
|
||
psutil.Process(actor_pid).kill() # sigkill
|
||
wait_for_condition(lambda: not psutil.pid_exists(child_pid), retry_interval_ms=100)
|
||
with pytest.raises(psutil.NoSuchProcess):
|
||
logger.info(get_process_info(child_pid))
|
||
|
||
|
||
@pytest.mark.skipif(
|
||
sys.platform != "linux",
|
||
reason="Orphan process killing only works on Linux.",
|
||
)
|
||
def test_background_child_survives_while_actor_alive_then_killed_with_pg_cleanup(
|
||
enable_pg_cleanup, shutdown_only
|
||
):
|
||
ray.init()
|
||
# Spawn a background child that remains in the same PG as the actor.
|
||
b = BedMaker.remote()
|
||
child_pid = ray.get(b.make_sleeper.remote())
|
||
actor_pid = ray.get(b.my_pid.remote())
|
||
|
||
# The background child remains alive while the actor is alive.
|
||
time.sleep(1)
|
||
assert psutil.pid_exists(child_pid)
|
||
|
||
# After the actor is killed, PG cleanup should terminate the background child.
|
||
psutil.Process(actor_pid).kill()
|
||
wait_for_condition(lambda: not psutil.pid_exists(child_pid), retry_interval_ms=100)
|
||
with pytest.raises(psutil.NoSuchProcess):
|
||
logger.info(get_process_info(child_pid))
|
||
|
||
|
||
@pytest.mark.skipif(
|
||
sys.platform != "linux",
|
||
reason="Orphan process killing only works on Linux.",
|
||
)
|
||
def test_detached_setsido_escape_with_pg_cleanup(enable_pg_cleanup, shutdown_only):
|
||
ray.init()
|
||
|
||
@ray.remote
|
||
class A:
|
||
def spawn_detached(self):
|
||
# Detach into a new session (escape worker PG); sleep long.
|
||
return subprocess.Popen(
|
||
[sys.executable, "-c", "import os,time; os.setsid(); time.sleep(1000)"]
|
||
).pid
|
||
|
||
def pid(self):
|
||
return os.getpid()
|
||
|
||
a = A.remote()
|
||
child_pid = ray.get(a.spawn_detached.remote())
|
||
actor_pid = ray.get(a.pid.remote())
|
||
psutil.Process(actor_pid).kill()
|
||
time.sleep(1)
|
||
# Detached child should still be alive (escaped PG cleanup).
|
||
assert psutil.pid_exists(child_pid)
|
||
|
||
|
||
@pytest.mark.skipif(
|
||
sys.platform != "linux" and sys.platform != "darwin",
|
||
reason="Process‑group cleanup is POSIX‑only (Linux/macOS).",
|
||
)
|
||
def test_nested_subprocess_cleanup_with_pg_cleanup(enable_pg_cleanup, shutdown_only):
|
||
"""
|
||
Test that a subprocess spawned by another subprocess is cleaned up when the actor
|
||
is killed.
|
||
"""
|
||
ray.init()
|
||
|
||
@ray.remote
|
||
class NestedSpawner:
|
||
def spawn_nested(self):
|
||
# Create a subprocess that spawns another subprocess.
|
||
proc = subprocess.Popen(
|
||
[
|
||
sys.executable,
|
||
"-c",
|
||
"import subprocess; "
|
||
"subprocess.Popen(['sleep', '150']); "
|
||
"import time; time.sleep(100)",
|
||
],
|
||
text=True,
|
||
)
|
||
child_pid = proc.pid
|
||
# Wait until the subprocess is running.
|
||
wait_for_condition(
|
||
lambda: psutil.pid_exists(child_pid), retry_interval_ms=100
|
||
)
|
||
wait_for_condition(
|
||
lambda: len(psutil.Process(child_pid).children()) > 0,
|
||
retry_interval_ms=100,
|
||
)
|
||
grandchild_pid = psutil.Process(child_pid).children()[0].pid
|
||
return proc.pid, grandchild_pid
|
||
|
||
actor = NestedSpawner.remote()
|
||
child_pid, grandchild_pid = ray.get(actor.spawn_nested.remote())
|
||
|
||
# Both child and grandchild should be alive while the actor is alive.
|
||
assert psutil.pid_exists(child_pid)
|
||
assert psutil.pid_exists(grandchild_pid)
|
||
|
||
del actor
|
||
wait_for_condition(lambda: not psutil.pid_exists(child_pid), retry_interval_ms=100)
|
||
wait_for_condition(
|
||
lambda: not psutil.pid_exists(grandchild_pid), retry_interval_ms=100
|
||
)
|
||
|
||
|
||
if __name__ == "__main__":
|
||
sys.exit(pytest.main(["-v", __file__]))
|