[Dist profiling] Fix flaky tests

[rohan-varma]

Stack from ghstack:

• [Dist profiling] Fix flaky tests #56963 [Dist profiling] Fix flaky tests

Tests such as #50840 and #56690 have been flaky due to the profiling part.

What was happening is that in these tests, we occasionally had inaccurate (off by one) counts of the expected no. of distributed collectives events expected to be profiled.

It was challenging to find the root cause of this as the issue was extremely rare. Digging into profiler's parse_legacy_records, I verified the correct no. of events were alway returned. However, when constructing EventList after parsing records, sometimes distributed collective events were getting filtered out leading to a reduction in their count.

The reason they were getting filtered out is because _remove_dup_nodes filters out parent/child events that appear to be the same event based on some heuristics. However these distributed collectives were incorrectly getting picked up as parent/child events since their times can overlap due to their async nature. To fix, add a blocklist and bypass the filtering if it is a distributed collective.

Testplan:

To reproduce:

1. SSH into circleCI instance that runs distributed_test.py with 3 GPUs
2. Modify test_send_recv_autograd_profiler to run 100+ send/recv calls. Without this patch the issue reproduces, verify that this patch fixes the issue.
   Was able to reproduce the issue more often by doing 100+ send/recv calls in _test_send_recv. Verified this fix on CI machine by running the test thousands of times.

An alternative fix is to mark these distributed collectives as async via the is_async flag (because work/wait() paradigm is basically async). This might be a bit cleaner as it also bypasses the parent/child heuristics.

Differential Revision: D28013540

[facebook-github-bot]

💊 CI failures summary and remediations

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

---

• 1/1 failures introduced in this PR

---

1 failure not recognized by patterns:

Job	Step	Action
binary_linux_libtorch_3_7m_cpu_devtoolset7_shared-with-deps_build	Build	🔁 rerun

---

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.

Click here to manually regenerate this comment.

[ilia-cher]

thanks, though it would be def. better to mark these events as async

We have logic that specifically exempts async events from parent-child analysis. Since these events are not marked as async, we recognize them as parent-child and the rest follows.

So, i think the right fix is not to add these kinds of hacky exceptions (it's also fragile in a sense that this list of exemptions needs to be maintained), but to address the root cause.

[rohan-varma]

@ilia-cher Sure, that approach is probably cleaner. Although, currently we check is_async by checking if the record function start/end callbacks have different thread IDs. For profiling these distributed collectives that is not the case (its only true for RPC). So I guess we would still have some sort of list needed to mark these events as async?

[rohan-varma]

Closing in favor of a new approach per offline discussion