chore: import upstream snapshot with attribution

python/ray/tests/test_actor_bounded_threads.py

@@ -0,0 +1,184 @@
+import logging
+import os
+import sys
+from collections import Counter
+from typing import Dict
+
+import pytest
+
+import ray
+from ray._private.test_utils import wait_for_dashboard_agent_available
+
+logger = logging.getLogger(__name__)
+
+
+def my_threads() -> Dict[str, int]:
+    """
+    Returns [(thread_id, thread_name)]
+    """
+    pid = os.getpid()
+    threads = Counter()
+    proc_dir = f"/proc/{pid}/task"
+
+    for tid_entry in os.listdir(proc_dir):
+        comm_path = os.path.join(proc_dir, tid_entry, "comm")
+
+        if os.path.exists(comm_path):
+            with open(comm_path, "r") as comm_file:
+                thread_name = comm_file.read().strip()
+                threads[thread_name] += 1
+    return threads
+
+
+# Tests a lot of workers sending tasks to an actor, the number of threads for that
+# actor should not infinitely go up.
+
+
+# These therads are from third party code, and may start any time, we can't control
+# them. So we allow them to be any number.
+KNOWN_THREADS = {
+    "grpc_global_tim",  # grpc global timer
+    "grpcpp_sync_ser",  # grpc
+    "jemalloc_bg_thd",  # jemalloc background thread
+}
+
+
+def assert_threads_are_bounded(
+    prev_threads: Dict[str, int], now_threads: Dict[str, int]
+):
+    """
+    Asserts that the threads did not grow unexpected.
+    Rule: For each (thread_name, count) in now_threads, it must either be <= the number
+    in prev_threads, or in KNOWN_THREADS.
+    """
+    for thread_name, count in now_threads.items():
+        if thread_name not in KNOWN_THREADS:
+            target = prev_threads.get(thread_name, 0)
+            assert count <= target, (
+                f"{thread_name} grows unexpectedly: "
+                f"expected <= {target}, got {count}. "
+                f"prev {prev_threads}, now: {now_threads}"
+            )
+
+
+# Spawns a lot of workers, each making 1 call to A.
+@ray.remote
+def fibonacci(a, i):
+    if i < 2:
+        return 1
+    f1 = fibonacci.remote(a, i - 1)
+    f2 = fibonacci.remote(a, i - 2)
+    return ray.get(a.add.remote(f1, f2))
+
+
+@pytest.mark.skipif(sys.platform != "linux", reason="procfs only works on linux.")
+def test_threaded_actor_have_bounded_num_of_threads(ray_start_cluster):
+    cluster = ray_start_cluster
+    cluster.add_node()
+    cluster.wait_for_nodes()
+    wait_for_dashboard_agent_available(cluster)
+    ray.init(address=cluster.address)
+
+    @ray.remote
+    class A:
+        def get_my_threads(self):
+            return my_threads()
+
+        def add(self, i, j):
+            return i + j
+
+    a = A.options(max_concurrency=2).remote()
+
+    prev_threads = ray.get(a.get_my_threads.remote())
+
+    assert ray.get(fibonacci.remote(a, 1)) == 1
+    now_threads = ray.get(a.get_my_threads.remote())
+    assert_threads_are_bounded(prev_threads, now_threads)
+
+    # Creates a lot of workers sending to actor
+    assert ray.get(fibonacci.remote(a, 10)) == 89
+    now_threads = ray.get(a.get_my_threads.remote())
+    assert_threads_are_bounded(prev_threads, now_threads)
+
+
+@pytest.mark.skipif(sys.platform != "linux", reason="procfs only works on linux.")
+def test_async_actor_have_bounded_num_of_threads(ray_start_cluster):
+    cluster = ray_start_cluster
+    cluster.add_node()
+    cluster.wait_for_nodes()
+    wait_for_dashboard_agent_available(cluster)
+    ray.init(address=cluster.address)
+
+    @ray.remote
+    class A:
+        async def get_my_threads(self):
+            return my_threads()
+
+        async def add(self, i, j):
+            return i + j
+
+    a = A.options(max_concurrency=2).remote()
+
+    prev_threads = ray.get(a.get_my_threads.remote())
+
+    assert ray.get(fibonacci.remote(a, 1)) == 1
+    now_threads = ray.get(a.get_my_threads.remote())
+    assert_threads_are_bounded(prev_threads, now_threads)
+
+    # Creates a lot of workers sending to actor
+    assert ray.get(fibonacci.remote(a, 10)) == 89
+    now_threads = ray.get(a.get_my_threads.remote())
+    assert_threads_are_bounded(prev_threads, now_threads)
+
+
+@pytest.mark.skipif(sys.platform != "linux", reason="procfs only works on linux.")
+def test_async_actor_cg_have_bounded_num_of_threads(ray_start_cluster):
+    cluster = ray_start_cluster
+    cluster.add_node()
+    cluster.wait_for_nodes()
+    wait_for_dashboard_agent_available(cluster)
+    ray.init(address=cluster.address)
+
+    @ray.remote(concurrency_groups={"io": 2, "compute": 4})
+    class A:
+        async def get_my_threads(self):
+            return my_threads()
+
+        @ray.method(concurrency_group="io")
+        async def io_add(self, i, j):
+            return i + j
+
+        @ray.method(concurrency_group="compute")
+        async def compute_add(self, i, j):
+            return i + j
+
+        async def default_add(self, i, j):
+            return i + j
+
+    # Spawns a lot of workers, each making 1 call to A.
+    @ray.remote
+    def fibonacci_cg(a, i):
+        if i < 2:
+            return 1
+        f1 = fibonacci_cg.remote(a, i - 1)
+        f2 = fibonacci_cg.remote(a, i - 2)
+        assert ray.get(a.io_add.remote(1, 2)) == 3
+        assert ray.get(a.compute_add.remote(4, 5)) == 9
+        return ray.get(a.default_add.remote(f1, f2))
+
+    a = A.options(max_concurrency=2).remote()
+
+    prev_threads = ray.get(a.get_my_threads.remote())
+
+    assert ray.get(fibonacci_cg.remote(a, 1)) == 1
+    now_threads = ray.get(a.get_my_threads.remote())
+    assert_threads_are_bounded(prev_threads, now_threads)
+
+    # Creates a lot of workers sending to actor
+    assert ray.get(fibonacci_cg.remote(a, 10)) == 89
+    now_threads = ray.get(a.get_my_threads.remote())
+    assert_threads_are_bounded(prev_threads, now_threads)
+
+
+if __name__ == "__main__":
+    sys.exit(pytest.main(["-sv", __file__]))