@@ -5,8 +5,50 @@ namespace torch {
55namespace distributed {
66namespace rpc {
77
8+ namespace {
9+ // WorkerInfo needs to be registered exactly once. Since the op registration
10+ // happens in libtorch_python we wrap the class registration in a helper to make
11+ // sure that if there's multiple copies of Python such as used in torch::deploy
12+ // we only ever register it once.
13+ static std::once_flag workerInfoFlag;
14+ static c10::optional<torch::class_<WorkerInfo>> workerInfo;
15+ } // namespace
16+
17+ RegisterWorkerInfoOnce::RegisterWorkerInfoOnce () {
18+ std::call_once (workerInfoFlag, []() {
19+ workerInfo = torch::class_<WorkerInfo>(" dist_rpc" " WorkerInfo"
20+ .def (torch::init<std::string, int64_t >());
21+ });
22+ }
23+
824constexpr size_t WorkerInfo::MAX_NAME_LEN;
925
26+ WorkerInfo::WorkerInfo (std::string name, int64_t id)
27+ : WorkerInfo(std::move(name), (worker_id_t )id) {
28+ TORCH_CHECK (
29+ id <= std::numeric_limits<worker_id_t >::max (),
30+ " RPC worker id "
31+ id,
32+ " out of bound of int16_t."
33+ }
34+
35+ WorkerInfo::WorkerInfo (std::string name, worker_id_t id)
36+ : name_(std::move(name)), id_(id) {
37+ bool validSize = name_.length () < MAX_NAME_LEN && name_.length () > 0 ;
38+ bool validChar =
39+ std::find_if (name_.begin (), name_.end (), [](char c) {
40+ return !(std::isalnum (c) || c == ' -' ' _' ' :'
41+ }) == name_.end ();
42+ TORCH_CHECK (
43+ validSize && validChar,
44+ " Worker name must match ^[A-Za-z0-9-_:]*$, "
45+ " and must be non-empty and shorter than "
46+ MAX_NAME_LEN,
47+ " chars, "
48+ " but got "
49+ name_);
50+ }
51+
1052// Large Time Duration for waiting on the condition variable until the map is
1153// population. Cannot use
1254// std::chrono::time_point<std::chrono::steady_clock>::max() due to a known
0 commit comments