/
githubmirror
/
ColossalAI
Обзор
Документация
Войти
/
githubmirror
/
ColossalAI
Код
Запросы
0
Пакеты
0
Релизы
0
Аналитика
Безопасность
main
tests/test_checkpoint_io/utils.py
21 строка
607 B
Hongxin Liu
[misc] update pre-commit and run all files (#4752)
19 сен 2023, 09:20
Не верифицирован
19 сен 2023, 09:20
079bf3c
Код
Авторство
О чём код?
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()