U Jc=%@sUddlmZddlZddlZddlZddlZddlmZmZddl m Z ddl Z ddl m ZeeZdaeZddZere jsederddlmZdd lmZmZmZmZm 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/m0Z0m1Z1m2Z2m3Z3m4Z4m5Z5m6Z6m7Z7m8Z8m9Z9m:Z:m;Z;mZ>m?Z?m@Z@d d l>TddlAZAddlBm mCZDd d l?mEZEd dlFmGZGd dlHmIZIeeeeJeJfddfeKd<dddZLddZMeEjNdddddfddZOe>jPddZQdS)) timedeltaN) GeneratorTuple)urlparsecCs ttjdS)N _rpc_init)hasattrtorch_Cr r B/tmp/pip-unpacked-wheel-gikjz4vx/torch/distributed/rpc/__init__.py is_availablesr z*Failed to initialize torch.distributed.rpc)Store)!_disable_jit_rref_pickle_enable_jit_rref_pickle'_disable_server_process_global_profiler&_enable_server_process_global_profiler_set_and_start_rpc_agent_reset_current_rpc_agent)_delete_all_user_and_unforked_owner_rrefs_destroy_rref_context_set_profiler_node_id_is_current_rpc_agent_set_rref_context_get_debug_info_cleanup_python_rpc_handler_invoke_rpc_builtin_invoke_rpc_python_udf_invoke_rpc_torchscript_invoke_remote_builtin_invoke_remote_python_udf_invoke_remote_torchscript_set_rpc_timeout_get_current_rpc_agentget_rpc_timeoutenable_gil_profilingRpcBackendOptions _TensorPipeRpcBackendOptionsBaseRpcAgentPyRRefTensorPipeAgentRemoteProfilerManager WorkerInfo_DEFAULT_INIT_METHOD_DEFAULT_NUM_WORKER_THREADS_UNSET_RPC_TIMEOUT_DEFAULT_RPC_TIMEOUT_SEC)apibackend_registry functions)*) BackendType)TensorPipeRpcBackendOptions)_server_process_global_profilerendezvous_iteratorc Csjtjd|dk r(t|tjs(td|dk rBt|tsBtd|dkr|dk rtD] }t|tt |rV|}qqVtd||tj krt d|dtj d|d |dkrtj }|dkrt |}|st ||}n t j|j||d att\}}}|t|jd t$t td t|}td 7aW5QRXt|t|t||||||dS)ak Initializes RPC primitives such as the local RPC agent and distributed autograd, which immediately makes the current process ready to send and receive RPCs. Args: name (str): a globally unique name of this node. (e.g., ``Trainer3``, ``ParameterServer2``, ``Master``, ``Worker1``) Name can only contain number, alphabet, underscore, colon, and/or dash, and must be shorter than 128 characters. backend (BackendType, optional): The type of RPC backend implementation. Supported values is ``BackendType.TENSORPIPE`` (the default). See :ref:`rpc-backends` for more information. rank (int): a globally unique id/rank of this node. world_size (int): The number of workers in the group. rpc_backend_options (RpcBackendOptions, optional): The options passed to the RpcAgent constructor. It must be an agent-specific subclass of :class:`~torch.distributed.rpc.RpcBackendOptions` and contains agent-specific initialization configurations. By default, for all agents, it sets the default timeout to 60 seconds and performs the rendezvous with an underlying process group initialized using ``init_method = "env://"``, meaning that environment variables ``MASTER_ADDR`` and ``MASTER_PORT`` need to be set properly. See :ref:`rpc-backends` for more information and find which options are available. ztorch.distributed.init_rpcNz0Argument backend must be a member of BackendTypezEArgument rpc_backend_options must be an instance of RpcBackendOptionsz$Could not infer backend for options zORPC was initialized with no explicit backend but with options corresponding to z9, hence that backend will be used instead of the default z(. To silence this warning pass `backend=z ` explicitly.)rank world_size)secondsz rpc_prefix_{}r/)rr Z_log_api_usage_once isinstancer1r4 TypeErrorr$typeZconstruct_rpc_backend_options TENSORPIPEloggerwarningdistZ_create_store_from_optionsZ rendezvousZ init_methodr7nextZ set_timeoutrZ rpc_timeout_init_counter_lockZ PrefixStorestrformat _init_counter dist_autograd_initr_init_rpc_backend)namebackendr9r:rpc_backend_optionsZcandidate_backendstore_r r r init_rpcPsj#     rPc Csb|tj|tj|t|tj|tjtdf|ti}| D](\}}t ||s4t d ||t|q4dS)Nz.Argument {} must be of type {} but got type {}) r1r4rBr rEnumbersIntegralr>r$itemsr< RuntimeErrorrF) rLrNrKr9r:rMZ type_mappingargZarg_typer r r _validate_rpc_argss,   rVcCsDt||||||tr tdtj||||||d}t|dS)NzRPC is already initialized)rNrKr9r:rM)rVrrTr1Z init_backendr0Z_init_rpc_states)rLrNrKr9r:rMZ rpc_agentr r r rJs  rJcCs*t}|t|t|S)N)rupdater0r!Zget_debug_inforH_get_debug_info)infor r r rXsrX)Nr8NN)Rdatetimerloggingos threadingwarningstypingrr urllib.parserrZtorch.distributedZ distributedrB getLogger__name__r@rGLockrDr r rrTZtorch._C._distributed_c10dr Ztorch._C._distributed_rpcrrrrrrrrrrrrrrrrrrr r!r"r#r$r%r&r'r(r)r*r+r,r-r.r0r1r2rQZtorch.distributed.autogradZautogradrHr4optionsr5Zserver_process_global_profilerr6int__annotations__rPrVr?rJZ_require_initializedrXr r r r sR     $    v