Completely remove FutureMessage type

[mrshenli]

Stack from ghstack:

• Apply clang-format to rpc cpp files #50236 Apply clang-format to rpc cpp files
• Completely remove FutureMessage type #50029 Completely remove FutureMessage type
• Completely Remove FutureMessage from RPC agents #50028 Completely Remove FutureMessage from RPC agents
• Completely Remove FutureMessage from RPC cpp tests #50027 Completely Remove FutureMessage from RPC cpp tests
• Remove FutureMessage from RPC request callback logic #50026 Remove FutureMessage from RPC request callback logic
• Completely remove FutureMessage from FaultyProcessGroupAgent #50025 Completely remove FutureMessage from FaultyProcessGroupAgent
• Remove FutureMessage from sender TensorPipeAgent #50024 Remove FutureMessage from sender TensorPipeAgent
• Remove FutureMessage from sender ProcessGroupAgent #50023 Remove FutureMessage from sender ProcessGroupAgent
• Completely remove FutureMessage from distributed autograd #50020 Completely remove FutureMessage from distributed autograd
• Completely remove FutureMessage from RPC TorchScript implementations #50005 Completely remove FutureMessage from RPC TorchScript implementations
• Completely remove FutureMessage from RRef Implementations #50004 Completely remove FutureMessage from RRef Implementations
• Replace FutureMessage with ivalue::Future in RpcAgent retry logic #49995 Replace FutureMessage with ivalue::Future in RpcAgent retry logic
• Replace FutureMessage with ivalue::Future in RRefContext #49960 Replace FutureMessage with ivalue::Future in RRefContext
• Replace FutureMessage with ivalue::Future in distributed/autograd/utils.* #49927 Replace FutureMessage with ivalue::Future in distributed/autograd/utils.*
• Let RpcAgent::send() return JitFuture #49906 Let RpcAgent::send() return JitFuture

Differential Revision: D25774915

[facebook-github-bot]

💊 CI failures summary and remediations

As of commit fccd346 (more details on the Dr. CI page):

---

💚 💚 Looks good so far! There are no failures yet. 💚 💚

---

This comment was automatically generated by Dr. CI (expand for details).

Follow this link to opt-out of these comments for your Pull Requests.

Please report bugs/suggestions to the (internal) Dr. CI Users group.

This comment has been revised 12 times.

[lw]

Awesome work! Thanks so much for this undertaking!

And thanks for such a well split stack, it made reviewing it a pleasure!

I left some minor comments sprinkled around, but I'm listing here some broader thoughts that apply to the whole stack. They can all be addressed as follow-up commits.

• Does it make sense to change the addCallback method of ivalue::Future so that it passes a reference to itself as first argument to the callback? This would remove all those times we create, capture and then lock a weak_ptr. This is already what we do in the Python wrapper for ivalue::Future, so it also brings a benefit of API alignment.
• There's some mix-up of shared_ptr and intrusive_ptr. I wonder if we should try to consolidate on the latter?
• JitFuture is an alias to ivalue::Future, which I guess was introduced for clarity to better separate it from FutureMessage. Does it make sense now to remove that alias and avoid the extra level of indirection that may hurt readability?
• Futures now are created with type at::AnyClassType::get(), although they then only contain Message instances. I know nothing about JIT but is there a way to use a narrower type, both for type safety and readability?

[mrshenli]

@lw thanks a lot for reviewing!

Does it make sense to change the addCallback method of ivalue::Future so that it passes a reference to itself as first argument to the callback? This would remove all those times we create, capture and then lock a weak_ptr. This is already what we do in the Python wrapper for ivalue::Future, so it also brings a benefit of API alignment.

Yep, this makes sense to me and I can add that. cc @wanchaol please let us know if there are concerns.

There's some mix-up of shared_ptr and intrusive_ptr. I wonder if we should try to consolidate on the latter?

Yep, this is on my TODO list

JitFuture is an alias to ivalue::Future, which I guess was introduced for clarity to better separate it from FutureMessage. Does it make sense now to remove that alias and avoid the extra level of indirection that may hurt readability?

Yep, also added this to my TODOs.

Futures now are created with type at::AnyClassType::get(), although they then only contain Message instances. I know nothing about JIT but is there a way to use a narrower type, both for type safety and readability?

Good point. @wanchaol if we want more explicit IValue type on Message, does it mean we need to make Message an IValue type. Or is there any lighter solutions for custom classes?

[mrshenli]

Micro benchmark results:

Process Group Agent (After):

test/distributed/rpc/test_process_group_agent.py::ProcessGroupRpcTestWithSpawn::test_stress_heavy_rpc 
Rank 3 finished testing 20 times in 0.09874725341796875 seconds.
Rank 2 finished testing 20 times in 0.11173415184020996 seconds.                                                          
Rank 1 finished testing 20 times in 0.11910843849182129 seconds.                                                           
Rank 0 finished testing 20 times in 0.11921072006225586 seconds. 

