Diff Coverage

    cached_res = planner.get_cached() if hasattr(planner, 'get_cached') else None
    if cached_res:
        # Get final plan and metadata from cache
        logger.info("Hit final plan and metadata cache.")
        final_plan, metadata = cached_res.final_plan, cached_res.metadata
    else:
        # First time generating a plan and metadata.
        final_plan, metadata = generate_final_plan_and_metadata()

    """
    # Copy the state_dict to cpu memory for async save
    staged = build_staged_state_dict(state_dict)

    persist_completion: Future[Metadata] = Future()
    result_queue: mp.Queue = mp.Queue(maxsize=1)
    proc = _create_persist_process(
        result_queue,
        staged,
        checkpoint_id,
        storage_writer,

        use_gloo,
    )

    # After the async save is completed, the user callback will be executed.
    def async_callback():
        if callback is not None:
            callback()

    # Deliberately not a daemon: it has to outlive the training loop long enough to resolve
    # persist_completion and run callback, and killing it early would not let the process exit
    # any sooner because the persist child is joined at interpreter exit regardless.

        target=resolve_async_persist_result,
        args=(proc, result_queue, persist_completion, async_callback),
        name="AsyncCheckpointPersistJoin",
    )
    proc.start()
    join_thread.start()
    ret = AsyncSaveResponse(persist_completion=persist_completion)
    return ret


def load(
    state_dict: dict[str, Any],

    use_collectives = False if no_dist else use_collectives

    # Check the DTensor broadcast args
    if broadcast_groups and not broadcast_from_minimum_rank:
        raise ValueError("If not use broadcast_from_minimum_rank, the broadcast_groups should be None.")

    # Set up storage reader
    if storage_reader is None:
        if checkpoint_id is None:

        """
        copy_func = cls.dispatch(obj)
        if copy_func:
            return copy_func(obj)
        logger.warning(
            "The user define object %s not implement copy method and register "
            "to DataCopier. So, the general copy method will be used, which may cause "
            "unexpected copy behavior and incorrect data storage. Please verify this on your own.",
            type(obj))
        return cls._copy_others(obj)

    @staticmethod
    def _copy_others(obj):
        obj_c = copy.deepcopy(obj)
        for name, value in obj_c.__dict__.items():
            if callable(value):
                continue
            setattr(obj_c, name, DataCopier.copy(value))
        return obj_c


def _copy_tensor_to_cpu(tensor: platform.Tensor) -> platform.Tensor:
    """Return a host-memory copy of a framework tensor, detached from autograd where applicable."""

    """Return a host-memory copy of a framework tensor, detached from autograd where applicable."""
    # ``to("cpu")`` is supported on both Torch and MindSpore tensor APIs used by HyperParallel.
    t = tensor.detach().clone() if tensor.is_cpu else tensor.detach().cpu()
    if hasattr(tensor, CHUNK_INFO):
        setattr(t, CHUNK_INFO, getattr(tensor, CHUNK_INFO))
    return t


@DataCopier.register(int, float, complex, bool, str, type(None), Enum, platform.tensor_dtype.dtype)

@DataCopier.register(np.ndarray)
def _copy_ndarray(obj):
    return obj.copy()


@DataCopier.register(platform.Tensor)
def _copy_tensor(obj):

@DataCopier.register(Namespace)
def _copy_namespace(obj):
    c = DataCopier.copy(vars(obj))
    return Namespace(**c)


@DataCopier.register(list)
def _copy_list(obj):

@DataCopier.register(tuple)
def _copy_tuple(obj):
    result = tuple(DataCopier.copy(item) for item in obj)
    return result


