[CI] Make MAX_IDLE_COUNT a feature of tests

ghstack-source-id: 2bf31df
Pull Request resolved: #2752

Author: Vincent Moens

.github/unittest/linux/scripts/run_all.sh

@@ -80,7 +80,9 @@ export DISPLAY=:0
 export SDL_VIDEODRIVER=dummy
 
 # legacy from bash scripts: remove?
-conda env config vars set MUJOCO_GL=$MUJOCO_GL PYOPENGL_PLATFORM=$MUJOCO_GL DISPLAY=:0 SDL_VIDEODRIVER=dummy LAZY_LEGACY_OP=False RL_LOGGING_LEVEL=DEBUG TOKENIZERS_PARALLELISM=true
+conda env config vars set \
+  MAX_IDLE_COUNT=1000 \
+  MUJOCO_GL=$MUJOCO_GL PYOPENGL_PLATFORM=$MUJOCO_GL DISPLAY=:0 SDL_VIDEODRIVER=dummy LAZY_LEGACY_OP=False RL_LOGGING_LEVEL=DEBUG TOKENIZERS_PARALLELISM=true
 
 pip3 install pip --upgrade
 pip install virtualenv

.github/unittest/linux_distributed/scripts/setup_env.sh

@@ -64,6 +64,7 @@ fi
 
 export MUJOCO_GL=$PRIVATE_MUJOCO_GL
 conda env config vars set MUJOCO_PY_MUJOCO_PATH=$root_dir/.mujoco/mujoco210 \
+  MAX_IDLE_COUNT=1000 \
   DISPLAY=unix:0.0 \
   MJLIB_PATH=$root_dir/.mujoco/mujoco-2.1.1/lib/libmujoco.so.2.1.1 \
   LD_LIBRARY_PATH=$LD_LIBRARY_PATH:$root_dir/.mujoco/mujoco210/bin \

.github/unittest/linux_libs/scripts_d4rl/setup_env.sh

@@ -85,6 +85,7 @@ fi
 
 export MUJOCO_GL=$PRIVATE_MUJOCO_GL
 conda env config vars set \
+  MAX_IDLE_COUNT=1000 \
   MUJOCO_PY_MUJOCO_PATH=$root_dir/.mujoco/mujoco200_linux \
   DISPLAY=unix:0.0 \
   MJLIB_PATH=$root_dir/.mujoco/mujoco200_linux/bin/libmujoco200.so \

.github/unittest/linux_libs/scripts_gym/setup_env.sh

@@ -71,6 +71,7 @@ cat "${this_dir}/environment.yml"
 
 export MUJOCO_GL=egl
 conda env config vars set \
+  MAX_IDLE_COUNT=1000 \
   MUJOCO_GL=egl \
   SDL_VIDEODRIVER=dummy \
   DISPLAY=unix:0.0 \

.github/unittest/linux_libs/scripts_habitat/run_test.sh

@@ -10,7 +10,9 @@ conda activate ./env
 # https://stackoverflow.com/questions/72540359/glibcxx-3-4-30-not-found-for-librosa-in-conda-virtual-environment-after-tryin
 #conda install -y -c conda-forge gcc=12.1.0
 conda install -y -c conda-forge libstdcxx-ng=12
-conda env config vars set LD_PRELOAD=$LD_PRELOAD:$STDC_LOC TOKENIZERS_PARALLELISM=true
+conda env config vars set \
+  MAX_IDLE_COUNT=1000 \
+  LD_PRELOAD=$LD_PRELOAD:$STDC_LOC TOKENIZERS_PARALLELISM=true
 
 ## find libstdc
 STDC_LOC=$(find conda/ -name "libstdc++.so.6" | head -1)
@@ -36,7 +38,9 @@ export MKL_THREADING_LAYER=GNU
 #wget https://github.com/openai/mujoco-py/blob/master/vendor/10_nvidia.json
 #mv 10_nvidia.json /usr/share/glvnd/egl_vendor.d/10_nvidia.json
 
