Method test_serialize

axlearn/common/array_serialization_test.py:237–292

(
        self, arrays: list[list[int]], max_concurrent_gb: int, expect_max_concurrent_gb: int
    )

        ),
    )
    def test_serialize(
        self, arrays: list[list[int]], max_concurrent_gb: int, expect_max_concurrent_gb: int
    ):
        arrays = [
            mock.Mock(
                addressable_shards=[
                    mock.Mock(replica_id=0, **{"data.nbytes": int(shard * 10**9), "data.shape": ()})
                    for shard in array
                ],
                nbytes=int(sum(array) * 10**9),
                dtype=jax.numpy.bfloat16,
            )
            for array in arrays
        ]
        tensorstore_specs = [{} for _ in range(len(arrays))]
        expect_max_concurrent_bytes = int(expect_max_concurrent_gb * 10**9)

        concurrent_bytes = 0

        class FakeTs:
            def __getitem__(self, *_):
                return self

            async def write(self, data: jax.Array, **_):
                await asyncio.sleep(0.1)
                nonlocal concurrent_bytes
                concurrent_bytes -= data.nbytes

        async def open_patch(*_, **__):
            return FakeTs()

        async def _copy_to_host_patch(shard_infos: list[_ShardInfo]):
            nonlocal concurrent_bytes
            for info in shard_infos:
                concurrent_bytes += info.data.nbytes
            # In-flight bytes should be lower than the expected max bytes
            self.assertLessEqual(concurrent_bytes, expect_max_concurrent_bytes)

        manager = BoundedDataShardedAsyncCheckpointManager(max_concurrent_gb=max_concurrent_gb)
        with (
            mock.patch(
                f"{array_serialization.__name__}._num_replicas_per_shard",
                lambda *args: defaultdict(lambda: 1),
            ),
            mock.patch(f"{array_serialization.__name__}._slices_to_tuple", lambda *_: 1),
            mock.patch(
                f"{array_serialization.__name__}._slice_shard_and_copy_to_host", _copy_to_host_patch
            ),
            mock.patch(
                f"{array_serialization.__name__}.serialization._get_metadata", lambda *_: {}
            ),
            mock.patch(f"{array_serialization.__name__}.ts.open", open_patch),
            mock.patch(f"{array_serialization.__name__}.ts.Spec", mock.MagicMock()),
        ):
            manager.serialize(arrays, tensorstore_specs, on_commit_callback=lambda: None)
            manager.wait_until_finished()

    @parameterized.product(