Repository: tensorflow
Issue: [BUG] race condition in local rendezvous
Labels: awaiting review,type:bug,comp:core,awaiting PR merge,TF 2.12

### Issue type

Bug

### Have you reproduced the bug with TensorFlow Nightly?

No

### TensorFlow version

2.10, 2.12, doesn't really matter

### Custom code

Yes

### OS platform and distribution

ubuntu 20.04

### Python version

3.8, 3.9

### Bazel version

5.3.0

### GCC/compiler version

9.4

### Current behavior?

The story is kinda bit long and it took us months to debug this issue. I'll try to keep it short.

### Background and what the problem is
We run the [DLRM from NV DLE](https://github.com/NVIDIA/DeepLearningExamples/tree/3f82b7f982a1d9a9c5392bd110a5b93dfa7e80eb/Tensorflow2/Recommendation/DLRM) in our CI daily test. 

Very rare, like once every few weeks, the daily test report such error.

`
W tensorflow/core/framework/op_kernel.cc:1874] OP_REQUIRES failed at strided_slice_op.cc:112 : INVALID_ARGUMENT: Expected begin, end, and strides to be 1D equal size tensors, but got shapes [65536], [1], and [1] instead.
`

The[ strided slice op validator](https://github.com/tensorflow/tensorflow/blob/3a029b19c9c156cd68cab671b5ce95bde839f15e/tensorflow/core/util/strided_slice_op.cc#L214) isn't happy about the input parameters.

The problem is like a ghost. It appears once every 1 or 2 weeks which makes it very hard to debug. But one thing is for sure is that it's not a random dram bit flip caused by cosmic rays because the error systoms is stable: one of (begin, end, stride) tensor's shape is incorrectly set to [65535] instead of [1]

### Evidence

I'll skip the lengthy debug process and jump to the last step.

![image](https://github.com/tensorflow/tensorflow/assets/8800468/d73ddf6d-7cc9-4a10-90f8-1e4d3885c91b)
![image](https://github.com/tensorflow/tensorflow/assets/8800468/0b99c252-59b0-490d-8321-924adb70be72)


We observed a malfunctioned SEND/RECV pair.

The item tensor in SENDOP has a shape of [1] but during the local rendezvous process the recevied tensor got a shape of [65536]. This tensor later flows to strided slice op and triggers the grumpy validator.

![image](https://github.com/tensorflow/tensorflow/assets/8800468/7626ef2f-0de4-4eb0-8c25-3cfd9618b359)

We grep the key hash (red highlighted 1632...9087) in the debug log.  There are two threads: 66787 and 63997. Most of the time the rendezvous runs in a SEND-RECV pattern within the same thread. 

But in the green rectangle you can see two threads interleaved and a SEND-SEND-RECV-RECV pattern is observed. Soon after this the bomb exploded.

### Analysis

It seems the root cause is a bug in the rendezvous mechanism. Two unrelated operations generated the same communication Rendezvous key. 

Assuming there are 2 threads, when the CPU load is not heavy, in most cases the scheduling order is

- THREAD A: SEND(KEY, VALUE_A)

- THREAD A: RECV(KEY)

- THREAD B: SEND(KEY, VALUE_B)

- THREAD B: RECV(KEY)

In this way, everyone will be fine even if the keys are the same. It is equivalent to time-division multiplexing of the same KEY.

When the CPU load becomes heavy (which is the case in our daily test scenario), thread scheduling becomes unpredictable. It is possible that such a pattern may occur

- THREAD A: SEND(KEY, VALUE_A)

- THREAD B: SEND(KEY, VALUE_B)

- THREAD A: RECV(KEY)

- THREAD B: RECV(KEY)

This will cause THREAD A to incorrectly receive the data sent by THREAD B.


### Related python code

1. Thread A is running the normal [training loop.](https://github.com/NVIDIA/DeepLearningExamples/blob/3f82b7f982a1d9a9c5392bd110a5b93dfa7e80eb/Tensorflow2/Recommendation/DLRM/main.py#L281)
2. Thread B is created for the RawBinaryDataset class. The pre-processing creates an [asynchronous thread pool.](https://github.com/NVIDIA/DeepLearningExamples/blob/3f82b7f982a1d9a9c5392bd110a5b93dfa7e80eb/Tensorflow2/Recommendation/DLRM/split_binary_dataset.py#L138)
3. Thread A is running[ x = x[self.begin_idx:self.end_idx]](https://github.com/NVIDIA/DeepLearningExamples/blob/3f82b7f982a1d9a9c5392bd110a5b93dfa7e80eb/Tensorflow2/Recommendation/DLRM/model.py#L58). Translated into StridedSlice operation by eager runtime. begin_idx, end_idx and the implicit stride are all constant tensors on CPU, and StridedSlice is an device operator. So these three tensors need to be sent to the device side.
4. Thread B is running[ tensor = tf.expand_dims(tensor, axis=1). ](https://github.com/NVIDIA/DeepLearningExamples/blob/3f82b7f982a1d9a9c5392bd110a5b93dfa7e80eb/Tensorflow2/Recommendation/DLRM/split_binary_dataset.py#L204). Execute expand_dims. Similar to thread A, the expand_dims operator runs on the device side while the input is on CPU, and the tensor needs to be sent to the device. The shape of this tensor is [65536].
5. Due to the defect of Rendezous::CreateKey, the keys generated by these two operators in the eager runtime are exactly the same.
6. Multi-threading + HASH KEY collision + SEND-SEND-RECV-RECV, all three together, BOOOM.

### FIX

![image](https://github.com/tensorflow/tensorflow/assets/8800468/f50ed76a-9cef-4bf7-87b2-e4a298236dd9)

The key string is 
/job:localhost/replica:0/task:0/device:CPU:0;ea74dbce35f0ab7e;/job:localhost/replica:0/task:0/device:MLU:0;edge_2_input;0:0

The rendezvous key consists 5 parts.
- src_device: /job:localhost/replica:0/task:0/device:CPU:0;
- src_incarnation: ea74dbce35f0ab7e;
- dst_device: /job:localhost/replica:0/task:0/device:MLU:0;
- name: edge_2_input;
- frame_inter: 0:0

In this case,  4 out of 5 (src_device, src_incarnation, dst_device_frame_iter) are naturely indistinguishable.

Unless we add new field in the key, the only field we can play with is `name`.

The creation of `edge_2_input` involes another few tons of code.

![image](https://github.com/tensorflow/tensorflow/assets/8800468/b2d85cd7-1e47-492b-be50-1e3d63f13d55)

A quick fix for my current problem is simply add `dst_name` in the `tensor_name_attr` during graph partition. The names will become `edge_2_input_strided_slice` and `edge_2_input_expand_dims` thus my problem is solved.

```diff
diff --git a/tensorflow/core/graph/graph_partition.cc b/tensorflow/core/graph/graph_partition.cc
index a4f09383c63..57a4e919526 100644
--- a/tensorflow/core/graph/graph_partition.cc
+++ b/tensorflow/core/graph/graph_partition.cc
@@ -1147,7 +1147,8 @@ Status Partition(const PartitionOptions& opts, Graph* g,
         tensor_name_attr = opts.get_tensor_name_attr(edge);
       } else {
         tensor_name_attr =
-            strings::StrCat("edge_", edge->id(), "_", edge->src()->name());
+            strings::StrCat("edge_", edge->id(), "_", edge->src()->name(),
+            "_", edge->dst()->name());
       }
```

But a more approriate and generic fix might be to have a unique src node name. The source name will be like input_xxxxxx and input_yyyyyy. But this sounds like a fundamental change and I'm not sure if this would break too many things. And I'm not sure about the right place to make the change, like manipulating the input name a little bit [here](https://github.com/tensorflow/tensorflow/blob/master/tensorflow/core/common_runtime/eager/execute.cc#L988).

I'd like to hear your opinion on this problem and I'd like to file a PR if you could point me the right place to apply the fix.

Cheers
Hengwen

### Standalone code to reproduce the issue

A minimum case derived from the [DLRM from NV DLE](https://github.com/NVIDIA/DeepLearningExamples/tree/3f82b7f982a1d9a9c5392bd110a5b93dfa7e80eb/Tensorflow2/Recommendation/DLRM) 

Click run all and you can reproduce the error in a few seconds.

https://colab.research.google.com/drive/1CUiyuG2Aob-Fj8xllrmbx2w3cNBXFYqU?usp=sharing