Estimates effective bytes read from storage per device, accounting for chunk overhead. TensorStore reads whole chunks even if only a partial chunk is needed, so actual bytes read can exceed the raw tensor shard size.
(
spec: dict[str, Any],
shape: tuple,
dtype: typing.DTypeLike,
sharding: jax.sharding.Sharding,
)
| 738 | """ |
| 739 | if len(cpu_devices) > 1: |
| 740 | logging.info("Creating CPU mesh from TPU mesh: %s", tpu_mesh) |
| 741 | cpu_mesh = colocated_python.colocated_cpu_devices(tpu_mesh) |
| 742 | logging.info("CPU Mesh: %s", cpu_mesh) |
| 743 | |
| 744 | return [ |
| 745 | ( |
| 746 | jax.sharding.NamedSharding(cpu_mesh, sharding.spec) |
| 747 | if isinstance(sharding, jax.sharding.NamedSharding) |
| 748 | else jax.sharding.NamedSharding(cpu_mesh, jax.sharding.PartitionSpec()) |
| 749 | ) |
| 750 | for sharding in tpu_shardings |
| 751 | ] |
| 752 | else: |
| 753 | # Single device - use simple single device sharding |
| 754 | return [jax.sharding.SingleDeviceSharding(cpu_devices[0]) for _ in tpu_shardings] |
| 755 | |
| 756 | |
| 757 | async def _effective_bytes_per_device( |
| 758 | spec: dict[str, Any], |
| 759 | shape: tuple, |
| 760 | dtype: typing.DTypeLike, |
| 761 | sharding: jax.sharding.Sharding, |
| 762 | ) -> int: |
| 763 | """Estimates effective bytes read from storage per device, accounting for chunk overhead. |
| 764 | |
| 765 | TensorStore reads whole chunks even if only a partial chunk is needed, so actual bytes |
| 766 | read can exceed the raw tensor shard size. |
| 767 | """ |
| 768 | t = await ts.open(ts.Spec(spec), open=True, context=serialization.TS_CONTEXT) |
| 769 | shard_shape = sharding.shard_shape(shape) |
| 770 | raw_bytes = math.prod(shard_shape) * np.dtype(dtype).itemsize |
| 771 |
no test coverage detected