From 9e68b77796c1478973f5236adfc09f3399721659 Mon Sep 17 00:00:00 2001 From: mvindiola1 Date: Wed, 28 Oct 2020 17:23:06 -0400 Subject: [PATCH] [RLLIB] Wait for remote_workers to finish closing environments before terminating (#11476) --- rllib/evaluation/worker_set.py | 13 +++++++++---- 1 file changed, 9 insertions(+), 4 deletions(-) diff --git a/rllib/evaluation/worker_set.py b/rllib/evaluation/worker_set.py index 5204cfe0a..9172e5f4a 100644 --- a/rllib/evaluation/worker_set.py +++ b/rllib/evaluation/worker_set.py @@ -144,10 +144,15 @@ class WorkerSet: def stop(self) -> None: """Stop all rollout workers.""" - self.local_worker().stop() - for w in self.remote_workers(): - w.stop.remote() - w.__ray_terminate__.remote() + try: + self.local_worker().stop() + tids = [w.stop.remote() for w in self.remote_workers()] + ray.get(tids) + except Exception: + logger.exception("Failed to stop workers") + finally: + for w in self.remote_workers(): + w.__ray_terminate__.remote() @DeveloperAPI def foreach_worker(self, func: Callable[[RolloutWorker], T]) -> List[T]: