diff --git a/python/ray/actor.py b/python/ray/actor.py index 86b87e7d5..6dfd62ff3 100644 --- a/python/ray/actor.py +++ b/python/ray/actor.py @@ -5,6 +5,8 @@ from __future__ import print_function import copy import hashlib import inspect +import os +import signal import traceback import ray.cloudpickle as pickle @@ -757,8 +759,10 @@ def make_actor(cls, num_cpus, num_gpus, resources, actor_method_cpus, # this is so that when the worker kills itself below, the local # scheduler won't push an error message to the driver. worker.local_scheduler_client.disconnect() - import os - os._exit(0) + # Kill the process group. We will get the SIGTERM and + # eventually call sys.exit(0) in worker.py as a result of this. + os.killpg(os.getpgid(os.getpid()), signal.SIGTERM) + assert False, "This process should have terminated." def __ray_save_checkpoint__(self): if hasattr(self, "__ray_save__"): diff --git a/python/ray/rllib/test/test_env_with_subprocess.py b/python/ray/rllib/test/test_env_with_subprocess.py new file mode 100644 index 000000000..44e833a67 --- /dev/null +++ b/python/ray/rllib/test/test_env_with_subprocess.py @@ -0,0 +1,77 @@ +"""Tests that envs clean up after themselves on agent exit.""" + +from __future__ import absolute_import +from __future__ import division +from __future__ import print_function + +from gym.spaces import Discrete +import atexit +import gym +import os +import subprocess +import tempfile +import time + +import ray +from ray.tune import run_experiments +from ray.tune.registry import register_env + +# Dummy command to run as a subprocess with a unique name +UNIQUE_CMD = "sleep {}".format(str(time.time())) +_, UNIQUE_FILE_0 = tempfile.mkstemp("test_env_with_subprocess") +_, UNIQUE_FILE_1 = tempfile.mkstemp("test_env_with_subprocess") + + +class EnvWithSubprocess(gym.Env): + """Our env that spawns a subprocess.""" + + def __init__(self, config): + self.action_space = Discrete(2) + self.observation_space = Discrete(2) + # Subprocess that should be cleaned up + self.subproc = subprocess.Popen(UNIQUE_CMD, shell=True) + # Exit handler should be called + if config.worker_index == 0: + atexit.register(lambda: os.unlink(UNIQUE_FILE_0)) + else: + atexit.register(lambda: os.unlink(UNIQUE_FILE_1)) + + def reset(self): + return [0] + + def step(self, action): + return [0], 0, True, {} + + +def leaked_processes(): + """Returns whether any subprocesses were leaked.""" + result = subprocess.check_output( + "ps aux | grep '{}' | grep -v grep || true".format(UNIQUE_CMD), + shell=True) + return result + + +if __name__ == "__main__": + register_env("subproc", lambda config: EnvWithSubprocess(config)) + ray.init() + assert os.path.exists(UNIQUE_FILE_0) + assert os.path.exists(UNIQUE_FILE_1) + assert not leaked_processes() + run_experiments({ + "demo": { + "run": "PG", + "env": "subproc", + "num_samples": 1, + "config": { + "num_workers": 1, + }, + "stop": { + "training_iteration": 1 + }, + }, + }) + leaked = leaked_processes() + assert not leaked, "LEAKED PROCESSES: {}".format(leaked) + assert not os.path.exists(UNIQUE_FILE_0), "atexit handler not called" + assert not os.path.exists(UNIQUE_FILE_1), "atexit handler not called" + print("OK") diff --git a/python/ray/worker.py b/python/ray/worker.py index 105332a72..5ed52f407 100644 --- a/python/ray/worker.py +++ b/python/ray/worker.py @@ -902,6 +902,10 @@ class Worker(object): self.actor_id = task.actor_creation_id().id() class_id = arguments[0] + # Set the process group id. This ensures that child processes spawned + # by this actor can be killed with this actor easily on termination. + os.setpgid(os.getpid(), os.getpid()) + key = b"ActorClass:" + class_id # Wait for the actor class key to have been imported by the import @@ -945,12 +949,14 @@ class Worker(object): } if task.actor_id().id() == NIL_ACTOR_ID: title = "ray_worker:{}()".format(function_name) + next_title = "ray_worker" else: actor = self.actors[task.actor_id().id()] title = "ray_{}:{}()".format(actor.__class__.__name__, function_name) + next_title = "ray_{}".format(actor.__class__.__name__) with profiling.profile("task", extra_data=extra_data, worker=self): - with _changeproctitle(title): + with _changeproctitle(title, next_title): self._process_task(task, execution_info) # Reset the state fields so the next task can run. with self.state_lock: @@ -2163,11 +2169,10 @@ def disconnect(worker=global_worker): @contextmanager -def _changeproctitle(title): - old_title = setproctitle.getproctitle() +def _changeproctitle(title, next_title): setproctitle.setproctitle(title) yield - setproctitle.setproctitle(old_title) + setproctitle.setproctitle(next_title) def _try_to_compute_deterministic_class_id(cls, depth=5): diff --git a/test/jenkins_tests/run_multi_node_tests.sh b/test/jenkins_tests/run_multi_node_tests.sh index 77875d59c..205071abc 100755 --- a/test/jenkins_tests/run_multi_node_tests.sh +++ b/test/jenkins_tests/run_multi_node_tests.sh @@ -261,6 +261,9 @@ docker run --rm --shm-size=10G --memory=10G $DOCKER_SHA \ docker run --rm --shm-size=10G --memory=10G $DOCKER_SHA \ python /ray/python/ray/rllib/test/test_supported_spaces.py +docker run --rm --shm-size=10G --memory=10G $DOCKER_SHA \ + python /ray/python/ray/rllib/test/test_env_with_subprocess.py + docker run --rm --shm-size=10G --memory=10G $DOCKER_SHA \ /ray/python/ray/rllib/test/test_rollout.sh