chore: import upstream snapshot with attribution

python/ray/tests/test_stress_sharded.py

@@ -0,0 +1,89 @@
+import sys
+
+import numpy as np
+import pytest
+
+import ray
+
+
+@pytest.fixture(params=[1])
+def ray_start_sharded(request):
+    # Start the Ray processes.
+    ray.init(
+        object_store_memory=int(0.5 * 10**9),
+        num_cpus=10,
+    )
+
+    yield None
+
+    # The code after the yield will run as teardown code.
+    ray.shutdown()
+
+
+def test_submitting_many_tasks(ray_start_sharded):
+    @ray.remote
+    def f(x):
+        return 1
+
+    def g(n):
+        x = 1
+        for i in range(n):
+            x = f.remote(x)
+        return x
+
+    ray.get([g(100) for _ in range(100)])
+    assert ray._private.services.remaining_processes_alive()
+
+
+def test_submitting_many_actors_to_one(ray_start_sharded):
+    @ray.remote
+    class Actor:
+        def __init__(self):
+            pass
+
+        def ping(self):
+            return
+
+    @ray.remote
+    class Worker:
+        def __init__(self, actor):
+            self.actor = actor
+
+        def ping(self):
+            return ray.get(self.actor.ping.remote())
+
+    a = Actor.remote()
+    workers = [Worker.remote(a) for _ in range(10)]
+    for _ in range(10):
+        out = ray.get([w.ping.remote() for w in workers])
+        assert out == [None for _ in workers]
+
+
+def test_getting_and_putting(ray_start_sharded):
+    for n in range(8):
+        x = np.zeros(10**n)
+
+        for _ in range(100):
+            ray.put(x)
+
+        x_id = ray.put(x)
+        for _ in range(1000):
+            ray.get(x_id)
+
+    assert ray._private.services.remaining_processes_alive()
+
+
+def test_getting_many_objects(ray_start_sharded):
+    @ray.remote
+    def f():
+        return 1
+
+    n = 10**4  # TODO(pcm): replace by 10 ** 5 once this is faster.
+    lst = ray.get([f.remote() for _ in range(n)])
+    assert lst == n * [1]
+
+    assert ray._private.services.remaining_processes_alive()
+
+
+if __name__ == "__main__":
+    sys.exit(pytest.main(["-sv", __file__]))