Exclude time of waiting in queue from gloo communication prof… (#61342)
Summary:
Background:
The gloo communication implementation is as follow:
1. Construct communication workers and push them into a queue.
2. Initialize a thread pool and each thread run a loop to get worker from the queue and execute it.
Issue:
The recorded profiling time span start from the worker construction and end at finish. So it will include the time of worker waiting in the queue and will result in multiple gloo communication time span overlapping with each other in a same thread in the timeline:
This is because when next work is waiting in the queue, the last work is not finished.
Solution:
This PR delays the profiling start time of gloo communication from worker construction to worker is really executed, so the profiling span will not include the time of waiting in queue. Implementation as follow:
1. Firstly, disable the original record function by specifying 'nullptr' to 'profilingTitle' argument of ProcessGroup::Work
2. Construct a 'recordFunctionBeforeCallback_' and 'recordFunctionEndCallback_' and save it as member of the worker.
3. When the worker is executed, invoke the 'recordFunctionBeforeCallback_'.
4. The 'recordFunctionEndCallback_' will be invoked at finish as before.
After this modification, the gloo profiling span in timeline will not overlap with each other:
Pull Request resolved: https://github.com/pytorch/pytorch/pull/61342
Reviewed By: albanD
Differential Revision: D29811656
Pulled By: gdankel
fbshipit-source-id: ff07e8906d90f21a072049998400b4a48791e441
