[PJRT:GPU] Add setting for mocked number of hosts per slice

[jaro-sevcik]

With the existing enable_mock_nccl setting it is impossible to warm up compilation cache when there are multiple processes per node. This is because the cache key includes topology and GPU topology contains information about number of slices and number of hosts per slice. The current mocking of topologies always sets num_hosts_per_slice to 1. However, if you have multiple GPUs on a node and run a process-per-GPU then num_hosts_per_slice must be set to the number of GPUs.

This patch allows setting num_hosts_per_slice explicitly when creating the GPU client.

[cheshire]

Looks reasonable to me, but I think @hawkinsp should LGTM as well.

[penpornk]

@hawkinsp Could you please help take a look?

[sgerrard]

@hawkinsp A friendly nudge to take a look, your review has been requested for this PR.

[jaro-sevcik]

@cheshire I have updated the patch to accept a parameter of the form mock_gpu_topology=4x2. Is this what you had in mind?

[jaro-sevcik]

Are you asking how it's going to be exposed in JAX? Here is a draft PR: https://github.com/google/jax/pull/23374

Or were you asking for some JAX scripts that make use of that option?

[dimitar-asenov]

@hawkinsp Could you please take a look at this PR?

[jaro-sevcik]

@cheshire Could you take another look, please?