[fix] Add FaultyTensorpipeAgent For TestCase-pytorch-AtomGit

Mtorch_npu/csrc/distributed/CMakeLists.txt+1-1

	@@ -3,7 +3,7 @@ if (DEFINED BUILD_LIBTORCH)
3	# Exclude Python binding files when building libtorch	3	# Exclude Python binding files when building libtorch
4	list(REMOVE_ITEM _DIST_SRCS "${CMAKE_CURRENT_SOURCE_DIR}/Init.cpp")	4	list(REMOVE_ITEM _DIST_SRCS "${CMAKE_CURRENT_SOURCE_DIR}/Init.cpp")
5	else()	5	else()
6	- FILE(GLOB _DIST_SRCS .cpp rpc/.cpp symm_mem/*.cpp)	6	+ FILE(GLOB _DIST_SRCS .cpp rpc/.cpp rpc/testing/.cpp symm_mem/.cpp)
7	endif()	7	endif()
8		8
9	LIST(APPEND DIST_SRCS ${_DIST_SRCS})	9	LIST(APPEND DIST_SRCS ${_DIST_SRCS})

Mtorch_npu/csrc/distributed/rpc/init.cpp+91-0

	@@ -18,6 +18,7 @@
18	#include <torch/types.h>	18	#include <torch/types.h>
19		19
20	#include "torch_npu/csrc/distributed/rpc/tensorpipe_agent.h"	20	#include "torch_npu/csrc/distributed/rpc/tensorpipe_agent.h"
		21	+#include "torch_npu/csrc/distributed/rpc/testing/faulty_tensorpipe_agent.h"
21		22
22	namespace torch_npu {	23	namespace torch_npu {
23	namespace distributed {	24	namespace distributed {
	@@ -84,6 +85,96 @@ PyObject rpc_npu_init(PyObject _unused, PyObject *noargs)
84	.def_readonly("is_static_group", &TensorPipeAgent::isStaticGroup_)	85	.def_readonly("is_static_group", &TensorPipeAgent::isStaticGroup_)
85	.def_property_readonly("store", &TensorPipeAgent::getStore);	86	.def_property_readonly("store", &TensorPipeAgent::getStore);
86		87
		88	+ shared_ptr_class_<FaultyTensorPipeRpcBackendOptions>(
		89	+ module,
		90	+ "FaultyTensorPipeRpcBackendOptions",
		91	+ rpc_module.attr("_TensorPipeRpcBackendOptionsBase"))
		92	+ .def(
		93	+ py::init<
		94	+ int,
		95	+ float,
		96	+ std::string,
		97	+ std::vector<std::string>,
		98	+ std::unordered_map<std::string, float>,
		99	+ int>(),
		100	+ py::arg("num_worker_threads"),
		101	+ py::arg("rpc_timeout"),
		102	+ py::arg("init_method"),
		103	+ py::arg("messages_to_fail"),
		104	+ py::arg("messages_to_delay"),
		105	+ py::arg("num_fail_sends"))
		106	+ .def_readwrite(
		107	+ "num_worker_threads", &TensorPipeRpcBackendOptions::numWorkerThreads)
		108	+ .def_readwrite(
		109	+ "messages_to_fail",
		110	+ &FaultyTensorPipeRpcBackendOptions::messagesToFail)
		111	+ .def_readwrite(
		112	+ "messages_to_delay",
		113	+ &FaultyTensorPipeRpcBackendOptions::messagesToDelay)
		114	+ .def_readwrite(
		115	+ "num_fail_sends", &FaultyTensorPipeRpcBackendOptions::numFailSends);
		116	+
		117	+ shared_ptr_class_<FaultyTensorPipeAgent>(
		118	+ module, "FaultyTensorPipeAgent", module.attr("TensorPipeAgent"))
		119	+ .def(
		120	+ py::init(
		121	+ [](const c10::intrusive_ptr<::c10d::Store> &store,
		122	+ std::string name,
		123	+ worker_id_t rank,
		124	+ c10::optional<int> world_size,
		125	+ FaultyTensorPipeRpcBackendOptions opts,
		126	+ std::unordered_map<std::string, DeviceMap> reverse_device_maps,
		127	+ std::vector<c10::Device> devices) {
		128	+ return std::shared_ptr<FaultyTensorPipeAgent>(
		129	+ new FaultyTensorPipeAgent(
		130	+ store,
		131	+ std::move(name),
		132	+ rank,
		133	+ world_size,
		134	+ std::move(opts),
		135	+ std::move(reverse_device_maps),
		136	+ std::move(devices),
		137	+ std::make_unique<RequestCallbackImpl>()),
		138	+ torch::impl::destroy_without_gil<FaultyTensorPipeAgent>);
		139	+ }),
		140	+ py::arg("store"),
		141	+ py::arg("name"),
		142	+ py::arg("rank"),
		143	+ py::arg("world_size"),
		144	+ py::arg("opts"),
		145	+ py::arg("reverse_device_maps"),
		146	+ py::arg("devices"))
		147	+ .def(
		148	+ "join",
		149	+ &TensorPipeAgent::join,
		150	+ py::call_guard<py::gil_scoped_release>(),
		151	+ py::arg("shutdown") = false,
		152	+ py::arg("timeout") = 0)
		153	+ .def(
		154	+ "shutdown",
		155	+ &TensorPipeAgent::shutdown,
		156	+ py::call_guard<py::gil_scoped_release>())
		157	+ .def(
		158	+ "get_worker_info",
		159	+ (const WorkerInfo &(TensorPipeAgent::*)(void) const) &
		160	+ RpcAgent::getWorkerInfo,
		161	+ py::call_guard<py::gil_scoped_release>())
		162	+ .def(
		163	+ "get_worker_info",
		164	+ (const WorkerInfo &(TensorPipeAgent::*)(const std::string &) const) &
		165	+ TensorPipeAgent::getWorkerInfo,
		166	+ py::call_guard<py::gil_scoped_release>())
		167	+ .def(
		168	+ "get_worker_info",
		169	+ (const WorkerInfo &(TensorPipeAgent::*)(worker_id_t id) const) &
		170	+ TensorPipeAgent::getWorkerInfo,
		171	+ py::call_guard<py::gil_scoped_release>())
		172	+ .def(
		173	+ "get_worker_infos",
		174	+ (std::vector<WorkerInfo>(TensorPipeAgent::*)() const) &
		175	+ TensorPipeAgent::getWorkerInfos,
		176	+ py::call_guard<py::gil_scoped_release>());
		177	+
87	Py_RETURN_TRUE;	178	Py_RETURN_TRUE;
88	}	179	}
89		180

Mtorch_npu/csrc/distributed/rpc/tensorpipe_agent.cpp+3-1

	@@ -9,6 +9,7 @@
9		9
10	#include <c10/core/StreamGuard.h>	10	#include <c10/core/StreamGuard.h>
11	#include <c10/util/irange.h>	11	#include <c10/util/irange.h>
		12	+#include <fmt/format.h>
12	#include <torch/csrc/distributed/rpc/agent_utils.h>	13	#include <torch/csrc/distributed/rpc/agent_utils.h>
13	#include <torch/csrc/distributed/rpc/utils.h>	14	#include <torch/csrc/distributed/rpc/utils.h>
14	#include <unistd.h>	15	#include <unistd.h>
	@@ -925,7 +926,8 @@ void TensorPipeAgent::pollTimeoutRpcs()
925	// outside the lock to prevent potential lock-order-inversions by callbacks	926	// outside the lock to prevent potential lock-order-inversions by callbacks
926	// triggered by the setError call.	927	// triggered by the setError call.
927	for (auto &timeoutMetadata : timedOutFutures) {	928	for (auto &timeoutMetadata : timedOutFutures) {
928	- std::string errorMsg = "";	929	+ std::string errorMsg =
		930	+ fmt::format(kRpcTimeoutErrorStr, timeoutMetadata.timeout.count());
929	auto err = makeRPCError(errorMsg, RPCErrorType::TIMEOUT);	931	auto err = makeRPCError(errorMsg, RPCErrorType::TIMEOUT);
930	markFutureWithError(std::move(timeoutMetadata.responseFuture), std::move(err));	932	markFutureWithError(std::move(timeoutMetadata.responseFuture), std::move(err));
931	}	933	}

Mtorch_npu/csrc/distributed/rpc/tensorpipe_agent.h+1-0

	@@ -44,6 +44,7 @@ using torch::distributed::rpc::collectNames;
44	using torch::distributed::rpc::createExceptionResponse;	44	using torch::distributed::rpc::createExceptionResponse;
45	using torch::distributed::rpc::DeviceMap;	45	using torch::distributed::rpc::DeviceMap;
46	using torch::distributed::rpc::JitFuture;	46	using torch::distributed::rpc::JitFuture;
		47	+using torch::distributed::rpc::kRpcTimeoutErrorStr;
47	using torch::distributed::rpc::kSecToMsConversion;	48	using torch::distributed::rpc::kSecToMsConversion;
48	using torch::distributed::rpc::kUnsetRpcTimeout;	49	using torch::distributed::rpc::kUnsetRpcTimeout;
49	using torch::distributed::rpc::makeRPCError;	50	using torch::distributed::rpc::makeRPCError;

Atorch_npu/csrc/distributed/rpc/testing/faulty_tensorpipe_agent.cpp+142-0

@@ -0,0 +1,142 @@
	1	+#ifdef USE_RPC_FRAMEWORK
	2	+
	3	+#include <torch_npu/csrc/distributed/rpc/testing/faulty_tensorpipe_agent.h>
	4	+#include <torch/csrc/distributed/rpc/utils.h>
	5	+
	6	+namespace torch_npu {
	7	+namespace distributed {
	8	+namespace rpc {
	9	+
	10	+static std::string fromVecToString(const std::vector<char>& vec) {
	11	+ return std::string(vec.begin(), vec.end());
	12	+}
	13	+
	14	+FaultyTensorPipeAgent::FaultyTensorPipeAgent(
	15	+ const c10::intrusive_ptr<::c10d::Store>& store,
	16	+ std::string selfName,
	17	+ worker_id_t selfId,
	18	+ c10::optional<int> worldSize,
	19	+ FaultyTensorPipeRpcBackendOptions opts,
	20	+ std::unordered_map<std::string, DeviceMap> reverseDeviceMaps,
	21	+ std::vector<c10::Device> devices,
	22	+ std::unique_ptr<RequestCallback> callback)
	23	+ : TensorPipeAgent(
	24	+ store,
	25	+ std::move(selfName),
	26	+ selfId,
	27	+ worldSize,
	28	+ static_cast<TensorPipeRpcBackendOptions>(opts),
	29	+ std::move(reverseDeviceMaps),
	30	+ std::move(devices),
	31	+ std::move(callback)),
	32	+ numFailSends_(opts.numFailSends),
	33	+ messageTypesToFail_(parseMessagesToFailInput(
	34	+ std::move(opts.messagesToFail))),
	35	+ messageTypesToDelay_(parseMessagesToDelay(
	36	+ std::move(opts.messagesToDelay))) {}
	37	+
	38	+std::vector<MessageType> FaultyTensorPipeAgent::parseMessagesToFailInput(
	39	+ const std::vector<std::string>& messagesToFail) const {
	40	+ std::vector<MessageType> messageTypesToFail;
	41	+ messageTypesToFail.reserve(messagesToFail.size());
	42	+ for (const auto& msgString : messagesToFail) {
	43	+ messageTypesToFail.push_back(messageStringToType(msgString));
	44	+ }
	45	+ return messageTypesToFail;
	46	+}
	47	+
	48	+std::unordered_map<MessageType, float, std::hash<int>> FaultyTensorPipeAgent::
	49	+ parseMessagesToDelay(const std::unordered_map<std::string, float>&
	50	+ messageTypesToDelay) const {
	51	+ std::unordered_map<MessageType, float, std::hash<int>> delayMessages;
	52	+ for (const auto& messagePair : messageTypesToDelay) {
	53	+ float delay = messagePair.second;
	54	+ TORCH_CHECK(
	55	+ delay >= 0,
	56	+ "Delays passed to FaultyTensorPipeAgent must be non-negative.")
	57	+ delayMessages.insert({messageStringToType(messagePair.first), delay});
	58	+ }
	59	+ return delayMessages;
	60	+}
	61	+
	62	+c10::intrusive_ptr<JitFuture> FaultyTensorPipeAgent::send(
	63	+ const WorkerInfo& to,
	64	+ c10::intrusive_ptr<Message> message,
	65	+ const float rpcTimeoutSeconds,
	66	+ const DeviceMap& /* unused */) {
	67	+ if (!shouldFailMessage(message->type())) {
	68	+ return TensorPipeAgent::send(to, std::move(message), rpcTimeoutSeconds);
	69	+ }
	70	+
	71	+ const auto key = fromVecToString(message->payload());
	72	+ std::unique_lock<std::mutex> lock(failMapMutex_);
	73	+ auto it = failMessageCountMap_.find(key);
	74	+ if (it == failMessageCountMap_.end()) {
	75	+ failMessageCountMap_[key] = 0;
	76	+ }
	77	+ if (failMessageCountMap_[key] < numFailSends_) {
	78	+ failMessageCountMap_[key]++;
	79	+ lock.unlock();
	80	+ auto jitFuture = c10::make_intrusive<JitFuture>(at::AnyClassType::get());
	81	+ jitFuture->setError(std::make_exception_ptr(std::runtime_error(makeRPCError(
	82	+ c10::str("Send attempt failed intentionally for ", key),
	83	+ RPCErrorType::INTENTIONAL_FAILURE))));
	84	+ return jitFuture;
	85	+ } else {
	86	+ lock.unlock();
	87	+ return TensorPipeAgent::send(to, std::move(message), rpcTimeoutSeconds);
	88	+ }
	89	+}
	90	+
	91	+void FaultyTensorPipeAgent::pipeWrite(
	92	+ const std::shared_ptr<tensorpipe_npu::Pipe>& pipe,
	93	+ c10::intrusive_ptr<Message> rpcMessage,
	94	+ std::vector<c10::Device>&& devices,
	95	+ std::vector<c10::Stream> streams,
	96	+ std::function<void(const tensorpipe_npu::Error&)> fn) noexcept {
	97	+ float msgDelay = getDelayForMessage(rpcMessage->type());
	98	+ if (msgDelay != 0) {
	99	+ std::this_thread::sleep_for(std::chrono::milliseconds(
	100	+ static_cast<int>(msgDelay * kSecToMsConversion)));
	101	+ }
	102	+ TensorPipeAgent::pipeWrite(pipe, rpcMessage, std::move(devices), streams, fn);
	103	+}
	104	+
	105	+bool FaultyTensorPipeAgent::shouldFailMessage(MessageType type) const {
	106	+ return (
	107	+ std::find(messageTypesToFail_.begin(), messageTypesToFail_.end(), type) !=
	108	+ messageTypesToFail_.end());
	109	+}
	110	+
	111	+float FaultyTensorPipeAgent::getDelayForMessage(MessageType type) const {
	112	+ const auto& it = messageTypesToDelay_.find(type);
	113	+ return it == messageTypesToDelay_.end() ? 0 : it->second;
	114	+}
	115	+
	116	+MessageType FaultyTensorPipeAgent::messageStringToType(
	117	+ const std::string& messageString) const {
	118	+ static std::unordered_map<std::string, MessageType> msgMap = {
	119	+ {"RREF_FORK_REQUEST", MessageType::RREF_FORK_REQUEST},
	120	+ {"RREF_CHILD_ACCEPT", MessageType::RREF_CHILD_ACCEPT},
	121	+ {"RREF_USER_DELETE", MessageType::RREF_USER_DELETE},
	122	+ {"CLEANUP_AUTOGRAD_CONTEXT_REQ",
	123	+ MessageType::CLEANUP_AUTOGRAD_CONTEXT_REQ},
	124	+ {"PYTHON_REMOTE_CALL", MessageType::PYTHON_REMOTE_CALL},
	125	+ {"SCRIPT_REMOTE_CALL", MessageType::SCRIPT_REMOTE_CALL},
	126	+ {"PYTHON_CALL", MessageType::PYTHON_CALL},
	127	+ {"SCRIPT_CALL", MessageType::SCRIPT_CALL},
	128	+ {"PYTHON_RREF_FETCH_CALL", MessageType::PYTHON_RREF_FETCH_CALL},
	129	+ {"SCRIPT_RREF_FETCH_CALL", MessageType::SCRIPT_RREF_FETCH_CALL}};
	130	+ const auto& it = msgMap.find(messageString);
	131	+ TORCH_CHECK(
	132	+ it != msgMap.end(),
	133	+ "No mapping to rpc::MessageType exists for ",
	134	+ messageString);
	135	+ return it->second;
	136	+}
	137	+
	138	+} // namespace rpc
	139	+} // namespace distributed
	140	+} // namespace torch_npu
	141	+
	142	+#endif

