Diff Coverage

    checkpoint_id = Path(checkpoint_id) if isinstance(checkpoint_id, str) else checkpoint_id

    # Determine if we're in distributed mode
    use_collectives = False if no_dist else use_collectives
    use_storage_comm = False if not use_collectives else use_storage_comm

    # Set up storage writer
    if storage_writer is None:
        if checkpoint_id is None:

    else:
        if checkpoint_id:
            storage_writer.initialize_writer(checkpoint_id)
        else:
            checkpoint_id = getattr(storage_writer, "checkpoint_dir", None)

    if use_storage_comm and checkpoint_id is None:
        raise ValueError(
            "Coordinating through the storage needs a checkpoint directory to exchange the "
            "plans and write results in: pass checkpoint_id, or a storage_writer exposing "
            "checkpoint_dir."
        )

        rank=rank,
        use_collectives=use_collectives
    )

    @dcp_timer_decorator
    def generate_final_plan_and_metadata():
        # Build local plan
        local_plan = planner.build_local_plan()
        local_plan = storage_writer.optimize_local_plan(local_plan)

        local_plan = planner.build_local_plan()
        local_plan = storage_writer.optimize_local_plan(local_plan)

        #Gather all local plans and build global plan
        all_local_plans = gather_all_results_from_storage(checkpoint_id,
                                                          is_coordinator,
                                                          world_size,
                                                          local_plan,
                                                          FileType.LOCAL_PLAN,

                                                          local_plan,
                                                          FileType.LOCAL_PLAN,
                                                          rank) if use_storage_comm \
            else all_gather_object(local_plan, world_size, use_collectives)
        global_plans, global_metadata = planner.build_global_plan(all_local_plans)
        global_plans = storage_writer.optimize_global_plan(global_plans)
        # Select central plan for current rank
        if use_collectives and world_size > 1 and global_plans:
            central_plan = global_plans[rank]

        else:
            central_plan = local_plan

        # Finalize and cache plan
        finalized_plan = planner.finalize_plan(central_plan)
        # Add final plan and metadata to the cache
        if hasattr(planner, 'cache_result'):
            planner.cache_result(finalized_plan, global_metadata)
        return finalized_plan, global_metadata

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

    # Write data
    write_results = storage_writer.execute_write(final_plan, planner)

    # Write data
    write_results = storage_writer.execute_write(final_plan, planner)

    # Finalize checkpoint
    all_write_results = gather_all_results_from_storage(checkpoint_id,
                                                        is_coordinator,
                                                        world_size,
                                                        write_results,
                                                        FileType.STORAGE_DATA,

    Raises:
        AssertionError: If gloo communication is requested without MASTER_ADDR or MASTER_PORT.
    """
    world_size = platform.get_world_size()
    need_comm = use_collectives and not no_dist and world_size > 1
    if not need_comm or not use_gloo:
        use_storage_comm = need_comm and not use_gloo
        if use_storage_comm:
            logger.info("use storage communication for dcp async save.")
        return mp.Process(
            target=execute_async_persist,
            args=(
                result_queue,
                staged,

            ),
            name="AsyncCheckpointPersist",
        )

    logger.info("use Gloo communication for dcp async save.")
    rank = platform.get_rank()
    master_addr = os.environ.get("MASTER_ADDR", None)
    master_port = os.environ.get("MASTER_PORT", None)
    if master_addr is None:
        raise AssertionError("Async DCP needs MASTER_ADDR to use prefix store")
    if master_port is None:
        raise AssertionError("Async DCP needs MASTER_PORT to use prefix store")
    master_port = int(master_port)
    global _async_save_count
    proc = mp.Process(
        target=execute_async_persist_with_gloo,
        args=(
            result_queue,
            staged,

            master_port,
        ),
        name=f"AsyncCheckpointPersistGloo_{_async_save_count}",
    )
    _async_save_count += 1
    return proc


@dcp_timer_decorator
def async_save(

    """
    # 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],

    # Determine if we're in distributed mode
    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:

    local_plan = planner.build_local_plan()
    local_plan = storage_reader.optimize_local_plan(local_plan)

    # Gather all local plans and build global plan
    all_local_plans = all_gather_object(local_plan, world_size, use_collectives)
    global_plans = planner.build_global_plan(all_local_plans)
    global_plans = storage_reader.optimize_global_plan(global_plans)

    # Select central plan for current rank

    # Finalize plan
    final_plan = planner.finalize_plan(central_plan)

    # Execute read
    storage_reader.execute_read(final_plan, planner, broadcast_groups)

        Returns:
            Optional[Callable]: The handler, or None when no registered type matches.
        """
        if type(obj) in cls._registry:
            return cls._registry[type(obj)]

        best_type = None
        for r_type in cls._registry:
            if not isinstance(obj, r_type):
                continue
            # Unrelated matches (neither one a subclass of the other) keep the first found.
            if best_type is None or issubclass(r_type, best_type):
                best_type = r_type
        return cls._registry[best_type] if best_type is not None else None

    @classmethod
    def copy(cls, obj):
        """

        Returns:
            Any: Copy of ``obj``, with tensor leaves residing in host memory.
        """
        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."""

def _copy_tensor_to_cpu(tensor: platform.Tensor) -> platform.Tensor:
    """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)
def _copy_const(obj):
    return obj


@DataCopier.register(bytes, bytearray)
def _copy_bytes(obj):
    return bytes(obj)


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


@DataCopier.register(platform.Tensor)
def _copy_tensor(obj):
    result = _copy_tensor_to_cpu(obj)
    return result


@DataCopier.register(DTensor)
def _copy_dtensor(obj):

@DataCopier.register(DTensor)
def _copy_dtensor(obj):
    """Stage a DTensor by copying its local shard to host memory, keeping mesh and placements."""
    staged_local = _copy_tensor_to_cpu(obj.to_local())
    # ``shape`` is required for RaggedShard: an unevenly-sharded DTensor cannot
    # recover its logical global shape from the local shard alone.
    result = DTensor.from_local(
        staged_local,
        obj.device_mesh,
        obj.placements,
        shape=tuple(obj.shape),

        obj.device_mesh,
        obj.placements,
        shape=tuple(obj.shape),
    )
    return result


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


@DataCopier.register(list)
def _copy_list(obj):
    result = [DataCopier.copy(item) for item in obj]
    return result


@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):
    result = {}
    for k, v in obj.items():
        result[DataCopier.copy(k)] = DataCopier.copy(v)
    return result


class AsyncPersistStatus(Enum):
    """Queue payload status from :func:`_async_persist_worker` to the parent join thread."""

    persist_completion: Future[Metadata]

    def get_result(self, timeout: int = None) -> Optional[Metadata]:
        return self.persist_completion.result(timeout=timeout)


@dcp_timer_decorator
def build_staged_state_dict(state_dict: dict[str, Any]) -> dict[str, Any]:

    Returns:
        dict[str, Any]: New dict with identical nesting and keys; tensor leaves are staging copies.
    """
    return DataCopier.copy(state_dict)


def cleanup_and_reinit_process_group(
    master_addr: str,

    """
    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: Training master port; the child connects to that store as a client.
    """
    # 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,

        persist_future: Future[None],
        async_callback: Callable[[], None],
) -> None:
    """Join persist ``proc`` and complete ``persist_future`` (runs on a background thread)."""
    result = None
    while result is None:
        try:
            result = result_queue.get(timeout=_PERSIST_RESULT_POLL_SECONDS)
        except queue.Empty:
            if proc.is_alive():
                continue
            # The child may have put its result and exited while the wait above was
            # timing out, so look once more before giving up on it.
            try:
                result = result_queue.get_nowait()
            except queue.Empty:
                pass
            break

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

    if result is None:
        persist_future.set_exception(
            RuntimeError(
                f"async_persist process exited with code {proc.exitcode} and no result on queue"
            )
        )
        return

    status, payload = result
    if status == AsyncPersistStatus.SUCCESS:
        try:
            async_callback()
            StandardSavePlanner.cached_save_result.update(payload[1])
        except Exception:  # pylint: disable=broad-except
            persist_future.set_exception(RuntimeError(traceback.format_exc()))
        else:
            persist_future.set_result(payload[0])
    elif status == AsyncPersistStatus.FAILURE:
        persist_future.set_exception(RuntimeError(payload))
    else:
        persist_future.set_exception(
            RuntimeError(f"async_persist queue returned unexpected status: {status!r}")
        )

    Raises:
        TimeoutError: If some files are still unreadable after 1800 seconds.
    """
    results = [None] * len(file_batch)
    file_to_read = file_batch
    first_file_id = file_batch[0]

    # Read files in a round-robin manner,
    # If a file has not been generated yet, skip it and proceed to the next file.
    # Until all files have been read.
    timeout = 1800
    deadline = time.monotonic() + timeout
    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 = []
        for file_id in file_to_read:
            ret = load_file(checkpoint_id, file_type, file_id)
            if ret is None:
                file_not_read.append(file_id)
            else:
                results[file_id - first_file_id] = ret
        file_to_read = tuple(file_not_read)

    return results


@dcp_timer_decorator
def parallel_read_files(

    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):
    """

    """
    Split the total file list into multiple batches
    based on the number of parallel threads
    """
    n = len(file_list)
    max_workers = min(max_workers, n)
    if max_workers == 0:
        raise ValueError("Max workers number for parallel read batch file is 0!")
    batch_size = (n + max_workers - 1) // max_workers
    # Split the file list into max_workers sublists.
    batches = [
        file_list[i * batch_size: (i + 1) * batch_size]
        for i in range(max_workers)
    ]
    batches = [b for b in batches if b]
    return batches, max_workers


@dcp_timer_decorator
def gather_all_results_from_storage(

    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:
    """

    Returns:
        str: Path of that rank's file.
    """
    return f"{checkpoint_id}/{file_type.name}_{file_id}.pkl"


def load_file(checkpoint_id: str, file_type: FileType, file_id: int):
    """

    Reads files from the storage.
    Currently, this function is mainly used to
    asynchronously save local plan and storage data files.
    """
    file_path = assemble_file_path(checkpoint_id, file_type, file_id)

    try:
        with open(file_path, "rb") as f:
            flag_len = len(_COMPLETE_FLAG)
            f.seek(0, os.SEEK_END)
            size = f.tell()
            # determining whether the file length is greater than the flag length
            if size < flag_len:
                return None
            f.seek(size - flag_len)
            # determining whether the file has a write complete flag
            if f.read(flag_len) != _COMPLETE_FLAG:
                return None
            f.seek(0)
            result = pickle.load(f)
        if file_type == FileType.STORAGE_DATA:
            os.remove(file_path)
        return result
    except FileNotFoundError:
        return None



@dcp_timer_decorator

    Write files to the storage.
    Currently, this function is mainly used to
    asynchronously save local plan and storage data files.
    """
    final_path = assemble_file_path(checkpoint_id, file_type, file_id)

    # To avoid read-write conflicts, the completion flag is appended last: a reader that
    # opens the file mid-write finds no flag at the end and retries on its next round.
    with open(final_path, "wb") as f:
        pickle.dump(content, f, protocol=pickle.HIGHEST_PROTOCOL)
        f.write(_COMPLETE_FLAG)

    logger.info("write file type: %s, file path: %s", file_type.name, final_path)

    _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):
    """

            is_coordinator (bool): Whether this rank is the coordinator.
            **kwargs: Additional keyword arguments (e.g., rank, use_collectives).
        """
        self.is_coordinator = is_coordinator
        self.rank = kwargs.get("rank") if "rank" in kwargs else platform.get_rank()
        self.use_collectives = kwargs.get("use_collectives", True)

    def optimize_local_plan(self, plan: SavePlan) -> SavePlan:
        """

    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(
        path: str,

            tensor_slices = tuple(
                slice(int(off), int(off) + int(length))
                for off, length in zip(req.storage_offsets, req.lengths)
            )
            tensor = tensor_file.get_slice(tensor_key)[tensor_slices]
            _validate_and_copy_tensor(req, tensor, planner)


def _load_ms_tensor_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):
    """

    def __init__(self, checkpoint_dir: Union[Path, str]):
        self.checkpoint_dir = Path(checkpoint_dir) if isinstance(checkpoint_dir, str) else checkpoint_dir
        # Cached storage layout: MetadataIndex -> StorageInfo (torch-aligned)
        self.storage_data: dict[MetadataIndex, StorageInfo] = {}
        self.rank: int = 0
        self.is_coordinator: bool = False
        self.broadcast_from_minimum_rank = False

    def initialize_reader(self, checkpoint_id: Optional[Union[Path, str]] = None) -> None:
        """
        Initialize storage reader with new checkpoint directory.

        if not metadata_file.exists():
            raise FileNotFoundError(f"Metadata file not found: {metadata_file}")
        with open(metadata_file, "rb") as f:
            metadata = _MetadataUnpickler(f).load()
        return metadata

    def configure_reader(self, metadata: Metadata, is_coordinator: bool, **kwargs) -> None:
        """Configure storage reader."""

        self.storage_data = getattr(metadata, "storage_data", None)
        self.is_coordinator = is_coordinator
        # Do not evaluate get_rank() when an offline caller supplies rank. The
        # default process group is intentionally not initialized by converters.
        self.rank = kwargs["rank"] if "rank" in kwargs else platform.get_rank()
        self.broadcast_from_minimum_rank = kwargs.get("broadcast_from_minimum_rank", self.broadcast_from_minimum_rank)

    def optimize_local_plan(self, plan: LoadPlan) -> LoadPlan:
        """
        Optimize local plan.

            else:
                # 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

            not self._enable_plan_caching
            or self._cached_plans_key not in StandardSavePlanner.cached_save_result
        ):
            return None
        return StandardSavePlanner.cached_save_result[self._cached_plans_key]

    def cache_result(self, final_plan: SavePlan, metadata: Metadata) -> None:
        """Store finalized plan and metadata in the class-level planner cache."""
        if not self._enable_plan_caching:

    def cache_result(self, final_plan: SavePlan, metadata: Metadata) -> None:
        """Store finalized plan and metadata in the class-level planner cache."""
        if not self._enable_plan_caching:
            return
        StandardSavePlanner.cached_save_result[self._cached_plans_key] = CachedSaveResult(
            final_plan=final_plan,
            metadata=metadata,
        )

        self.metadata: Optional[Metadata] = None
        self.is_coordinator: bool = False
        self.rank: int = 0
        self.allow_partial_load = allow_partial_load
        self.broadcast_from_minimum_rank: bool = broadcast_from_minimum_rank
        self.flatten_state_dict: bool = True

    def configure_planner(self, state_dict: dict[str, Any], metadata: Metadata, **kwargs) -> None:
        """

        self.state_dict = state_dict
        self.metadata = metadata
        self.is_coordinator = kwargs.get("is_coordinator", False)
        self.rank = kwargs.get("rank", 0)
        self.broadcast_from_minimum_rank = kwargs.get("broadcast_from_minimum_rank", self.broadcast_from_minimum_rank)
        self.flatten_state_dict = kwargs.get("flatten_state_dict", True)
        self.original_state_dict = state_dict
        if self.flatten_state_dict:
            state_dict, self.name_mapping = flatten_state_dict(state_dict)

        Raises:
            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 len(group_ranks) > 1:
            setattr(tensor, BROADCAST_INFO, BroadcastInfo(group_ranks, load_rank))
        return self.rank == load_rank

    def _rank_owns_dtensor_shard(self, obj: Any) -> bool:
        """
        Check whether the current rank appears in the rank list of a DTensor layout.

        Returns:
            bool: False only when the layout has a rank list the current rank is absent from.
        """
        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:
        """
        Build local load plan.

                if obj_size is None or md.size != tuple(obj_size):
                    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)
                requests += create_read_items_for_chunk_list(fqn, md, local_chunks)
            else:

    with ``HP_LOG_CONFIG=DCP:INFO``.
    """
    @wraps(func)
    def wrapper(*args, **kwargs):
        try:
            rank_id = platform.get_rank()
        except ValueError:
            # No process group yet (offline converters, single-process tools).
            rank_id = 0
        logger.info("[rank=%d] >>> func %s start exec", rank_id, func.__name__)
        start_time = time.time()
        result = func(*args, **kwargs)
        end_time = time.time()
        execution_time = end_time - start_time
        logger.info("[rank=%d] >>> func %s cost %.4f seconds", rank_id, func.__name__, execution_time)
        return result

    return wrapper

def check_path(path: Union[Path, str]) -> None:

    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(

    Returns:
        list[Any]: List of all objects from all ranks.
    """
    if use_collectives and world_size > 1:
        all_objects = [None] * world_size
        platform.all_gather_object(all_objects, local_object)
        return all_objects
    return [local_object]


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

    Raises:
        ValueError: If the broadcast info attached to an entry has an unexpected type.
    """
    missing_groups_ranks = {}
    for obj in state_dict.values():
        broadcast_info = getattr(obj, BROADCAST_INFO, None)
        if broadcast_info is None:
            continue
        if not isinstance(broadcast_info, BroadcastInfo):
            raise ValueError(f"The broadcast info attached to tensor must be of type {BroadcastInfo}.")
        group_ranks, src_rank = tuple(broadcast_info.group_ranks), broadcast_info.src_rank
        if group_ranks in groups:
            platform.broadcast(
                obj.to_local().detach() if isinstance(obj, DTensor) else obj.detach(),
                src_rank,
                groups[group_ranks]
            )
            delattr(obj, BROADCAST_INFO)
        else:
            missing_groups_ranks.setdefault(group_ranks, []).append(obj)
    return missing_groups_ranks


def _destroy_groups(groups: dict[tuple, Any]) -> None:
    """

    Args:
        groups (dict[tuple, Any]): Groups to release, keyed by their rank tuple.
    """
    current_rank = platform.get_rank()
    # Destroying a group is collective over its members, so keep the same deterministic order
    # the groups were created in.
    for group_ranks in sorted(groups):
        if current_rank not in group_ranks:
            continue
        try:
            platform.destroy_process_group(groups[group_ranks])
        except Exception as e:  # pylint: disable=broad-except
            logger.warning("Failed to destroy the broadcast group %s: %s", group_ranks, e)


def _create_groups_and_broadcast(missing_groups_ranks: dict[tuple, list[Any]]) -> None:
    """

    Args:
        missing_groups_ranks (dict[tuple, list[Any]]): The entries waiting for a group, keyed
            by the rank tuple of the group they need.
    """
    logger.warning("There are missing groups %s. Then all gather the missing groups on each rank and "
                   "create them one by one, which will increase some time consumption.",
                   missing_groups_ranks)
    # all gather all missing groups, create them and broadcast the tensors
    all_missing_groups_ranks = all_gather_object(tuple(missing_groups_ranks.keys()),
                                                 platform.get_world_size(),
                                                 use_collectives=True)
    final_missing_groups_ranks = set(g for sub in all_missing_groups_ranks for g in sub)

    # The groups are only needed for the broadcasts right below, so create them without
    # registering them in the global group cache. Iterate in a deterministic order:
    # creating a group is a collective call and every rank must issue them in the same order.
    new_groups = {}
    for group_ranks in sorted(final_missing_groups_ranks):
        new_groups[group_ranks] = platform.new_group(group_ranks)
    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]
            )
            delattr(tensor, BROADCAST_INFO)

    # Only reached once every broadcast went through: releasing a group whose collectives just
    # failed can block instead of raising, which would turn a clean error into a hang.
    _destroy_groups(new_groups)


@dcp_timer_decorator
def broadcast_loaded_tensors(

        groups (dict): The Communication groups for broadcast.
    """
    # A caller that enabled broadcasting without pre-building groups lands in the
    # missing-groups path below, which creates them on demand.
    missing_groups_ranks = _broadcast_within_existing_groups(state_dict, groups or {})

    # if no group missing, return
    if not missing_groups_ranks:
        return

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