Function save

python/oneflow/framework/check_point_v2.py:684–741

r"""Save an object to a directory. Args: obj: The object to be saved path_or_buffer: a file-like object (has to implement write and flush) or a string or os.PathLike object containing a file name global_dst_rank (int, optional): The destination rank for

(
    obj: Any,
    path_or_buffer: FILE_LIKE,
    global_dst_rank: Optional[int] = None,
    save_as_external_data: bool = False,
)

def save(
    obj: Any,
    path_or_buffer: FILE_LIKE,
    global_dst_rank: Optional[int] = None,
    save_as_external_data: bool = False,
) -> None:
    r"""Save an object to a directory.

    Args:
        obj: The object to be saved
        path_or_buffer: a file-like object (has to implement write and flush) or a string or
           os.PathLike object containing a file name
        global_dst_rank (int, optional): The destination rank for
            saving global tensors. When specified, whole tensors
            will be saved by the process whose rank ==
            global_src_rank, while other processes will not do any
            disk I/O.
        save_as_external_data (bool): useful only if path_or_buffer is a string or
           os.PathLike object containing a file name
    """
    if isinstance(path_or_buffer, str):
        path_or_buffer = Path(path_or_buffer)

    if isinstance(obj, graph_util.Graph):
        if not _is_path(path_or_buffer):
            raise ValueError(
                "path_or_buffer must be the type of {`str`, `pathlib.Path`} while obj is Graph"
            )
        _save_graph(obj, path_or_buffer)
        return

    # this `path` is only used for `ContextData` and is set to empty when `path_or_buffer` is IO[bytes] or BinaryIO
    path: Path = Path(path_or_buffer if _is_path(path_or_buffer) else "")
    obj = {"protocol_version": PROTOCOL_VERSION, ONEFLOW_MAGIC_KEY: None, "data": obj}

    with tensor_pickling_context(path, global_dst_rank, None, save_as_external_data):
        pickled_bytes = pickle.dumps(obj)

    if _is_path(path_or_buffer) and save_as_external_data:
        path_or_buffer.mkdir(exist_ok=True)
        path_or_buffer = path_or_buffer / PICKLE_FILENAME

    def write_file():
        with _open_file_like(path_or_buffer, "wb") as f:
            f.write(pickled_bytes)

    if global_dst_rank is not None:
        assert isinstance(
            global_dst_rank, int
        ), f"global_dst_rank expected type int, but got {type(global_dst_rank)}."
        assert (
            global_dst_rank >= 0 and global_dst_rank < flow.env.get_world_size()
        ), f"out of range (expected to be in range of [0, {flow.env.get_world_size()}), but got {global_dst_rank})."
        if flow.env.get_rank() == global_dst_rank:
            write_file()
    else:
        # global_dst_rank is None
        write_file()