test/distributed/rpc/test_process_group_agent.py::ProcessGroupRpcTestWithSpawn::test_stress_heavy_rpc_torchscript
Rank 1 finished testing 20 times in 0.028783082962036133 seconds.
Rank 3 finished testing 20 times in 0.030823945999145508 seconds.
Rank 0 finished testing 20 times in 0.031084299087524414 seconds.
Rank 2 finished testing 20 times in 0.032097578048706055 seconds.

test/distributed/rpc/test_process_group_agent.py::ProcessGroupRpcTestWithSpawn::test_stress_light_rpc
Rank 3 finished testing 1000 times in 0.4663074016571045 seconds.
Rank 2 finished testing 1000 times in 0.5232148170471191 seconds.
Rank 1 finished testing 1000 times in 0.5746622085571289 seconds.
Rank 0 finished testing 1000 times in 0.5750923156738281 seconds.

Process Group Agent (Before):

test/distributed/rpc/test_process_group_agent.py::ProcessGroupRpcTestWithSpawn::test_stress_heavy_rpc
Rank 2 finished testing 20 times in 0.12031865119934082 seconds. 
Rank 1 finished testing 20 times in 0.12074995040893555 seconds.
Rank 0 finished testing 20 times in 0.12139439582824707 seconds.
Rank 3 finished testing 20 times in 0.12353968620300293 seconds. 

test/distributed/rpc/test_process_group_agent.py::ProcessGroupRpcTestWithSpawn::test_stress_heavy_rpc_torchscript
Rank 3 finished testing 20 times in 0.035836219787597656 seconds.
Rank 1 finished testing 20 times in 0.03717660903930664 seconds.
Rank 0 finished testing 20 times in 0.03726530075073242 seconds.
Rank 2 finished testing 20 times in 0.038858652114868164 seconds.

test/distributed/rpc/test_process_group_agent.py::ProcessGroupRpcTestWithSpawn::test_stress_light_rpc
Rank 3 finished testing 1000 times in 0.530799150466919 seconds.
Rank 2 finished testing 1000 times in 0.5381615161895752 seconds.
Rank 1 finished testing 1000 times in 0.5415563583374023 seconds.
Rank 0 finished testing 1000 times in 0.5458028316497803 seconds.

TensorPipe Agent (After):

test/distributed/rpc/test_tensorpipe_agent.py::TensorPipeRpcTestWithSpawn::test_stress_heavy_rpc 
Rank 1 finished testing 20 times in 0.11632132530212402 seconds.
Rank 0 finished testing 20 times in 0.12914705276489258 seconds.
Rank 3 finished testing 20 times in 0.13401508331298828 seconds.
Rank 2 finished testing 20 times in 0.13787078857421875 seconds.

test/distributed/rpc/test_tensorpipe_agent.py::TensorPipeRpcTestWithSpawn::test_stress_heavy_rpc_torchscript 
Rank 0 finished testing 20 times in 0.032578468322753906 seconds.
Rank 3 finished testing 20 times in 0.035418033599853516 seconds.
Rank 1 finished testing 20 times in 0.039852142333984375 seconds.
Rank 2 finished testing 20 times in 0.04088330268859863 seconds.

test/distributed/rpc/test_tensorpipe_agent.py::TensorPipeRpcTestWithSpawn::test_stress_light_rpc
Rank 3 finished testing 1000 times in 0.1573326587677002 seconds.
Rank 2 finished testing 1000 times in 0.2936713695526123 seconds.
Rank 1 finished testing 1000 times in 0.2966134548187256 seconds.
Rank 0 finished testing 1000 times in 0.2967543601989746 seconds.

TensorPipe Agent (Before):

test/distributed/rpc/test_tensorpipe_agent.py::TensorPipeRpcTestWithSpawn::test_stress_heavy_rpc
Rank 1 finished testing 20 times in 0.12186193466186523 seconds.
Rank 0 finished testing 20 times in 0.12729144096374512 seconds.
Rank 2 finished testing 20 times in 0.1270914077758789 seconds.
Rank 3 finished testing 20 times in 0.13142681121826172 seconds.

test/distributed/rpc/test_tensorpipe_agent.py::TensorPipeRpcTestWithSpawn::test_stress_heavy_rpc_torchscript
Rank 1 finished testing 20 times in 0.04239463806152344 seconds.
Rank 3 finished testing 20 times in 0.04551196098327637 seconds.
Rank 2 finished testing 20 times in 0.04797840118408203 seconds.
Rank 0 finished testing 20 times in 0.06020164489746094 seconds.

test/distributed/rpc/test_tensorpipe_agent.py::TensorPipeRpcTestWithSpawn::test_stress_light_rpc
Rank 1 finished testing 1000 times in 0.26510000228881836 seconds.
Rank 3 finished testing 1000 times in 0.3550393581390381 seconds.
Rank 2 finished testing 1000 times in 0.35509538650512695 seconds.
Rank 0 finished testing 1000 times in 0.36699509620666504 seconds.

[mrshenli]

ci-all tests in #50216

[facebook-github-bot]

@mrshenli merged this pull request in c480eeb.