Update

[ghstack-poisoned]

Author: vmoens

torchrl/collectors/_constants.py

@@ -12,7 +12,6 @@
 import torch
 from torch import multiprocessing as mp
 
-from torchrl._utils import WEIGHT_SYNC_TIMEOUT
 from torchrl.envs.utils import ExplorationType
 
 try:
@@ -40,6 +39,9 @@ def cudagraph_mark_step_begin():
 _TIMEOUT = 1.0
 INSTANTIATE_TIMEOUT = 20
 _MIN_TIMEOUT = 1e-3  # should be several orders of magnitude inferior wrt time spent collecting a trajectory
+# Timeout for weight synchronization during collector init.
+# Increase this when using many collectors across different CUDA devices.
+WEIGHT_SYNC_TIMEOUT = float(os.environ.get("TORCHRL_WEIGHT_SYNC_TIMEOUT", 120.0))
 # MAX_IDLE_COUNT is the maximum number of times a Dataloader worker can timeout with his queue.
 _MAX_IDLE_COUNT = int(os.environ.get("MAX_IDLE_COUNT", torch.iinfo(torch.int64).max))
 

torchrl/weight_update/_shared.py

@@ -10,7 +10,8 @@
 
 from torch import multiprocessing as mp, nn
 
-from torchrl._utils import logger as torchrl_logger, WEIGHT_SYNC_TIMEOUT
+from torchrl._utils import logger as torchrl_logger
+from torchrl.collectors._constants import WEIGHT_SYNC_TIMEOUT
 
 from torchrl.weight_update.utils import _resolve_model
 from torchrl.weight_update.weight_sync_schemes import (