Use tfrt cpu client

[darisoy]

• Add the CPU_ASYNC_CLIENT env var
• Change PjRtComputationClient to use GetTfrtCpuClient when testing multiple processes with a CPU
• Add unit tests

[will-cromar]

LGTM. Thanks!

[JackCaoG]

error seems real

test (__main__.TestParallelTensorMNIST) ... 2022-08-17 02:23:06.009032: E  498776 ./tensorflow/compiler/xla/service/collective_ops_utils.h:224] This thread has been waiting for 5000ms for and may be stuck: participant AllReduceParticipantData{buffers=[{element_count=250},{element_count=10},{element_count=5000},{element_count=20},{element_count=16000},{element_count=50},{element_count=500},{element_count=10},{element_count=1}], rendezvous_key=RendezvousKey{run_id=RunId: 564, global_devices=[0,1,2,3], num_local_participants=4, collective_op_kind=cross_replica, op_id=246}, device_ordinal=0, stream=(nil)} waiting for all participants to arrive at rendezvous RendezvousKey{run_id=RunId: 564, global_devices=[0,1,2,3], num_local_participants=4, collective_op_kind=cross_replica, op_id=246}
2022-08-17 02:23:06.021322: E  498773 ./tensorflow/compiler/xla/service/collective_ops_utils.h:224] This thread has been waiting for 5000ms for and may be stuck: participant AllReduceParticipantData{buffers=[{element_count=250},{element_count=10},{element_count=5000},{element_count=20},{element_count=16000},{element_count=50},{element_count=500},{element_count=10},{element_count=1}], rendezvous_key=RendezvousKey{run_id=RunId: 566, global_devices=[0,1,2,3], num_local_participants=4, collective_op_kind=cross_replica, op_id=247}, device_ordinal=2, stream=(nil)} waiting for all participants to arrive at rendezvous RendezvousKey{run_id=RunId: 566, global_devices=[0,1,2,3], num_local_participants=4, collective_op_kind=cross_replica, op_id=247}
2022-08-17 02:23:06.037646: E  498775 ./tensorflow/compiler/xla/service/collective_ops_utils.h:224] This thread has been waiting for 5000ms for and may be stuck: participant AllReduceParticipantData{buffers=[{element_count=250},{element_count=10},{element_count=5000},{element_count=20},{element_count=16000},{element_count=50},{element_count=500},{element_count=10},{element_count=1}], rendezvous_key=RendezvousKey{run_id=RunId: 568, global_devices=[0,1,2,3], num_local_participants=4, collective_op_kind=cross_replica, op_id=248}, device_ordinal=3, stream=(nil)} waiting for all participants to arrive at rendezvous RendezvousKey{run_id=RunId: 568, global_devices=[0,1,2,3], num_local_participants=4, collective_op_kind=cross_replica, op_id=248}
2022-08-17 02:23:06.045691: E  498774 ./tensorflow/compiler/xla/service/collective_ops_utils.h:224] This thread has been waiting for 5000ms for and may be stuck: participant AllReduceParticipantData{buffers=[{element_count=250},{element_count=10},{element_count=5000},{element_count=20},{element_count=16000},{element_count=50},{element_count=500},{element_count=10},{element_count=1}], rendezvous_key=RendezvousKey{run_id=RunId: 570, global_devices=[0,1,2,3], num_local_participants=4, collective_op_kind=cross_replica, op_id=249}, device_ordinal=1, stream=(nil)} waiting for all participants to arrive at rendezvous RendezvousKey{run_id=RunId: 570, global_devices=[0,1,2,3], num_local_participants=4, collective_op_kind=cross_replica, op_id=249}

[will-cromar]

Collectives hang forever on CPU with multiple devices. This test passed before because we ignored CPU_NUM_DEVICES. @darisoy is investigating if there is an easy fix. Otherwise, we can implement a workaround to go back to testing on one CPU device.

[will-cromar]

Thanks!