Jax + Ray[Core]

[GJBoth]

What happened + What you expected to happen

I'm trying to run Jax inside a ray actor, but since Jax is multi threaded it's throwing errors due to ray's use of os.fork:

RuntimeWarning: os.fork() was called. os.fork() is incompatible with multithreaded code, and JAX is multithreaded, so this will likely lead to a deadlock.

Is there some way to use 'spawn' or 'forkserver' strategy, or am I missing something else?

Versions / Dependencies

ray 2.9.3
Jax 0.4.25

Reproduction script

import ray
import jax.random as jr

@ray.remote
class RemoteEnv(object):
    def __init__(self):
        key = jr.PRNGKey(42)
        x = jr.normal(key, (1, )

Issue Severity

High: It blocks me from completing my task.

[jjyao]

Assigning P0 assumption basic Jax + Ray doesn't work. If it turns out not to be the case, we can downgrade.

[rynewang]

can't repro. I made a jax script to create random matrices, and dot product, and had no problems.

import ray
import jax.random as jr
import jax.numpy as jnp

@ray.remote
class RemoteEnv(object):
    def __init__(self):
        self.key = jr.PRNGKey(42)

    def create_matrix(self, shape):
        new_key, self.key = jr.split(self.key)
        return jr.normal(new_key, shape)

    def dot(self, mat1, mat2):
        result = jnp.dot(mat1, mat2)
        return result

import numpy as np
ray.init()
remote_env = RemoteEnv.remote()
mat1_ref = remote_env.create_matrix.remote((20,50))
mat2_ref = remote_env.create_matrix.remote((50, 60))
result_ref = remote_env.dot.remote(mat1_ref, mat2_ref)
mat1, mat2, result = ray.get([mat1_ref, mat2_ref, result_ref])

for m in (mat1, mat2, result):
    print(f"{m.shape}")
    print(jnp.sum(m))

result:

(20, 50)
-36.139606
(50, 60)
37.13751
(20, 60)
-41.57849

can you share a repro script which can print that warning?

[GJBoth]

Hi @rynewang I just tried with your example and it prints the warning for me. Just to clarify: the code does run - but I'm not sure how to check for deadlocks.

I tried on two machines (x86 ubuntu and m1 osx) and both gave the error - I also tried with both python 3.10 and 3.11 and again, errors on both. Are you sure you're matching my jax / ray versions? EDIT: I'm running Jax on CPU here - maybe that explains the difference?

An old (but related) issue: google/jax#1805

[rynewang]

Repro'd the warnings on my laptop.

It turns out to be the process creations for the ray clusters (raylets, gcs, ...). It should not have any "deadlocks" with Jax since we will execve Ray binaries right after the fork.

For the warning, you can move ray.init() up before any jax imports and it should work. Like this: (see line 2)

import ray
ray.init()

import jax.random as jr
import jax.numpy as jnp

@ray.remote
class RemoteEnv(object):
    def __init__(self):
        self.key = jr.PRNGKey(42)

    def create_matrix(self, shape):
        new_key, self.key = jr.split(self.key)
        return jr.normal(new_key, shape)

    def dot(self, mat1, mat2):
        result = jnp.dot(mat1, mat2)
        return result

import numpy as np

remote_env = RemoteEnv.remote()
mat1_ref = remote_env.create_matrix.remote((20,50))
mat2_ref = remote_env.create_matrix.remote((50, 60))
result_ref = remote_env.dot.remote(mat1_ref, mat2_ref)
mat1, mat2, result = ray.get([mat1_ref, mat2_ref, result_ref])

for m in (mat1, mat2, result):
    print(f"{m.shape}")
    print(jnp.sum(m))

[GJBoth]

Great, thanks for the quick help! Given the label I'll keep the issue open, but feel free to close it if I'm wrong :-)