[Feature] Checkpointing (pytorch#549)

* init * Update HABITAT.md * init * [Feature] checkpointing RewardNormalization (pytorch#550) * amend * amend * [Feature] checkpointing ReplayBuffer (pytorch#553) * init * amend * amend * amend * [Feature] checkpointing SelectKeys (pytorch#551) * [Feature] checkpointing SubSampler (pytorch#555) * tests * typo * permission * pip install * egl * egl * habitat.utils.gym_definitions * time to impact: 5 mins * hope * amend * amend * LD_PRELOAD * amend * amend * amend * amend * amend * amend * amend * revert yum * amend * amend * amend * amend * amend * LD_LIBRARY_PATH * amend * amend * amend * amend * amend * amend * amend * amend * amend * amend * amend * amend * amend * amend * amend * amend * amend * amend * amend * amend * amend * amend * amend * amend * --force-reinstall * amend * amend * removing habitat baselines * amend * VP suggestions * Checkpointing -- TorchSnapshot (pytorch#577) * init * amend * temp * amend * amend * bf * fix tests * lock bugfix * install snapshot branch * amend * amend * amend * amend * amend * amend * bf * coverage xml -i * amend * amend
ai4co · Oct 28, 2022 · 86380c8 · 86380c8
1 parent fae718a
commit 86380c8
Show file tree

Hide file tree

Showing 13 changed files with 1,259 additions and 228 deletions.
diff --git a/.circleci/unittest/linux/scripts/install.sh b/.circleci/unittest/linux/scripts/install.sh
@@ -36,5 +36,8 @@ fi
 # smoke test
 python -c "import functorch"
 
+# install snapshot
+pip install git+https://github.com/pytorch/torchsnapshot
+
 printf "* Installing torchrl\n"
 python setup.py develop
diff --git a/.circleci/unittest/linux/scripts/run_test.sh b/.circleci/unittest/linux/scripts/run_test.sh
@@ -27,4 +27,4 @@ export MKL_THREADING_LAYER=GNU
 coverage run -m pytest test/smoke_test.py -v --durations 20
 coverage run -m pytest test/smoke_test_deps.py -v --durations 20 -k 'test_gym or test_dm_control_pixels or test_dm_control or test_tb'
 coverage run -m pytest --instafail -v --durations 20
-coverage xml
+coverage xml -i
diff --git a/.circleci/unittest/linux_optdeps/scripts/run_test.sh b/.circleci/unittest/linux_optdeps/scripts/run_test.sh
@@ -11,7 +11,8 @@ python -m torch.utils.collect_env
 git config --global --add safe.directory '*'
 root_dir="$(git rev-parse --show-toplevel)"
 export MKL_THREADING_LAYER=GNU
+export CKPT_BACKEND=torch
 
 #MUJOCO_GL=glfw pytest --cov=torchrl --junitxml=test-results/junit.xml -v --durations 20
 MUJOCO_GL=egl coverage run -m pytest --instafail -v --durations 20
-coverage xml
+coverage xml -i
diff --git a/.circleci/unittest/linux_stable/scripts/run_test.sh b/.circleci/unittest/linux_stable/scripts/run_test.sh
@@ -17,8 +17,9 @@ lib_dir="${env_dir}/lib"
 # solves ImportError: /lib64/libstdc++.so.6: version `GLIBCXX_3.4.21' not found
 export LD_LIBRARY_PATH=$LD_LIBRARY_PATH:$lib_dir
 export MKL_THREADING_LAYER=GNU
+export CKPT_BACKEND=torch
 
 coverage run -m pytest test/smoke_test.py -v --durations 20
 coverage run -m pytest test/smoke_test_deps.py -v --durations 20 -k 'test_gym or test_dm_control_pixels or test_dm_control or test_tb'
 coverage run -m pytest --instafail -v --durations 20
-coverage xml
+coverage xml -i
diff --git a/setup.py b/setup.py
@@ -222,6 +222,9 @@ def _main(argv):
                 "hydra-core>=1.1",
                 "hydra-submitit-launcher",
             ],
+            "checkpointing": [
+                "torchinductor",
+            ],
         },
         zip_safe=False,
         classifiers=[