[pytorch] Support process_group_agent "sending to itself"

[jjlilley]

Stack from ghstack:

• [pytorch] process_group_agent optimizations #29324 [pytorch] process_group_agent optimizations
• [pytorch] Support process_group_agent "sending to itself" #29253 [pytorch] Support process_group_agent "sending to itself"

Some operations can be simpler if a worker can send an rpc to itself.
The main reason for not doing previous was that Gloo doesn't support
self-sending.

That said, this changes the process_group_agent to skip the assert
check, and simply enqueue the rpc message in its receiving queue.

Differential Revision: D18339715

[pietern]

LGTM, but will defer to @mrshenli for the approval.

[mrshenli]

This is awesome! Thank you!

[mrshenli]

there is a RpcAgent::getWorkerInfo() for this.

[jjlilley]

I can change that. I was just copying the code a few lines below, later in this function. :)

[mrshenli]

This is capture by value? Does this prevent the local var serializedPayload being freed before the recv is done?

[jjlilley]

serializedPayload is a string*, so we capture the pointer by value.
The goal is for the deleter can delete the buffer when the torch::Tensor is done with it.
The local var won't clean itself up.

In the other torch::from_blob() cases in this file, the tensor buffer doesn't need to be persisted beyond the immediate function scope. This case is different because the Tensor's use is delayed until the other listen thread picks it up.

Maybe I'll go ahead and make it a unique_ptr<>, to make the code a bit more exception-safe, and to clarify the lifetime.

[jjlilley]

I went ahead and uploaded a version with a smart pointer to make memory lifetime clearer in this func, and safer with exceptions.

That said, I had to use shared_ptr<> instead of unique_ptr<> because Pytorch is supposed to avoid c++14 code (TIL there are some c++11 issues with unique_ptr<> lambda capture that were solved in c++14)

[pietern]

@jjlilley Thanks to @smessmer, we can use C++14 since 2 weeks (see #28443).

Capturing a std::unique_ptr is possible only if you don't convert it to std::function, because that requires the wrapped function to be copy-constructible. And of course, if you capture a std::unique_ptr, it won't be copy-constructible.

[mrshenli]

Do we then want to wait for #27864?

Nov 07 01:02:43 /var/lib/jenkins/workspace/torch/csrc/distributed/rpc/process_group_agent.cpp: In lambda function:
Nov 07 01:02:43 /var/lib/jenkins/workspace/torch/csrc/distributed/rpc/process_group_agent.cpp:316:20: error: lambda capture initializers only available with -std=c++14 or -std=gnu++14 [-Werror]
Nov 07 01:02:43                    [payload = std::move(payload)](void*) {

[jjlilley]

actually, if I create a unique_ptr<> and release the raw pointer into the deleter lambda captures, I think we both:

1. are still fairly exception-safe, and mostly avoid raw pointers
2. self-document the memory lifecycle reasonably
   I'll upload a version that does this.

[mrshenli]

Shall we add a test for rpc.remote as well?

[jjlilley]

• Uploading another try to deal with c++14 vs c++11 lambda issue. I think this one will pass. :)

• re: rpc.remote - there's still additional work for this, see python_functions.cpp:142, and kind of distinct code to change.
  (I tried trivially uncommenting out the asserts and adding a test case, it seems best deferred to another PR)

[mrshenli]

I tried trivially uncommenting out the asserts and adding a test case, it seems best deferred to another PR.

I see, that makes sense. Yes, besides commenting that assert out, at least the following line should create an owner RRef instead of a user RRef, if calling to self.

[mrshenli]

LGTM! Thanks!

Would I be correct if I assume that we will get rid of passing raw pointers to the deleter when #27864 is in?

[jjlilley]

Thanks for the review!
And happy to change the lambda capture after we transition to c++14.

[facebook-github-bot]

This pull request has been merged in 2cd4f86.