<commit_before># Copyright 2022 The Orbax 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
#
#     http://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.

"""PyTreeCheckpointHandler class.

Implementation of CheckpointHandler interface.
"""

import asyncio
import dataclasses
import functools
import os
import re
from typing import Any, List, MutableMapping, Optional, Tuple, Union, cast

from etils import epath
import flax
from flax import traverse_util
import jax
from jax.experimental import multihost_utils
from jax.experimental import pjit
from jax.experimental.gda_serialization import serialization

from jax.experimental.global_device_array import GlobalDeviceArray
from jax.experimental.maps import Mesh
import jax.numpy as jnp
import numpy as np
from orbax.checkpoint import lazy_array
from orbax.checkpoint import transform_utils
from orbax.checkpoint import utils
from orbax.checkpoint.async_checkpoint_handler import AsyncCheckpointHandler
from orbax.checkpoint.future import Future
import tensorstore as ts

# TODO(orbax-dev): Remove this when jax>=0.3.19
# pylint: disable=g-import-not-at-top,bare-except
try:
  from jax import sharding
  from jax import make_array_from_callback
except ImportError:
  from jax.experimental import sharding  # pytype: disable=import-error
  from jax.experimental.array import make_array_from_callback  # pytype: disable=import-error
# pylint: enable=g-import-not-at-top,bare-except

PyTree = type(jax.tree_util.tree_structure(None))
ArrayOrScalarType = Union[int, float, np.number, np.ndarray, jnp.ndarray,
                          GlobalDeviceArray, jax.Array]
ArrayType = Union[np.ndarray, jnp.ndarray, GlobalDeviceArray, jax.Array]
StructureLeaf = Union[str, dict, ts.Spec, ArrayOrScalarType]

_FLAX_CHECKPOINT_FILE = 'checkpoint'

utils.register_ts_spec_for_serialization()


@dataclasses.dataclass
class ParamInfo:
  """Information describing a parameter in a PyTree.

  name: name of the parameter.
  tspec: Tensorstore spec in JSON format.
  """
  name: Optional[str]
  tspec: Optional[MutableMapping[str, Any]]


@dataclasses.dataclass
class SaveArgs:
  """Extra arguments that can be provided for saving.

  use_flax: if true, saves the given parameter using flax.serialization to a
    unified checkpoint file. Must be false if the given array value is a GDA.
  dtype: if provided, casts the parameter to the given dtype before saving.
    Note that the parameter must be compatible with the given type (e.g.
    jnp.bfloat16 is not compatible with np.ndarray).
  """
  use_flax: bool = False
  dtype: Optional[jnp.dtype] = None


# TODO(b/233807751): Raise an error if mesh/mesh_axes are None.
@dataclasses.dataclass
class RestoreArgs:
  """Extra arguments that can be provided for restoration.

  as_gda: if true, restores the given paramater as a GlobalDeviceArray
    regardless of how it was saved. Mesh and mesh_axes are required.
  as_jax_array: if true, restores the given paramater as a JAX Array type
    regardless of how it was saved. Mesh and mesh_axes are required.
  mesh: the device mesh that the array should be restored as. If restoring as
    GDA, cannot be None.
  mesh_axes: the mesh_axes that the array should be restored as. If restoring as
    GDA, cannot be None.
  global_shapes: the global shape that the array should be restored into. If not
    provided, the shape will be restored as written.
  lazy: if True, restores using LazyArray. The actual read operation will not be
    performed until `get` is called for the restored LazyArray
  dtype: if provided, casts the parameter to the given dtype after restoring.
    Note that the parameter must be compatible with the given type (e.g.
    jnp.bfloat16 is not compatible with np.ndarray).
  """
  as_gda: bool = True
  as_jax_array: bool = False
  mesh: Optional[Mesh] = None
  mesh_axes: Optional[pjit.PartitionSpec] = None
  global_shape: Optional[Tuple[int]] = None
  lazy: bool = False
  dtype: Optional[jnp.dtype] = None


