code: def test_memory_consumption(self):
      global_mesh = jtu.create_global_mesh((2, 4), ("x", "y"))
      inp_shape = (2_048, 4_096)
      pspec = P("x", "y")
      num = math.prod(inp_shape)
      sharding = NamedSharding(global_mesh, pspec)
      src = jax.numpy.arange(num, dtype=np.int32).reshape(
          inp_shape
      )  # 8e9
      inp = array.make_array_from_callback(
          inp_shape, sharding, lambda idx: src[idx]
      )
      ckpt_dir = pathlib.Path(self.create_tempdir("memprof").full_path)
      tspec = serialization.get_tensorstore_spec(str(ckpt_dir))
  
      manager = serialization.GlobalAsyncCheckpointManager()
      manager.serialize(
          [inp],
          [tspec],
          on_commit_callback=partial(
              self._on_commit_callback, ckpt_dir, ckpt_dir
          ),
      )
      manager.wait_until_finished()
  
      async def deserialize_with_byte_limit():
          r = await serialization.async_deserialize(
              sharding,
              tspec,
              inp_shape,
              byte_limiter=serialization._LimitInFlightBytes(4_200_000),
          )
          r.block_until_ready()
  
      tm.start()
      asyncio.run(deserialize_with_byte_limit())
      unused_current, peak = tm.get_traced_memory()
      # NB: some padding + tensorstore overhead. It should always be
      # less than array size (2048 * 4096 * 4 = 32M)
      self.assertLess(peak, 10_000_000)
      deserialize_wo_limit = serialization.async_deserialize(
          sharding, tspec, inp_shape
      )
      tm.clear_traces()
      # NB: call block_until_ready() is important here and above
      # because otherwise this leads to racing condition and segfault with
      # tensorstore attempting to dealloc using tracemalloc which is already
      # destroyed.
      asyncio.run(deserialize_wo_limit).block_until_ready()
  
      unused_current, peak = tm.get_traced_memory()
      # We load entire array in memory here.
      self.assertGreater(peak, 30_000_000)
      tm.stop()
  
imports: import asyncio
  import math
  from functools import partial
  import os
  import pathlib
  import tracemalloc
  from absl.testing import absltest
  from absl.testing import parameterized
  import jax
  from jax._src import test_util
  from jax import config
  from jax._src import array
  from jax.sharding import NamedSharding, GSPMDSharding
  from jax.sharding import PartitionSpec
  from jax.experimental.array_serialization import serialization
  import numpy
  import tensorstore
  import unittest
label_desc: Use image node_id 2 for calling the CheckpointTest obj's underlying member method code with example usage: obj.test_memory_consumption() without return types