-conda env config vars set MAGNUM_LOG=quiet HABITAT_SIM_LOG=quiet TOKENIZERS_PARALLELISM=true
+conda env config vars set \
+  MAX_IDLE_COUNT=1000 \
+  MAGNUM_LOG=quiet HABITAT_SIM_LOG=quiet TOKENIZERS_PARALLELISM=true
 
 conda deactivate && conda activate ./env
 

.github/unittest/linux_libs/scripts_habitat/setup_env.sh

@@ -41,7 +41,9 @@ fi
 conda activate "${env_dir}"
 
 # set debug variables
-conda env config vars set MAGNUM_LOG=debug HABITAT_SIM_LOG=debug TOKENIZERS_PARALLELISM=true
+conda env config vars set \
+  MAX_IDLE_COUNT=1000 \
+  MAGNUM_LOG=debug HABITAT_SIM_LOG=debug TOKENIZERS_PARALLELISM=true
 conda deactivate && conda activate "${env_dir}"
 
 pip3 install "cython<3"

.github/unittest/linux_libs/scripts_robohive/setup_env.sh

@@ -61,6 +61,7 @@ cat "${this_dir}/environment.yml"
 
 export MUJOCO_GL=egl
 conda env config vars set \
+  MAX_IDLE_COUNT=1000 \
   MUJOCO_GL=egl \
   SDL_VIDEODRIVER=dummy \
   DISPLAY=unix:0.0 \

.github/unittest/linux_olddeps/scripts_gym_0_13/setup_env.sh

@@ -77,6 +77,7 @@ cat "${this_dir}/environment.yml"
 
 export MUJOCO_GL=egl
 conda env config vars set \
+  MAX_IDLE_COUNT=1000 \
   MUJOCO_GL=egl \
   SDL_VIDEODRIVER=dummy \
   DISPLAY=unix:0.0 \

.github/unittest/linux_sota/scripts/run_all.sh

@@ -77,7 +77,9 @@ export PYOPENGL_PLATFORM=egl
 export LAZY_LEGACY_OP=False
 export COMPOSITE_LP_AGGREGATE=0
 
-conda env config vars set MUJOCO_PY_MUJOCO_PATH=$root_dir/.mujoco/mujoco210 \
+conda env config vars set \
+  MAX_IDLE_COUNT=1000 \
+  MUJOCO_PY_MUJOCO_PATH=$root_dir/.mujoco/mujoco210 \
   DISPLAY=unix:0.0 \
   LD_LIBRARY_PATH=$LD_LIBRARY_PATH:$root_dir/.mujoco/mujoco210/bin \
   SDL_VIDEODRIVER=dummy \
@@ -101,7 +103,9 @@ pip install git+https://github.com/Farama-Foundation/d4rl@master#egg=d4rl
 conda install -y -c conda-forge libstdcxx-ng=12
 ## find libstdc
 STDC_LOC=$(find conda/ -name "libstdc++.so.6" | head -1)
-conda env config vars set LD_PRELOAD=${root_dir}/$STDC_LOC TOKENIZERS_PARALLELISM=true
+conda env config vars set \
+  MAX_IDLE_COUNT=1000 \
+  LD_PRELOAD=${root_dir}/$STDC_LOC TOKENIZERS_PARALLELISM=true
 
 # compile mujoco-py (bc it's done at runtime for whatever reason someone thought it was a good idea)
 python -c """import gym;import d4rl"""

torchrl/collectors/collectors.py

@@ -81,7 +81,7 @@ def cudagraph_mark_step_begin():
 INSTANTIATE_TIMEOUT = 20
 _MIN_TIMEOUT = 1e-3  # should be several orders of magnitude inferior wrt time spent collecting a trajectory
 # 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", 1000))
+_MAX_IDLE_COUNT = int(os.environ.get("MAX_IDLE_COUNT", float("inf")))
 
 DEFAULT_EXPLORATION_TYPE: ExplorationType = ExplorationType.RANDOM