Speed up client_test with sharding.
Speed up client_test with sharding.

PiperOrigin-RevId: 313359020
Change-Id: I8ec2ab94104f6ea4c77437d32cb3c6a0b72636c8

deepmind/reverb,deepmind/reverb,deepmind/reverb,deepmind/reverb