diff --git a/.github/workflows/build-wheels-windows.yml b/.github/workflows/build-wheels-windows.yml index 1d47449568..556d805c64 100644 --- a/.github/workflows/build-wheels-windows.yml +++ b/.github/workflows/build-wheels-windows.yml @@ -37,7 +37,7 @@ jobs: post-script: "python packaging/wheel/relocate.py" smoke-test-script: test/smoke_test.py package-name: torchrl - name: pytorch/rl + name: ${{ matrix.repository }} uses: pytorch/test-infra/.github/workflows/build_wheels_windows.yml@main with: repository: ${{ matrix.repository }} diff --git a/.github/workflows/wheels-legacy.yml b/.github/workflows/wheels-legacy.yml index 25a42b8043..998b9eba4b 100644 --- a/.github/workflows/wheels-legacy.yml +++ b/.github/workflows/wheels-legacy.yml @@ -19,7 +19,7 @@ jobs: runs-on: windows-latest strategy: matrix: - python_version: [["3.9", "3.9"], ["3.10", "3.10.3"], ["3.11", "3.11"], ["3.12", "3.12"]] + python_version: [["3.9", "3.9"], ["3.10", "3.10.3"], ["3.11", "3.11"], ["3.12", "3.12"]] steps: - name: Setup Python uses: actions/setup-python@v2 @@ -37,12 +37,12 @@ jobs: python3 -mpip install wheel TORCHRL_BUILD_VERSION=0.5.0 python3 setup.py bdist_wheel - name: Upload wheel for the test-wheel job - uses: actions/upload-artifact@v2 + uses: actions/upload-artifact@v3 with: name: torchrl-win-${{ matrix.python_version[0] }}.whl path: dist/torchrl-*.whl - name: Upload wheel for download - uses: actions/upload-artifact@v2 + uses: actions/upload-artifact@v3 with: name: torchrl-batch.whl path: dist/*.whl @@ -77,7 +77,7 @@ jobs: run: | python3 -mpip install numpy pytest pytest-cov codecov unittest-xml-reporting pillow>=4.1.1 scipy av networkx expecttest pyyaml - name: Download built wheels - uses: actions/download-artifact@v2 + uses: actions/download-artifact@v3 with: name: torchrl-win-${{ matrix.python_version }}.whl path: wheels diff --git a/test/test_rb.py b/test/test_rb.py index af7140b398..3d8db28553 100644 --- a/test/test_rb.py +++ b/test/test_rb.py @@ -26,6 +26,7 @@ assert_allclose_td, is_tensor_collection, is_tensorclass, + LazyStackedTensorDict, tensorclass, TensorDict, TensorDictBase, @@ -715,6 +716,20 @@ def test_storage_state_dict(self, storage_in, storage_out, init_out, backend): s = new_replay_buffer.sample() assert (s.exclude("index") == 1).all() + @pytest.mark.parametrize("storage_type", [LazyMemmapStorage, LazyTensorStorage]) + def test_extend_lazystack(self, storage_type): + + rb = ReplayBuffer( + storage=storage_type(6), + batch_size=2, + ) + td1 = TensorDict(a=torch.rand(5, 4, 8), batch_size=5) + td2 = TensorDict(a=torch.rand(5, 3, 8), batch_size=5) + ltd = LazyStackedTensorDict(td1, td2, stack_dim=1) + rb.extend(ltd) + rb.sample(3) + assert len(rb) == 5 + @pytest.mark.parametrize("device_data", get_default_devices()) @pytest.mark.parametrize("storage_type", [LazyMemmapStorage, LazyTensorStorage]) @pytest.mark.parametrize("data_type", ["tensor", "tc", "td", "pytree"]) diff --git a/torchrl/data/replay_buffers/storages.py b/torchrl/data/replay_buffers/storages.py index 20b2169cc8..f047d5b2d2 100644 --- a/torchrl/data/replay_buffers/storages.py +++ b/torchrl/data/replay_buffers/storages.py @@ -5,6 +5,8 @@ from __future__ import annotations import abc + +import logging import os import textwrap import warnings @@ -1116,16 +1118,17 @@ def max_size_along_dim0(data_shape): out = data.clone().to(self.device) out = out.expand(max_size_along_dim0(data.shape)) out = out.memmap_like(prefix=self.scratch_dir, existsok=self.existsok) - for key, tensor in sorted( - out.items(include_nested=True, leaves_only=True), key=str - ): - try: - filesize = os.path.getsize(tensor.filename) / 1024 / 1024 - torchrl_logger.debug( - f"\t{key}: {tensor.filename}, {filesize} Mb of storage (size: {tensor.shape})." - ) - except (AttributeError, RuntimeError): - pass + if torchrl_logger.getEffectiveLevel() == logging.DEBUG: + for key, tensor in sorted( + out.items(include_nested=True, leaves_only=True), key=str + ): + try: + filesize = os.path.getsize(tensor.filename) / 1024 / 1024 + torchrl_logger.debug( + f"\t{key}: {tensor.filename}, {filesize} Mb of storage (size: {tensor.shape})." + ) + except (AttributeError, RuntimeError): + pass else: out = _init_pytree(self.scratch_dir, max_size_along_dim0, data) self._storage = out