1from typing import Optional 2 3import torch.distributed as dist 4 5from . import config 6 7 8_COMPILE_PG: Optional[dist.ProcessGroup] = None 9 10 11def get_compile_pg() -> Optional[dist.ProcessGroup]: 12 if ( 13 config.enable_compiler_collectives 14 and dist.is_available() 15 and dist.is_initialized() 16 ): 17 global _COMPILE_PG 18 if _COMPILE_PG is None: 19 # , timeout=datetime.timedelta(seconds=2) 20 _COMPILE_PG = dist.distributed_c10d._new_group_with_tag( 21 pg_tag="pt2_compile_pg" 22 ) 23 return _COMPILE_PG 24 25 return None 26