@DataCopier.register(dict)
def _copy_dict(obj):

    """
    Cleanup hccl process group and reinitialize gloo process group.
    """
    # torch is imported lazily so this module stays importable on a MindSpore-only install.
    import torch.distributed as dist  # pylint: disable=import-outside-toplevel
    import torch.distributed.distributed_c10d as c10d  # pylint: disable=import-outside-toplevel

    # Resetting the private c10d state is the only way to drop the process group inherited
    # from the parent process, so the accesses below are deliberate.
    # pylint: disable=protected-access

    # from the parent process, so the accesses below are deliberate.
    # pylint: disable=protected-access

    # Step 1: clear C++ ProcessGroupRegistry (avoids "already registered" error).
    if hasattr(c10d, "_unregister_all_process_groups"):
        try:
            c10d._unregister_all_process_groups()
        except Exception:  # pylint: disable=broad-except
            pass

    # Step 2: reset Python-layer _world to an uninitialised state.
    if hasattr(c10d, "_World"):
        try:
            c10d._world = c10d._World()
        except Exception:  # pylint: disable=broad-except
            pass
    # Also reset the module-level init-method string (present in all versions).
    if hasattr(c10d, "_default_pg_init_method"):
        c10d._default_pg_init_method = None
    for attr in ("_default_pg", "_pg_map", "_pg_names", "_pg_group_ranks",
                 "_pg_backend_config", "_group_count", "_tags_to_pg",
                 "_pg_to_tag"):
        if not hasattr(c10d, attr):
            continue
        old_val = getattr(c10d, attr)
        if isinstance(old_val, dict):
            setattr(c10d, attr, {})
        elif isinstance(old_val, int):
            setattr(c10d, attr, 0)
        else:
            setattr(c10d, attr, None)

    # Step 3: build a fresh TCPStore on the dedicated gloo port and init.
    # rank-0 child is the store master; all others connect as clients.
    store = dist.PrefixStore(
        prefix=mp.current_process().name,
        store=dist.TCPStore(
            host_name=master_addr,
            port=master_port

        ),
    )

    # Step 4: init new process group with backend "gloo".
    dist.init_process_group(
        backend="gloo",
        store=store,
        rank=rank,
        world_size=world_size,

        master_port: Original training master port (informational, not used here).
    """
    # Lazy imports: torch keeps this module framework-agnostic, and ``api`` imports this
    # module at load time, so importing it here is what breaks the import cycle.
    import torch.distributed as dist  # pylint: disable=import-outside-toplevel
    from hyper_parallel.core.distributed_checkpoint.api import _save_impl  # pylint: disable=import-outside-toplevel
    try:
        # Clear the original communication group information and establish the gloo communication.
        cleanup_and_reinit_process_group(
            master_addr,
            master_port,
            rank,
            world_size

            master_port,
            rank,
            world_size
        )
        meta = _save_impl(
            staged,
            checkpoint_id=checkpoint_id,
            storage_writer=storage_writer,
            planner=planner,

            storage_writer=storage_writer,
            planner=planner,
        )
        # The plan cache is also returned and can be reused for the next dcp async save.
        result_queue.put((AsyncPersistStatus.SUCCESS, (meta, StandardSavePlanner.cached_save_result)))
    except Exception:  # pylint: disable=broad-except
        result_queue.put((AsyncPersistStatus.FAILURE, traceback.format_exc()))
    finally:
        # Clean up the gloo group we created; ignore errors (process is exiting).
        if dist.is_initialized():
            dist.destroy_process_group()


@dcp_timer_decorator
def execute_async_persist(

        use_collectives: Total world size (inherited from parent).
        use_storage_comm: TCP rendezvous address for the gloo store.
    """
    # ``api`` imports this module at load time, so this import has to stay function local.
    from hyper_parallel.core.distributed_checkpoint.api import _save_impl  # pylint: disable=import-outside-toplevel
    try:
        meta = _save_impl(
            staged,
            checkpoint_id=checkpoint_id,
            storage_writer=storage_writer,
            planner=planner,

            use_collectives=use_collectives,
            use_storage_comm=use_storage_comm
        )
        # The plan cache is also returned and can be reused for the next dcp async save.
        result_queue.put((AsyncPersistStatus.SUCCESS, (meta, StandardSavePlanner.cached_save_result)))
    except Exception:  # pylint: disable=broad-except
        result_queue.put((AsyncPersistStatus.FAILURE, traceback.format_exc()))


def resolve_async_persist_result(
        proc: mp.Process,

            break

    proc.join()
    if persist_future.done():
        return

    if result is None:
        persist_future.set_exception(
            RuntimeError(

    while True:
        if not file_to_read:
            break
        if time.monotonic() > deadline:
            raise TimeoutError(
                f"After waiting for 1800 seconds, "
                f"the {file_type.name} files: {file_to_read} still have not been read completely."
            )
        file_not_read = []

    Split the total file list into multiple batches
    based on the number of parallel threads,
    then perform multi-threaded parallel reading to improve read efficiency.
    """
    if not file_list:
        return []

    batches, max_workers = construct_file_batches(file_list, max_workers)
    logger.info(
        "parallel read files, file type: %s, file count: %d, max read workers: %d",
        file_type.name, len(file_list), max_workers
    )

        file_type.name, len(file_list), max_workers
    )

    # Use the thread pool to read data concurrently, improving the read efficiency.
    with ThreadPoolExecutor(max_workers=max_workers) as executor:
        all_results = tuple(
            executor.map(lambda b: batch_read_worker(b, checkpoint_id, file_type), batches)
        )
    return [result for batch_results in all_results for result in batch_results]


def construct_file_batches(file_list: tuple[int], max_workers: int) -> (list[tuple[int]], int):
    """

    In the first step, each rank writes its own data.
    In the second step, for local plan, each rank reads the file written by all ranks.
    For storage data, only rank 0 reads the file written by all ranks.
    """
    logger.info("gather all results from storage, file type: %s", file_type.name)
    write_file(checkpoint_id, file_type, file_id, results)
    if not is_coordinator and file_type != FileType.LOCAL_PLAN:
        return []

    # For LOCAL_PLAN, each node has 16 rank processes reading simultaneously.
    if file_type == FileType.LOCAL_PLAN:
        max_workers = min(16, max(os.cpu_count() // 32, 1))
    else:
        max_workers = min(64, max(os.cpu_count() // 2, 1))
    file_list = tuple(range(0, world_size))
    gathered = parallel_read_files(file_list,
                                   checkpoint_id,
                                   file_type,
                                   max_workers=max_workers)
    return gathered


def assemble_file_path(checkpoint_id: str, file_type: FileType, file_id: int) -> str:
    """

    _FOREIGN_MODULE_MAP = {}

    def find_class(self, module: str, name: str) -> Any:
        if module in _MetadataUnpickler._FOREIGN_MODULE_MAP:
            module = _MetadataUnpickler._FOREIGN_MODULE_MAP[module]
        return super().find_class(module, name)


class FileSystemWriter(StorageWriter):

    if hasattr(target_tensor, "copy_"):
        # for torch and ms, call 'copy_' to copy data to target tensor
        target_tensor.copy_(tensor)
    else:
        target_tensor[...] = tensor
    planner.apply_tensor(req, target_tensor)


def _load_bytes_file(

    """
    if platform.platform_type == PlatformType.PYTORCH:
        _load_torch_tensor_file(path, reqs, planner, storage_data)
    else:
        _load_ms_tensor_file(path, reqs, planner, storage_data)


class FileSystemReader(StorageReader):
    """

                # TENSOR: one safetensors file per rank
                _load_tensor_file(path, reqs, planner, self.storage_data)

        if self.broadcast_from_minimum_rank:
            broadcast_loaded_tensors(planner.state_dict, broadcast_groups)

                    size = item.tensor_data['size']

                    # Validate consistency across ranks
                    if fqn in fqn_to_chunks and (fqn_to_properties[fqn] != properties or fqn_to_size[fqn] != size):
                        raise ValueError(f"The {fqn} in different rank has different properties and size, "
                                         f"properties: {fqn_to_properties[fqn]} != {properties}, "
                                         f"size: or {fqn_to_size[fqn]} != {size}.")

                    # Initialize FQN entry if not exists

            ValueError: If the chunk info has an unexpected type, or if the shard group is
                empty, or if the current rank is not a member of the shard group.
        """
        if isinstance(tensor, DTensor):
            group_ranks = infer_same_shard_ranks_for_dtensor(tensor)
        elif hasattr(tensor, CHUNK_INFO):
            if not isinstance(getattr(tensor, CHUNK_INFO), ChunkInfo):
                raise ValueError(f"The chunk info attached to tensor must be of type {ChunkInfo}")
            group_ranks = getattr(tensor, CHUNK_INFO).replica_rank_list
            if group_ranks is None:
                return True
        else:
            return True

        if not group_ranks:
            raise ValueError("The tensor must be distributed on at least one rank.")
        if self.rank not in group_ranks:
            raise ValueError(f"Current rank {self.rank} is not in the same shard group {group_ranks}.")

        load_rank = min(group_ranks)

        if not isinstance(obj, DTensor):
            return True
        layout = getattr(obj, "layout", None)
        if layout is None:
            return True
        rank_list = getattr(layout, "rank_list", None) if layout else None
        if rank_list is None:
            rank_list = getattr(layout, "_rank_list", None)
        if rank_list is None:
            return True
        return self.rank in rank_list

    def build_local_plan(self) -> LoadPlan:
        """

                    raise ValueError(
                        f"Size mismatch between saved {md.size} and current: {obj_size} for {fqn}",
                    )
                if not self._rank_owns_dtensor_shard(obj):
                    continue
                if self.broadcast_from_minimum_rank and not self.should_load_shard(obj):
                    continue
                # Both DTensor and platform.Tensor: create local chunks and read items
                local_chunks = create_chunk_list_for_tensor(obj)

    Returns:
        tuple[int, ...]: Sorted global-rank tuples; each tuple is one
            same-shard group (length 1 when the shard is unique to one rank).
    """
    current_rank = platform.get_rank()
    layout = dtensor.layout
    if layout is None:
        return (current_rank,)

    mesh_shape = layout.mesh_shape
    tensor_map = layout.tensor_map
    rank_list = layout.rank_list

    if mesh_shape is None or tensor_map is None or rank_list is None:
        return (current_rank,)

    if current_rank not in rank_list:
        return (current_rank,)

    n_mesh_dims = len(mesh_shape)
    inner_rank_id = rank_list.index(current_rank)

    def dev_id_list_from_rank(inner_id: int) -> list[int]:
        dev_id_list = [0] * n_mesh_dims
        temp = inner_id
        for i in range(n_mesh_dims - 1, -1, -1):
            dev_id_list[i] = temp % mesh_shape[i]
            temp //= mesh_shape[i]
        return dev_id_list

    def compute_shard_key(dev_id_list: list[int]) -> tuple:
        key_parts = []
        for mapping in tensor_map:
            if isinstance(mapping, int):
                mapping = (mapping,) if mapping != -1 else ()
            elif not isinstance(mapping, tuple):
                mapping = (mapping,)
            if not mapping:
                continue
            shard_id = 0
            coef = 1
            for dim in reversed(mapping):
                if dim == -1:
                    continue
                shard_id += dev_id_list[-dim - 1] * coef
                coef *= mesh_shape[-dim - 1]
            key_parts.append(shard_id)
        return tuple(key_parts)

    current_dev_id_list = dev_id_list_from_rank(inner_rank_id)
    current_shard_key = compute_shard_key(current_dev_id_list)

    same_shard_ranks = []
    for idx, global_rank in enumerate(rank_list):
        if idx == inner_rank_id:
            same_shard_ranks.append(global_rank)
            continue
        other_dev_id_list = dev_id_list_from_rank(idx)
        if compute_shard_key(other_dev_id_list) == current_shard_key:
            same_shard_ranks.append(global_rank)

    return tuple(same_shard_ranks)


@dcp_timer_decorator
def all_gather_object(

    for group_ranks, tensors in missing_groups_ranks.items():
        for tensor in tensors:
            broadcast_info = getattr(tensor, BROADCAST_INFO, None)
            if broadcast_info is None:
                continue
            platform.broadcast(
                tensor.to_local().detach() if isinstance(tensor, DTensor) else tensor.detach(),
                broadcast_info.src_rank,
                new_groups[group_ranks]

        Raises:
            NotImplementedError: MindSpore support is not yet implemented.
        """
        raise NotImplementedError(
            "new_group is not yet supported on MindSpore"
        )

    def _create_group(self, rank_list, pg_options: Any = None):

        Returns:
            The newly created communication group.
        """
        raise NotImplementedError("Platform subclasses must implement new_group")

    def create_group(self, rank_list):
        """Create or retrieve a communication group with the specified ranks.

        Returns:
            ProcessGroup: The newly created process group.
        """
        return dist.new_group(ranks=list(rank_list))

    def _create_group(self, rank_list):
        normalized_rank_list = tuple(sorted(rank_list))
        world_rank_list = tuple(range(self.get_world_size()))