Coverage for / home / jenkins / .local / lib / python3.10 / site-packages / hyper_parallel / core / distributed_checkpoint / async_persist.py: 74%
271 statements
« prev ^ index » next coverage.py v7.13.1, created at 2026-10-25 05:45 +0800
1# Copyright 2026 Huawei Technologies Co., Ltd. All rights reserved.
2#
3# Licensed under the Apache License, Version 2.0 (the "License");
4# you may not use this file except in compliance with the License.
5# You may obtain a copy of the License at
6#
7# http://www.apache.org/licenses/LICENSE-2.0
8#
9# Unless required by applicable law or agreed to in writing, software
10# distributed under the License is distributed on an "AS IS" BASIS,
11# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
12# See the License for the specific language governing permissions and
13# limitations under the License.
14# ============================================================================
15"""Tensor staging and persist for asynchronous distributed checkpoint save.
17Staging materializes a host-side snapshot of ``state_dict`` so ``async_save`` can
18run persistence (plan + I/O) on a copy while training continues on the original tensors.
19"""
21import os
22import copy
23import pickle
24import dataclasses
25import queue
26import time
27import traceback
28import multiprocessing as mp
29from pathlib import Path
30from enum import Enum, auto
31from argparse import Namespace
32from collections.abc import Callable
33from typing import Any, Optional, Union
34from concurrent.futures import Future, ThreadPoolExecutor
36import numpy as np
37import torch
38import torch.distributed as dist
39import torch.distributed.distributed_c10d as c10d
41from hyper_parallel.core.dtensor.dtensor import DTensor
42from hyper_parallel.core.distributed_checkpoint.metadata import (
43 Metadata,
44 CHUNK_INFO,
45)
46from hyper_parallel.core.distributed_checkpoint.planner import (
47 SavePlan,
48 SavePlanner,
49)
50from hyper_parallel.core.distributed_checkpoint.storage import (
51 StorageWriter,
52 WriteResult,
53)
54from hyper_parallel.core.distributed_checkpoint.standard_planner import (
55 StandardSavePlanner,
56)
57from hyper_parallel.core.distributed_checkpoint.utils import (
58 dcp_timer_decorator,
59 logger,
60)
62_COMPLETE_FLAG = b"\x00__PICKLE_WRITE_COMPLETE__\x00"
64# The join wait sits in the *training* process for as long as the child persists,
65# so it blocks on the queue with a timeout rather than spinning.
66_PERSIST_RESULT_POLL_SECONDS = 0.1
69class DataCopier:
70 """Type dispatched copy of ``state_dict`` values into host memory.
72 Copy handlers are registered per type through :meth:`register`; :meth:`copy` resolves the
73 handler of an object by exact type first and by ``isinstance`` afterwards, falling back to
74 a generic recursive deep copy for objects no handler was registered for.
75 """
77 _registry = {}
79 @classmethod
80 def register(cls, *types: type) -> Callable:
81 """Register the decorated function as the copy handler of every type given.
83 Args:
84 *types (type): Types the handler is responsible for.
86 Returns:
87 Callable: Decorator recording the handler and returning it unchanged.
88 """
89 def decorator(func: Callable) -> Callable:
90 """Record ``func`` against every registered type and hand it back."""
91 for t in types:
92 cls._registry[t] = func
93 return func
95 return decorator
97 @classmethod
98 def dispatch(cls, obj: Any) -> Optional[Callable]:
99 """
100 Resolve the copy handler of ``obj``: its exact type first, its most derived
101 registered base class afterwards.
103 Registration order must not decide the handler. ``DTensor`` is registered after
104 ``torch.Tensor`` and is a subclass of it, so scanning in registration order would
105 hand a ``DTensor`` subclass to the plain tensor handler, staging it as a bare local
106 tensor with its mesh and placements dropped.
108 Args:
109 obj (Any): Object to find a copy handler for.
111 Returns:
112 Optional[Callable]: The handler, or None when no registered type matches.
113 """
114 if type(obj) in cls._registry:
115 return cls._registry[type(obj)]
117 best_type = None
118 for r_type in cls._registry:
119 if not isinstance(obj, r_type):
120 continue
121 # Unrelated matches (neither one a subclass of the other) keep the first found.
122 if best_type is None or issubclass(r_type, best_type):
123 best_type = r_type
124 return cls._registry[best_type] if best_type is not None else None
126 @classmethod
127 def copy(cls, obj: Any) -> Any:
128 """
129 Copy ``obj`` with its registered handler, or with a generic deep copy as fallback.
131 Args:
132 obj (Any): Object to copy, of any type held by a state dict.
134 Returns:
135 Any: Copy of ``obj``, with tensor leaves residing in host memory.
136 """
137 copy_func = cls.dispatch(obj)
138 if copy_func:
139 return copy_func(obj)
140 logger.warning(
141 "The user define object %s not implement copy method and register "
142 "to DataCopier. So, the general copy method will be used, which may cause "
143 "unexpected copy behavior and incorrect data storage. Please verify this on your own.",
144 type(obj))
145 return cls._copy_others(obj)
147 @staticmethod
148 def _copy_others(obj):
149 obj_c = copy.deepcopy(obj)
150 for name, value in obj_c.__dict__.items():
151 if callable(value):
152 continue
153 setattr(obj_c, name, DataCopier.copy(value))
154 return obj_c
157def _copy_tensor_to_cpu(tensor: torch.Tensor) -> torch.Tensor:
158 """Return a host-memory copy of a tensor, detached from autograd."""
159 t = tensor.detach().to("cpu", copy=True)
160 if hasattr(tensor, CHUNK_INFO):
161 setattr(t, CHUNK_INFO, getattr(tensor, CHUNK_INFO))
162 return t
165@DataCopier.register(int, float, complex, bool, str, type(None), Enum, torch.dtype)
166def _copy_const(obj):
167 return obj
170@DataCopier.register(bytes, bytearray)
171def _copy_bytes(obj):
172 return bytes(obj)
175@DataCopier.register(np.ndarray)
176def _copy_ndarray(obj):
177 return obj.copy()
180@DataCopier.register(torch.Tensor)
181def _copy_tensor(obj):
182 result = _copy_tensor_to_cpu(obj)
183 return result
186@DataCopier.register(DTensor)
187def _copy_dtensor(obj):
188 """Stage a DTensor by copying its local shard to host memory, keeping mesh and placements."""
189 staged_local = _copy_tensor_to_cpu(obj.to_local())
190 # ``shape`` is required for RaggedShard: an unevenly-sharded DTensor cannot
191 # recover its logical global shape from the local shard alone.
192 result = DTensor.from_local(
193 staged_local,
194 obj.device_mesh,
195 obj.placements,
196 shape=tuple(obj.shape),
197 )
198 return result
201@DataCopier.register(Namespace)
202def _copy_namespace(obj):
203 c = DataCopier.copy(vars(obj))
204 return Namespace(**c)
207@DataCopier.register(list)
208def _copy_list(obj):
209 result = [DataCopier.copy(item) for item in obj]
210 return result
213@DataCopier.register(tuple)
214def _copy_tuple(obj):
215 result = tuple(DataCopier.copy(item) for item in obj)
216 return result
219@DataCopier.register(dict)
220def _copy_dict(obj):
221 result = {}
222 for k, v in obj.items():
223 result[DataCopier.copy(k)] = DataCopier.copy(v)
224 return result
227class AsyncPersistStatus(Enum):
228 """Queue payload status from :func:`_async_persist_worker` to the parent join thread."""
230 SUCCESS = auto()
231 FAILURE = auto()
234class FileType(Enum):
236 LOCAL_PLAN = auto()
237 STORAGE_DATA = auto()
240@dataclasses.dataclass
241class AsyncSaveResponse:
242 """Result of :func:`async_save`.
244 Host staging runs synchronously before :func:`async_save` returns; only checkpoint
245 **persistence** is asynchronous. ``persist_completion`` completes when the child
246 process finishes :func:`_save_impl` (plan, collectives, disk I/O) and supplies
247 :class:`Metadata`.
248 """
250 persist_completion: Future[Metadata]
252 def get_result(self, timeout: int = None) -> Optional[Metadata]:
253 """Wait for the save to be persisted and return its metadata.
255 Args:
256 timeout (int): Seconds to wait. Default ``None``, which waits indefinitely.
258 Returns:
259 Optional[Metadata]: Metadata of the persisted checkpoint.
260 """
261 return self.persist_completion.result(timeout=timeout)
264@dcp_timer_decorator
265def build_staged_state_dict(state_dict: dict[str, Any]) -> dict[str, Any]:
266 """
267 Build a deep structural copy of ``state_dict`` with tensor / DTensor data
268 copied to host memory so the original can be mutated while checkpoint I/O runs.
270 Uses the same flattening rules as :class:`StandardSavePlanner`.
272 Args:
273 state_dict (dict[str, Any]): Nested or flat training state dict.
275 Returns:
276 dict[str, Any]: New dict with identical nesting and keys; tensor leaves are staging copies.
277 """
278 return DataCopier.copy(state_dict)
281def cleanup_and_reinit_process_group(
282 master_addr: str,
283 master_port: int,
284 rank: int,
285 world_size: int,
286) -> None:
287 """
288 Cleanup hccl process group and reinitialize gloo process group.
290 Args:
291 master_addr (str): Address of the rendezvous the new group meets at.
292 master_port (int): Port of that rendezvous.
293 rank (int): Rank of this process in the new group.
294 world_size (int): Number of processes in the new group.
295 """
296 # Resetting the private c10d state is the only way to drop the process group inherited
297 # from the parent process, so the accesses below are deliberate.
298 # pylint: disable=protected-access
300 # Step 1: clear C++ ProcessGroupRegistry (avoids "already registered" error).
301 if hasattr(c10d, "_unregister_all_process_groups"):
302 try:
303 c10d._unregister_all_process_groups()
304 except Exception: # pylint: disable=broad-except
305 pass
307 # Step 2: reset Python-layer _world to an uninitialised state.
308 if hasattr(c10d, "_World"):
309 try:
310 c10d._world = c10d._World()
311 except Exception: # pylint: disable=broad-except
312 pass
313 # Also reset the module-level init-method string (present in all versions).
314 if hasattr(c10d, "_default_pg_init_method"):
315 c10d._default_pg_init_method = None
316 for attr in ("_default_pg", "_pg_map", "_pg_names", "_pg_group_ranks",
317 "_pg_backend_config", "_group_count", "_tags_to_pg",
318 "_pg_to_tag"):
319 if not hasattr(c10d, attr):
320 continue
321 old_val = getattr(c10d, attr)
322 if isinstance(old_val, dict):
323 setattr(c10d, attr, {})
324 elif isinstance(old_val, int):
325 setattr(c10d, attr, 0)
326 else:
327 setattr(c10d, attr, None)
329 # Step 3: build a fresh TCPStore on the dedicated gloo port and init.
330 # rank-0 child is the store master; all others connect as clients.
331 store = dist.PrefixStore(
332 prefix=mp.current_process().name,
333 store=dist.TCPStore(
334 host_name=master_addr,
335 port=master_port
336 ),
337 )
339 # Step 4: init new process group with backend "gloo".
340 dist.init_process_group(
341 backend="gloo",
342 store=store,
343 rank=rank,
344 world_size=world_size,
345 )
348@dcp_timer_decorator
349def execute_async_persist_with_gloo(
350 result_queue: mp.Queue,
351 staged: dict[str, Any],
352 checkpoint_id: Optional[Union[Path, str]],
353 storage_writer: Optional[StorageWriter],
354 planner: Optional[SavePlanner],
355 rank: int,
356 world_size: int,
357 master_addr: str,
358 master_port: int,
359) -> None:
360 """Child-process entry for gloo-based async save.
361 Reinitialises a **CPU-only gloo process group** on a dedicated port so that
362 all collective communication (all_gather_object for plan exchange,
363 all_gather_object for write_results) can run inside the child process,
364 exactly mirroring what :func:`_save_impl` does synchronously.
365 The parent process only stages tensors to CPU; all planning, collectives,
366 disk I/O and metadata finalisation happen here.
367 Design rationale
368 ----------------
369 * The training group uses NCCL (or hccl on Ascend) and lives in C++ global
370 state that is invalidated across ``fork``. We never touch it.
371 * A fresh ``gloo`` group is initialised via a ``TCPStore`` client on the training
372 rendezvous itself (``master_addr:master_port``), namespaced by a ``PrefixStore``
373 keyed on this process name. The name carries a per-``async_save`` counter, so each
374 save gets its own key space on the same store; every rank must therefore reach the
375 same number of gloo saves, or the prefixes diverge and the rendezvous hangs.
376 * After :func:`_save_impl` completes the gloo group is destroyed so resources
377 are released cleanly.
379 Args:
380 result_queue: IPC queue for returning status / metadata to parent join thread.
381 staged: CPU-resident staged state dict (output of build_staged_state_dict).
382 checkpoint_id: Checkpoint directory path.
383 storage_writer: Pre-constructed StorageWriter (or None → FileSystemWriter).
384 planner: Pre-constructed SavePlanner (or None → StandardSavePlanner).
385 rank: Global rank of this process (inherited from parent).
386 world_size: Total world size (inherited from parent).
387 master_addr: TCP rendezvous address for the gloo store.
388 master_port: Training master port; the child connects to that store as a client.
389 """
390 # ``api`` imports this module at load time, so this import stays function local: that is
391 # what keeps the two from importing each other at load time.
392 from hyper_parallel.core.distributed_checkpoint.api import _save_impl # pylint: disable=import-outside-toplevel
393 try:
394 # Clear the original communication group information and establish the gloo communication.
395 cleanup_and_reinit_process_group(
396 master_addr,
397 master_port,
398 rank,
399 world_size
400 )
401 meta = _save_impl(
402 staged,
403 checkpoint_id=checkpoint_id,
404 storage_writer=storage_writer,
405 planner=planner,
406 )
407 # The plan cache is also returned and can be reused for the next dcp async save.
408 result_queue.put((AsyncPersistStatus.SUCCESS, (meta, StandardSavePlanner.cached_save_result)))
409 except Exception: # pylint: disable=broad-except
410 result_queue.put((AsyncPersistStatus.FAILURE, traceback.format_exc()))
411 finally:
412 # Clean up the gloo group we created; ignore errors (process is exiting).
413 if dist.is_initialized():
414 dist.destroy_process_group()
417@dcp_timer_decorator
418def execute_async_persist(
419 result_queue: mp.Queue,
420 staged: dict[str, Any],
421 checkpoint_id: Optional[Union[Path, str]],
422 storage_writer: Optional[StorageWriter],
423 planner: Optional[SavePlanner],
424 no_dist: bool,
425 use_collectives: bool,
426 use_storage_comm: bool,
427) -> None:
428 """
429 Perform asynchronous flushing to disk,
430 which is not suitable for collective communication,
431 or use the file system for collective communication.
433 Args:
434 result_queue: IPC queue for returning status / metadata to parent join thread.
435 staged: CPU-resident staged state dict (output of build_staged_state_dict).
436 checkpoint_id: Checkpoint directory path.
437 storage_writer: Pre-constructed StorageWriter (or None → FileSystemWriter).
438 planner: Pre-constructed SavePlanner (or None → StandardSavePlanner).
439 no_dist: Whether it is a distributed scenario or a standalone scenario.
440 use_collectives: Total world size (inherited from parent).
441 use_storage_comm: TCP rendezvous address for the gloo store.
442 """
443 # ``api`` imports this module at load time, so this import has to stay function local.
444 from hyper_parallel.core.distributed_checkpoint.api import _save_impl # pylint: disable=import-outside-toplevel
445 try:
446 meta = _save_impl(
447 staged,
448 checkpoint_id=checkpoint_id,
449 storage_writer=storage_writer,
450 planner=planner,
451 no_dist=no_dist,
452 use_collectives=use_collectives,
453 use_storage_comm=use_storage_comm
454 )
455 # The plan cache is also returned and can be reused for the next dcp async save.
456 result_queue.put((AsyncPersistStatus.SUCCESS, (meta, StandardSavePlanner.cached_save_result)))
457 except Exception: # pylint: disable=broad-except
458 result_queue.put((AsyncPersistStatus.FAILURE, traceback.format_exc()))
461def resolve_async_persist_result(
462 proc: mp.Process,
463 result_queue: mp.Queue,
464 persist_future: Future[None],
465 async_callback: Callable[[], None],
466) -> None:
467 """Join persist ``proc`` and complete ``persist_future`` (runs on a background thread)."""
468 result = None
469 while result is None:
470 try:
471 result = result_queue.get(timeout=_PERSIST_RESULT_POLL_SECONDS)
472 except queue.Empty:
473 if proc.is_alive():
474 continue
475 # The child may have put its result and exited while the wait above was
476 # timing out, so look once more before giving up on it.
477 try:
478 result = result_queue.get_nowait()
479 except queue.Empty:
480 pass
481 break
483 proc.join()
484 if persist_future.done():
485 return
487 if result is None:
488 persist_future.set_exception(
489 RuntimeError(
490 f"async_persist process exited with code {proc.exitcode} and no result on queue"
491 )
492 )
493 return
495 status, payload = result
496 if status == AsyncPersistStatus.SUCCESS:
497 try:
498 async_callback()
499 StandardSavePlanner.cached_save_result.update(payload[1])
500 except Exception as exc: # pylint: disable=broad-except
501 persist_future.set_exception(exc)
502 else:
503 persist_future.set_result(payload[0])
504 elif status == AsyncPersistStatus.FAILURE:
505 # The child's exception object cannot cross the process boundary, so what
506 # arrives is its formatted traceback and a RuntimeError carries the text.
507 persist_future.set_exception(RuntimeError(payload))
508 else:
509 persist_future.set_exception(
510 RuntimeError(f"async_persist queue returned unexpected status: {status!r}")
511 )
514def batch_read_worker(file_batch: tuple[int], checkpoint_id: str, file_type: FileType) -> list:
515 """
516 Read one batch of files, waiting for the files that the other ranks have not written yet.
518 Args:
519 file_batch (tuple[int]): Contiguous file ids this worker is responsible for.
520 checkpoint_id (str): Checkpoint directory holding the files.
521 file_type (FileType): Which kind of file to read, local plan or storage data.
523 Returns:
524 list: Content of every file of ``file_batch``, in the order of the batch.
526 Raises:
527 TimeoutError: If some files are still unreadable after 1800 seconds.
528 """
529 results = [None] * len(file_batch)
530 file_to_read = file_batch
531 first_file_id = file_batch[0]
533 # Read files in a round-robin manner,
534 # If a file has not been generated yet, skip it and proceed to the next file.
535 # Until all files have been read.
536 timeout = 1800
537 deadline = time.monotonic() + timeout
538 while True:
539 if not file_to_read:
540 break
541 if time.monotonic() > deadline:
542 raise TimeoutError(
543 f"After waiting for 1800 seconds, "
544 f"the {file_type.name} files: {file_to_read} still have not been read completely."
545 )
546 file_not_read = []
547 for file_id in file_to_read:
548 ret = load_file(checkpoint_id, file_type, file_id)
549 if ret is None:
550 file_not_read.append(file_id)
551 else:
552 results[file_id - first_file_id] = ret
553 file_to_read = tuple(file_not_read)
555 return results
558@dcp_timer_decorator
559def parallel_read_files(
560 file_list: tuple[int],
561 checkpoint_id: str,
562 file_type: FileType,
563 max_workers: int = 32
564) -> list:
565 """
566 Split the total file list into multiple batches
567 based on the number of parallel threads,
568 then perform multi-threaded parallel reading to improve read efficiency.
569 """
570 if not file_list:
571 return []
573 batches, max_workers = construct_file_batches(file_list, max_workers)
574 logger.info(
575 "parallel read files, file type: %s, file count: %d, max read workers: %d",
576 file_type.name, len(file_list), max_workers
577 )
579 # Use the thread pool to read data concurrently, improving the read efficiency.
580 with ThreadPoolExecutor(max_workers=max_workers) as executor:
581 all_results = tuple(
582 executor.map(lambda b: batch_read_worker(b, checkpoint_id, file_type), batches)
583 )
584 return [result for batch_results in all_results for result in batch_results]
587def construct_file_batches(file_list: tuple[int], max_workers: int) -> (list[tuple[int]], int):
588 """
589 Split the total file list into multiple batches
590 based on the number of parallel threads
591 """
592 n = len(file_list)
593 max_workers = min(max_workers, n)
594 if max_workers == 0:
595 raise ValueError("Max workers number for parallel read batch file is 0!")
596 batch_size = (n + max_workers - 1) // max_workers
597 # Split the file list into max_workers sublists.
598 batches = [
599 file_list[i * batch_size: (i + 1) * batch_size]
600 for i in range(max_workers)
601 ]
602 batches = [b for b in batches if b]
603 return batches, max_workers
606@dcp_timer_decorator
607def gather_all_results_from_storage(
608 checkpoint_id: str,
609 is_coordinator: bool,
610 world_size: int,
611 results: Union[SavePlan, WriteResult],
612 file_type: FileType,
613 file_id: int
614) -> list:
615 """
616 All node data is collected from the storage,
617 which is classified into local plan and storage data.
618 In the first step, each rank writes its own data.
619 In the second step, for local plan, each rank reads the file written by all ranks.
620 For storage data, only rank 0 reads the file written by all ranks.
621 """
622 logger.info("gather all results from storage, file type: %s", file_type.name)
623 write_file(checkpoint_id, file_type, file_id, results)
624 if not is_coordinator and file_type != FileType.LOCAL_PLAN:
625 return []
627 # For LOCAL_PLAN, each node has 16 rank processes reading simultaneously.
628 if file_type == FileType.LOCAL_PLAN:
629 max_workers = min(16, max(os.cpu_count() // 32, 1))
630 else:
631 max_workers = min(64, max(os.cpu_count() // 2, 1))
632 file_list = tuple(range(0, world_size))
633 gathered = parallel_read_files(file_list,
634 checkpoint_id,
635 file_type,
636 max_workers=max_workers)
637 return gathered
640def assemble_file_path(checkpoint_id: str, file_type: FileType, file_id: int) -> str:
641 """
642 Assemble file path to exchanges its results through.
644 Args:
645 checkpoint_id (str): Checkpoint directory holding the exchange files.
646 file_type (FileType): Which kind of file, local plan or storage data.
647 file_id (int): Rank that owns the file.
649 Returns:
650 str: Path of that rank's file.
651 """
652 return f"{checkpoint_id}/{file_type.name}_{file_id}.pkl"
655def _safe_remove_file(file_path: str) -> None:
656 """
657 Delete a file that has been read, without letting the delete fail the load.
659 What it held has already been read, so a file that will not go is a leftover worth a
660 warning rather than a reason to fail a save that has its data. One attempt is enough:
661 what refuses here - a read-only mount, a directory this rank may not write in - does
662 not clear by asking again.
664 Args:
665 file_path (str): File to delete.
666 """
667 try:
668 os.remove(file_path)
669 except Exception as e: # pylint: disable=broad-except
670 logger.warning(
671 "Could not remove the STORAGE_DATA file %s, because the error: %s", file_path, e
672 )
675def load_file(checkpoint_id: str, file_type: FileType, file_id: int) -> Any:
676 """
677 Reads files from the storage.
678 Currently, this function is mainly used to
679 asynchronously save local plan and storage data files.
681 Args:
682 checkpoint_id (str): Checkpoint directory holding the file.
683 file_type (FileType): Which kind of file to read, local plan or storage data.
684 file_id (int): Id of the file within that kind.
686 Returns:
687 Any: The unpickled content of the file.
688 """
689 file_path = assemble_file_path(checkpoint_id, file_type, file_id)
691 try:
692 with open(file_path, "rb") as f:
693 flag_len = len(_COMPLETE_FLAG)
694 f.seek(0, os.SEEK_END)
695 size = f.tell()
696 # determining whether the file length is greater than the flag length
697 if size < flag_len:
698 return None
699 f.seek(size - flag_len)
700 # determining whether the file has a write complete flag
701 if f.read(flag_len) != _COMPLETE_FLAG:
702 return None
703 f.seek(0)
704 result = pickle.load(f)
705 except FileNotFoundError:
706 # The rank that owns this file has not written it yet; the caller comes back round.
707 return None
709 # Delete only the STORAGE_DATA file. Do not delete the LOCAL_PALN file,
710 # as it is generated only during the initial save.
711 if file_type == FileType.STORAGE_DATA:
712 _safe_remove_file(file_path)
713 return result
716@dcp_timer_decorator
717def write_file(
718 checkpoint_id: str,
719 file_type: FileType,
720 file_id: int,
721 content: Union[SavePlan, WriteResult],
722) -> None:
723 """
724 Write files to the storage.
725 Currently, this function is mainly used to
726 asynchronously save local plan and storage data files.
728 Args:
729 checkpoint_id (str): Checkpoint directory to write into.
730 file_type (FileType): Which kind of file to write, local plan or storage data.
731 file_id (int): Id of the file within that kind.
732 content (Union[SavePlan, WriteResult]): What to pickle into it.
733 """
734 final_path = assemble_file_path(checkpoint_id, file_type, file_id)
736 # To avoid read-write conflicts, the completion flag is appended last: a reader that
737 # opens the file mid-write finds no flag at the end and retries on its next round.
738 with open(final_path, "wb") as f:
739 pickle.dump(content, f, protocol=pickle.HIGHEST_PROTOCOL)
740 f.write(_COMPLETE_FLAG)
742 logger.info("write file type: %s, file path: %s", file_type.name, final_path)