From 97333cff14f2742dd76fe503dd35c0d11e6c06a0 Mon Sep 17 00:00:00 2001 From: =?UTF-8?q?Nathan=20Lichtl=C3=A9?= Date: Fri, 10 Jul 2020 02:43:56 +0200 Subject: [PATCH] add --multi_node flag --- examples/train.py | 7 ++++++- 1 file changed, 6 insertions(+), 1 deletion(-) diff --git a/examples/train.py b/examples/train.py index d062fd39a..3bdd751e6 100644 --- a/examples/train.py +++ b/examples/train.py @@ -84,6 +84,9 @@ def parse_args(args): parser.add_argument( '--checkpoint_path', type=str, default=None, help='Directory with checkpoint to restore training from.') + parser.add_argument('--multi_node', action='store_true', + help='Set to true if this will be run in cluster mode.' + 'Relevant for rllib') return parser.parse_known_args(args)[0] @@ -350,7 +353,9 @@ def train_rllib(submodule, flags): def trial_str_creator(trial): return "{}_{}".format(trial.trainable_name, trial.experiment_tag) - if flags.local_mode: + if flags.multi_node: + ray.init(redis_address='localhost:6379') + elif flags.local_mode: ray.init(local_mode=True) else: ray.init()