chore: import upstream snapshot with attribution

python/ray/tests/test_component_failures.py

@@ -0,0 +1,191 @@
+import os
+import signal
+import sys
+import time
+
+import numpy as np
+import pytest
+
+import ray
+from ray._common.test_utils import SignalActor
+from ray._private.test_utils import run_string_as_driver_nonblocking
+
+SIGKILL = signal.SIGKILL if sys.platform != "win32" else signal.SIGTERM
+
+
+# This test checks that when a worker dies in the middle of a get, the raylet will not die.
+def test_dying_worker_get(ray_start_2_cpus):
+    @ray.remote
+    def wait_on_signal(signal_1, signal_2):
+        ray.get(signal_1.send.remote())
+        ray.get(signal_2.wait.remote())
+        return np.ones(200 * 1024, dtype=np.uint8)
+
+    @ray.remote
+    def get_worker_pid():
+        return os.getpid()
+
+    signal_1 = SignalActor.remote()
+    signal_2 = SignalActor.remote()
+
+    x_id = wait_on_signal.remote(signal_1, signal_2)
+    ray.get(signal_1.wait.remote())
+    # Get the PID of the other worker.
+    worker_pid = ray.get(get_worker_pid.remote())
+
+    @ray.remote
+    def f(id_in_a_list):
+        ray.get(id_in_a_list[0])
+
+    # Have the worker wait in a get call.
+    result_id = f.remote([x_id])
+    time.sleep(1)
+
+    # Make sure the task hasn't finished.
+    ready_ids, _ = ray.wait([result_id], timeout=0)
+    assert len(ready_ids) == 0
+
+    # Kill the worker.
+    os.kill(worker_pid, SIGKILL)
+    time.sleep(0.1)
+
+    # Make sure the sleep task hasn't finished.
+    ready_ids, _ = ray.wait([x_id], timeout=0)
+    assert len(ready_ids) == 0
+
+    # So that we attempt to notify the worker that the object is available.
+    ray.get(signal_2.send.remote())
+    ray.get(x_id)
+    time.sleep(0.1)
+
+    # Make sure that nothing has died.
+    assert ray._private.services.remaining_processes_alive()
+
+
+# This test checks that when a driver dies in the middle of a get, the raylet will not die.
+def test_dying_driver_get(ray_start_regular):
+    # Start the Ray processes.
+    address_info = ray_start_regular
+
+    @ray.remote
+    def wait_on_signal(signal):
+        ray.get(signal.wait.remote())
+        return np.ones(200 * 1024, dtype=np.uint8)
+
+    signal = SignalActor.remote()
+    x_id = wait_on_signal.remote(signal)
+
+    driver = """
+import ray
+ray.init("{}")
+ray.get(ray.ObjectRef(ray._common.utils.hex_to_binary("{}")))
+""".format(
+        address_info["address"], x_id.hex()
+    )
+
+    p = run_string_as_driver_nonblocking(driver)
+    # Make sure the driver is running.
+    time.sleep(1)
+    assert p.poll() is None
+
+    # Kill the driver process.
+    p.kill()
+    p.wait()
+    time.sleep(0.1)
+
+    # Make sure the original task hasn't finished.
+    ready_ids, _ = ray.wait([x_id], timeout=0)
+    assert len(ready_ids) == 0
+    # So that we attempt to notify the worker that the object is available.
+    ray.get(signal.send.remote())
+    ray.get(x_id)
+    time.sleep(0.1)
+
+    # Make sure that nothing has died.
+    assert ray._private.services.remaining_processes_alive()
+
+
+# This test checks that when a worker dies in the middle of a wait, the raylet will not die.
+def test_dying_worker_wait(ray_start_2_cpus):
+    @ray.remote
+    def wait_on_signal(signal):
+        ray.get(signal.wait.remote())
+        return np.ones(200 * 1024, dtype=np.uint8)
+
+    @ray.remote
+    def get_pid():
+        return os.getpid()
+
+    signal = SignalActor.remote()
+    x_id = wait_on_signal.remote(signal)
+    # Get the PID of the worker that block_in_wait will run on (sleep a little
+    # to make sure that wait_on_signal has already started).
+    time.sleep(0.1)
+    worker_pid = ray.get(get_pid.remote())
+
+    @ray.remote
+    def block_in_wait(object_ref_in_list):
+        ray.wait(object_ref_in_list)
+
+    # Have the worker wait in a wait call.
+    block_in_wait.remote([x_id])
+    time.sleep(0.1)
+
+    # Kill the worker.
+    os.kill(worker_pid, SIGKILL)
+    time.sleep(0.1)
+
+    # So that we attempt to notify the worker that the object is available.
+    ray.get(signal.send.remote())
+    ray.get(x_id)
+    time.sleep(0.1)
+
+    # Make sure that nothing has died.
+    assert ray._private.services.remaining_processes_alive()
+
+
+# This test checks that when a driver dies in the middle of a wait, the raylet will not die.
+def test_dying_driver_wait(ray_start_regular):
+    # Start the Ray processes.
+    address_info = ray_start_regular
+
+    @ray.remote
+    def wait_on_signal(signal):
+        ray.get(signal.wait.remote())
+        return np.ones(200 * 1024, dtype=np.uint8)
+
+    signal = SignalActor.remote()
+    x_id = wait_on_signal.remote(signal)
+
+    driver = """
+import ray
+ray.init("{}")
+ray.wait([ray.ObjectRef(ray._common.utils.hex_to_binary("{}"))])
+""".format(
+        address_info["address"], x_id.hex()
+    )
+
+    p = run_string_as_driver_nonblocking(driver)
+    # Make sure the driver is running.
+    time.sleep(1)
+    assert p.poll() is None
+
+    # Kill the driver process.
+    p.kill()
+    p.wait()
+    time.sleep(0.1)
+
+    # Make sure the original task hasn't finished.
+    ready_ids, _ = ray.wait([x_id], timeout=0)
+    assert len(ready_ids) == 0
+    # So that we attempt to notify the worker that the object is available.
+    ray.get(signal.send.remote())
+    ray.get(x_id)
+    time.sleep(0.1)
+
+    # Make sure that nothing has died.
+    assert ray._private.services.remaining_processes_alive()
+
+
+if __name__ == "__main__":
+    sys.exit(pytest.main(["-sv", __file__]))