[WIP] jax pytree data serialization

PiperOrigin-RevId: 668858634

jax/experimental/array_serialization/__init__.py

@@ -11,3 +11,8 @@
 # WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
 # See the License for the specific language governing permissions and
 # limitations under the License.
+
+from jax.experimental.array_serialization.serialization import (
+ GlobalAsyncCheckpointManager)
+from jax.experimental.array_serialization.pytree_serialization import (
+ save, load, load_pytree, nonblocking_load, nonblocking_save)

jax/experimental/array_serialization/asyncio_utils.py

@@ -0,0 +1,73 @@
+# Copyright 2021 The JAX Authors.
+#
+# Licensed under the Apache License, Version 2.0 (the "License");
+# you may not use this file except in compliance with the License.
+# You may obtain a copy of the License at
+#
+# https://www.apache.org/licenses/LICENSE-2.0
+#
+# Unless required by applicable law or agreed to in writing, software
+# distributed under the License is distributed on an "AS IS" BASIS,
+# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+# See the License for the specific language governing permissions and
+# limitations under the License.
+
+import asyncio
+import functools
+import logging
+from concurrent.futures import ThreadPoolExecutor
+
+logger = logging.getLogger(__name__)
+
+_PARALLEL_THREAD_POOL_EXECUTOR = ThreadPoolExecutor(max_workers=64)
+_ORDERED_THREAD_EXECUTOR = ThreadPoolExecutor(max_workers=1)
+
+# Lifted from T5X.
+class _LimitInFlightBytes:
+ """Limits in-flight bytes when reading/writing checkpoints per process."""
+
+ def __init__(self, num_bytes):
+ self._max_bytes = num_bytes
+ self._available_bytes = num_bytes
+ self._cv = asyncio.Condition(lock=asyncio.Lock())
+
+ async def wait_for_bytes(self, requested_bytes):
+ if requested_bytes > self._max_bytes:
+ self._max_bytes = requested_bytes
+ logger.warning("Requested more bytes than we reserved space for: %d > %d"
+ ". Increasing the limit to %d.", requested_bytes,
+ self._max_bytes, self._max_bytes)
+ async with self._cv:
+ await self._cv.wait_for(lambda: self._available_bytes > requested_bytes)
+ self._available_bytes -= requested_bytes
+ assert self._available_bytes >= 0
+
+ async def release_bytes(self, requested_bytes):
+ async with self._cv:
+ self._available_bytes += requested_bytes
+ assert self._available_bytes <= self._max_bytes
+ self._cv.notify_all()
+
+
+def _maybe_run_async_sync(name, async_fn, ordered_execution: bool = False):
+ """Run async routine synchronously irrespective of the current environment.
+
+ Args:
+ name: The name of the function.
+ async_fn: The function to run.
+ ordered_execution: If True, the function will be run in an ordered sequence
+ Otherwise, it will be run in a separate thread pool.
+ Returns:
+ The result of the function async_fn or raises an exception.
+ """
+ thread_pool_executor = (_ORDERED_THREAD_EXECUTOR if ordered_execution
+ else _PARALLEL_THREAD_POOL_EXECUTOR)
+
+ def wrapped_fn(*args, **kw):
+ return thread_pool_executor.submit(
+ lambda: asyncio.run(async_fn(*args, **kw))).result()
+
+ functools.update_wrapper(wrapper=wrapped_fn, wrapped=async_fn)
+ wrapped_fn.__name__ = name
+ wrapped_fn.__qualname__ = name
+ return wrapped_fn