fix: 修复 pipeline 模块中的打包逻辑缺陷并完善测试覆盖

2026-03-30 12:22:37 +08:00
parent 07ef471aa6
commit 71887bb4bb
11 changed files with 1186 additions and 39 deletions
@@ -0,0 +1,257 @@
+"""单元测试：pipeline.packing 模块中的 SequencePacker 类"""
+
+import pytest
+import torch
+from pipeline.packing import SequencePacker
+
+
+class TestSequencePacker:
+    """SequencePacker 类的测试套件"""
+
+    def test_normal_packing(self):
+        """测试正常打包场景：多个序列正确打包成固定长度的包"""
+        packer = SequencePacker(pack_size=10, pad_value=0)
+        
+        sequences = [
+            torch.tensor([1, 2, 3], dtype=torch.int32),
+            torch.tensor([4, 5], dtype=torch.int32),
+            torch.tensor([6, 7, 8, 9], dtype=torch.int32),
+        ]
+        
+        packages = packer.pack(sequences)
+        
+        # 验证至少有包输出
+        assert len(packages) >= 1
+        
+        # 验证每个包的长度是正确的
+        for pkg in packages:
+            assert pkg.shape == (10,)
+            # 验证填充值
+            # 检查所有非零元素都在前几个位置，或者包是满的
+            non_zero_count = (pkg != 0).sum().item()
+            # 非零元素的数量应该等于原始序列元素的总和
+            total_elements = sum(s.numel() for s in sequences)
+            # 由于打包，第一个包包含3+2=5个元素，第二个包包含4个元素
+            # 第一个包应该包含前两个序列
+            pkg1 = packages[0]
+            # 序列[1,2,3]和[4,5]按长度降序排序后是[1,2,3]在前，然后[4,5]
+            # 但排序是原地修改...等等，我们已经修复了使用sorted()
+            # 所以排序后的顺序是[6,7,8,9], [1,2,3], [4,5]
+            # 第一个包包含[6,7,8,9]和部分[1,2,3] = 4+3=7，剩余3个位置放[4,5]
+            # 所以第一个包应该是[6,7,8,9,1,2,3,4,5,0]
+        
+        # 简化测试：验证打包后的张量包含所有原始数据
+        all_values = []
+        for pkg in packages:
+            non_zero = pkg[pkg != 0].tolist()
+            all_values.extend(non_zero)
+        
+        # 检查所有原始数据是否都被包含
+        original_values = [1, 2, 3, 4, 5, 6, 7, 8, 9]
+        for val in original_values:
+            assert val in all_values, f"Value {val} not found in packages"
+
+    def test_empty_list_input(self):
+        """测试空列表输入"""
+        packer = SequencePacker(pack_size=10)
+        
+        packages = packer.pack([])
+        
+        assert packages == []
+        
+        # 验证内部状态已正确初始化
+        assert packer._current_pack is not None
+        assert packer._current_pos == 0
+
+    def test_single_sequence_input(self):
+        """测试单个序列输入"""
+        packer = SequencePacker(pack_size=10, pad_value=-1)
+        
+        sequences = [torch.tensor([1, 2, 3], dtype=torch.int32)]
+        
+        packages = packer.pack(sequences)
+        
+        assert len(packages) == 1
+        pkg = packages[0]
+        assert pkg.shape == (10,)
+        assert pkg[:3].tolist() == [1, 2, 3]
+        assert pkg[3:].tolist() == [-1] * 7
+
+    def test_truncate_long_sequence(self, caplog):
+        """测试超长序列截断，验证警告日志是否触发"""
+        packer = SequencePacker(pack_size=5, pad_value=0)
+        
+        sequences = [
+            torch.tensor([1, 2, 3, 4, 5, 6, 7, 8], dtype=torch.int32),  # 长度8，超过pack_size=5
+        ]
+        
+        packages = packer.pack(sequences)
+        
+        assert len(packages) == 1
+        pkg = packages[0]
+        assert pkg.shape == (5,)
+        assert pkg.tolist() == [1, 2, 3, 4, 5]  # 只保留前5个元素
+        
+        # 验证警告日志已触发
+        assert "truncating" in caplog.text.lower() or "exceeds" in caplog.text.lower()
+
+    def test_padding_value(self):
+        """测试填充值正确应用"""
+        packer = SequencePacker(pack_size=8, pad_value=99)
+        
+        sequences = [
+            torch.tensor([1, 2], dtype=torch.int32),
+            torch.tensor([3], dtype=torch.int32),
+        ]
+        
+        packages = packer.pack(sequences)
+        
+        assert len(packages) == 1
+        pkg = packages[0]
+        
+        # 前3个元素是数据
+        assert pkg[:3].tolist() == [1, 2, 3]
+        # 后5个元素是填充值
+        assert pkg[3:].tolist() == [99] * 5
+
+    def test_different_dtypes(self):
+        """测试支持不同 dtype (int32, int64, float32)"""
+        # int32
+        packer_int32 = SequencePacker(pack_size=10, pad_value=0, dtype=torch.int32)
+        sequences_int32 = [torch.tensor([1, 2, 3], dtype=torch.int32)]
+        packages_int32 = packer_int32.pack(sequences_int32)
+        assert packages_int32[0].dtype == torch.int32
+
+        # int64
+        packer_int64 = SequencePacker(pack_size=10, pad_value=0, dtype=torch.int64)
+        sequences_int64 = [torch.tensor([1, 2, 3], dtype=torch.int64)]
+        packages_int64 = packer_int64.pack(sequences_int64)
+        assert packages_int64[0].dtype == torch.int64
+
+        # float32
+        packer_float32 = SequencePacker(pack_size=10, pad_value=0.0, dtype=torch.float32)
+        sequences_float32 = [torch.tensor([1.0, 2.0, 3.0], dtype=torch.float32)]
+        packages_float32 = packer_float32.pack(sequences_float32)
+        assert packages_float32[0].dtype == torch.float32
+
+    def test_non_1d_tensor_raises_error(self):
+        """测试非1D张量是否抛出异常"""
+        packer = SequencePacker(pack_size=10)
+        
+        # 2D 张量应该抛出异常
+        sequences_2d = [torch.tensor([[1, 2], [3, 4]])]  # shape: (2, 2)
+        with pytest.raises(ValueError, match="Expected 1D tensor"):
+            packer.pack(sequences_2d)
+        
+        # 0D 张量 (标量) 应该抛出异常
+        sequences_0d = [torch.tensor(5)]  # shape: ()
+        with pytest.raises(ValueError, match="Expected 1D tensor"):
+            packer.pack(sequences_0d)
+        
+        # 3D 张量应该抛出异常
+        sequences_3d = [torch.tensor([[[1, 2]]])]  # shape: (1, 1, 2)
+        with pytest.raises(ValueError, match="Expected 1D tensor"):
+            packer.pack(sequences_3d)
+
+    def test_reset_method(self):
+        """测试 reset() 方法是否正确重置内部状态"""
+        packer = SequencePacker(pack_size=10, pad_value=0)
+        
+        # 第一次打包
+        sequences1 = [torch.tensor([1, 2, 3], dtype=torch.int32)]
+        packer.pack(sequences1)
+        
+        # 验证内部状态已更新
+        assert packer._current_pos == 0
+        assert packer._current_pack is None  # 最后一个包已发送，设置为None
+        
+        # 重置
+        packer.reset()
+        
+        # 验证重置后的状态
+        assert packer._current_pos == 0
+        assert packer._current_pack is not None
+        assert packer._current_pack.shape == (10,)
+        assert packer._current_pack.tolist() == [0] * 10
+        
+        # 验证重置后可以继续正常使用
+        sequences2 = [torch.tensor([4, 5, 6], dtype=torch.int32)]
+        packages = packer.pack(sequences2)
+        
+        assert len(packages) == 1
+        assert packages[0][:3].tolist() == [4, 5, 6]
+
+    def test_input_list_not_modified(self):
+        """测试输入列表是否未被修改（使用 sorted 而非 sort）"""
+        packer = SequencePacker(pack_size=10)
+        
+        # 创建原始序列列表（故意不按长度排序）
+        original_sequences = [
+            torch.tensor([3], dtype=torch.int32),   # 长度1
+            torch.tensor([1, 2], dtype=torch.int32), # 长度2
+            torch.tensor([4, 5, 6, 7], dtype=torch.int32), # 长度4
+        ]
+        
+        # 保存原始顺序的字符串表示
+        original_repr = [seq.tolist() for seq in original_sequences]
+        
+        # 打包
+        packer.pack(original_sequences)
+        
+        # 验证输入列表未被修改
+        current_repr = [seq.tolist() for seq in original_sequences]
+        assert current_repr == original_repr, "输入列表被修改了，应该使用 sorted() 而非 sort()"
+
+    def test_exact_pack_size_fit(self):
+        """测试序列长度恰好等于 pack_size 的情况"""
+        packer = SequencePacker(pack_size=5, pad_value=0)
+        
+        sequences = [
+            torch.tensor([1, 2, 3, 4, 5], dtype=torch.int32),
+            torch.tensor([6, 7, 8, 9, 10], dtype=torch.int32),
+        ]
+        
+        packages = packer.pack(sequences)
+        
+        # 每个序列恰好占满一个包
+        assert len(packages) == 2
+        
+        assert packages[0].tolist() == [1, 2, 3, 4, 5]
+        assert packages[1].tolist() == [6, 7, 8, 9, 10]
+
+    def test_multiple_packs_full_utilization(self):
+        """测试多个包的高效利用"""
+        packer = SequencePacker(pack_size=10, pad_value=-1)
+        
+        # 创建多个小序列，确保高效打包
+        sequences = [
+            torch.tensor([1], dtype=torch.int32),
+            torch.tensor([2], dtype=torch.int32),
+            torch.tensor([3], dtype=torch.int32),
+            torch.tensor([4], dtype=torch.int32),
+            torch.tensor([5], dtype=torch.int32),
+            torch.tensor([6], dtype=torch.int32),
+            torch.tensor([7], dtype=torch.int32),
+            torch.tensor([8], dtype=torch.int32),
+            torch.tensor([9], dtype=torch.int32),
+            torch.tensor([10], dtype=torch.int32),
+            torch.tensor([11], dtype=torch.int32),
+        ]
+        
+        packages = packer.pack(sequences)
+        
+        # 前10个序列打包成一个包，最后一个序列单独一个包
+        assert len(packages) == 2
+        assert packages[0].tolist() == list(range(1, 11))
+        assert packages[1].tolist() == [11] + [-1] * 9
+
+    def test_dtype_mismatch_warning(self, caplog):
+        """测试 dtype 不匹配时的警告"""
+        packer = SequencePacker(pack_size=10, pad_value=0, dtype=torch.int32)
+        
+        sequences = [torch.tensor([1, 2, 3], dtype=torch.int64)]
+        
+        packages = packer.pack(sequences)
+        
+        # 应该触发 dtype 不匹配警告
+        assert "dtype" in caplog.text.lower() or "converted" in caplog.text.lower()