Reduce iota_test execution time: shared pjrt client + automatic shard number

[mfrancepillois]

Motivation

The iota_test was very slow on AMD targets (compared to NVDIA) for 2 main reasons:

• The number of shards was fixed to 50, which led to resource contention and unnecessary context switches when running on node that does not have that amount of GPU available. => We therefore defined the shard number automatically based on the number of available GPU. On H100 node this reduces the execution time of iota from about 700 seconds to ~200 seconds.
• The pjrt client was destroyed and recreated for each of the 4500 tests that make up the iota_test. This task in ROCm, which calls hipCtxDestroy() + hipCtxCreate(), was ~40× slower than with CUDA (see table below). => We therefore modified the code to create only 1 pjrt client per process and share it among the tests within that process.

Phase	H100	AMD MI300X
Previous client teardown + new client init (pre-BFC log)	~35ms total	~963ms total
BFC allocator re-setup (8 GPUs)	~0.3ms	~0.1ms
Per-test GPU lifecycle cost	35ms	1009ms

This PR re-enables iota_test in AMD CIs

[i-chaochen]

this PR should be against to the main branch instead of 0.9.2, isn't?

[mfrancepillois]

this PR should be against to the main branch instead of 0.9.2, isn't?

Another PR has been open against the main rocm branch: #852
This PR can therefore be closed.