Atorch_npu/csrc/distributed/rpc/testing/faulty_tensorpipe_agent.h+91-0

@@ -0,0 +1,91 @@
	1	+#ifdef USE_RPC_FRAMEWORK
	2	+
	3	+#pragma once
	4	+
	5	+#include <torch/csrc/distributed/rpc/message.h>
	6	+#include <torch_npu/csrc/distributed/rpc/tensorpipe_agent.h>
	7	+
	8	+namespace torch_npu {
	9	+namespace distributed {
	10	+namespace rpc {
	11	+
	12	+struct FaultyTensorPipeRpcBackendOptions : public TensorPipeRpcBackendOptions {
	13	+ FaultyTensorPipeRpcBackendOptions(
	14	+ int num_worker_threads,
	15	+ float rpc_timeout,
	16	+ std::string init_method,
	17	+ std::vector<std::string> messages_to_fail,
	18	+ std::unordered_map<std::string, float> messages_to_delay,
	19	+ int num_fail_sends = 0)
	20	+ : TensorPipeRpcBackendOptions(
	21	+ num_worker_threads,
	22	+ std::optional<std::vector<std::string>>(),
	23	+ std::optional<std::vector<std::string>>(),
	24	+ rpc_timeout,
	25	+ std::move(init_method)),
	26	+ messagesToFail(std::move(messages_to_fail)),
	27	+ messagesToDelay(std::move(messages_to_delay)),
	28	+ numFailSends(num_fail_sends) {
	29	+ TORCH_CHECK(numFailSends >= 0, "numFailSends should be non-negative");
	30	+ }
	31	+
	32	+ std::vector<std::string> messagesToFail;
	33	+ std::unordered_map<std::string, float> messagesToDelay;
	34	+ int numFailSends;
	35	+};
	36	+
	37	+class FaultyTensorPipeAgent : public TensorPipeAgent {
	38	+ public:
	39	+ FaultyTensorPipeAgent(
	40	+ const c10::intrusive_ptr<::c10d::Store>& store,
	41	+ std::string selfName,
	42	+ worker_id_t selfId,
	43	+ c10::optional<int> worldSize,
	44	+ FaultyTensorPipeRpcBackendOptions opts,
	45	+ std::unordered_map<std::string, DeviceMap> reverseDeviceMaps,
	46	+ std::vector<c10::Device> devices,
	47	+ std::unique_ptr<RequestCallback> callback);
	48	+
	49	+ c10::intrusive_ptr<JitFuture> send(
	50	+ const WorkerInfo& to,
	51	+ c10::intrusive_ptr<Message> message,
	52	+ const float rpcTimeoutSeconds = kUnsetRpcTimeout,
	53	+ const DeviceMap& deviceMap = {}) override;
	54	+
	55	+ void pipeWrite(
	56	+ const std::shared_ptr<tensorpipe_npu::Pipe>& pipe,
	57	+ c10::intrusive_ptr<Message> rpcMessage,
	58	+ std::vector<c10::Device>&& devices,
	59	+ std::vector<c10::Stream> streams,
	60	+ std::function<void(const tensorpipe_npu::Error&)> fn) noexcept override;
	61	+
	62	+ protected:
	63	+ bool shouldFailMessage(MessageType type) const;
	64	+
	65	+ private:
	66	+ std::vector<MessageType> parseMessagesToFailInput(
	67	+ const std::vector<std::string>& messagesToFail) const;
	68	+
	69	+ float getDelayForMessage(MessageType type) const;
	70	+
	71	+ std::unordered_map<MessageType, float, std::hash<int>> parseMessagesToDelay(
	72	+ const std::unordered_map<std::string, float>& messageTypesToDelay) const;
	73	+
	74	+ const int numFailSends_;
	75	+
	76	+ const std::vector<MessageType> messageTypesToFail_;
	77	+
	78	+ std::unordered_map<MessageType, float, std::hash<int>> messageTypesToDelay_;
	79	+
	80	+ std::unordered_map<std::string, int> failMessageCountMap_;
	81	+
	82	+ std::mutex failMapMutex_;
	83	+
	84	+ MessageType messageStringToType(const std::string& messageString) const;
	85	+};
	86	+
	87	+} // namespace rpc
	88	+} // namespace distributed
	89	+} // namespace torch_npu
	90	+
	91	+#endif

