mirror of
https://github.com/pytorch/pytorch.git
synced 2025-10-20 21:14:14 +08:00
d55dc00f84
Pull Request resolved: https://github.com/pytorch/pytorch/pull/156321 Approved by: https://github.com/jingsh ghstack dependencies: #156313, #156314, #156315, #156316, #156317, #156319
43 lines
1.6 KiB
C++
43 lines
1.6 KiB
C++
#pragma once
|
|
|
|
#include <torch/csrc/distributed/c10d/PrefixStore.hpp>
|
|
#include <torch/csrc/distributed/rpc/utils.h>
|
|
|
|
namespace torch::distributed::rpc {
|
|
|
|
// All RPC peers should call into this function at the same time. Each peer
|
|
// provides its own id and name, and this function uses the given Store to
|
|
// gather global name-to-id mapping on all peers.
|
|
TORCH_API std::unordered_map<std::string, worker_id_t> collectNames(
|
|
::c10d::PrefixStore store,
|
|
const worker_id_t selfId,
|
|
const std::string& selfName,
|
|
const int worldSize);
|
|
|
|
// Ranks in dynamic RPC groups will initially call into this to establish the
|
|
// name-to-id mapping for the current peers in the group. The current rank will
|
|
// put its own worker info in the store and discover all the ranks that came
|
|
// before it. NOTE: This needs to be called with the Dynamic RPC group
|
|
// membership management token held.
|
|
TORCH_API std::unordered_map<std::string, worker_id_t> collectCurrentNames(
|
|
::c10d::PrefixStore store,
|
|
const worker_id_t selfId,
|
|
const std::string& selfName);
|
|
|
|
// Remove name from Store, used in dynamic RPC groups.
|
|
// NOTE: This needs to be called with the Dynamic RPC group
|
|
// membership management token held.
|
|
TORCH_API void removeCurrentName(
|
|
::c10d::PrefixStore store,
|
|
const worker_id_t selfId,
|
|
const std::string& selfName);
|
|
|
|
// This performs a synchronization of all call counts by using store.
|
|
// All RPC peers wait for others to join to exit at the same time.
|
|
TORCH_API int syncCallCount(
|
|
::c10d::PrefixStore store,
|
|
const int worldSize,
|
|
int activeCalls = 0);
|
|
|
|
} // namespace torch::distributed::rpc
|