Kill actor child processes on shutdown (ray-project#3297)

* example * add env * test pg * change to test * add atexit test * Update rllib-env.rst * comment * revert unnecessary file * fix title when actor is idle * Update python/ray/actor.py Co-Authored-By: ericl <ekhliang@gmail.com>

Author: ericl

python/ray/actor.py

@@ -5,6 +5,8 @@
 import copy
 import hashlib
 import inspect
+import os
+import signal
 import traceback
 
 import ray.cloudpickle as pickle
@@ -757,8 +759,10 @@ def __ray_terminate__(self):
  # 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__"):

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")

python/ray/worker.py

@@ -902,6 +902,10 @@ def _become_actor(self, task):
  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 @@ def _wait_for_and_process_task(self, task):
  }
  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):

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