Kill actor child processes on shutdown (#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>

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

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 @@ 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):