chore: import upstream snapshot with attribution

python/ray/tests/test_network_failure_e2e.py

@@ -0,0 +1,319 @@
+import json
+import sys
+import threading
+from time import sleep
+
+import pytest
+
+from ray._common.test_utils import wait_for_condition
+from ray.tests.conftest_docker import *  # noqa
+from ray.tests.conftest_docker import gen_head_node, gen_worker_node
+
+SLEEP_TASK_SCRIPTS = """
+import ray
+ray.init()
+@ray.remote(max_retries=-1)
+def f():
+    import time
+    time.sleep(10000)
+ray.get([f.remote() for _ in range(2)])
+"""
+
+head = gen_head_node(
+    {
+        "RAY_grpc_keepalive_time_ms": "1000",
+        "RAY_grpc_client_keepalive_time_ms": "1000",
+        "RAY_grpc_client_keepalive_timeout_ms": "1000",
+        "RAY_health_check_initial_delay_ms": "1000",
+        "RAY_health_check_period_ms": "1000",
+        "RAY_health_check_timeout_ms": "1000",
+        "RAY_health_check_failure_threshold": "2",
+    }
+)
+
+worker = gen_worker_node(
+    envs={
+        "RAY_grpc_keepalive_time_ms": "1000",
+        "RAY_grpc_client_keepalive_time_ms": "1000",
+        "RAY_grpc_client_keepalive_timeout_ms": "1000",
+        "RAY_health_check_initial_delay_ms": "1000",
+        "RAY_health_check_period_ms": "1000",
+        "RAY_health_check_timeout_ms": "1000",
+        "RAY_health_check_failure_threshold": "2",
+    },
+    num_cpus=8,
+)
+
+
+def test_network_task_submit(head, worker, gcs_network):
+    network = gcs_network
+    # https://docker-py.readthedocs.io/en/stable/containers.html#docker.models.containers.Container.exec_run
+    head.exec_run(
+        cmd=f"python -c '{SLEEP_TASK_SCRIPTS}'",
+        detach=True,
+        environment=[
+            "RAY_grpc_client_keepalive_time_ms=1000",
+            "RAY_grpc_client_keepalive_timeout_ms=1000",
+        ],
+    )
+
+    def check_task_running(n=None):
+        output = head.exec_run(cmd="ray list tasks --format json")
+        if output.exit_code == 0:
+            tasks_json = json.loads(output.output)
+            print("tasks_json:", json.dumps(tasks_json, indent=2))
+            if n is not None and n != len(tasks_json):
+                return False
+            return all([task["state"] == "RUNNING" for task in tasks_json])
+        return False
+
+    # list_task make sure all tasks are running
+    wait_for_condition(lambda: check_task_running(2))
+
+    # Previously grpc client will only send 2 ping frames when there is no
+    # data/header frame to be sent.
+    # keepalive interval is 1s. So after 3s it wouldn't send anything and
+    # failed the test previously.
+    sleep(3)
+
+    # partition the network between head and worker
+    # https://docker-py.readthedocs.io/en/stable/networks.html#docker.models.networks.Network.disconnect
+    network.disconnect(worker.name)
+    print("Disconnected network")
+
+    def check_dead_node():
+        output = head.exec_run(cmd="ray list nodes --format json")
+        if output.exit_code == 0:
+            nodes_json = json.loads(output.output)
+            print("nodes_json:", json.dumps(nodes_json, indent=2))
+            for node in nodes_json:
+                if node["state"] == "DEAD" and not node["is_head_node"]:
+                    return True
+        return False
+
+    wait_for_condition(check_dead_node)
+    print("observed node died")
+
+    # Previously under network partition, the tasks would stay in RUNNING state
+    # and hanging forever.
+    # We write this test to check that.
+    def check_task_not_running():
+        output = head.exec_run(cmd="ray list tasks --format json")
+        if output.exit_code == 0:
+            tasks_json = json.loads(output.output)
+            print("tasks_json:", json.dumps(tasks_json, indent=2))
+            return all([task["state"] != "RUNNING" for task in tasks_json])
+        return False
+
+    # we set num_cpus=0 for head node.
+    # which ensures no task was scheduled on the head node.
+    wait_for_condition(check_task_not_running)
+
+    # After the fix, we should observe that the tasks are not running.
+    # `ray list tasks` would show two FAILED and
+    # two PENDING_NODE_ASSIGNMENT states.
+
+    def check_task_pending(n=0):
+        output = head.exec_run(cmd="ray list tasks --format json")
+        if output.exit_code == 0:
+            tasks_json = json.loads(output.output)
+            print("tasks_json:", json.dumps(tasks_json, indent=2))
+            return n == sum(
+                [task["state"] == "PENDING_NODE_ASSIGNMENT" for task in tasks_json]
+            )
+        return False
+
+    wait_for_condition(lambda: check_task_pending(2))
+
+
+head2 = gen_head_node(
+    {
+        "RAY_grpc_keepalive_time_ms": "1000",
+        "RAY_grpc_client_keepalive_time_ms": "1000",
+        "RAY_grpc_client_keepalive_timeout_ms": "1000",
+        "RAY_health_check_initial_delay_ms": "1000",
+        "RAY_health_check_period_ms": "1000",
+        "RAY_health_check_timeout_ms": "100000",
+        "RAY_health_check_failure_threshold": "20",
+    }
+)
+
+worker2 = gen_worker_node(
+    envs={
+        "RAY_grpc_keepalive_time_ms": "1000",
+        "RAY_grpc_client_keepalive_time_ms": "1000",
+        "RAY_grpc_client_keepalive_timeout_ms": "1000",
+        "RAY_health_check_initial_delay_ms": "1000",
+        "RAY_health_check_period_ms": "1000",
+        "RAY_health_check_timeout_ms": "100000",
+        "RAY_health_check_failure_threshold": "20",
+    },
+    num_cpus=2,
+)
+
+
+def test_transient_network_error(head2, worker2, gcs_network):
+    # Test to make sure the head node and worker node
+    # connection can be recovered from transient network error.
+    network = gcs_network
+
+    check_two_nodes = """
+import ray
+from ray._common.test_utils import wait_for_condition
+
+ray.init()
+wait_for_condition(lambda: len(ray.nodes()) == 2)
+"""
+    result = head2.exec_run(cmd=f"python -c '{check_two_nodes}'")
+    assert result.exit_code == 0, result.output.decode("utf-8")
+
+    # Simulate transient network error
+    worker_ip = worker2._container.attrs["NetworkSettings"]["Networks"][network.name][
+        "IPAddress"
+    ]
+    network.disconnect(worker2.name, force=True)
+    sleep(2)
+    network.connect(worker2.name, ipv4_address=worker_ip)
+
+    # Make sure the connection is recovered by scheduling
+    # an actor.
+    check_actor_scheduling = """
+import ray
+from ray._common.test_utils import wait_for_condition
+
+ray.init()
+
+@ray.remote(num_cpus=1)
+class Actor:
+    def ping(self):
+        return 1
+
+actor = Actor.remote()
+ray.get(actor.ping.remote())
+wait_for_condition(lambda: ray.available_resources()["CPU"] == 1.0)
+"""
+    result = head2.exec_run(cmd=f"python -c '{check_actor_scheduling}'")
+    assert result.exit_code == 0, result.output.decode("utf-8")
+
+
+head3 = gen_head_node(
+    {
+        "RAY_grpc_keepalive_time_ms": "1000",
+        "RAY_grpc_client_keepalive_time_ms": "1000",
+        "RAY_grpc_client_keepalive_timeout_ms": "1000",
+        "RAY_health_check_initial_delay_ms": "1000",
+        "RAY_health_check_period_ms": "1000",
+        "RAY_health_check_timeout_ms": "100000",
+        "RAY_health_check_failure_threshold": "20",
+    }
+)
+
+worker3 = gen_worker_node(
+    envs={
+        "RAY_grpc_keepalive_time_ms": "1000",
+        "RAY_grpc_client_keepalive_time_ms": "1000",
+        "RAY_grpc_client_keepalive_timeout_ms": "1000",
+        "RAY_health_check_initial_delay_ms": "1000",
+        "RAY_health_check_period_ms": "1000",
+        "RAY_health_check_timeout_ms": "100000",
+        "RAY_health_check_failure_threshold": "20",
+    },
+    num_cpus=2,
+)
+
+
+def test_async_actor_task_retry(head3, worker3, gcs_network):
+    # Test that if transient network error happens
+    # after an async actor task is submitted and being executed,
+    # a secon attempt will be submitted and executed after the
+    # first attempt finishes.
+    network = gcs_network
+
+    driver = """
+import asyncio
+import ray
+from ray.util.state import list_tasks
+
+ray.init(namespace="test")
+
+@ray.remote(num_cpus=0.1, name="counter", lifetime="detached")
+class Counter:
+  def __init__(self):
+    self.count = 0
+
+  def inc(self):
+    self.count = self.count + 1
+    return self.count
+
+  @ray.method(max_task_retries=-1)
+  def get(self):
+    return self.count
+
+@ray.remote(num_cpus=0.1, max_task_retries=-1)
+class AsyncActor:
+  def __init__(self, counter):
+    self.counter = counter
+
+  async def run(self):
+    count = await self.counter.get.remote()
+    if count == 0:
+      # first attempt
+      await self.counter.inc.remote()
+      while len(list_tasks(
+            filters=[("name", "=", "AsyncActor.run")])) < 2:
+        # wait for second attempt to be made
+        await asyncio.sleep(1)
+      # wait until the second attempt reaches the actor
+      await asyncio.sleep(2)
+      await self.counter.inc.remote()
+      return "first"
+    else:
+      # second attempt
+      # make sure second attempt only runs
+      # after first attempt finishes
+      assert count == 2
+      return "second"
+
+counter = Counter.remote()
+async_actor = AsyncActor.remote(counter)
+assert ray.get(async_actor.run.remote()) == "second"
+"""
+
+    check_async_actor_run_is_called = """
+import ray
+from ray._common.test_utils import wait_for_condition
+ray.init(namespace="test")
+
+wait_for_condition(lambda: ray.get_actor("counter") is not None)
+counter = ray.get_actor("counter")
+wait_for_condition(lambda: ray.get(counter.get.remote()) == 1)
+"""
+
+    def inject_transient_network_failure():
+        try:
+            result = head3.exec_run(
+                cmd=f"python -c '{check_async_actor_run_is_called}'"
+            )
+            assert result.exit_code == 0, result.output.decode("utf-8")
+
+            worker_ip = worker3._container.attrs["NetworkSettings"]["Networks"][
+                network.name
+            ]["IPAddress"]
+            network.disconnect(worker3.name, force=True)
+            sleep(2)
+            network.connect(worker3.name, ipv4_address=worker_ip)
+        except Exception as e:
+            print(f"Network failure injection failed {e}")
+
+    t = threading.Thread(target=inject_transient_network_failure, daemon=True)
+    t.start()
+
+    result = head3.exec_run(
+        cmd=f"python -c '{driver}'",
+    )
+    assert result.exit_code == 0, result.output.decode("utf-8")
+
+
+if __name__ == "__main__":
+
+    sys.exit(pytest.main(["-sv", __file__]))