def _validate_save_args(args: SaveArgs, enable_flax: bool):
  if args.use_flax and not enable_flax:
    raise ValueError(
        'Saving with flax disabled for this handler, but saving with flax was requested for at least one parameter.'
    )


def _validate_restore_args(args: RestoreArgs):
  if args.mesh is None or args.mesh_axes is None:
    raise ValueError(
        'Sharding of GlobalDeviceArray/Array cannot be None. Provide `mesh` and `mesh_axes`.'
    )


async def _create_param_save_dir(param_info: ParamInfo):
  tspec = param_info.tspec
  if tspec is None:
    return
  path = tspec['kvstore']['path']
  if jax.process_index() == 0:
    await utils.async_makedirs(epath.Path(path))


async def _maybe_deserialize(args, value, info):
  """Deserialize from tensorstore or return value if already deserialized."""
  if value is None or (isinstance(value, dict) and not value):
    return {}
  if isinstance(value, (ts.Spec, utils.Leaf)):
    result = lazy_array.LazyAwaitableArray.from_tensor_store_spec(
        ts.Spec(info.tspec),
        get_fn=lambda: _deserialize_array(args, info),<commit_msg>Add validation for restore args preventing improper use of GDA vs. jax.Array arguments.<commit_after># Copyright 2022 The Orbax 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
#
#     http://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.

"""PyTreeCheckpointHandler class.

Implementation of CheckpointHandler interface.
"""

import asyncio
import dataclasses
import functools
import os
import re
from typing import Any, List, MutableMapping, Optional, Tuple, Union, cast

from etils import epath
import flax
from flax import traverse_util
import jax
from jax.experimental import multihost_utils
from jax.experimental import pjit
from jax.experimental.gda_serialization import serialization

from jax.experimental.global_device_array import GlobalDeviceArray
from jax.experimental.maps import Mesh
import jax.numpy as jnp
import numpy as np
from orbax.checkpoint import lazy_array
from orbax.checkpoint import transform_utils
from orbax.checkpoint import utils
from orbax.checkpoint.async_checkpoint_handler import AsyncCheckpointHandler
from orbax.checkpoint.future import Future
import tensorstore as ts

# TODO(orbax-dev): Remove this when jax>=0.3.19
# pylint: disable=g-import-not-at-top,bare-except
try:
  from jax import sharding
  from jax import make_array_from_callback
except ImportError:
  from jax.experimental import sharding  # pytype: disable=import-error
  from jax.experimental.array import make_array_from_callback  # pytype: disable=import-error
# pylint: enable=g-import-not-at-top,bare-except

PyTree = type(jax.tree_util.tree_structure(None))
ArrayOrScalarType = Union[int, float, np.number, np.ndarray, jnp.ndarray,
                          GlobalDeviceArray, jax.Array]
ArrayType = Union[np.ndarray, jnp.ndarray, GlobalDeviceArray, jax.Array]
StructureLeaf = Union[str, dict, ts.Spec, ArrayOrScalarType]

_FLAX_CHECKPOINT_FILE = 'checkpoint'

utils.register_ts_spec_for_serialization()


@dataclasses.dataclass
class ParamInfo:
  """Information describing a parameter in a PyTree.

  name: name of the parameter.
  tspec: Tensorstore spec in JSON format.
  """
  name: Optional[str]
  tspec: Optional[MutableMapping[str, Any]]


@dataclasses.dataclass
class SaveArgs:
  """Extra arguments that can be provided for saving.

  use_flax: if true, saves the given parameter using flax.serialization to a
    unified checkpoint file. Must be false if the given array value is a GDA.
  dtype: if provided, casts the parameter to the given dtype before saving.
    Note that the parameter must be compatible with the given type (e.g.
    jnp.bfloat16 is not compatible with np.ndarray).
  """
  use_flax: bool = False
  dtype: Optional[jnp.dtype] = None


# TODO(b/233807751): Raise an error if mesh/mesh_axes are None.
@dataclasses.dataclass
class RestoreArgs:
  """Extra arguments that can be provided for restoration.

  as_gda: if true, restores the given paramater as a GlobalDeviceArray
    regardless of how it was saved. Mesh and mesh_axes are required.
  as_jax_array: if true, restores the given paramater as a JAX Array type
    regardless of how it was saved. Mesh and mesh_axes are required.
  mesh: the device mesh that the array should be restored as. If restoring as
    GDA, cannot be None.
  mesh_axes: the mesh_axes that the array should be restored as. If restoring as
    GDA, cannot be None.
  global_shapes: the global shape that the array should be restored into. If not
    provided, the shape will be restored as written.
  lazy: if True, restores using LazyArray. The actual read operation will not be
    performed until `get` is called for the restored LazyArray
  dtype: if provided, casts the parameter to the given dtype after restoring.
    Note that the parameter must be compatible with the given type (e.g.
    jnp.bfloat16 is not compatible with np.ndarray).
  """
  as_gda: bool = True
  as_jax_array: bool = False
  mesh: Optional[Mesh] = None
  mesh_axes: Optional[pjit.PartitionSpec] = None
  global_shape: Optional[Tuple[int]] = None
  lazy: bool = False
  dtype: Optional[jnp.dtype] = None


def _validate_save_args(args: SaveArgs, enable_flax: bool):
  if args.use_flax and not enable_flax:
    raise ValueError(
        'Saving with flax disabled for this handler, but saving with flax was requested for at least one parameter.'
    )


def _validate_restore_args(args: RestoreArgs):
  """Validates RestoreArgs object."""
  if args.mesh is None or args.mesh_axes is None:
    raise ValueError(
        'Sharding of GlobalDeviceArray/Array cannot be None. Provide `mesh` and `mesh_axes`.'
    )
  if args.as_gda and not jax.config.jax_parallel_functions_output_gda:
    raise ValueError('Requested restoration as GDA, but GDA is not enabled.')
  if args.as_jax_array and not jax.config.jax_array:
    raise ValueError(
        'Requested restoration as JAX Array, but jax.Array is not enabled.')
  if jax.config.jax_parallel_functions_output_gda and jax.config.jax_array:
    raise ValueError('Cannot enable both GDA and jax.Array.')


async def _create_param_save_dir(param_info: ParamInfo):
  tspec = param_info.tspec
  if tspec is None:
    return
  path = tspec['kvstore']['path']
  if jax.process_index() == 0:
    await utils.async_makedirs(epath.Path(path))


async def _maybe_deserialize(args, value, info):
  """Deserialize from tensorstore or return value if already deserialized."""
  if value is None or (isinstance(value, dict) and not value):
    return {}
  if isinstance(value, (ts.Spec, utils.Leaf)):
    result = lazy_array.LazyAwaitableArray.from_tensor_store_spec(
        ts.Spec(info.tspec),
        get_fn=lambda: _deserialize_array(args, info),
{'fraction_non_alphanumeric': 0.0708559606930437, 'fraction_numerical': 0.004482372209292302, 'mean_word_length': 4.341620626151013, 'pattern_counts': {'":': 0, '<': 3, '<?xml version=': 0, '>': 5, 'https://': 0, 'lorem ipsum': 0, 'www.': 2, 'xml': 0}, 'pii_count': 0, 'substrings_counts': 0, 'word_list_counts': {'cursed_substrings.json': 0, 'profanity_word_list.json': 0, 'sexual_word_list.json': 0, 'zh_pornsignals.json': 0}}
{'dir': 'git-commits-cleaned', 'id': '3424962', 'n_tokens_mistral': 3506, 'n_tokens_neox': 3439, 'n_words': 1313}