mirror of
https://github.com/hpcaitech/ColossalAI.git
synced 2025-04-28 11:45:23 +00:00
079bf3cb26
* [misc] update pre-commit * [misc] run pre-commit * [misc] remove useless configuration files * [misc] ignore cuda for clang-format
22 lines
607 B
Python
22 lines
607 B
Python
import tempfile
|
|
from contextlib import contextmanager, nullcontext
|
|
from typing import Iterator
|
|
|
|
import torch.distributed as dist
|
|
|
|
|
|
@contextmanager
|
|
def shared_tempdir() -> Iterator[str]:
|
|
"""
|
|
A temporary directory that is shared across all processes.
|
|
"""
|
|
ctx_fn = tempfile.TemporaryDirectory if dist.get_rank() == 0 else nullcontext
|
|
with ctx_fn() as tempdir:
|
|
try:
|
|
obj = [tempdir]
|
|
dist.broadcast_object_list(obj, src=0)
|
|
tempdir = obj[0] # use the same directory on all ranks
|
|
yield tempdir
|
|
finally:
|
|
dist.barrier()
|