o
    "i                     @   s  d dl Z d dlZd dlmZ d dlZd dlZd dlZdefddZe r-ej	
 s-edej	jZej	jZej	jZej	jZe rd dlmZmZmZmZmZmZmZmZmZmZm Z!m"Z"m#Z#m$Z$m%Z%m&Z&m'Z'm(Z(m)Z)m*Z*m+Z+m,Z,m-Z-m.Z.m/Z/ G dd	 d	ej0Z1dd
e2fddZ3ej4dkrd dlm5Z5m6Z6 ddl7T ddl7m8Z8m9Z9m:Z:m;Z;m<Z<m=Z=m>Z>m?Z? ddl@m@Z@mAZAmBZB ddlCmDZD ddlEmFZFmGZG e,  dS G dd dZHeHejId _dS )    N)Enumreturnc                   C   s   t tjdS )a  
    Return ``True`` if the distributed package is available.

    Otherwise,
    ``torch.distributed`` does not expose any other APIs. Currently,
    ``torch.distributed`` is available on Linux, MacOS and Windows. Set
    ``USE_DISTRIBUTED=1`` to enable it when building PyTorch from source.
    Currently, the default value is ``USE_DISTRIBUTED=1`` for Linux and Windows,
    ``USE_DISTRIBUTED=0`` for MacOS.
    
_c10d_init)hasattrtorch_C r   r   X/home/ubuntu/SoloSpeech/.venv/lib/python3.10/site-packages/torch/distributed/__init__.pyis_available
   s   r
   z&Failed to initialize torch.distributed)Store	FileStoreTCPStoreProcessGroupBackendPrefixStoreReducerLoggerBuiltinCommHookType
GradBucketWork_DEFAULT_FIRST_BUCKET_BYTES_register_comm_hook_register_builtin_comm_hook_broadcast_coalesced"_compute_bucket_assignment_by_size_verify_params_across_processes_test_python_store
DebugLevelget_debug_levelset_debug_levelset_debug_level_from_env_make_nccl_premul_sum_ControlCollectives_StoreCollectivesc                   @   s   e Zd ZdZdd ZdS )_DistributedPdbz
        Supports using PDB from inside a multiprocessing child process.

        Usage:
        _DistributedPdb().set_trace()
        c                 O   s@   t j}ztdt _tjj| g|R i | W |t _d S |t _w )Nz
/dev/stdin)sysstdinopenpdbPdbinteraction)selfargskwargs_stdinr   r   r	   r*   E   s
   
z_DistributedPdb.interactionN)__name__
__module____qualname____doc__r*   r   r   r   r	   r$   >   s    r$   rankc              	   C   sv   t  | krt }|d|  d |  tj }tj }tjd zt	  W tj| ~dS tj| ~w )z
        Set a breakpoint, but only on a single rank.  All other ranks will wait for you to be
        done with the breakpoint before continuing.

        Args:
            rank (int): Which rank to break on.  Default: ``0``
        zS
!!! ATTENTION !!!

Type 'up' to get to the frame that called dist.breakpoint(rank=z)
FN)
get_rankr$   message	set_tracer   r   _meta_in_tls_dispatch_include_DisableTorchDispatch!_set_meta_in_tls_dispatch_includebarrier)r3   r(   meta_in_tlsguardr   r   r	   
breakpointM   s"   


r=   win32)	HashStore_round_robin_process_groups   )*)_all_gather_base_reduce_scatter_base_create_process_group_wrapper_rank_not_in_group_coalescing_manager_CoalescingManager_get_process_group_nameget_node_local_rank)
rendezvous_create_store_from_optionsregister_rendezvous_handler)_remote_device)init_device_mesh
DeviceMeshc                   @   s   e Zd ZdS )_ProcessGroupStubN)r/   r0   r1   r   r   r   r	   rQ      s    rQ   ztorch.distributed)r   )Josr%   enumr   r(   ior   boolr
   r   r   RuntimeError
_DistError	DistError_DistBackendErrorDistBackendError_DistNetworkErrorDistNetworkError_DistStoreErrorDistStoreErrortorch._C._distributed_c10dr   r   r   r   r   _Backendr   r   r   r   r   r   _Workr   r   r   r   r   r   r   r   r   r   r    r!   r"   r#   r)   r$   intr=   platformr?   r@   distributed_c10drC   rD   rE   rF   rG   rH   rI   rJ   rK   rL   rM   remote_devicerN   device_meshrO   rP   rQ   modulesr   r   r   r	   <module>   s6   l
(
	