chore: import upstream snapshot with attribution

python/ray/tests/test_basic_2.py

@@ -0,0 +1,787 @@
+# coding: utf-8
+import array
+import logging
+import os
+import random
+import subprocess
+import sys
+import tempfile
+import threading
+import time
+from unittest.mock import MagicMock, patch
+
+import pytest
+
+from ray._private.ray_constants import KV_NAMESPACE_FUNCTION_TABLE
+from ray._private.test_utils import client_test_enabled
+from ray.cluster_utils import Cluster, cluster_not_supported
+from ray.exceptions import GetTimeoutError, RayTaskError
+from ray.tests.client_test_utils import create_remote_signal_actor
+
+if client_test_enabled():
+    from ray.util.client import ray
+else:
+    import ray
+
+logger = logging.getLogger(__name__)
+
+
+def test_variable_number_of_args(shutdown_only):
+    ray.init(num_cpus=1)
+
+    @ray.remote
+    def varargs_fct1(*a):
+        return " ".join(map(str, a))
+
+    @ray.remote
+    def varargs_fct2(a, *b):
+        return " ".join(map(str, b))
+
+    x = varargs_fct1.remote(0, 1, 2)
+    assert ray.get(x) == "0 1 2"
+    x = varargs_fct2.remote(0, 1, 2)
+    assert ray.get(x) == "1 2"
+
+    @ray.remote
+    def f1(*args):
+        return args
+
+    @ray.remote
+    def f2(x, y, *args):
+        return x, y, args
+
+    assert ray.get(f1.remote()) == ()
+    assert ray.get(f1.remote(1)) == (1,)
+    assert ray.get(f1.remote(1, 2, 3)) == (1, 2, 3)
+    with pytest.raises(TypeError):
+        f2.remote()
+    with pytest.raises(TypeError):
+        f2.remote(1)
+    assert ray.get(f2.remote(1, 2)) == (1, 2, ())
+    assert ray.get(f2.remote(1, 2, 3)) == (1, 2, (3,))
+    assert ray.get(f2.remote(1, 2, 3, 4)) == (1, 2, (3, 4))
+
+    def testNoArgs(self):
+        @ray.remote
+        def no_op():
+            pass
+
+        self.ray_start()
+
+        ray.get(no_op.remote())
+
+
+def test_defining_remote_functions(shutdown_only):
+    ray.init(num_cpus=3)
+
+    # Test that we can close over plain old data.
+    data = [
+        (1, 2, "a"),
+        [0.0, 1.0, 1 << 62],
+        1 << 60,
+        {"a": bytes(3)},
+    ]
+
+    @ray.remote
+    def g():
+        return data
+
+    ray.get(g.remote())
+
+    # Test that we can close over modules.
+    @ray.remote
+    def h():
+        return array.array("d", [1.0, 2.0, 3.0])
+
+    assert ray.get(h.remote()) == array.array("d", [1.0, 2.0, 3.0])
+
+    @ray.remote
+    def j():
+        return time.time()
+
+    ray.get(j.remote())
+
+    # Test that we can define remote functions that call other remote
+    # functions.
+    @ray.remote
+    def k(x):
+        return x + 1
+
+    @ray.remote
+    def k2(x):
+        return ray.get(k.remote(x))
+
+    @ray.remote
+    def m(x):
+        return ray.get(k2.remote(x))
+
+    assert ray.get(k.remote(1)) == 2
+    assert ray.get(k2.remote(1)) == 2
+    assert ray.get(m.remote(1)) == 2
+
+
+def test_redefining_remote_functions(shutdown_only):
+    ray.init(num_cpus=1)
+
+    # Test that we can define a remote function in the shell.
+    @ray.remote
+    def f(x):
+        return x + 1
+
+    assert ray.get(f.remote(0)) == 1
+
+    # Test that we can redefine the remote function.
+    @ray.remote
+    def f(x):
+        return x + 10
+
+    while True:
+        val = ray.get(f.remote(0))
+        assert val in [1, 10]
+        if val == 10:
+            break
+        else:
+            logger.info("Still using old definition of f, trying again.")
+
+    # Check that we can redefine functions even when the remote function source
+    # doesn't change (see https://github.com/ray-project/ray/issues/6130).
+    @ray.remote
+    def g():
+        return nonexistent()
+
+    with pytest.raises(RayTaskError, match="nonexistent"):
+        ray.get(g.remote())
+
+    def nonexistent():
+        return 1
+
+    # Redefine the function and make sure it succeeds.
+    @ray.remote
+    def g():
+        return nonexistent()
+
+    assert ray.get(g.remote()) == 1
+
+    # Check the same thing but when the redefined function is inside of another
+    # task.
+    @ray.remote
+    def h(i):
+        @ray.remote
+        def j():
+            return i
+
+        return j.remote()
+
+    for i in range(20):
+        assert ray.get(ray.get(h.remote(i))) == i
+
+
+def test_call_matrix(shutdown_only):
+    ray.init(object_store_memory=1000 * 1024 * 1024)
+
+    @ray.remote
+    class Actor:
+        def small_value(self):
+            return 0
+
+        def large_value(self):
+            return bytes(80 * 1024 * 1024)
+
+        def echo(self, x):
+            if isinstance(x, list):
+                x = ray.get(x[0])
+            return x
+
+    @ray.remote
+    def small_value():
+        return 0
+
+    @ray.remote
+    def large_value():
+        return bytes(80 * 1024 * 1024)
+
+    @ray.remote
+    def echo(x):
+        if isinstance(x, list):
+            x = ray.get(x[0])
+        return x
+
+    def check(source_actor, dest_actor, is_large, out_of_band):
+        print(
+            "CHECKING",
+            "actor" if source_actor else "task",
+            "to",
+            "actor" if dest_actor else "task",
+            "large_object" if is_large else "small_object",
+            "out_of_band" if out_of_band else "in_band",
+        )
+        if source_actor:
+            a = Actor.remote()
+            if is_large:
+                x_id = a.large_value.remote()
+            else:
+                x_id = a.small_value.remote()
+        else:
+            if is_large:
+                x_id = large_value.remote()
+            else:
+                x_id = small_value.remote()
+        if out_of_band:
+            x_id = [x_id]
+        if dest_actor:
+            b = Actor.remote()
+            x = ray.get(b.echo.remote(x_id))
+        else:
+            x = ray.get(echo.remote(x_id))
+        if is_large:
+            assert isinstance(x, bytes)
+        else:
+            assert isinstance(x, int)
+
+    for is_large in [False, True]:
+        for source_actor in [False, True]:
+            for dest_actor in [False, True]:
+                for out_of_band in [False, True]:
+                    check(source_actor, dest_actor, is_large, out_of_band)
+
+
+def test_actor_call_order(shutdown_only):
+    ray.init(num_cpus=4)
+
+    @ray.remote
+    def small_value():
+        time.sleep(0.01 * random.randint(0, 10))
+        return 0
+
+    @ray.remote
+    class Actor:
+        def __init__(self):
+            self.count = 0
+
+        def inc(self, count, dependency):
+            assert count == self.count
+            self.count += 1
+            return count
+
+    a = Actor.remote()
+    assert ray.get([a.inc.remote(i, small_value.remote()) for i in range(100)]) == list(
+        range(100)
+    )
+
+
+def test_actor_pass_by_ref_order_optimization(shutdown_only):
+    ray.init(num_cpus=4)
+
+    @ray.remote
+    class Actor:
+        def __init__(self):
+            pass
+
+        def f(self, x):
+            pass
+
+    a = Actor.remote()
+
+    @ray.remote
+    def fast_value():
+        print("fast value")
+        pass
+
+    @ray.remote
+    def slow_value():
+        print("start sleep")
+        time.sleep(30)
+
+    @ray.remote
+    def runner(f):
+        print("runner", a, f)
+        return ray.get(a.f.remote(f.remote()))
+
+    runner.remote(slow_value)
+    time.sleep(1)
+    x2 = runner.remote(fast_value)
+    start = time.time()
+    ray.get(x2)
+    delta = time.time() - start
+    assert delta < 10, "did not skip slow value"
+
+
+@pytest.mark.parametrize(
+    "ray_start_cluster",
+    [
+        {
+            "num_cpus": 1,
+            "num_nodes": 1,
+        },
+        {
+            "num_cpus": 1,
+            "num_nodes": 2,
+        },
+    ],
+    indirect=True,
+)
+def test_call_chain(ray_start_cluster):
+    @ray.remote
+    def g(x):
+        return x + 1
+
+    x = 0
+    for _ in range(100):
+        x = g.remote(x)
+    assert ray.get(x) == 100
+
+
+@pytest.mark.xfail(cluster_not_supported, reason="cluster not supported")
+@pytest.mark.skipif(client_test_enabled(), reason="init issue")
+def test_system_config_when_connecting(ray_start_cluster):
+    config = {"object_timeout_milliseconds": 200}
+    cluster = Cluster()
+    cluster.add_node(_system_config=config, object_store_memory=100 * 1024 * 1024)
+    cluster.wait_for_nodes()
+
+    # Specifying _system_config when connecting to a cluster is disallowed.
+    with pytest.raises(ValueError):
+        ray.init(address=cluster.address, _system_config=config)
+
+    # Check that the config was picked up (object pinning is disabled).
+    ray.init(address=cluster.address)
+    obj_ref = ray.put(bytes(40 * 1024 * 1024))
+
+    for _ in range(5):
+        put_ref = ray.put(bytes(40 * 1024 * 1024))
+    del put_ref
+
+    ray.get(obj_ref)
+
+
+def test_get_multiple(ray_start_regular_shared):
+    object_refs = [ray.put(i) for i in range(10)]
+    assert ray.get(object_refs) == list(range(10))
+
+    # Get a random choice of object refs with duplicates.
+    indices = [random.choice(range(10)) for i in range(5)]
+    indices += indices
+    results = ray.get([object_refs[i] for i in indices])
+    assert results == indices
+
+
+def test_get_with_timeout(ray_start_regular_shared):
+    SignalActor = create_remote_signal_actor(ray)
+    signal = SignalActor.remote()
+
+    # Check that get() returns early if object is ready.
+    start = time.time()
+    ray.get(signal.wait.remote(should_wait=False), timeout=30)
+    assert time.time() - start < 30
+
+    # Check that get() raises a TimeoutError after the timeout if the object
+    # is not ready yet.
+    result_id = signal.wait.remote()
+    with pytest.raises(GetTimeoutError):
+        ray.get(result_id, timeout=0.1)
+
+    assert issubclass(GetTimeoutError, TimeoutError)
+    with pytest.raises(TimeoutError):
+        ray.get(result_id, timeout=0.1)
+
+    # timeout of 0 should raise an error
+    with pytest.raises(GetTimeoutError):
+        ray.get(result_id, timeout=0)
+
+    # Check that a subsequent get() returns early.
+    ray.get(signal.send.remote())
+    start = time.time()
+    ray.get(result_id, timeout=30)
+    assert time.time() - start < 30
+
+
+# https://github.com/ray-project/ray/issues/6329
+def test_call_actors_indirect_through_tasks(ray_start_regular_shared):
+    @ray.remote
+    class Counter:
+        def __init__(self, value):
+            self.value = int(value)
+
+        def increase(self, delta):
+            self.value += int(delta)
+            return self.value
+
+    @ray.remote
+    def foo(object):
+        return ray.get(object.increase.remote(1))
+
+    @ray.remote
+    def bar(object):
+        return ray.get(object.increase.remote(1))
+
+    @ray.remote
+    def zoo(object):
+        return ray.get(object[0].increase.remote(1))
+
+    c = Counter.remote(0)
+    for _ in range(0, 100):
+        ray.get(foo.remote(c))
+        ray.get(bar.remote(c))
+        ray.get(zoo.remote([c]))
+
+
+def test_inline_arg_memory_corruption(ray_start_regular_shared):
+    @ray.remote
+    def f():
+        return bytes(1000)
+
+    @ray.remote
+    class Actor:
+        def __init__(self):
+            self.z = []
+
+        def add(self, x):
+            self.z.append(x)
+            for prev in self.z:
+                assert sum(prev) == 0, ("memory corruption detected", prev)
+
+    a = Actor.remote()
+    for i in range(100):
+        ray.get(a.add.remote(f.remote()))
+
+
+@pytest.mark.skipif(client_test_enabled(), reason="internal api")
+def test_skip_plasma(ray_start_regular_shared):
+    @ray.remote
+    class Actor:
+        def __init__(self):
+            pass
+
+        def f(self, x):
+            return x * 2
+
+    a = Actor.remote()
+    obj_ref = a.f.remote(1)
+    # it is not stored in plasma
+    assert not ray._private.worker.global_worker.core_worker.object_exists(obj_ref)
+    assert ray.get(obj_ref) == 2
+
+
+@pytest.mark.skipif(client_test_enabled(), reason="internal api")
+def test_actor_large_objects(ray_start_regular_shared):
+    @ray.remote
+    class Actor:
+        def __init__(self):
+            pass
+
+        def f(self):
+            time.sleep(1)
+            return bytes(80000000)
+
+    a = Actor.remote()
+    obj_ref = a.f.remote()
+    assert not ray._private.worker.global_worker.core_worker.object_exists(obj_ref)
+    done, _ = ray.wait([obj_ref])
+    assert len(done) == 1
+    assert ray._private.worker.global_worker.core_worker.object_exists(obj_ref)
+    assert isinstance(ray.get(obj_ref), bytes)
+
+
+def test_actor_pass_by_ref(ray_start_regular_shared):
+    @ray.remote
+    class Actor:
+        def __init__(self):
+            pass
+
+        def f(self, x):
+            return x * 2
+
+    @ray.remote
+    def f(x):
+        return x
+
+    @ray.remote
+    def error():
+        sys.exit(0)
+
+    a = Actor.remote()
+    assert ray.get(a.f.remote(f.remote(1))) == 2
+
+    fut = [a.f.remote(f.remote(i)) for i in range(100)]
+    assert ray.get(fut) == [i * 2 for i in range(100)]
+
+    # propagates errors for pass by ref
+    with pytest.raises(Exception):
+        ray.get(a.f.remote(error.remote()))
+
+
+def test_actor_recursive(ray_start_regular_shared):
+    @ray.remote
+    class Actor:
+        def __init__(self, delegate=None):
+            self.delegate = delegate
+
+        def f(self, x):
+            if self.delegate:
+                return ray.get(self.delegate.f.remote(x))
+            return x * 2
+
+    a = Actor.remote()
+    b = Actor.remote(a)
+    c = Actor.remote(b)
+
+    result = ray.get([c.f.remote(i) for i in range(100)])
+    assert result == [x * 2 for x in range(100)]
+
+    result, _ = ray.wait([c.f.remote(i) for i in range(100)], num_returns=100)
+    result = ray.get(result)
+    assert result == [x * 2 for x in range(100)]
+
+
+def test_actor_concurrent(ray_start_regular_shared):
+    @ray.remote
+    class Batcher:
+        def __init__(self):
+            self.batch = []
+            self.event = threading.Event()
+
+        def add(self, x):
+            self.batch.append(x)
+            if len(self.batch) >= 3:
+                self.event.set()
+            else:
+                self.event.wait()
+            return sorted(self.batch)
+
+    a = Batcher.options(max_concurrency=3).remote()
+    x1 = a.add.remote(1)
+    x2 = a.add.remote(2)
+    x3 = a.add.remote(3)
+    r1 = ray.get(x1)
+    r2 = ray.get(x2)
+    r3 = ray.get(x3)
+    assert r1 == [1, 2, 3]
+    assert r1 == r2 == r3
+
+
+def test_actor_max_concurrency(ray_start_regular_shared):
+    """
+    Test that an actor of max_concurrency=N should only run
+    N tasks at most concurrently.
+    """
+    CONCURRENCY = 3
+
+    @ray.remote
+    class ConcurentActor:
+        def __init__(self):
+            self.threads = set()
+
+        def call(self):
+            # Record the current thread that runs this function.
+            self.threads.add(threading.get_ident())
+
+        def get_num_threads(self):
+            return len(self.threads)
+
+    @ray.remote
+    def call(actor):
+        for _ in range(CONCURRENCY * 100):
+            ray.get(actor.call.remote())
+        return
+
+    actor = ConcurentActor.options(max_concurrency=CONCURRENCY).remote()
+    # Start many concurrent tasks that will call the actor many times.
+    ray.get([call.remote(actor) for _ in range(CONCURRENCY * 10)])
+
+    # Check that the number of threads shouldn't be greater than CONCURRENCY.
+    assert ray.get(actor.get_num_threads.remote()) <= CONCURRENCY
+
+
+def test_duplicate_args(ray_start_regular_shared):
+    @ray.remote
+    def f(arg1, arg2, arg1_duplicate, kwarg1=None, kwarg2=None, kwarg1_duplicate=None):
+        assert arg1 == kwarg1
+        assert arg1 != arg2
+        assert arg1 == arg1_duplicate
+        assert kwarg1 != kwarg2
+        assert kwarg1 == kwarg1_duplicate
+
+    # Test by-value arguments.
+    arg1 = [1]
+    arg2 = [2]
+    ray.get(f.remote(arg1, arg2, arg1, kwarg1=arg1, kwarg2=arg2, kwarg1_duplicate=arg1))
+
+    # Test by-reference arguments.
+    arg1 = ray.put([1])
+    arg2 = ray.put([2])
+    ray.get(f.remote(arg1, arg2, arg1, kwarg1=arg1, kwarg2=arg2, kwarg1_duplicate=arg1))
+
+    # Test by-reference arguments on an actor task.
+    @ray.remote
+    class Actor:
+        def f(
+            self,
+            arg1,
+            arg2,
+            arg1_duplicate,
+            kwarg1=None,
+            kwarg2=None,
+            kwarg1_duplicate=None,
+        ):
+            assert arg1 == kwarg1
+            assert arg1 != arg2
+            assert arg1 == arg1_duplicate
+            assert kwarg1 != kwarg2
+            assert kwarg1 == kwarg1_duplicate
+
+    actor = Actor.remote()
+    ray.get(
+        actor.f.remote(
+            arg1, arg2, arg1, kwarg1=arg1, kwarg2=arg2, kwarg1_duplicate=arg1
+        )
+    )
+
+
+@pytest.mark.skipif(client_test_enabled(), reason="internal api")
+def test_get_correct_node_ip():
+    with patch("ray._private.worker") as worker_mock:
+        node_mock = MagicMock()
+        node_mock.node_ip_address = "10.0.0.111"
+        worker_mock._global_node = node_mock
+        found_ip = ray.util.get_node_ip_address()
+        assert found_ip == "10.0.0.111"
+
+
+def test_load_code_from_local(ray_start_regular_shared):
+    # This case writes a driver python file to a temporary directory.
+    #
+    # The driver starts a cluster with
+    # `ray.init(ray.job_config.JobConfig(code_search_path=<path list>))`,
+    # then creates a nested actor. The actor will be loaded from code in
+    # worker.
+    #
+    # This tests the following two cases when :
+    # 1) Load a nested class.
+    # 2) Load a class defined in the `__main__` module.
+    code_test = """
+import os
+import ray
+
+class A:
+    @ray.remote
+    class B:
+        def get(self):
+            return "OK"
+
+if __name__ == "__main__":
+    current_path = os.path.dirname(__file__)
+    job_config = ray.job_config.JobConfig(code_search_path=[current_path])
+    ray.init({}, job_config=job_config)
+    b = A.B.remote()
+    print(ray.get(b.get.remote()))
+"""
+
+    # Test code search path contains space.
+    with tempfile.TemporaryDirectory(suffix="a b") as tmpdir:
+        test_driver = os.path.join(tmpdir, "test_load_code_from_local.py")
+        with open(test_driver, "w") as f:
+            f.write(code_test.format(repr(ray_start_regular_shared["address"])))
+
+        # Ray's handling of sys.path does not work with PYTHONSAFEPATH.
+        env = os.environ.copy()
+        if env.get("PYTHONSAFEPATH", "") != "":
+            env["PYTHONSAFEPATH"] = ""  # Set to empty string to disable.
+        output = subprocess.check_output([sys.executable, test_driver], env=env)
+        assert b"OK" in output, f"Output has no 'OK': {output.decode()}"
+
+
+def test_load_from_local_function(ray_start_regular_shared):
+    code_test = """
+import ray
+import os
+
+@ray.remote
+def wrapped_func():
+    return "OK"
+def normal_func():
+    return "OK"
+if __name__ == "__main__":
+    current_path = os.path.dirname(__file__)
+    job_config = ray.job_config.JobConfig(code_search_path=[current_path])
+    ray.init({}, job_config=job_config)
+    res1 = wrapped_func.remote()
+    wrapped_normal_func = ray.remote(normal_func)
+    res2 = wrapped_normal_func.remote()
+    s1 = ray.get(res1)
+    s2 = ray.get(res2)
+    print(s1+"+"+s2)
+"""
+
+    # Test code search path contains space.
+    with tempfile.TemporaryDirectory(suffix="a b") as tmpdir:
+        test_driver = os.path.join(tmpdir, "test_load_code_from_local.py")
+        with open(test_driver, "w") as f:
+            f.write(code_test.format(repr(ray_start_regular_shared["address"])))
+        env = os.environ.copy()
+        if env.get("PYTHONSAFEPATH", "") != "":
+            env["PYTHONSAFEPATH"] = ""  # Set to empty string to disable.
+        output = subprocess.check_output([sys.executable, test_driver], env=env)
+        assert b"OK+OK" in output, f"Output has no 'OK+OK': {output.decode()}"
+
+
+@pytest.mark.skipif(
+    client_test_enabled(), reason="JobConfig doesn't work in client mode"
+)
+def test_use_dynamic_function_and_class():
+    # Test use dynamically defined functions
+    # and classes for remote tasks and actors.
+    # See https://github.com/ray-project/ray/issues/12834.
+    ray.shutdown()
+    current_path = os.path.dirname(__file__)
+    job_config = ray.job_config.JobConfig(code_search_path=[current_path])
+    ray.init(job_config=job_config)
+
+    def foo1():
+        @ray.remote
+        def foo2():
+            return "OK"
+
+        return foo2
+
+    @ray.remote
+    class Foo:
+        @ray.method(num_returns=1)
+        def foo(self):
+            return "OK"
+
+    f = foo1()
+    assert ray.get(f.remote()) == "OK"
+    # Check whether the dynamic function is exported to GCS.
+    # Note, the key format should be kept
+    # the same as in `FunctionActorManager.export`.
+    key_func = (
+        b"RemoteFunction:"
+        + ray._private.worker.global_worker.current_job_id.hex().encode()
+        + b":"
+        + f._function_descriptor.function_id.binary()
+    )
+    assert ray._private.worker.global_worker.gcs_client.internal_kv_exists(
+        key_func, KV_NAMESPACE_FUNCTION_TABLE
+    )
+    foo_actor = Foo.remote()
+
+    assert ray.get(foo_actor.foo.remote()) == "OK"
+    # Check whether the dynamic class is exported to GCS.
+    # Note, the key format should be kept
+    # the same as in `FunctionActorManager.export_actor_class`.
+    key_cls = (
+        b"ActorClass:"
+        + ray._private.worker.global_worker.current_job_id.hex().encode()
+        + b":"
+        + foo_actor._ray_actor_creation_function_descriptor.function_id.binary()
+    )
+    assert ray._private.worker.global_worker.gcs_client.internal_kv_exists(
+        key_cls, namespace=KV_NAMESPACE_FUNCTION_TABLE
+    )
+
+
+if __name__ == "__main__":
+    import pytest
+
+    # Skip test_basic_2_client_mode for now- the test suite is breaking.
+    sys.exit(pytest.main(["-sv", __file__]))