TransferQueue 0.1.1.dev0__py3-none-any.whl

tests/test_samplers.py

@@ -0,0 +1,492 @@
+# Copyright 2025 The TransferQueue Team
+#
+# Licensed under the Apache License, Version 2.0 (the "License");
+# you may not use this file except in compliance with the License.
+# You may obtain a copy of the License at
+#
+#     http://www.apache.org/licenses/LICENSE-2.0
+#
+# Unless required by applicable law or agreed to in writing, software
+# distributed under the License is distributed on an "AS IS" BASIS,
+# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+# See the License for the specific language governing permissions and
+# limitations under the License.
+
+"""Unit tests for TransferQueue samplers."""
+
+from typing import Any
+
+import pytest
+
+from transfer_queue.sampler import BaseSampler
+from transfer_queue.sampler.grpo_group_n_sampler import GRPOGroupNSampler
+from transfer_queue.sampler.sequential_sampler import SequentialSampler
+
+
+class TestBaseSampler:
+    """Test cases for BaseSampler abstract class."""
+
+    def test_base_sampler_is_abstract(self):
+        """Test that BaseSampler cannot be instantiated directly."""
+        with pytest.raises(TypeError) as exc_info:
+            BaseSampler()
+
+        assert "Can't instantiate abstract class" in str(exc_info.value)
+        assert "sample" in str(exc_info.value)
+
+    def test_base_sampler_has_abstract_methods(self):
+        """Test that BaseSampler defines abstract methods."""
+        assert hasattr(BaseSampler, "sample")
+        assert getattr(BaseSampler.sample, "__isabstractmethod__", False)
+
+    def test_base_sampler_has_call_method(self):
+        """Test that BaseSampler has __call__ method."""
+        assert callable(BaseSampler)
+
+    def test_base_sampler_initialization_states(self):
+        """Test BaseSampler initialization sets _states correctly."""
+
+        # Create a concrete implementation for testing
+        class TestSampler(BaseSampler):
+            def sample(self, ready_indexes: list[int], batch_size: int, **kwargs: Any) -> tuple[list[int], list[int]]:
+                return ready_indexes[:batch_size], ready_indexes[:batch_size]
+
+        sampler = TestSampler()
+        assert hasattr(sampler, "_states")
+        assert sampler._states == {}
+
+
+class TestSequentialSampler:
+    """Test cases for SequentialSampler."""
+
+    def test_sequential_sampler_initialization(self):
+        """Test SequentialSampler initialization."""
+        sampler = SequentialSampler()
+        assert isinstance(sampler, BaseSampler)
+        assert hasattr(sampler, "_states")
+        assert sampler._states == {}
+
+    def test_sequential_sampler_basic_functionality(self):
+        """Test basic sampling functionality."""
+        sampler = SequentialSampler()
+        ready_indexes = [0, 1, 2, 3, 4, 5]
+        batch_size = 3
+
+        sampled, consumed = sampler.sample(ready_indexes, batch_size)
+
+        assert sampled == [0, 1, 2]
+        assert consumed == [0, 1, 2]
+        assert len(sampled) == batch_size
+        assert len(consumed) == batch_size
+
+    def test_sequential_sampler_empty_ready_indexes(self):
+        """Test behavior with empty ready indexes."""
+        sampler = SequentialSampler()
+        ready_indexes = []
+        batch_size = 3
+
+        sampled, consumed = sampler.sample(ready_indexes, batch_size)
+
+        assert sampled == []
+        assert consumed == []
+
+    def test_sequential_sampler_batch_size_larger_than_ready(self):
+        """Test behavior when batch_size > len(ready_indexes)."""
+        sampler = SequentialSampler()
+        ready_indexes = [0, 1]
+        batch_size = 5
+
+        sampled, consumed = sampler.sample(ready_indexes, batch_size)
+
+        assert sampled == [0, 1]
+        assert consumed == [0, 1]
+        assert len(sampled) == len(ready_indexes)
+
+    def test_sequential_sampler_zero_batch_size(self):
+        """Test behavior with zero batch size."""
+        sampler = SequentialSampler()
+        ready_indexes = [0, 1, 2, 3]
+        batch_size = 0
+
+        sampled, consumed = sampler.sample(ready_indexes, batch_size)
+
+        assert sampled == []
+        assert consumed == []
+
+    def test_sequential_sampler_negative_batch_size(self):
+        """Test behavior with negative batch size."""
+        sampler = SequentialSampler()
+        ready_indexes = [0, 1, 2, 3]
+        batch_size = -1
+
+        sampled, consumed = sampler.sample(ready_indexes, batch_size)
+
+        # Python slicing with negative numbers should work as expected
+        expected = ready_indexes[:batch_size]  # This gives [0, 1, 2] for -1
+        assert sampled == expected
+        assert consumed == expected
+
+    def test_sequential_sampler_non_sequential_indexes(self):
+        """Test behavior with non-sequential ready indexes."""
+        sampler = SequentialSampler()
+        ready_indexes = [10, 5, 15, 20, 8]
+        batch_size = 3
+
+        sampled, consumed = sampler.sample(ready_indexes, batch_size)
+
+        assert sampled == [10, 5, 15]
+        assert consumed == [10, 5, 15]
+
+    def test_sequential_sampler_duplicate_indexes(self):
+        """Test behavior with duplicate indexes."""
+        sampler = SequentialSampler()
+        ready_indexes = [0, 1, 0, 2, 1, 3]
+        batch_size = 4
+
+        sampled, consumed = sampler.sample(ready_indexes, batch_size)
+
+        assert sampled == [0, 1, 0, 2]
+        assert consumed == [0, 1, 0, 2]
+
+    def test_sequential_sampler_call_method(self):
+        """Test that __call__ method works correctly."""
+        sampler = SequentialSampler()
+        ready_indexes = [0, 1, 2, 3]
+        batch_size = 2
+
+        sampled, consumed = sampler(ready_indexes, batch_size)
+
+        assert sampled == [0, 1]
+        assert consumed == [0, 1]
+
+    def test_sequential_sampler_with_extra_kwargs(self):
+        """Test that SequentialSampler accepts extra kwargs but ignores them."""
+        sampler = SequentialSampler()
+        ready_indexes = [0, 1, 2, 3]
+        batch_size = 2
+
+        # SequentialSampler should accept extra kwargs but ignore them
+        sampled, consumed = sampler.sample(ready_indexes, batch_size, extra_param="ignored")
+
+        assert sampled == [0, 1]
+        assert consumed == [0, 1]
+
+
+class TestGRPOGroupNSampler:
+    """Test cases for GRPOGroupNSampler."""
+
+    def test_grpo_sampler_initialization(self):
+        """Test GRPOGroupNSampler initialization."""
+        sampler = GRPOGroupNSampler()
+        assert isinstance(sampler, BaseSampler)
+        assert hasattr(sampler, "_states")
+        assert sampler._states == {}
+
+    def test_grpo_sampler_basic_functionality(self):
+        """Test basic grouped sampling functionality."""
+        sampler = GRPOGroupNSampler()
+        ready_indexes = [0, 1, 2, 3, 4, 5, 6, 7]  # 8 indexes
+        batch_size = 8
+        n_samples_per_prompt = 4  # 2 groups of 4
+
+        sampled, consumed = sampler.sample(ready_indexes, batch_size, n_samples_per_prompt)
+
+        assert sampled == [0, 1, 2, 3, 4, 5, 6, 7]
+        assert consumed == [0, 1, 2, 3, 4, 5, 6, 7]
+        assert len(sampled) == batch_size
+        assert len(consumed) == batch_size
+
+    def test_grpo_sampler_partial_batch(self):
+        """Test partial batch sampling."""
+        sampler = GRPOGroupNSampler()
+        ready_indexes = [0, 1, 2, 3, 4, 5, 6, 7, 8, 9, 10, 11]  # 12 indexes
+        batch_size = 8  # Want 8 samples total
+        n_samples_per_prompt = 4  # 2 groups of 4
+
+        sampled, consumed = sampler.sample(ready_indexes, batch_size, n_samples_per_prompt)
+
+        assert sampled == [0, 1, 2, 3, 4, 5, 6, 7]
+        assert consumed == [0, 1, 2, 3, 4, 5, 6, 7]
+        assert len(sampled) == batch_size
+        assert len(consumed) == batch_size
+
+    def test_grpo_sampler_different_group_sizes(self):
+        """Test different n_samples_per_prompt values."""
+        sampler = GRPOGroupNSampler()
+        ready_indexes = [0, 1, 2, 3, 4, 5, 6, 7, 8, 9, 10, 11, 12, 13, 14, 15]
+
+        # Test with 2 samples per prompt (8 groups)
+        sampled, consumed = sampler.sample(ready_indexes, 8, n_samples_per_prompt=2)
+        assert sampled == [0, 1, 2, 3, 4, 5, 6, 7]
+        assert consumed == [0, 1, 2, 3, 4, 5, 6, 7]
+
+        # Test with 8 samples per prompt (2 groups)
+        sampled, consumed = sampler.sample(ready_indexes, 8, n_samples_per_prompt=8)
+        assert sampled == [0, 1, 2, 3, 4, 5, 6, 7]
+        assert consumed == [0, 1, 2, 3, 4, 5, 6, 7]
+
+    def test_grpo_sampler_batch_size_divisibility(self):
+        """Test that batch_size must be divisible by n_samples_per_prompt."""
+        sampler = GRPOGroupNSampler()
+        ready_indexes = [0, 1, 2, 3, 4, 5, 6, 7]  # 8 indexes, sufficient for batch_size=7
+        batch_size = 7
+        n_samples_per_prompt = 4
+
+        with pytest.raises(ValueError) as exc_info:
+            sampler.sample(ready_indexes, batch_size, n_samples_per_prompt)
+
+        assert "must be a multiple of n_samples_per_prompt" in str(exc_info.value)
+
+    def test_grpo_sampler_insufficient_ready_indexes(self):
+        """Test behavior when not enough ready indexes are available."""
+        sampler = GRPOGroupNSampler()
+        ready_indexes = [0, 1, 2, 3]  # Only 4 indexes, but need 8 for 2 groups of 4
+        batch_size = 8
+        n_samples_per_prompt = 4
+
+        # Should return empty lists when insufficient complete groups
+        sampled, consumed = sampler.sample(ready_indexes, batch_size, n_samples_per_prompt)
+        assert sampled == []
+        assert consumed == []
+
+    def test_grpo_sampler_exact_multiple_available(self):
+        """Test when ready_indexes length is exactly a multiple of n_samples_per_prompt."""
+        sampler = GRPOGroupNSampler()
+        ready_indexes = [0, 1, 2, 3, 4, 5, 6, 7]  # 8 indexes
+        batch_size = 8
+        n_samples_per_prompt = 4
+
+        sampled, consumed = sampler.sample(ready_indexes, batch_size, n_samples_per_prompt)
+
+        assert sampled == [0, 1, 2, 3, 4, 5, 6, 7]
+        assert consumed == [0, 1, 2, 3, 4, 5, 6, 7]
+
+    def test_grpo_sampler_zero_batch_size(self):
+        """Test behavior with zero batch size."""
+        sampler = GRPOGroupNSampler()
+        ready_indexes = [0, 1, 2, 3]
+        batch_size = 0
+        n_samples_per_prompt = 2
+
+        sampled, consumed = sampler.sample(ready_indexes, batch_size, n_samples_per_prompt)
+
+        assert sampled == []
+        assert consumed == []
+
+    def test_grpo_sampler_single_sample_per_prompt(self):
+        """Test with n_samples_per_prompt = 1."""
+        sampler = GRPOGroupNSampler()
+        ready_indexes = [0, 1, 2, 3, 4, 5]
+        batch_size = 3
+        n_samples_per_prompt = 1
+
+        sampled, consumed = sampler.sample(ready_indexes, batch_size, n_samples_per_prompt)
+
+        assert sampled == [0, 1, 2]
+        assert consumed == [0, 1, 2]
+
+    def test_grpo_sampler_large_group_size(self):
+        """Test with large n_samples_per_prompt."""
+        sampler = GRPOGroupNSampler()
+        ready_indexes = list(range(20))  # 20 indexes
+        batch_size = 20
+        n_samples_per_prompt = 10
+
+        sampled, consumed = sampler.sample(ready_indexes, batch_size, n_samples_per_prompt)
+
+        assert sampled == list(range(20))
+        assert consumed == list(range(20))
+
+    def test_grpo_sampler_call_method(self):
+        """Test that __call__ method works correctly."""
+        sampler = GRPOGroupNSampler()
+        ready_indexes = [0, 1, 2, 3, 4, 5, 6, 7]
+        batch_size = 4
+        n_samples_per_prompt = 2
+
+        sampled, consumed = sampler(ready_indexes, batch_size, n_samples_per_prompt=n_samples_per_prompt)
+
+        assert sampled == [0, 1, 2, 3]
+        assert consumed == [0, 1, 2, 3]
+
+    def test_grpo_sampler_parameter_order_independence(self):
+        """Test that parameter order doesn't matter when using kwargs."""
+        sampler = GRPOGroupNSampler()
+        ready_indexes = [0, 1, 2, 3, 4, 5, 6, 7]
+
+        # Try different parameter orders
+        sampled1, consumed1 = sampler.sample(n_samples_per_prompt=4, batch_size=8, ready_indexes=ready_indexes)
+
+        sampled2, consumed2 = sampler.sample(batch_size=8, ready_indexes=ready_indexes, n_samples_per_prompt=4)
+
+        assert sampled1 == sampled2
+        assert consumed1 == consumed2
+
+    def test_grpo_sampler_with_extra_kwargs(self):
+        """Test that GRPOGroupNSampler accepts extra kwargs but ignores them."""
+        sampler = GRPOGroupNSampler()
+        ready_indexes = [0, 1, 2, 3, 4, 5, 6, 7]
+        batch_size = 8
+        n_samples_per_prompt = 4
+
+        # GRPOGroupNSampler should accept extra kwargs but ignore them
+        sampled, consumed = sampler.sample(
+            ready_indexes, batch_size, n_samples_per_prompt, extra_param="ignored", another_param=42
+        )
+
+        assert sampled == [0, 1, 2, 3, 4, 5, 6, 7]
+        assert consumed == [0, 1, 2, 3, 4, 5, 6, 7]
+
+    def test_grpo_sampler_non_sequential_indexes(self):
+        """Test with non-sequential ready indexes that get sorted."""
+        sampler = GRPOGroupNSampler()
+        ready_indexes = [3, 4, 5, 6, 9, 10, 11, 12]  # Non-sequential order but has consecutive groups after sorting
+        batch_size = 8
+        n_samples_per_prompt = 4
+
+        sampled, consumed = sampler.sample(ready_indexes, batch_size, n_samples_per_prompt)
+
+        # Should find consecutive groups after sorting: [3,4,5,6] and [9,10,11,12]
+        expected = [3, 4, 5, 6, 9, 10, 11, 12]
+        assert sampled == expected
+        assert consumed == expected
+
+    def test_grpo_sampler_invalid_n_samples_per_prompt(self):
+        """Test behavior with invalid n_samples_per_prompt values."""
+        sampler = GRPOGroupNSampler()
+        ready_indexes = [0, 1, 2, 3, 4, 5, 6, 7]
+        batch_size = 8
+
+        # Test zero n_samples_per_prompt
+        with pytest.raises(ValueError) as exc_info:
+            sampler.sample(ready_indexes, batch_size, n_samples_per_prompt=0)
+        assert "must be positive" in str(exc_info.value)
+
+        # Test negative n_samples_per_prompt
+        with pytest.raises(ValueError) as exc_info:
+            sampler.sample(ready_indexes, batch_size, n_samples_per_prompt=-2)
+        assert "must be positive" in str(exc_info.value)
+
+    def test_grpo_sampler_no_complete_groups(self):
+        """Test behavior when no complete groups are available."""
+        sampler = GRPOGroupNSampler()
+        ready_indexes = [0, 1, 3, 4, 6, 7]  # No consecutive groups of size 3
+        batch_size = 6
+        n_samples_per_prompt = 3
+
+        # Should return empty lists when no complete groups found
+        sampled, consumed = sampler.sample(ready_indexes, batch_size, n_samples_per_prompt)
+        assert sampled == []
+        assert consumed == []
+
+    def test_grpo_sampler_mixed_groups(self):
+        """Test behavior with mixed complete and incomplete groups."""
+        sampler = GRPOGroupNSampler()
+        ready_indexes = [0, 1, 3, 4, 5, 6, 7, 9, 10, 11]  # Mixed groups
+        batch_size = 6
+        n_samples_per_prompt = 3
+
+        # Should find the complete groups [3,4,5] and [9,10,11]
+        sampled, consumed = sampler.sample(ready_indexes, batch_size, n_samples_per_prompt)
+        assert sampled == [3, 4, 5, 9, 10, 11]
+        assert consumed == [3, 4, 5, 9, 10, 11]
+
+    def test_grpo_sampler_sorting_functionality(self):
+        """Test that ready_indexes are properly sorted before group detection."""
+        sampler = GRPOGroupNSampler()
+        ready_indexes = [10, 11, 12, 5, 6, 7, 8, 9]  # Out of order but contains consecutive groups
+        batch_size = 8
+        n_samples_per_prompt = 4
+
+        sampled, consumed = sampler.sample(ready_indexes, batch_size, n_samples_per_prompt)
+
+        # After sorting: [5,6,7,8,9,10,11,12], should find [5,6,7,8] and [9,10,11,12]
+        expected = [5, 6, 7, 8, 9, 10, 11, 12]
+        assert sampled == expected
+        assert consumed == expected
+
+    def test_grpo_sampler_insufficient_groups(self):
+        """Test behavior when requesting more groups than available."""
+        sampler = GRPOGroupNSampler()
+        ready_indexes = [0, 1, 2, 3, 4, 5, 6, 7, 8, 9, 10, 11, 12, 13, 14, 15]  # 4 groups of 4
+        batch_size = 12  # Requesting 3 groups of 4 - this should work
+        n_samples_per_prompt = 4
+
+        # This should actually work fine since we have 4 groups and request 3
+        sampled, consumed = sampler.sample(ready_indexes, batch_size, n_samples_per_prompt)
+        assert len(sampled) == 12
+        assert len(consumed) == 12
+
+        # Now test requesting more than available
+        batch_size = 20  # Requesting 5 groups of 4, but only have 4
+        sampled, consumed = sampler.sample(ready_indexes, batch_size, n_samples_per_prompt)
+
+        # Should return empty lists when requesting more complete groups than available
+        assert sampled == []
+        assert consumed == []
+
+
+class TestSamplerIntegration:
+    """Integration tests for samplers."""
+
+    def test_samplers_implement_base_interface(self):
+        """Test that all samplers properly implement BaseSampler interface."""
+        samplers = [SequentialSampler(), GRPOGroupNSampler()]
+
+        for sampler in samplers:
+            # Test that they are instances of BaseSampler
+            assert isinstance(sampler, BaseSampler)
+
+            # Test that they have the required methods
+            assert hasattr(sampler, "sample")
+            assert callable(sampler.sample)
+            assert callable(sampler)
+            assert callable(sampler.__call__)
+
+    def test_samplers_return_consistent_types(self):
+        """Test that all samplers return consistent tuple types."""
+        samplers = [(SequentialSampler(), {}), (GRPOGroupNSampler(), {"n_samples_per_prompt": 2})]
+
+        ready_indexes = [0, 1, 2, 3, 4, 5, 6, 7]
+        batch_size = 4
+
+        for sampler, kwargs in samplers:
+            sampled, consumed = sampler.sample(ready_indexes, batch_size, **kwargs)
+
+            # Check return types
+            assert isinstance(sampled, list)
+            assert isinstance(consumed, list)
+            assert isinstance(sampled[0], int) if sampled else True
+            assert isinstance(consumed[0], int) if consumed else True
+
+            # Check return value consistency
+            assert len(sampled) <= batch_size
+            assert len(sampled) == len(consumed)
+
+    def test_samplers_handle_edge_cases_consistently(self):
+        """Test that samplers handle edge cases consistently."""
+        samplers = [(SequentialSampler(), {}), (GRPOGroupNSampler(), {"n_samples_per_prompt": 2})]
+
+        # Test empty ready indexes
+        for sampler, kwargs in samplers:
+            try:
+                sampled, consumed = sampler.sample([], 0, **kwargs)
+                assert sampled == []
+                assert consumed == []
+            except Exception:
+                # GRPO sampler might fail with empty list, that's expected
+                pass
+
+        # Test zero batch size
+        for sampler, kwargs in samplers:
+            try:
+                sampled, consumed = sampler.sample([0, 1, 2, 3], 0, **kwargs)
+                assert sampled == []
+                assert consumed == []
+            except Exception:
+                # Some samplers might not handle zero batch size
+                pass
+
+
+if __name__ == "__main__":
+    pytest.main([__file__])

