Update

[ghstack-poisoned]

Author: kurtamohler

test/test_rb.py

@@ -116,6 +116,7 @@
 _has_gym = importlib.util.find_spec("gym") is not None
 _has_snapshot = importlib.util.find_spec("torchsnapshot") is not None
 _os_is_windows = sys.platform == "win32"
+TORCH_VERSION = version.parse(version.parse(torch.__version__).base_version)
 
 torch_2_3 = version.parse(
     ".".join([str(s) for s in version.parse(str(torch.__version__)).release])
@@ -404,14 +405,16 @@ def data_iter():
         ) if cond else contextlib.nullcontext():
             rb.extend(data2)
 
+    @pytest.mark.skipif(
+        TORCH_VERSION < version.parse("2.5.0"), reason="requires Torch >= 2.5.0"
+    )
+    # Compiling on Windows requires "cl" compiler to be installed.
+    # <https://github.com/pytorch/pytorch/blob/8231180147a096a703d8891756068c89365292e0/torch/_inductor/cpp_builder.py#L143>
+    # Our Windows CI jobs do not have "cl", so skip this test.
+    @pytest.mark.skipif(_os_is_windows, reason="windows tests do not support compile")
     def test_extend_sample_recompile(
         self, rb_type, sampler, writer, storage, size, datatype
     ):
-        if _os_is_windows:
-            # Compiling on Windows requires "cl" compiler to be installed.
-            # <https://github.com/pytorch/pytorch/blob/8231180147a096a703d8891756068c89365292e0/torch/_inductor/cpp_builder.py#L143>
-            # Our Windows CI jobs do not have "cl", so skip this test.
-            pytest.skip("This test does not support Windows.")
         if rb_type is not ReplayBuffer:
             pytest.skip(
                 "Only replay buffer of type 'ReplayBuffer' is currently supported."
@@ -429,7 +432,7 @@ def test_extend_sample_recompile(
         if datatype == "tensordict":
             pytest.skip("'tensordict' datatype is not currently supported.")
 
-        torch._dynamo.reset()
+        torch._dynamo.reset_code_caches()
 
         storage_size = 10 * size
         rb = self._get_rb(

torchrl/data/replay_buffers/storages.py

@@ -144,9 +144,19 @@ def load_state_dict(self, state_dict: Dict[str, Any]) -> None:
     def _empty(self):
         ...
 
-    # NOTE: This property is used to enable compiled Storages. A `len(self)`
-    # call can cause recompiles, but for some reason, wrapping the call in a
-    # `property` decorated function avoids the recompiles.
+    # NOTE: This property is used to enable compiled Storages. Calling
+    # `len(self)` on a TensorStorage should normally cause a graph break since
+    # it uses a `mp.Value`, and it does cause a break when the `len(self)` call
+    # happens within a method of TensorStorage itself. However, when the
+    # `len(self)` call happens in the Storage base class, for an unknown reason
+    # the compiler doesn't seem to recognize that there should be a graph break,
+    # and the lack of a break causes a recompile each time `len(self)` is called
+    # in this context. Also for an unknown reason, we can force the graph break
+    # to happen if we wrap the `len(self)` call with a `property`-decorated
+    # function. For another unknown reason, if we change
+    # `TensorStorage._len_value` from `mp.Value` to int, it seems like there
+    # should no longer be any need to recompile, but recompiles happen anyway.
+    # Ideally, this should all be investigated and understood in the future.
     @property
     def len(self):
         return len(self)