chore: import upstream snapshot with attribution

python/ray/tests/test_actor_resources.py

@@ -0,0 +1,682 @@
+import collections
+import os
+import sys
+import time
+
+import pytest
+
+import ray
+import ray.cluster_utils
+
+
+def test_actor_deletion_with_gpus(shutdown_only):
+    ray.init(num_cpus=1, num_gpus=1, object_store_memory=int(150 * 1024 * 1024))
+
+    # When an actor that uses a GPU exits, make sure that the GPU resources
+    # are released.
+
+    @ray.remote(num_gpus=1)
+    class Actor:
+        def getpid(self):
+            return os.getpid()
+
+    for _ in range(5):
+        # If we can successfully create an actor, that means that enough
+        # GPU resources are available.
+        a = Actor.remote()
+        ray.get(a.getpid.remote())
+
+
+def test_actor_state(ray_start_regular):
+    @ray.remote
+    class Counter:
+        def __init__(self):
+            self.value = 0
+
+        def increase(self):
+            self.value += 1
+
+        def value(self):
+            return self.value
+
+    c1 = Counter.remote()
+    c1.increase.remote()
+    assert ray.get(c1.value.remote()) == 1
+
+    c2 = Counter.remote()
+    c2.increase.remote()
+    c2.increase.remote()
+    assert ray.get(c2.value.remote()) == 2
+
+
+def test_actor_class_methods(ray_start_regular):
+    class Foo:
+        x = 2
+
+        @classmethod
+        def as_remote(cls):
+            return ray.remote(cls)
+
+        @classmethod
+        def f(cls):
+            return cls.x
+
+        @classmethod
+        def g(cls, y):
+            return cls.x + y
+
+        def echo(self, value):
+            return value
+
+    a = Foo.as_remote().remote()
+    assert ray.get(a.echo.remote(2)) == 2
+    assert ray.get(a.f.remote()) == 2
+    assert ray.get(a.g.remote(2)) == 4
+
+
+@pytest.mark.skipif(sys.platform == "win32", reason="Failing on Windows.")
+def test_actor_gpus(ray_start_cluster):
+    cluster = ray_start_cluster
+    num_nodes = 3
+    num_gpus_per_raylet = 4
+    for i in range(num_nodes):
+        cluster.add_node(
+            num_cpus=10 * num_gpus_per_raylet, num_gpus=num_gpus_per_raylet
+        )
+    ray.init(address=cluster.address)
+
+    @ray.remote(num_gpus=1)
+    class Actor1:
+        def __init__(self):
+            self.gpu_ids = ray.get_gpu_ids()
+
+        def get_location_and_ids(self):
+            assert ray.get_gpu_ids() == self.gpu_ids
+            return (
+                ray._private.worker.global_worker.node.unique_id,
+                tuple(self.gpu_ids),
+            )
+
+    # Create one actor per GPU.
+    actors = [Actor1.remote() for _ in range(num_nodes * num_gpus_per_raylet)]
+    # Make sure that no two actors are assigned to the same GPU.
+    locations_and_ids = ray.get(
+        [actor.get_location_and_ids.remote() for actor in actors]
+    )
+    node_names = {location for location, gpu_id in locations_and_ids}
+    assert len(node_names) == num_nodes
+    location_actor_combinations = []
+    for node_name in node_names:
+        for gpu_id in range(num_gpus_per_raylet):
+            location_actor_combinations.append((node_name, (gpu_id,)))
+    assert set(locations_and_ids) == set(location_actor_combinations)
+
+    # Creating a new actor should fail because all of the GPUs are being
+    # used.
+    a = Actor1.remote()
+    ready_ids, _ = ray.wait([a.get_location_and_ids.remote()], timeout=0.01)
+    assert ready_ids == []
+
+
+def test_actor_multiple_gpus(ray_start_cluster):
+    cluster = ray_start_cluster
+    num_nodes = 3
+    num_gpus_per_raylet = 5
+    for i in range(num_nodes):
+        cluster.add_node(
+            num_cpus=10 * num_gpus_per_raylet, num_gpus=num_gpus_per_raylet
+        )
+    ray.init(address=cluster.address)
+
+    @ray.remote(num_gpus=2)
+    class Actor1:
+        def __init__(self):
+            self.gpu_ids = ray.get_gpu_ids()
+
+        def get_location_and_ids(self):
+            assert ray.get_gpu_ids() == self.gpu_ids
+            return (
+                ray._private.worker.global_worker.node.unique_id,
+                tuple(self.gpu_ids),
+            )
+
+    # Create some actors.
+    actors1 = [Actor1.remote() for _ in range(num_nodes * 2)]
+    # Make sure that no two actors are assigned to the same GPU.
+    locations_and_ids = ray.get(
+        [actor.get_location_and_ids.remote() for actor in actors1]
+    )
+    node_names = {location for location, gpu_id in locations_and_ids}
+    assert len(node_names) == num_nodes
+
+    # Keep track of which GPU IDs are being used for each location.
+    gpus_in_use = {node_name: [] for node_name in node_names}
+    for location, gpu_ids in locations_and_ids:
+        gpus_in_use[location].extend(gpu_ids)
+    for node_name in node_names:
+        assert len(set(gpus_in_use[node_name])) == 4
+
+    # Creating a new actor should fail because all of the GPUs are being
+    # used.
+    a = Actor1.remote()
+    ready_ids, _ = ray.wait([a.get_location_and_ids.remote()], timeout=0.01)
+    assert ready_ids == []
+
+    # We should be able to create more actors that use only a single GPU.
+    @ray.remote(num_gpus=1)
+    class Actor2:
+        def __init__(self):
+            self.gpu_ids = ray.get_gpu_ids()
+
+        def get_location_and_ids(self):
+            return (
+                ray._private.worker.global_worker.node.unique_id,
+                tuple(self.gpu_ids),
+            )
+
+    # Create some actors.
+    actors2 = [Actor2.remote() for _ in range(num_nodes)]
+    # Make sure that no two actors are assigned to the same GPU.
+    locations_and_ids = ray.get(
+        [actor.get_location_and_ids.remote() for actor in actors2]
+    )
+    names = {location for location, gpu_id in locations_and_ids}
+    assert node_names == names
+    for location, gpu_ids in locations_and_ids:
+        gpus_in_use[location].extend(gpu_ids)
+    for node_name in node_names:
+        assert len(gpus_in_use[node_name]) == 5
+        assert set(gpus_in_use[node_name]) == set(range(5))
+
+    # Creating a new actor should fail because all of the GPUs are being
+    # used.
+    a = Actor2.remote()
+    ready_ids, _ = ray.wait([a.get_location_and_ids.remote()], timeout=0.01)
+    assert ready_ids == []
+
+
+@pytest.mark.skipif(sys.platform == "win32", reason="Very flaky.")
+def test_actor_different_numbers_of_gpus(ray_start_cluster):
+    # Test that we can create actors on two nodes that have different
+    # numbers of GPUs.
+    cluster = ray_start_cluster
+    cluster.add_node(num_cpus=10, num_gpus=0)
+    cluster.add_node(num_cpus=10, num_gpus=5)
+    cluster.add_node(num_cpus=10, num_gpus=10)
+    ray.init(address=cluster.address)
+
+    @ray.remote(num_gpus=1)
+    class Actor1:
+        def __init__(self):
+            self.gpu_ids = ray.get_gpu_ids()
+
+        def get_location_and_ids(self):
+            return (
+                ray._private.worker.global_worker.node.unique_id,
+                tuple(self.gpu_ids),
+            )
+
+    # Create some actors.
+    actors = [Actor1.remote() for _ in range(0 + 5 + 10)]
+    # Make sure that no two actors are assigned to the same GPU.
+    locations_and_ids = ray.get(
+        [actor.get_location_and_ids.remote() for actor in actors]
+    )
+    node_names = {location for location, gpu_id in locations_and_ids}
+    assert len(node_names) == 2
+    for node_name in node_names:
+        node_gpu_ids = [
+            gpu_id for location, gpu_id in locations_and_ids if location == node_name
+        ]
+        assert len(node_gpu_ids) in [5, 10]
+        assert set(node_gpu_ids) == {(i,) for i in range(len(node_gpu_ids))}
+
+    # Creating a new actor should fail because all of the GPUs are being
+    # used.
+    a = Actor1.remote()
+    ready_ids, _ = ray.wait([a.get_location_and_ids.remote()], timeout=0.01)
+    assert ready_ids == []
+
+
+def test_actor_multiple_gpus_from_multiple_tasks(ray_start_cluster):
+    cluster = ray_start_cluster
+    num_nodes = 3
+    num_gpus_per_raylet = 2
+    for i in range(num_nodes):
+        cluster.add_node(
+            num_cpus=4 * num_gpus_per_raylet,
+            num_gpus=num_gpus_per_raylet,
+            _system_config={"health_check_failure_threshold": 100} if i == 0 else {},
+        )
+    ray.init(address=cluster.address)
+
+    @ray.remote
+    def create_actors(i, n):
+        @ray.remote(num_gpus=1)
+        class Actor:
+            def __init__(self, i, j):
+                self.gpu_ids = ray.get_gpu_ids()
+
+            def get_location_and_ids(self):
+                return (
+                    (ray._private.worker.global_worker.node.unique_id),
+                    tuple(self.gpu_ids),
+                )
+
+            def sleep(self):
+                time.sleep(100)
+
+        # Create n actors.
+        actors = []
+        for j in range(n):
+            actors.append(Actor.remote(i, j))
+
+        locations = ray.get([actor.get_location_and_ids.remote() for actor in actors])
+
+        # Put each actor to sleep for a long time to prevent them from getting
+        # terminated.
+        for actor in actors:
+            actor.sleep.remote()
+
+        return locations
+
+    all_locations = ray.get(
+        [create_actors.remote(i, num_gpus_per_raylet) for i in range(num_nodes)]
+    )
+
+    # Make sure that no two actors are assigned to the same GPU.
+    node_names = {
+        location for locations in all_locations for location, gpu_id in locations
+    }
+    assert len(node_names) == num_nodes
+
+    # Keep track of which GPU IDs are being used for each location.
+    gpus_in_use = {node_name: [] for node_name in node_names}
+    for locations in all_locations:
+        for location, gpu_ids in locations:
+            gpus_in_use[location].extend(gpu_ids)
+    for node_name in node_names:
+        assert len(set(gpus_in_use[node_name])) == num_gpus_per_raylet
+
+    @ray.remote(num_gpus=1)
+    class Actor:
+        def __init__(self):
+            self.gpu_ids = ray.get_gpu_ids()
+
+        def get_location_and_ids(self):
+            return (
+                ray._private.worker.global_worker.node.unique_id,
+                tuple(self.gpu_ids),
+            )
+
+    # All the GPUs should be used up now.
+    a = Actor.remote()
+    ready_ids, _ = ray.wait([a.get_location_and_ids.remote()], timeout=0.01)
+    assert ready_ids == []
+
+
+def test_actors_and_tasks_with_gpus(ray_start_cluster):
+    cluster = ray_start_cluster
+    num_nodes = 3
+    num_gpus_per_raylet = 2
+    for i in range(num_nodes):
+        cluster.add_node(num_cpus=num_gpus_per_raylet, num_gpus=num_gpus_per_raylet)
+    ray.init(address=cluster.address)
+
+    def check_intervals_non_overlapping(list_of_intervals):
+        for i in range(len(list_of_intervals)):
+            for j in range(i):
+                first_interval = list_of_intervals[i]
+                second_interval = list_of_intervals[j]
+                # Check that list_of_intervals[i] and list_of_intervals[j]
+                # don't overlap.
+                assert first_interval[0] < first_interval[1]
+                assert second_interval[0] < second_interval[1]
+                intervals_nonoverlapping = (
+                    first_interval[1] <= second_interval[0]
+                    or second_interval[1] <= first_interval[0]
+                )
+                assert (
+                    intervals_nonoverlapping
+                ), "Intervals {} and {} are overlapping.".format(
+                    first_interval, second_interval
+                )
+
+    @ray.remote(num_gpus=1)
+    def f1():
+        t1 = time.time()
+        time.sleep(0.1)
+        t2 = time.time()
+        gpu_ids = ray.get_gpu_ids()
+        assert len(gpu_ids) == 1
+        assert gpu_ids[0] in range(num_gpus_per_raylet)
+        return (
+            ray._private.worker.global_worker.node.unique_id,
+            tuple(gpu_ids),
+            [t1, t2],
+        )
+
+    @ray.remote(num_gpus=2)
+    def f2():
+        t1 = time.time()
+        time.sleep(0.1)
+        t2 = time.time()
+        gpu_ids = ray.get_gpu_ids()
+        assert len(gpu_ids) == 2
+        assert gpu_ids[0] in range(num_gpus_per_raylet)
+        assert gpu_ids[1] in range(num_gpus_per_raylet)
+        return (
+            ray._private.worker.global_worker.node.unique_id,
+            tuple(gpu_ids),
+            [t1, t2],
+        )
+
+    @ray.remote(num_gpus=1)
+    class Actor1:
+        def __init__(self):
+            self.gpu_ids = ray.get_gpu_ids()
+            assert len(self.gpu_ids) == 1
+            assert self.gpu_ids[0] in range(num_gpus_per_raylet)
+
+        def get_location_and_ids(self):
+            assert ray.get_gpu_ids() == self.gpu_ids
+            return (
+                ray._private.worker.global_worker.node.unique_id,
+                tuple(self.gpu_ids),
+            )
+
+    def locations_to_intervals_for_many_tasks():
+        # Launch a bunch of GPU tasks.
+        locations_ids_and_intervals = ray.get(
+            [f1.remote() for _ in range(5 * num_nodes * num_gpus_per_raylet)]
+            + [f2.remote() for _ in range(5 * num_nodes * num_gpus_per_raylet)]
+            + [f1.remote() for _ in range(5 * num_nodes * num_gpus_per_raylet)]
+        )
+
+        locations_to_intervals = collections.defaultdict(lambda: [])
+        for location, gpu_ids, interval in locations_ids_and_intervals:
+            for gpu_id in gpu_ids:
+                locations_to_intervals[(location, gpu_id)].append(interval)
+        return locations_to_intervals
+
+    # Run a bunch of GPU tasks.
+    locations_to_intervals = locations_to_intervals_for_many_tasks()
+    # For each GPU, verify that the set of tasks that used this specific
+    # GPU did not overlap in time.
+    for locations in locations_to_intervals:
+        check_intervals_non_overlapping(locations_to_intervals[locations])
+
+    # Create an actor that uses a GPU.
+    a = Actor1.remote()
+    actor_location = ray.get(a.get_location_and_ids.remote())
+    actor_location = (actor_location[0], actor_location[1][0])
+    # This check makes sure that actor_location is formatted the same way
+    # that the keys of locations_to_intervals are formatted.
+    assert actor_location in locations_to_intervals
+
+    # Run a bunch of GPU tasks.
+    locations_to_intervals = locations_to_intervals_for_many_tasks()
+    # For each GPU, verify that the set of tasks that used this specific
+    # GPU did not overlap in time.
+    for locations in locations_to_intervals:
+        check_intervals_non_overlapping(locations_to_intervals[locations])
+    # Make sure that the actor's GPU was not used.
+    assert actor_location not in locations_to_intervals
+
+    # Create more actors to fill up all the GPUs.
+    more_actors = [Actor1.remote() for _ in range(num_nodes * num_gpus_per_raylet - 1)]
+    # Wait for the actors to finish being created.
+    ray.get([actor.get_location_and_ids.remote() for actor in more_actors])
+
+    # Now if we run some GPU tasks, they should not be scheduled.
+    results = [f1.remote() for _ in range(30)]
+    ready_ids, remaining_ids = ray.wait(results, timeout=1.0)
+    assert len(ready_ids) == 0
+
+
+def test_actors_and_tasks_with_gpus_version_two(shutdown_only):
+    # Create tasks and actors that both use GPUs and make sure that they
+    # are given different GPUs
+    num_gpus = 4
+
+    ray.init(
+        num_cpus=(num_gpus + 1),
+        num_gpus=num_gpus,
+        object_store_memory=int(150 * 1024 * 1024),
+    )
+
+    # The point of this actor is to record which GPU IDs have been seen. We
+    # can't just return them from the tasks, because the tasks don't return
+    # for a long time in order to make sure the GPU is not released
+    # prematurely.
+    @ray.remote
+    class RecordGPUs:
+        def __init__(self):
+            self.gpu_ids_seen = []
+            self.num_calls = 0
+
+        def add_ids(self, gpu_ids):
+            self.gpu_ids_seen += gpu_ids
+            self.num_calls += 1
+
+        def get_gpu_ids_and_calls(self):
+            return self.gpu_ids_seen, self.num_calls
+
+    @ray.remote(num_gpus=1)
+    def f(record_gpu_actor):
+        gpu_ids = ray.get_gpu_ids()
+        assert len(gpu_ids) == 1
+        record_gpu_actor.add_ids.remote(gpu_ids)
+        # Sleep for a long time so that the GPU never gets released. This task
+        # will be killed by ray.shutdown() before it actually finishes.
+        time.sleep(1000)
+
+    @ray.remote(num_gpus=1)
+    class Actor:
+        def __init__(self, record_gpu_actor):
+            self.gpu_ids = ray.get_gpu_ids()
+            assert len(self.gpu_ids) == 1
+            record_gpu_actor.add_ids.remote(self.gpu_ids)
+
+        def check_gpu_ids(self):
+            assert ray.get_gpu_ids() == self.gpu_ids
+
+    record_gpu_actor = RecordGPUs.remote()
+
+    actors = []
+    actor_results = []
+    for _ in range(num_gpus // 2):
+        f.remote(record_gpu_actor)
+        a = Actor.remote(record_gpu_actor)
+        actor_results.append(a.check_gpu_ids.remote())
+        # Prevent the actor handle from going out of scope so that its GPU
+        # resources don't get released.
+        actors.append(a)
+
+    # Make sure that the actor method calls succeeded.
+    ray.get(actor_results)
+
+    start_time = time.time()
+    while time.time() - start_time < 30:
+        seen_gpu_ids, num_calls = ray.get(
+            record_gpu_actor.get_gpu_ids_and_calls.remote()
+        )
+        if num_calls == num_gpus:
+            break
+    assert set(seen_gpu_ids) == set(range(num_gpus))
+
+
+def test_blocking_actor_task(shutdown_only):
+    ray.init(num_cpus=1, num_gpus=1, object_store_memory=int(150 * 1024 * 1024))
+
+    @ray.remote(num_gpus=1)
+    def f():
+        return 1
+
+    @ray.remote
+    class Foo:
+        def __init__(self):
+            pass
+
+        def blocking_method(self):
+            ray.get(f.remote())
+
+    # Make sure we can execute a blocking actor method even if there is
+    # only one CPU.
+    actor = Foo.remote()
+    ray.get(actor.blocking_method.remote())
+
+    @ray.remote(num_cpus=1)
+    class CPUFoo:
+        def __init__(self):
+            pass
+
+        def blocking_method(self):
+            ray.get(f.remote())
+
+    # Make sure that lifetime CPU resources are not released when actors
+    # block.
+    actor = CPUFoo.remote()
+    x_id = actor.blocking_method.remote()
+    ready_ids, remaining_ids = ray.wait([x_id], timeout=1.0)
+    assert ready_ids == []
+    assert remaining_ids == [x_id]
+
+    @ray.remote(num_gpus=1)
+    class GPUFoo:
+        def __init__(self):
+            pass
+
+        def blocking_method(self):
+            ray.get(f.remote())
+
+    # Make sure that GPU resources are not released when actors block.
+    actor = GPUFoo.remote()
+    x_id = actor.blocking_method.remote()
+    ready_ids, remaining_ids = ray.wait([x_id], timeout=1.0)
+    assert ready_ids == []
+    assert remaining_ids == [x_id]
+
+
+def test_lifetime_and_transient_resources(ray_start_regular):
+    # This actor acquires resources only when running methods.
+    @ray.remote
+    class Actor1:
+        def method(self):
+            pass
+
+    # This actor acquires resources for its lifetime.
+    @ray.remote(num_cpus=1)
+    class Actor2:
+        def method(self):
+            pass
+
+    actor1s = [Actor1.remote() for _ in range(10)]
+    ray.get([a.method.remote() for a in actor1s])
+
+    actor2s = [Actor2.remote() for _ in range(2)]
+    results = [a.method.remote() for a in actor2s]
+    ready_ids, remaining_ids = ray.wait(results, num_returns=len(results), timeout=5.0)
+    assert len(ready_ids) == 1
+
+
+def test_custom_label_placement(ray_start_cluster):
+    cluster = ray_start_cluster
+    custom_resource1_node = cluster.add_node(
+        num_cpus=2, resources={"CustomResource1": 2}
+    )
+    custom_resource2_node = cluster.add_node(
+        num_cpus=2, resources={"CustomResource2": 2}
+    )
+    ray.init(address=cluster.address)
+
+    @ray.remote(resources={"CustomResource1": 1})
+    class ResourceActor1:
+        def get_location(self):
+            return ray._private.worker.global_worker.node.unique_id
+
+    @ray.remote(resources={"CustomResource2": 1})
+    class ResourceActor2:
+        def get_location(self):
+            return ray._private.worker.global_worker.node.unique_id
+
+    # Create some actors.
+    actors1 = [ResourceActor1.remote() for _ in range(2)]
+    actors2 = [ResourceActor2.remote() for _ in range(2)]
+    locations1 = ray.get([a.get_location.remote() for a in actors1])
+    locations2 = ray.get([a.get_location.remote() for a in actors2])
+    for location in locations1:
+        assert location == custom_resource1_node.unique_id
+    for location in locations2:
+        assert location == custom_resource2_node.unique_id
+
+
+def test_creating_more_actors_than_resources(shutdown_only):
+    ray.init(num_cpus=10, num_gpus=2, resources={"CustomResource1": 1})
+
+    @ray.remote(num_gpus=1)
+    class ResourceActor1:
+        def method(self):
+            return ray.get_gpu_ids()[0]
+
+    @ray.remote(resources={"CustomResource1": 1})
+    class ResourceActor2:
+        def method(self):
+            pass
+
+    # Make sure the first two actors get created and the third one does
+    # not.
+    actor1 = ResourceActor1.remote()
+    result1 = actor1.method.remote()
+    ray.wait([result1])
+    actor2 = ResourceActor1.remote()
+    result2 = actor2.method.remote()
+    ray.wait([result2])
+    actor3 = ResourceActor1.remote()
+    result3 = actor3.method.remote()
+    ready_ids, _ = ray.wait([result3], timeout=0.2)
+    assert len(ready_ids) == 0
+
+    # By deleting actor1, we free up resources to create actor3.
+    del actor1
+
+    results = ray.get([result1, result2, result3])
+    assert results[0] == results[2]
+    assert set(results) == {0, 1}
+
+    # Make sure that when one actor goes out of scope a new actor is
+    # created because some resources have been freed up.
+    results = []
+    for _ in range(3):
+        actor = ResourceActor2.remote()
+        object_ref = actor.method.remote()
+        results.append(object_ref)
+        # Wait for the task to execute. We do this because otherwise it may
+        # be possible for the __ray_terminate__ task to execute before the
+        # method.
+        ray.wait([object_ref])
+
+    ray.get(results)
+
+
+def test_actor_cuda_visible_devices(shutdown_only):
+    """Test user can overwrite CUDA_VISIBLE_DEVICES
+    after the actor is created."""
+    ray.init(num_gpus=1)
+
+    @ray.remote(num_gpus=1)
+    class Actor:
+        def set_cuda_visible_devices(self, cuda_visible_devices):
+            os.environ["CUDA_VISIBLE_DEVICES"] = cuda_visible_devices
+
+        def get_cuda_visible_devices(self):
+            return os.environ["CUDA_VISIBLE_DEVICES"]
+
+    actor = Actor.remote()
+    assert ray.get(actor.get_cuda_visible_devices.remote()) == "0"
+    ray.get(actor.set_cuda_visible_devices.remote("0,1"))
+    assert ray.get(actor.get_cuda_visible_devices.remote()) == "0,1"
+
+
+if __name__ == "__main__":
+    sys.exit(pytest.main(["-sv", __file__]))