Mtorch_npu/distributed/rpc/backend_registry.py+59-1

	@@ -285,14 +285,72 @@ def _npu_tensorpipe_init_backend_handler(
285	return agent	285	return agent
286		286
287		287
		288	+def _faulty_tensorpipe_construct_rpc_backend_options_handler(
		289	+ rpc_timeout,
		290	+ init_method,
		291	+ num_worker_threads,
		292	+ messages_to_fail,
		293	+ messages_to_delay,
		294	+ num_fail_sends=0,
		295	+ **kwargs,
		296	+):
		297	+ from torch_npu._C._distributed_rpc import FaultyTensorPipeRpcBackendOptions
		298	+
		299	+ return FaultyTensorPipeRpcBackendOptions(
		300	+ num_worker_threads=num_worker_threads,
		301	+ rpc_timeout=rpc_timeout,
		302	+ init_method=init_method,
		303	+ messages_to_fail=messages_to_fail,
		304	+ messages_to_delay=messages_to_delay,
		305	+ num_fail_sends=num_fail_sends,
		306	+ )
		307	+
		308	+
		309	+def _faulty_tensorpipe_init_backend_handler(
		310	+ store, name, rank, world_size, rpc_backend_options
		311	+):
		312	+ from torch_npu._C._distributed_rpc import FaultyTensorPipeAgent, FaultyTensorPipeRpcBackendOptions
		313	+
		314	+ if not isinstance(store, dist.Store):
		315	+ raise TypeError(f"`store` must be a c10d::Store. {store}" + dist_error(ErrCode.TYPE))
		316	+
		317	+ if not isinstance(rpc_backend_options, FaultyTensorPipeRpcBackendOptions):
		318	+ raise TypeError(
		319	+ f"`rpc_backend_options` must be a `FaultyTensorPipeRpcBackendOptions`. {rpc_backend_options}" +
		320	+ dist_error(ErrCode.TYPE)
		321	+ )
		322	+
		323	+ _init_device_state(_get_privateuse1_backend_name())
		324	+
		325	+ agent = FaultyTensorPipeAgent(
		326	+ store,
		327	+ name,
		328	+ rank,
		329	+ world_size,
		330	+ rpc_backend_options,
		331	+ {},
		332	+ [],
		333	+ )
		334	+ api._init_rpc_states(agent)
		335	+
		336	+ return agent
		337	+
		338	+
288	def _rpc_backend_registry():	339	def _rpc_backend_registry():
289	if hasattr(torch_npu._C, "_rpc_npu_init"):	340	if hasattr(torch_npu._C, "_rpc_npu_init"):
290	torch_npu._C._rpc_npu_init()	341	torch_npu._C._rpc_npu_init()
		342	+
291	rpc.backend_registry.register_backend(	343	rpc.backend_registry.register_backend(
292	"NPU_TENSORPIPE",	344	"NPU_TENSORPIPE",
293	_npu_tensorpipe_construct_rpc_backend_options_handler,	345	_npu_tensorpipe_construct_rpc_backend_options_handler,
294	_npu_tensorpipe_init_backend_handler,	346	_npu_tensorpipe_init_backend_handler,
295	)	347	)
296		348
		349	+ rpc.backend_registry.register_backend(
		350	+ "NPU_FAULTY_TENSORPIPE",
		351	+ _faulty_tensorpipe_construct_rpc_backend_options_handler,
		352	+ _faulty_tensorpipe_init_backend_handler,
		353	+ )
		354	+
297	import torch.distributed.rpc as _rpc_module	355	import torch.distributed.rpc as _rpc_module
298	- _rpc_module.BackendType = rpc.backend_registry.BackendType	356	+ _rpc_module.BackendType = rpc.backend_registry.BackendType