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()

    join_thread = threading.Thread(
        target=resolve_async_persist_result,
        args=(proc, result_queue, persist_completion, async_callback),

        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 []

    file_list = tuple(range(0, world_size))
    gathered = parallel_read_files(file_list,
                                   checkpoint_id,
                                   file_type,
                                   max_workers=min(64, os.cpu_count() // 2))
    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 only rank of the group is not the current rank.
        """
        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}")
            if getattr(tensor, CHUNK_INFO).replica_rank_list is None:
                return True
            group_ranks = getattr(tensor, CHUNK_INFO).replica_rank_list
        else:
            return True

        if not group_ranks:
            raise ValueError("The tensor must be distributed on at least one rank.")
        if len(group_ranks) == 1 and self.rank != group_ranks[0]:
            raise ValueError(f"The same shard group only have one rank {group_ranks[0]}, "
                             f"but is not the current rank {self.rank}.")

        load_rank = min(group_ranks)
        if len(group_ranks) > 1:

        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()))