diff --git a/axlearn/cloud/gcp/examples/colocated_python_benchmark.py b/axlearn/cloud/gcp/examples/colocated_python_benchmark.py index d548bc25e..28a052a52 100644 --- a/axlearn/cloud/gcp/examples/colocated_python_benchmark.py +++ b/axlearn/cloud/gcp/examples/colocated_python_benchmark.py @@ -19,14 +19,9 @@ """ import argparse -import asyncio -import functools -import logging import os -import sys import time -from concurrent.futures import ThreadPoolExecutor -from contextlib import nullcontext +from contextlib import contextmanager from datetime import datetime from typing import Any, Dict, Optional, Sequence @@ -34,104 +29,31 @@ import jax.numpy as jnp import pathwaysutils # pytype: disable=import-error from jax._src.mesh import thread_resources -from jax.experimental import colocated_python, mesh_utils +from jax.experimental import mesh_utils from jax.experimental.array_serialization import serialization as array_serialization -from jax.experimental.array_serialization import tensorstore_impl from axlearn.common import utils -from axlearn.common.array_serialization import GlobalAsyncCheckpointManager, _async_deserialize +from axlearn.common.array_serialization import GlobalAsyncCheckpointManager from axlearn.common.checkpointer import parse_step_from_dir, read_index_file from axlearn.common.utils import TensorSpec, infer_mesh_shape +@contextmanager def maybe_profile(enabled: bool, profile_dir: Optional[str]): - """Return JAX profiler context if enabled, otherwise a no-op context manager. + """JAX profiler context if enabled, otherwise a no-op. Args: enabled: Whether profiling is enabled. profile_dir: Directory to save profiling results. - - Returns: - Context manager for profiling or no-op. """ if enabled: assert profile_dir is not None, "profile_dir must be set when profiling is enabled" - return jax.profiler.trace(profile_dir) - else: - # Return a no-op context manager - return nullcontext() - - -def _colocated_deserialize( - shardings: Sequence[jax.sharding.NamedSharding], - tensorstore_specs: Sequence[Dict[str, Any]], - global_shapes: Sequence[tuple], - dtypes: Sequence[jnp.dtype], -): - concurrent_bytes = 34359738368 * 6 # 32GB * 6 - cpu_devices = colocated_python.colocated_cpu_devices(jax.devices()) - print(f"{cpu_devices=}") - - if len(cpu_devices) > 1: - print(f"TPU Mesh: {thread_resources.env.physical_mesh}") - cpu_mesh = colocated_python.colocated_cpu_devices(thread_resources.env.physical_mesh) - print(f"CPU Mesh: {cpu_mesh}") - cpu_shardings = [ - jax.sharding.NamedSharding(cpu_mesh, sharding.spec) for sharding in shardings - ] - else: - cpu_shardings = [ - jax.sharding.SingleDeviceSharding(cpu_devices[0]) for sharding in shardings - ] - - def output_spec_fn(): - return [ - jax.ShapeDtypeStruct(shape=shape, dtype=dtype, sharding=sharding) - for shape, dtype, sharding in zip(global_shapes, dtypes, cpu_shardings) - ] - - @colocated_python.colocated_python - def run_deserializer(): - # Object should be created once per process. - # pylint: disable=protected-access - # print("Print statement inside colocated") - logging.info("Logging statement inside colocated") - sys.stderr.write("Stdder statement in colocated") - start_colocated_time = time.perf_counter() - byte_limiter = tensorstore_impl._LimitInFlightBytes(concurrent_bytes) - h2d_limiter = tensorstore_impl._LimitInFlightBytes(concurrent_bytes) - thread_pool = ThreadPoolExecutor(1) - multi_thread_pool = ThreadPoolExecutor(2) - - future_arrays = jax.tree.map( - functools.partial( - _async_deserialize, - byte_limiter=byte_limiter, - h2d_limiter=h2d_limiter, - single_thread_pool=thread_pool, - multi_thread_pool=multi_thread_pool, - ), - cpu_shardings, - tensorstore_specs, - global_shapes, - dtypes, - ) - - async def gather_func(): - return await asyncio.gather(*future_arrays) - - result = asyncio.run(gather_func()) - logging.info("Deserialize took %.2f seconds", time.perf_counter() - start_colocated_time) - return result - - run_deserializer = run_deserializer.specialize( - devices=cpu_devices, - out_specs_fn=output_spec_fn, - ) - - # Try running in the current event loop if one exists, otherwise create new one - result = run_deserializer() - return result + jax.profiler.start_trace(profile_dir) + try: + yield + finally: + if enabled: + jax.profiler.stop_trace() def create_mesh(mesh_shape=(1, 1, 1, 1, 1, 16, -1)): @@ -210,6 +132,13 @@ def create_checkpoint_spec_from_state(ckpt_dir: str, state_spec: dict): if not mesh.shape: raise RuntimeError("Checkpoint restoration must take place within the context of a Mesh") + # Track sharding statistics + sharded_bytes = 0 + replicated_bytes = 0 + per_shard_bytes = 0 + num_sharded = 0 + num_replicated = 0 + # Process each tensor in the state spec for path, value in utils.flatten_items(state_spec, separator="/"): if isinstance(value, TensorSpec): @@ -231,6 +160,40 @@ def create_checkpoint_spec_from_state(ckpt_dir: str, state_spec: dict): dtypes.append(dtype) shardings.append(sharding) + # Compute tensor size in bytes downloaded from GCS. + # Checkpoints are stored as fp32 (4 bytes per element). + element_size = 4 + tensor_bytes = 1 + for d in value.shape: + tensor_bytes *= d + tensor_bytes *= element_size + + if partition_spec == jax.sharding.PartitionSpec(): + replicated_bytes += tensor_bytes + num_replicated += 1 + else: + sharded_bytes += tensor_bytes + num_sharded += 1 + # Compute per-shard size by dividing by the number of shards + num_shards = 1 + for axis in partition_spec: + if axis is not None: + if isinstance(axis, tuple): + for a in axis: + num_shards *= mesh.shape[a] + else: + num_shards *= mesh.shape[axis] + per_shard_bytes += tensor_bytes // num_shards + + num_devices = len(mesh.devices.flat) + print(f"Sharding stats ({num_devices} devices):") + print( + f" Sharded: {num_sharded} tensors, {sharded_bytes / 10**9:.2f} GB " + f"(per-shard total: {per_shard_bytes / 10**9:.2f} GB)" + ) + print(f" Replicated: {num_replicated} tensors, {replicated_bytes / 10**9:.2f} GB") + print(f" Per-device total: {(per_shard_bytes + replicated_bytes) / 10**9:.2f} GB") + return tensorstore_specs, shardings, global_shapes, dtypes @@ -266,64 +229,32 @@ def cleanup_loaded_arrays(loaded_arrays: list) -> None: print("Cleanup complete.") -def load_model_default( +def load_model( tensorstore_specs: Sequence[Dict[str, Any]], shardings: Sequence[jax.sharding.NamedSharding], global_shapes: Sequence[tuple], dtypes: Sequence[jnp.dtype], ): - """Load model using default method (direct to TPU).""" - print("Preloading checkpoint to TPU HBM...") - start_time = time.perf_counter() + """Load model from checkpoint. + Args: + tensorstore_specs: TensorStore specifications for each array. + shardings: Target shardings for the restored arrays. + global_shapes: Global shapes for each array. + dtypes: Data types for each array. + + Returns: + List of restored JAX arrays. + """ manager = GlobalAsyncCheckpointManager() restored_values = manager.deserialize( shardings=shardings, tensorstore_specs=tensorstore_specs, global_shapes=global_shapes, dtypes=dtypes, - concurrent_gb=192, - ) - - preload_time = time.perf_counter() - start_time - print(f"Preload completed in {preload_time:.2f} seconds") - print(f"Preloaded {len(restored_values)} arrays") - - return restored_values - - -def load_model_colocated( - tensorstore_specs: Sequence[Dict[str, Any]], - shardings: Sequence[jax.sharding.NamedSharding], - global_shapes: Sequence[tuple], - dtypes: Sequence[jnp.dtype], -): - """Load model using colocated Python (CPU preload then transfer to TPU).""" - print("Preloading checkpoint to CPU memory...") - start_time = time.perf_counter() - - preloaded_values = _colocated_deserialize( - shardings=shardings, - tensorstore_specs=tensorstore_specs, - global_shapes=global_shapes, - dtypes=dtypes, + concurrent_gb=400, ) - # for x in preloaded_values: - # x.block_until_ready() - - preload_time = time.perf_counter() - start_time - print(f"Preload completed in {preload_time:.2f} seconds") - print(f"Preloaded {len(preloaded_values)} arrays") - - print("Transferring arrays to TPU...") - start_time = time.perf_counter() - - restored_values = [jax.device_put(x, s) for x, s in zip(preloaded_values, shardings)] - for x in restored_values: - x.block_until_ready() - - transfer_time = time.perf_counter() - start_time - print(f"Transfer completed in {transfer_time:.2f} seconds") + print(f"Loaded {len(restored_values)} arrays") return restored_values @@ -346,6 +277,12 @@ def main(): action="store_true", help="Enable JAX profiler (adds overhead, disable for accurate benchmarking)", ) + parser.add_argument( + "--num_iters", + type=int, + default=1, + help="Number of times to repeat the load benchmark (default: 1)", + ) args = parser.parse_args() # Disable persistent compilation cache for fair benchmarking @@ -359,13 +296,6 @@ def main(): print(f"JAX devices: {jax.devices()}") - # Select loading function and profile prefix based on method - if args.method == "colocated": - loader_fn = load_model_colocated - else: # args.method == "default" - loader_fn = load_model_default - print(f"--- Running {args.method} benchmark ---") - # Validate checkpoint path if not args.ckpt_path.startswith("gs://"): raise ValueError(f"Only GCS paths (gs://) are supported, got: {args.ckpt_path}") @@ -385,6 +315,8 @@ def main(): state_spec = create_state_spec_from_checkpoint(args.ckpt_path) print(f"Found {len(jax.tree_util.tree_leaves(state_spec))} tensors in checkpoint") + num_iterations = args.num_iters + print(f"--- Running {args.method} benchmark ({num_iterations} iterations) ---") loaded_values = None try: with create_mesh(): @@ -393,19 +325,28 @@ def main(): args.ckpt_path, state_spec ) + if args.method == "default": + os.environ["COLOCATED_PYTHON_DESERIALIZE"] = "0" + + loaded_values = None with maybe_profile(args.profile, profile_dir): - start_time = time.perf_counter() - loaded_values = loader_fn( - tensorstore_specs=tensorstore_specs, - shardings=shardings, - global_shapes=global_shapes, - dtypes=dtypes, - ) - print(f"✅ Successfully loaded model from {args.ckpt_path}") - print(f"Deserialize took {time.perf_counter() - start_time:.2f} seconds") - print(f" Total parameters: {sum(x.size for x in loaded_values):,}") + for i in range(num_iterations): + if loaded_values is not None: + del loaded_values + print(f"\n--- Iteration {i + 1}/{num_iterations} ---") + start_time = time.perf_counter() + loaded_values = load_model( + tensorstore_specs=tensorstore_specs, + shardings=shardings, + global_shapes=global_shapes, + dtypes=dtypes, + ) + elapsed = time.perf_counter() - start_time + print(f"✅ Successfully loaded model from {args.ckpt_path}") + print(f"Total time took {elapsed:.2f} seconds") + print(f" Total parameters: {sum(x.size for x in loaded_values):,}") finally: - # Always clean up, even if benchmark fails + # Always clean up, even if benchmark fails. if loaded_values is not None: cleanup_loaded_arrays(loaded_values) diff --git a/axlearn/common/array_serialization.py b/axlearn/common/array_serialization.py index db78cfa27..897e20c96 100644 --- a/axlearn/common/array_serialization.py +++ b/axlearn/common/array_serialization.py @@ -18,6 +18,7 @@ """ import asyncio import functools +import gc import math import os import threading @@ -35,11 +36,95 @@ from absl import logging from jax._src import array, typing from jax._src.layout import Format +from jax._src.mesh import thread_resources +from jax.experimental import colocated_python from jax.experimental.array_serialization import serialization from axlearn.common.utils import Tensor +class _ColocatedResourceManager: + """Manages TensorStore context and event loop on colocated Python workers. + + Encapsulates all sidecar-side state for checkpoint deserialization. + Created by _colocated_deserialize_setup and stored as an attribute on the + colocated_python module so it can be accessed from other colocated functions + without cloudpickle frozen-globals issues. + + Must be created before each deserialization to ensure fresh resources, + avoiding accumulation of stale resources across checkpoint reloads. + """ + + def __init__(self): + """Initialize TensorStore context and event loop on colocated Python workers.""" + self.ts_context = ts.Context(serialization.TS_CONTEXT.spec) + + # Creates a dedicated event loop with a dedicated executor pool. + self.event_loop = asyncio.new_event_loop() + self.event_loop.set_default_executor(futures.ThreadPoolExecutor(max_workers=os.cpu_count())) + + # Runs the event loop in a background thread. + self.loop_thread = threading.Thread(target=self.event_loop.run_forever, daemon=True) + self.loop_thread.start() + logging.info("_ColocatedResourceManager initialized on sidecar.") + + def __del__(self): + """Release TensorStore context and event loop on colocated Python workers.""" + if hasattr(self, "event_loop") and self.event_loop is not None: + # Stop the dedicated event loop. + self.event_loop.call_soon_threadsafe(self.event_loop.stop) + # Exit the background thread and release resources of the event loop. + self.loop_thread.join() + self.event_loop.close() + logging.info("_ColocatedResourceManager destroyed on sidecar.") + + +@colocated_python.colocated_python +def _colocated_deserialize_setup(dummy_array: jax.Array) -> jax.Array: + """Initialize TensorStore context and event loop on colocated Python workers. + + Must be called before each deserialization to ensure fresh resources, + avoiding accumulation of stale resources across checkpoint reloads. + """ + colocated_python.resource_manager = _ColocatedResourceManager() + return dummy_array + + +@colocated_python.colocated_python +def _colocated_deserialize_teardown(dummy_array: jax.Array) -> jax.Array: + """Release TensorStore context and event loop on colocated Python workers. + + Must be called after each deserialization to release resources. + """ + import tracemalloc # pylint: disable=import-outside-toplevel + from collections import Counter # pylint: disable=import-outside-toplevel + + if hasattr(colocated_python, "resource_manager"): + colocated_python.resource_manager = None + + gc.collect() + + # Log object counts by type to identify what's accumulating. + type_counts = Counter(type(obj).__name__ for obj in gc.get_objects()) + logging.info("Top 20 object types: %s", type_counts.most_common(20)) + + # Log top memory allocations by source location. + if not tracemalloc.is_tracing(): + tracemalloc.start() + else: + snapshot = tracemalloc.take_snapshot() + top_stats = snapshot.statistics("lineno") + for stat in top_stats[:20]: + logging.info("tracemalloc: %s", stat) + + logging.info( + "Teardown complete. live objects: %d, gc garbage: %d", + len(gc.get_objects()), + len(gc.garbage), + ) + return dummy_array + + @dataclass class _ShardInfo: """Stores information for a maybe sliced jax.Shard. @@ -211,6 +296,22 @@ def running_on_pathways(): return os.getenv("JAX_PLATFORMS") == "proxy" +def colocated_python_available() -> bool: + """Returns True if colocated Python CPU workers are available at runtime. + + Colocated Python requires: + 1. Running on Pathways (JAX_PLATFORMS == "proxy"). + 2. At least one colocated CPU device is present (i.e. the sidecar is running). + """ + if not running_on_pathways(): + return False + try: + cpu_devices = colocated_python.colocated_cpu_devices(jax.devices()) + return len(cpu_devices) > 0 + except Exception: # pylint: disable=broad-exception-caught + return False + + async def _slice_shard_and_copy_to_host(shard_infos: list[_ShardInfo]): """Slices each shard according to shard info and then copy the sliced result to host. @@ -402,8 +503,24 @@ async def _run_serializer( raise e -def _blocking_device_put(out: Tensor, layout: Format) -> Tensor: - return jax.block_until_ready(jax.device_put(out, layout)) +def _blocking_device_put( + tensor: Tensor, target: Union[Format, jax.sharding.Sharding], *, verbose: bool = False +) -> Tensor: + """Device put and block until ready. + + Args: + tensor: Array to transfer. + target: Either a Format (with layout + sharding) or Sharding. + verbose: If True, log execution time. + + Returns: + Transferred array. + """ + start = time.perf_counter() + result = jax.block_until_ready(jax.device_put(tensor, target)) + if verbose: + logging.info("device_put took %.3f seconds", time.perf_counter() - start) + return result async def _async_deserialize( @@ -412,10 +529,11 @@ async def _async_deserialize( global_shape: Optional[Sequence[int]], dtype: Optional[typing.DTypeLike], *, - h2d_limiter: serialization._LimitInFlightBytes, + h2d_limiter: Optional[serialization._LimitInFlightBytes], byte_limiter: serialization._LimitInFlightBytes, - single_thread_pool: ThreadPoolExecutor, - multi_thread_pool: ThreadPoolExecutor, + single_thread_pool: Optional[ThreadPoolExecutor], + multi_thread_pool: Optional[ThreadPoolExecutor], + ts_context: Optional[ts.Context] = None, ): """Modified from https://github.com/jax-ml/jax/blob/e7ec418eba9ada336f755613948cbdf4a9e97d59/jax/experimental/array_serialization/serialization.py#L345 @@ -432,6 +550,9 @@ async def _async_deserialize( in flight H2D is imposed. Note that Pathways checkpoint loading does not require h2d limiter since the H2D doesn't happen in the head node, and each worker has preemapped a chunk of host memory that is larger than the total device memory. + h2d_limiter, single_thread_pool, and multi_thread_pool are all optional. When omitted (e.g. + in the colocated Python path where a global pipeline limiter already bounds concurrency), + H2D is submitted to the default asyncio thread pool without per-shard gating. 4. Let user pass in a multi_thread_pool thread pool for ckpt loading, instead of letting async io to create a default pool, to make it more configurable. @@ -469,9 +590,17 @@ async def _async_deserialize( # - On AWS (or other non-GCP environments) accessing GCS, gcs_grpc may hit auth/network # issues due to cross-cloud constraints. So we enable this optimization on Pathways # which only runs on GCP for now. - context = serialization.TS_CONTEXT + if ts_context is not None: + logging.info("Using provided TensorStore context for deserialization.") + context = ts_context + else: + logging.info("Using global TensorStore context for deserialization.") + context = serialization.TS_CONTEXT if os.getenv("ENABLE_GCS_GRPC", "false") == "true": + logging.debug("gcs_grpc enabled") tensorstore_spec, context = use_gcs_grpc(tensorstore_spec) + else: + logging.debug("gcs_grpc not enabled") t = await ts.open( tensorstore_spec, @@ -528,10 +657,16 @@ async def cb(index: array.Index, device: jax.Device): dll, jax.sharding.SingleDeviceSharding(device, memory_kind=in_sharding.memory_kind) ) + if h2d_limiter is None: + # No per-shard H2D gating — caller (e.g. colocated Python path) relies on a global + # pipeline limiter to bound concurrency instead. + result = await loop.run_in_executor( + None, functools.partial(_blocking_device_put, verbose=True), out, layout + ) # Jax >= 0.6.2 changes the behavior of _LimitInFlightBytes, where wait_for_bytes no longer # throws an exception if requested_bytes > max_bytes # pylint: disable-next=protected-access - if out_size > h2d_limiter._max_bytes: + elif out_size > h2d_limiter._max_bytes: logging.log_first_n( logging.WARNING, "Tensor shard for tensor %s (padded size %d bytes) exceeded " @@ -569,6 +704,254 @@ async def cb(index: array.Index, device: jax.Device): ) +def _create_cpu_shardings( + cpu_devices: list, + tpu_shardings: Sequence[Union[jax.sharding.Sharding, Format]], + tpu_mesh: jax.sharding.Mesh, +) -> list[jax.sharding.Sharding]: + """Create CPU shardings that mirror the structure of TPU shardings. + + Args: + cpu_devices: List of CPU devices. + tpu_shardings: Target TPU shardings to mirror. + tpu_mesh: TPU mesh for creating multi-device CPU mesh. + + Returns: + List of CPU shardings matching the TPU sharding structure. + """ + if len(cpu_devices) > 1: + logging.info("Creating CPU mesh from TPU mesh: %s", tpu_mesh) + cpu_mesh = colocated_python.colocated_cpu_devices(tpu_mesh) + logging.info("CPU Mesh: %s", cpu_mesh) + + return [ + ( + jax.sharding.NamedSharding(cpu_mesh, sharding.spec) + if isinstance(sharding, jax.sharding.NamedSharding) + else jax.sharding.NamedSharding(cpu_mesh, jax.sharding.PartitionSpec()) + ) + for sharding in tpu_shardings + ] + else: + # Single device - use simple single device sharding + return [jax.sharding.SingleDeviceSharding(cpu_devices[0]) for _ in tpu_shardings] + + +async def _effective_bytes_per_device( + spec: dict[str, Any], + shape: tuple, + dtype: typing.DTypeLike, + sharding: jax.sharding.Sharding, +) -> int: + """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. + """ + t = await ts.open(ts.Spec(spec), open=True, context=serialization.TS_CONTEXT) + shard_shape = sharding.shard_shape(shape) + raw_bytes = math.prod(shard_shape) * np.dtype(dtype).itemsize + + read_chunk = t.chunk_layout.read_chunk + if read_chunk is None or read_chunk.shape is None: + return raw_bytes + + chunk_shape = read_chunk.shape + overhead_ratio = 1.0 + for dim_idx in range(len(shape)): + cs = chunk_shape[dim_idx] + needed = shard_shape[dim_idx] + if needed > 0: + num_chunks = (needed + cs - 1) // cs + overhead_ratio *= (num_chunks * cs) / needed + + return int(raw_bytes * overhead_ratio) + + +async def _run_colocated_deserialize( + shardings: Sequence[Union[jax.sharding.Sharding, Format]], + tensorstore_specs: Sequence[dict[str, Any]], + global_shapes: Sequence[tuple], + dtypes: Sequence[typing.DTypeLike], + *, + concurrent_bytes: int, + tpu_mesh: jax.sharding.Mesh, + pipeline_concurrent_bytes: int, + multi_thread_pool: ThreadPoolExecutor, +): + """Deserialize checkpoint with pipelined load to CPU then transfer to TPU. + + This approach uses colocated Python to load checkpoint data to CPU worker, then + transfers each array to TPU. A global limiter controls the entire pipeline (load + transfer) + to keep CPU memory usage bounded. + + Args: + shardings: Target TPU shardings for the restored arrays. + tensorstore_specs: TensorStore specifications for each array. + global_shapes: Global shapes for each array. + dtypes: Data types for each array. + concurrent_bytes: Maximum concurrent bytes for reading from storage (used by worker). + tpu_mesh: TPU mesh for creating CPU mesh. Should be captured from the main thread + where the mesh context is active. + pipeline_concurrent_bytes: Maximum concurrent bytes for the entire pipeline. + Controls peak CPU memory usage. + multi_thread_pool: Thread pool for blocking operations (block_until_ready, TPU transfer). + + Returns: + List of JAX arrays on TPU devices. + """ + cpu_devices = colocated_python.colocated_cpu_devices(jax.devices()) + logging.info("Colocated CPU devices: %s", cpu_devices) + + # Create CPU shardings matching the TPU sharding structure + cpu_shardings = _create_cpu_shardings(cpu_devices, shardings, tpu_mesh) + + # Pre-compute effective per-device bytes in parallel across all arrays. + effective_bytes = list( + await asyncio.gather( + *[ + _effective_bytes_per_device( + tensorstore_specs[i], + global_shapes[i], + dtypes[i], + cpu_sharding, + ) + for i, cpu_sharding in enumerate(cpu_shardings) + ] + ) + ) + + # Global limiter controls the entire pipeline (load + transfer) + # This limits total bytes in flight across all arrays + # pylint: disable=protected-access + global_limiter = serialization._LimitInFlightBytes(pipeline_concurrent_bytes) + + logging.info( + "Created global pipeline limiter with %.1fGB capacity", + pipeline_concurrent_bytes / (10**9), + ) + + # Capture all per-array data as lists in the closure so that load_to_cpu can be defined + # and specialized exactly once. colocated_python only supports jax.Array as inputs, so we + # pass a scalar index array and look up the per-array data inside the function. + # _get_specialized_func is cached on (FunctionInfo, Specialization): with a single function + # object the cache hits for all arrays that share the same output shape/dtype/sharding, + # reducing compilations from O(num_arrays) to O(num_unique_output_specs). + cpu_mesh = colocated_python.colocated_cpu_devices(tpu_mesh) + replicated_sharding = jax.sharding.NamedSharding(cpu_mesh, jax.sharding.PartitionSpec()) + + # Initialize fresh TensorStore context and event loop on the sidecar workers. + dummy_array = jax.device_put(jnp.array(0), replicated_sharding) + _colocated_deserialize_setup.specialize(devices=cpu_devices)(dummy_array) + + @colocated_python.colocated_python + def load_to_cpu(idx: jax.Array): + i = int(idx) + # pylint: disable=protected-access + byte_limiter = serialization._LimitInFlightBytes(concurrent_bytes) + + rm = colocated_python.resource_manager + return asyncio.run_coroutine_threadsafe( + _async_deserialize( + cpu_shardings[i], + tensorstore_specs[i], + global_shapes[i], + dtypes[i], + byte_limiter=byte_limiter, + h2d_limiter=None, + single_thread_pool=None, + multi_thread_pool=None, + ts_context=rm.ts_context, + ), + rm.event_loop, + ).result() + + # Pre-create replicated index arrays across all cpu_devices. colocated_python requires inputs + # to have one shard per device in the specialized device list. + idx_arrays = [jax.device_put(jnp.array(i), replicated_sharding) for i in range(len(shardings))] + + async def _load_and_transfer_one( + idx: int, + tpu_sharding: jax.sharding.Sharding, + dispatch_pool: ThreadPoolExecutor, + ): + """Load one array to CPU via colocated Python, then transfer to TPU.""" + # Specialize per array with explicit out_specs_fn to avoid stale cached output specs + # when arrays have different shapes. _get_specialized_func caches on + # (FunctionInfo, Specialization) so arrays sharing the same (shape, dtype, sharding) + # will reuse the compiled executable. + specialized = load_to_cpu.specialize( + devices=cpu_devices, + out_specs_fn=lambda _: jax.ShapeDtypeStruct( + shape=global_shapes[idx], dtype=dtypes[idx], sharding=cpu_shardings[idx] + ), + ) + loop = asyncio.get_running_loop() + # Each call dispatched from a separate thread enables concurrent sidecar execution. + cpu_array = await loop.run_in_executor(dispatch_pool, specialized, idx_arrays[idx]) + # block_until_ready is run in an executor so the event loop can continue + # dispatching other coroutines while waiting for this array to materialize. + await loop.run_in_executor(multi_thread_pool, cpu_array.block_until_ready) + tpu_array = await loop.run_in_executor( + multi_thread_pool, + _blocking_device_put, + cpu_array, + tpu_sharding, + ) + # Explicitly release the CPU buffer once H2D transfer is complete. + # This is for deterministic memory release within colocated-python sidecar container. + del cpu_array + return tpu_array + + async def _load_and_transfer_one_rate_limited( + idx: int, + tpu_sharding: jax.sharding.Sharding, + dispatch_pool: ThreadPoolExecutor, + ): + """Wrapper that applies global limiter to the entire load+transfer operation.""" + bytes_per_device = effective_bytes[idx] + # pylint: disable-next=protected-access + max_capacity = global_limiter._max_bytes + + # Clamp to limiter capacity so oversized arrays don't deadlock wait_for_bytes. + if bytes_per_device > max_capacity: + logging.warning( + "Array shard size per device (%.2f GB) exceeded " + "limiter capacity (%.2f GB). Clamping reservation.", + bytes_per_device / (1024**3), + max_capacity / (1024**3), + ) + bytes_to_reserve = min(bytes_per_device, max_capacity) + + await global_limiter.wait_for_bytes(bytes_to_reserve) + try: + return await _load_and_transfer_one(idx, tpu_sharding, dispatch_pool) + finally: + await global_limiter.release_bytes(bytes_to_reserve) + + # Dedicated pool for colocated_python dispatch. colocated_python only executes calls + # concurrently when dispatched from different threads — a separate pool ensures dispatch + # threads are never starved by blocking work (block_until_ready, TPU transfer) on + # multi_thread_pool. + with ThreadPoolExecutor(max_workers=min(len(shardings), 256)) as dispatch_pool: + tasks = [ + _load_and_transfer_one_rate_limited(idx, tpu_sharding, dispatch_pool) + for idx, tpu_sharding in enumerate(shardings) + ] + + logging.info("Starting pipelined colocated load and transfer...") + start_time = time.perf_counter() + tpu_arrays = await asyncio.gather(*tasks) + + total_time = time.perf_counter() - start_time + logging.info("Pipelined colocated deserialization completed in %.2f seconds", total_time) + + # Tear down the TensorStore context and event loop on colocated workers. + _colocated_deserialize_teardown.specialize(devices=cpu_devices)(dummy_array) + + return tpu_arrays + + # Reference: # https://github.com/google/orbax/blob/ebb3e6d75f9ccb52bf862f1740943a45b18f4dac/checkpoint/orbax/checkpoint/future.py#L49 class _ThreadRaisingException(threading.Thread): @@ -678,34 +1061,102 @@ def deserialize( dtypes: Optional[Sequence[typing.DTypeLike]] = None, concurrent_gb: int = 32, ): + """Deserialize arrays from TensorStore. + + Args: + shardings: Sharding specifications for each array. + tensorstore_specs: TensorStore specifications for each array. + global_shapes: Global shapes for each array. If None, uses shape from TensorStore. + dtypes: Data types for each array. If None, uses dtype from TensorStore. + concurrent_gb: Maximum concurrent GB for reading from storage. + + Returns: + List of deserialized JAX arrays. + + Environment variables: + COLOCATED_PYTHON_DESERIALIZE: Set to "0" or "false" to disable colocated Python + deserialization even when available. Defaults to enabled when available. + COLOCATED_PYTHON_PIPELINE_CONCURRENT_GB: Maximum concurrent GB in flight during + CPU to TPU transfer when using colocated Python. Defaults to 64. + """ self.wait_until_finished() start_time = time.perf_counter() + use_colocated_python = colocated_python_available() and ( + os.getenv("COLOCATED_PYTHON_DESERIALIZE", "").lower() not in ("0", "false") + ) + logging.info("use_colocated_python=%s", use_colocated_python) + concurrent_bytes = concurrent_gb * 10**9 - async def _run_deserializer(): - # Object should be created once per process. - # pylint: disable=protected-access - byte_limiter = serialization._LimitInFlightBytes(concurrent_bytes) - h2d_limiter = serialization._LimitInFlightBytes(_get_premapped_buffer_size()) - - future_arrays = jax.tree.map( - functools.partial( - _async_deserialize, - byte_limiter=byte_limiter, - h2d_limiter=h2d_limiter, - single_thread_pool=self._single_thread_pool, - multi_thread_pool=self._multi_thread_pool, - ), - shardings, - tensorstore_specs, - [None] * len(tensorstore_specs) if global_shapes is None else global_shapes, - [None] * len(tensorstore_specs) if dtypes is None else dtypes, + # Prepare global shapes and dtypes + global_shapes_list = ( + [None] * len(tensorstore_specs) + if global_shapes is None + else [tuple(s) for s in global_shapes] + ) + dtypes_list = [None] * len(tensorstore_specs) if dtypes is None else list(dtypes) + + # Create the appropriate coroutine based on mode + if use_colocated_python: + # Capture mesh from main thread where context is active + # (mesh is thread-local and won't be available in background event loop) + tpu_mesh = thread_resources.env.physical_mesh + pipeline_concurrent_gb = int(os.getenv("COLOCATED_PYTHON_PIPELINE_CONCURRENT_GB", "64")) + + # Resolve any None shapes from TensorStore metadata before entering the colocated + # path, since out_specs_fn requires concrete shapes. + resolved_shapes = [ + ( + tuple( + ts.open(ts.Spec(spec), open=True, context=serialization.TS_CONTEXT) + .result() + .shape + ) + if shape is None + else shape + ) + for spec, shape in zip(tensorstore_specs, global_shapes_list) + ] + + # Use colocated Python path + coro = _run_colocated_deserialize( + shardings=shardings, + tensorstore_specs=tensorstore_specs, + global_shapes=resolved_shapes, + dtypes=dtypes_list, + concurrent_bytes=concurrent_bytes, + tpu_mesh=tpu_mesh, + pipeline_concurrent_bytes=pipeline_concurrent_gb * 10**9, + multi_thread_pool=self._multi_thread_pool, ) - return await asyncio.gather(*future_arrays) + else: + # Use default deserialization path (direct to TPU) + async def _run_deserializer(): + # Object should be created once per process. + # pylint: disable=protected-access + byte_limiter = serialization._LimitInFlightBytes(concurrent_bytes) + h2d_limiter = serialization._LimitInFlightBytes(_get_premapped_buffer_size()) + + future_arrays = jax.tree.map( + functools.partial( + _async_deserialize, + byte_limiter=byte_limiter, + h2d_limiter=h2d_limiter, + single_thread_pool=self._single_thread_pool, + multi_thread_pool=self._multi_thread_pool, + ), + shardings, + tensorstore_specs, + global_shapes_list, + dtypes_list, + ) + return await asyncio.gather(*future_arrays) - fut = asyncio.run_coroutine_threadsafe(_run_deserializer(), self._loop) - result = fut.result() + coro = _run_deserializer() + + # Run the coroutine and get result + result = asyncio.run_coroutine_threadsafe(coro, self._loop).result() logging.info("deserialize took %.4f seconds.", time.perf_counter() - start_time) return result @@ -843,4 +1294,4 @@ def serialize( self._tensorstore_spec_log_fn(tensorstore_specs) logging.info("D2H during save took %fs. Starting async commit.", time.time() - start_t) - self._start_async_commit(on_commit_callback) + self._start_async_commit(on_commit_callback) \ No newline at end of file