diff --git a/python/ray/actor.py b/python/ray/actor.py index c52d7ba2a..b20da5f05 100644 --- a/python/ray/actor.py +++ b/python/ray/actor.py @@ -275,6 +275,13 @@ def actor(*args, **kwargs): def make_actor(Class, num_cpus, num_gpus): + # Modify the class to have an additional method that will be used for + # terminating the worker. + class Class(Class): + def __ray_terminate__(self): + import os + os._exit(0) + class_id = random_actor_class_id() # The list exported will have length 0 if the class has not been exported # yet, and length one if it has. This is just implementing a bool, but we @@ -389,6 +396,11 @@ def make_actor(Class, num_cpus, num_gpus): def __reduce__(self): raise Exception("Actor objects cannot be pickled.") + def __del__(self): + """Kill the worker that is running this actor.""" + actor_method_call(self._ray_actor_id, "__ray_terminate__", + self._ray_method_signatures["__ray_terminate__"]) + return NewClass diff --git a/test/actor_test.py b/test/actor_test.py index 65623b0d0..c250652f9 100644 --- a/test/actor_test.py +++ b/test/actor_test.py @@ -682,10 +682,10 @@ class ActorsWithGPUs(unittest.TestCase): tuple(self.gpu_ids)) # Create some actors. - actors = [Actor1.remote() for _ in range(num_local_schedulers * 2)] + actors1 = [Actor1.remote() for _ in range(num_local_schedulers * 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 actors]) + for actor in actors1]) node_names = set([location for location, gpu_id in locations_and_ids]) self.assertEqual(len(node_names), num_local_schedulers) @@ -711,10 +711,10 @@ class ActorsWithGPUs(unittest.TestCase): tuple(self.gpu_ids)) # Create some actors. - actors = [Actor2.remote() for _ in range(num_local_schedulers)] + actors2 = [Actor2.remote() for _ in range(num_local_schedulers)] # 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]) + for actor in actors2]) self.assertEqual(node_names, set([location for location, gpu_id in locations_and_ids])) for location, gpu_ids in locations_and_ids: @@ -962,10 +962,14 @@ class ActorsWithGPUs(unittest.TestCase): return self.gpu_ids[0] results = [] + actors = [] for _ in range(5): results.append(f.remote()) a = Actor.remote() results.append(a.get_gpu_id.remote()) + # Prevent the actor handle from going out of scope so that its GPU + # resources don't get released. + actors.append(a) gpu_ids = ray.get(results) self.assertEqual(set(gpu_ids), set(range(10)))