def check_if_run_distributed(cfg):
from easyrl import HOROVOD_AVAILABLE
if HOROVOD_AVAILABLE:
import horovod.torch as hvd
hvd.init()
if hvd.size() > 1:
cfg.distributed = True
if cfg.distributed and not HOROVOD_AVAILABLE:
logger.error('Horovod is not installed! Will not run in distributed training')
distributed = HOROVOD_AVAILABLE and cfg.distributed
cfg.distributed = distributed
if distributed:
gpu_id = hvd.local_rank() + cfg.gpu_shift
if cfg.gpus is not None:
gpu_id = cfg.gpus[gpu_id]
logger.info(f'Rank {hvd.local_rank()} GPU ID: {gpu_id}')
torch.cuda.set_device(gpu_id)
logger.info(f'Using Horovod for distributed training, number of processes:{hvd.size()}')
return distributed
def get_git_infos(path):
git_info = None
try:
repo = git.Repo(path)
try:
branch_name = repo.active_branch.name
except TypeError:
branch_name = '[DETACHED]'
git_info = dict(
directory=str(path),
code_diff=repo.git.diff(None),
code_diff_staged=repo.git.diff('--staged'),
commit_hash=repo.head.commit.hexsha,
branch_name=branch_name,
)
except git.exc.InvalidGitRepositoryError as e:
logger.error(f'Not a valid git repo: {path}')
except git.exc.NoSuchPathError as e:
logger.error(f'{path} does not exist.')
return git_info
def _subproc_worker(pipe, parent_pipe, env_fn_wrapper, obs_bufs, obs_shapes, obs_dtypes, keys):
"""
Control a single environment instance using IPC and
shared mem.
"""
def _write_obs(maybe_dict_obs):
flatdict = obs_to_dict(maybe_dict_obs)
for k in keys:
dst = obs_bufs[k].get_obj()
dst_np = np.frombuffer(dst,
dtype=obs_dtypes[k]).reshape(obs_shapes[k]) # pylint: disable=W0212
np.copyto(dst_np, flatdict[k])
env = env_fn_wrapper.x()
parent_pipe.close()
try:
while True:
cmd, data = pipe.recv()
if cmd == 'reset':
if data is None:
pipe.send(_write_obs(env.reset()))
else:
pipe.send(_write_obs(env.reset(**data)))
elif cmd == 'step':
obs, reward, done, info = env.step(data)
if done:
info['true_next_ob'] = obs
obs = env.reset()
pipe.send((_write_obs(obs), reward, done, info))
elif cmd == 'render':
pipe.send(env.render(mode='rgb_array'))
elif cmd == 'close':
pipe.send(None)
break
else:
raise RuntimeError('Got unrecognized cmd %s' % cmd)
except KeyboardInterrupt:
logger.error('ShmemVecEnv worker: got KeyboardInterrupt')
finally:
env.close()