tests/test_serial_utils_on_cpu.py

@@ -0,0 +1,202 @@
+# Copyright 2025 The TransferQueue Team
+#
+# Licensed under the Apache License, Version 2.0 (the "License");
+# you may not use this file except in compliance with the License.
+# You may obtain a copy of the License at
+#
+#     http://www.apache.org/licenses/LICENSE-2.0
+#
+# Unless required by applicable law or agreed to in writing, software
+# distributed under the License is distributed on an "AS IS" BASIS,
+# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+# See the License for the specific language governing permissions and
+# limitations under the License.
+
+import sys
+from pathlib import Path
+
+import numpy as np
+import pytest
+import tensordict
+import torch
+from tensordict import NonTensorData, NonTensorStack, TensorDict
+
+# Import your classes here
+parent_dir = Path(__file__).resolve().parent.parent
+sys.path.append(str(parent_dir))
+
+from transfer_queue.utils.serial_utils import MsgpackDecoder, MsgpackEncoder  # noqa: E402
+
+
+def get_tensordict(tensor_dict: dict[str, torch.Tensor | list], non_tensor_dict: dict = None) -> TensorDict:
+    if non_tensor_dict is None:
+        non_tensor_dict = {}
+
+    batch_size = None
+
+    for key, val in tensor_dict.items():
+        if isinstance(val, list):
+            for v in val:
+                assert not isinstance(v, torch.Tensor), (
+                    "Passing a list makes the data NonTensorStack, "
+                    "which doesn't support torch.Tensor. Please convert to numpy first"
+                )
+
+        assert isinstance(val, torch.Tensor | list)
+
+        if batch_size is None:
+            batch_size = len(val)
+        else:
+            assert len(val) == batch_size
+
+    if batch_size is None:
+        batch_size = []
+    else:
+        batch_size = [batch_size]
+
+    for key, val in non_tensor_dict.items():
+        assert key not in tensor_dict
+        tensor_dict[key] = NonTensorData(val)
+
+    return TensorDict(source=tensor_dict, batch_size=batch_size)
+
+
+@pytest.mark.parametrize(
+    "dtype",
+    [
+        torch.float16,
+        torch.bfloat16,
+        torch.float32,
+    ],
+)
+def test_tensor_serialization(dtype):
+    encoder = MsgpackEncoder()
+    decoder = MsgpackDecoder(torch.Tensor)
+
+    tensor = torch.randn(100, 10, dtype=dtype)
+    serialized = encoder.encode(tensor)
+    deserialized = decoder.decode(serialized)
+    assert torch.allclose(tensor, deserialized)
+
+    vocab_size = 128
+    a = torch.randint(low=0, high=vocab_size, size=(11,))
+    b = torch.randint(low=0, high=vocab_size, size=(13,))
+    input_ids = [a, b]
+    input_ids = torch.nested.as_nested_tensor(input_ids, layout=torch.jagged, dtype=dtype)
+
+    input_ids_serialized = encoder.encode(input_ids)
+    input_ids_deserialized = decoder.decode(input_ids_serialized)
+    for i in range(len(input_ids.unbind())):
+        assert torch.allclose(input_ids[0], input_ids_deserialized[0])
+
+
+def test_tensordict_serialization_with_nontensor():
+    encoder = MsgpackEncoder()
+    decoder = MsgpackDecoder(TensorDict)
+
+    obs = torch.randn(100, 10)
+    data1 = {"obs": obs, "act": torch.randn(100, 3), "data_sources": ["gsm8k"] * 100}
+    data1 = get_tensordict(tensor_dict=data1)
+
+    serialized = encoder.encode(data1)
+    deserialized = decoder.decode(serialized)
+
+    assert deserialized.keys() == data1.keys()
+    assert deserialized.batch_size[0] == 100
+    assert isinstance(deserialized.get("data_sources"), NonTensorStack)
+    for k, v in data1.items():
+        if isinstance(v, torch.Tensor):
+            assert torch.allclose(deserialized[k], v)
+        elif isinstance(v, NonTensorStack):
+            assert deserialized[k] == data1[k]
+
+
+def test_tensordict_serialization_with_images():
+    # each sample contains a sequence with multiple images of different sizes
+    vocab_size = 128
+    a = torch.randint(low=0, high=vocab_size, size=(11,))
+    b = torch.randint(low=0, high=vocab_size, size=(13,))
+    input_ids = [a, b]
+    input_ids = torch.nested.as_nested_tensor(input_ids, layout=torch.jagged)
+
+    a_images = [
+        torch.randint(low=0, high=255, size=(3, 256, 256), dtype=torch.uint8).numpy(),
+        torch.randint(low=0, high=255, size=(3, 128, 128), dtype=torch.uint8).numpy(),
+    ]
+    b_images = [
+        torch.randint(low=0, high=255, size=(3, 256, 256), dtype=torch.uint8).numpy(),
+        torch.randint(low=0, high=255, size=(3, 128, 128), dtype=torch.uint8).numpy(),
+        torch.randint(low=0, high=255, size=(3, 64, 64), dtype=torch.uint8).numpy(),
+    ]
+
+    images = [a_images, b_images]
+
+    data = get_tensordict({"input_ids": input_ids, "images": images})
+
+    encoder = MsgpackEncoder()
+    decoder = MsgpackDecoder(TensorDict)
+
+    serialized = encoder.encode(data)
+    deserialized = decoder.decode(serialized)
+
+    assert np.all(np.equal(deserialized[0]["images"][0], a_images[0]))
+    assert torch.all(torch.eq(deserialized[0]["input_ids"], a))
+
+
+# Copied from https://github.com/volcengine/verl/blob/33edd95e13c72b9494585765b5fedc679fd73923/tests/test_protocol_v2_on_cpu.py#L119
+def test_tensordict_with_packing():
+    vocab_size = 128
+    a = torch.randint(low=0, high=vocab_size, size=(11,))
+    b = torch.randint(low=0, high=vocab_size, size=(13,))
+    input_ids = [a, b]
+    input_ids = torch.nested.as_nested_tensor(input_ids, layout=torch.jagged)
+
+    data = get_tensordict({"input_ids": input_ids})
+    encoder = MsgpackEncoder()
+    decoder = MsgpackDecoder(TensorDict)
+    deserialized_data = decoder.decode(encoder.encode(data))
+
+    # test cu_seqlens
+    cu_seqlens = torch.tensor([0, 11, 24])
+    assert torch.all(torch.eq(cu_seqlens, deserialized_data["input_ids"].offsets()))
+
+    # test index
+    assert torch.all(torch.eq(deserialized_data["input_ids"][0], a))
+    assert torch.all(torch.eq(deserialized_data["input_ids"][1], b))
+
+    assert torch.all(torch.eq(deserialized_data[0]["input_ids"], a))
+    assert torch.all(torch.eq(deserialized_data[1]["input_ids"], b))
+
+    data_lst = deserialized_data.chunk(2)
+
+    assert torch.all(torch.eq(data_lst[0]["input_ids"][0], a))
+    assert torch.all(torch.eq(data_lst[1]["input_ids"][0], b))
+
+
+def test_nested_tensordict_serialization():
+    td1 = tensordict.TensorDict({"a": torch.randn(2, 3), "b": torch.randn(2, 4)}, batch_size=[2])
+
+    td2 = tensordict.TensorDict({"c": torch.randn(2, 5), "d": torch.randn(2, 6)}, batch_size=[2])
+
+    td = tensordict.TensorDict({"part1": td1, "part2": td2, "e": torch.randn(2, 7)}, batch_size=[2])
+
+    encoder = MsgpackEncoder()
+    decoder = MsgpackDecoder(TensorDict)
+    deserialized_td = decoder.decode(encoder.encode(td))
+
+    assert isinstance(deserialized_td, tensordict.TensorDict)
+    assert set(deserialized_td.keys()) == set(td.keys())
+    assert isinstance(deserialized_td["part1"], tensordict.TensorDict)
+    assert isinstance(deserialized_td["part2"], tensordict.TensorDict)
+
+    assert set(deserialized_td["part1"].keys()) == set(td1.keys())
+    assert set(deserialized_td["part2"].keys()) == set(td2.keys())
+
+    for key in td.keys():
+        if isinstance(td[key], tensordict.TensorDict):
+            for inner_key in td[key].keys():
+                assert torch.allclose(deserialized_td[key][inner_key], td[key][inner_key]), (
+                    f"Values for key '{key}.{inner_key}' do not match"
+                )
+        else:
+            assert torch.allclose(deserialized_td[key], td[key]), f"Values for key '{key}' do not match"