Skip to content

Commit

Permalink
ProcConcatVec close fix (#219)
Browse files Browse the repository at this point in the history
  • Loading branch information
elliottower committed Jul 7, 2023
2 parents 96fa00d + 4bc83fe commit 9b26d8c
Show file tree
Hide file tree
Showing 2 changed files with 26 additions and 19 deletions.
2 changes: 1 addition & 1 deletion supersuit/__init__.py
Original file line number Diff line number Diff line change
Expand Up @@ -65,4 +65,4 @@ def __getattr__(wrapper_name):
raise ImportError(f"cannot import name '{wrapper_name}' from 'supersuit'")


__version__ = "3.8.0"
__version__ = "3.8.1"
43 changes: 25 additions & 18 deletions supersuit/vector/multiproc_vec.py
Original file line number Diff line number Diff line change
@@ -1,5 +1,6 @@
import copy
import multiprocessing as mp
import time
import traceback

import gymnasium.vector
Expand Down Expand Up @@ -133,6 +134,8 @@ def __init__(
self.observation_space, self.shared_obs, n=self.num_envs
)

self.graceful_shutdown_timeout = 10

pipes = []
procs = []
for constr in vec_env_constrs:
Expand Down Expand Up @@ -219,13 +222,7 @@ def step(self, actions):
return self.step_wait()

def __del__(self):
for pipe in self.pipes:
try:
pipe.send("terminate")
except ConnectionError:
pass
for proc in self.procs:
proc.join()
self.close()

def render(self):
self.pipes[0].send("render")
Expand All @@ -239,17 +236,27 @@ def render(self):
return render_result

def close(self):
for pipe in self.pipes:
pipe.send("close")
for pipe in self.pipes:
try:
pipe.recv()
except EOFError:
raise RuntimeError(
"only one multiproccessing vector environment can open a window over the duration of a process"
)
except ConnectionError:
pass
try:
for pipe, proc in zip(self.pipes, self.procs):
if proc.is_alive():
pipe.send("close")
except OSError:
pass
else:
deadline = (
None
if self.graceful_shutdown_timeout is None
else time.monotonic() + self.graceful_shutdown_timeout
)
for proc in self.procs:
timeout = None if deadline is None else deadline - time.monotonic()
if timeout is not None and timeout <= 0:
break
proc.join(timeout)
for pipe, proc in zip(self.pipes, self.procs):
if proc.is_alive():
proc.kill()
pipe.close()

def env_is_wrapped(self, wrapper_class, indices=None):
for i, pipe in enumerate(self.pipes):
Expand Down

0 comments on commit 9b26d8c

Please sign in to comment.