Better environment shutdown (#2620)

* Wait for env process to exit before killing it

* don't propagate signals, better error logging

* set proc1 to None when done

* comments

ml-agents-envs/mlagents/envs/environment.py

 
 from .rpc_communicator import RpcCommunicator
 from sys import platform
+import signal
 
 logging.basicConfig(level=logging.INFO)
 logger = logging.getLogger("mlagents.envs")
         self.proc1 = (
             None
         )  # The process that is started. If None, no process was started
+        self.timeout_wait: int = timeout_wait
         self.communicator = self.get_communicator(worker_id, base_port, timeout_wait)
         self.worker_id = worker_id
 
                 subprocess_args += ["--port", str(self.port)]
                 subprocess_args += args
                 try:
-                    self.proc1 = subprocess.Popen(subprocess_args)
+                    self.proc1 = subprocess.Popen(
+                        subprocess_args,
+                        # start_new_session=True means that signals to the parent python process
+                        # (e.g. SIGINT from keyboard interrupt) will not be sent to the new process on POSIX platforms.
+                        # This is generally good since we want the environment to have a chance to shutdown,
+                        # but may be undesirable in come cases; if so, we'll add a command-line toggle.
+                        # Note that on Windows, the CTRL_C signal will still be sent.
+                        start_new_session=True,
+                    )
                 except PermissionError as perm:
                     # This is likely due to missing read or execute permissions on file.
                     raise UnityEnvironmentException(
         self._loaded = False
         self.communicator.close()
         if self.proc1 is not None:
-            self.proc1.kill()
+            # Wait a bit for the process to shutdown, but kill it if it takes too long
+            try:
+                self.proc1.wait(timeout=self.timeout_wait)
+                signal_name = self.returncode_to_signal_name(self.proc1.returncode)
+                signal_name = f" ({signal_name})" if signal_name else ""
+                return_info = f"Environment shut down with return code {self.proc1.returncode}{signal_name}."
+                logger.info(return_info)
+            except subprocess.TimeoutExpired:
+                logger.info("Environment timed out shutting down. Killing...")
+                self.proc1.kill()
+            # Set to None so we don't try to close multiple times.
+            self.proc1 = None
 
     @classmethod
     def _flatten(cls, arr: Any) -> List[float]:
         result = UnityInput()
         result.rl_input.CopyFrom(rl_input)
         return result
+
+    @staticmethod
+    def returncode_to_signal_name(returncode: int) -> Optional[str]:
+        """
+        Try to convert return codes into their corresponding signal name.
+        E.g. returncode_to_signal_name(-2) -> "SIGINT"
+        """
+        try:
+            # A negative value -N indicates that the child was terminated by signal N (POSIX only).
+            s = signal.Signals(-returncode)
+            return s.name
+        except Exception:
+            # Should generally be a ValueError, but catch everything just in case.
+            return None

ml-agents-envs/mlagents/envs/subprocess_env_manager.py

             elif cmd.name == "close":
                 break
     except (KeyboardInterrupt, UnityCommunicationException):
-        print("UnityEnvironment worker: environment stopping.")
+        logger.info(f"UnityEnvironment worker {worker_id}: environment stopping.")
         step_queue.put(EnvironmentResponse("env_close", worker_id, None))
     finally:
         # If this worker has put an item in the step queue that hasn't been processed by the EnvManager, the process
-        logger.debug(f"Worker {worker_id} closing.")
+        logger.debug(f"UnityEnvironment worker {worker_id} closing.")
-        logger.debug(f"Worker {worker_id} done.")
+        logger.debug(f"UnityEnvironment worker {worker_id} done.")
 
 
 class SubprocessEnvManager(EnvManager):

ml-agents-envs/mlagents/envs/tests/test_envs.py

     assert comm.has_been_closed
 
 
+def test_returncode_to_signal_name():
+    assert UnityEnvironment.returncode_to_signal_name(-2) == "SIGINT"
+    assert UnityEnvironment.returncode_to_signal_name(42) is None
+    assert UnityEnvironment.returncode_to_signal_name("SIGINT") is None
+
+
 if __name__ == "__main__":
     pytest.main()