the conflict between the jax after 0.4.16 and ray

[Jeasuiiya]

service_server = xla_client._xla.get_distributed_runtime_service(
            server_address, self.num_hosts, use_coordination_service=False)

When JAX upgrade to version 0.4.16, the option "--jax_coordination_service" will be removed. Once I attempt to use Ray for distributed training, the following situation will occur.
I think the problem is caused by the conflict between the coordination services of Ray and Jax.I want to set the use_coordination_service to False,but in the new version,this parameter will be set as true by default.

How can i solve this problem?HELPPPPP!!!!!!

[jakevdp]

Hello - thanks for the question! I'm not sure about compatibility between JAX and Ray. I'm not entirely familiar with Ray's execution model, but in most cases it's probably better to use JAX's native multi-device/sharding capabilities than to try to layer another multi-process or multi-thread solution on top of it. In particular, JAX will not work well with systems that launch multiple threads (see https://docs.jax.dev/en/latest/concurrency.html for some discussion).

You may find more help on this from the Ray community: it looks like there are a number of posts there where people are discussing compatibility between the projects (e.g. https://discuss.ray.io/t/interaction-between-ray-and-jax/7311, https://discuss.ray.io/t/ray-train-and-sharded-models-with-jax-possible-desirable/13128, ray-project/ray#44087).

Hope that hleps!

[Jeasuiiya]

Sorry, it's my fault. The problem should not be caused by ray. I'm currently checking what impact changing "--jax_coordination_service" to true will have. Thank you for your explanation!!