Jax behaviour with callbacks that use multiprocessing.

[miltonllera]

Let's say I have the following code, which performs some computations in staged-out code, but at some point part of the evaluation involves running code on the host:

import os
from time import time

import numpy as np
import jax
import jax.numpy as jnp
import multiprocessing as mp
from itertools import product
from collections import deque

os.environ["OPENBLAS_NUM_THREADS"] = "1"
os.environ["MKL_NUM_THREADS"] = "1"
os.environ["NUMEXPR_NUM_THREADS"] = "1"
os.environ["OMP_NUM_THREADS"] = "1"
os.environ["VECLIB_MAXIMUM_THREADS"] = "1"


class SnippetTimer:
    def __init__(self, snippet_description=None):
        self.start_time = 0
        self.snippet_description = snippet_description

    def __enter__(self):
        self.start_time = time()
        return self.start_time

    def __exit__(self, *args, **kwargs):
        total_time = time() - self.start_time
        if self.snippet_description is None:
            print(f"Total time taken: {total_time}")
        else:
            print(f"Total time taken for snippet '{self.snippet_description}': {total_time}")


N_CPUS = mp.cpu_count() - 1

INNER_SAMPLES = 100
OUTER_SAMPLES = 100

C, H, W = 5, 16, 16

MAX_VALUE = np.finfo(np.float32).max
DIRECTIONS = [(0, 1), (1, 0), (0, -1), (-1, 0)]


def batched_n_islands(int_maps):
    # Note that because the callbacks above are using vectorization, they will sync over all
    # vmaps applied. Thus there could be 1 or 2 leading dimenisions if we are either evaluating
    # a model or if this is called when running evolutionary training, respectively.
    if len(int_maps.shape) == 3:
        batch_shapes = int_maps.shape[:1]
    else:
        batch_shapes = int_maps.shape[:2]
        int_maps = int_maps.reshape((-1, *int_maps.shape[2:]))

    # with SnippetTimer("batched_n_islands"):
    with mp.Pool(N_CPUS) as pool:
        result = pool.imap(n_islands, int_maps, int(np.ceil(len(int_maps) / N_CPUS)))
        result = list(result)

    result = np.concatenate(result)

    return result.reshape(batch_shapes)[..., None]


def n_islands(int_map: np.ndarray, non_traversible_tiles=(0,)):
    h, w = int_map.shape

    int_map = ~np.isin(int_map, non_traversible_tiles)

    if np.all(int_map == 0):
        return np.asarray([h * w], dtype=np.float32)

    visited = np.zeros_like(int_map, dtype=bool)

    def in_bounds(pos):
        return 0 <= pos[0] < h and 0 <= pos[1] < w

    def visit(i, j):
        to_visit = deque()

        visited[i, j] = True
        to_visit.append((i, j))

        while len(to_visit) > 0:
            i, j = to_visit.popleft()
            for di, dj in DIRECTIONS:
                nb = i + di, j + dj
                if in_bounds(nb) and int_map[nb] and not visited[nb]:
                    to_visit.append(nb)
                    visited[nb] = True

    n_components = 0
    for i, j in product(range(h), range(w)):
        if int_map[i, j] and not visited[i, j]:
            visit(i, j)
            n_components += 1

    return np.asarray([n_components], dtype=np.float32)


if __name__ == "__main__":
    def produce_output(key):
        return jnp.argmax(jax.random.normal(key=key, shape=(C, H, W)), axis=0)

    @jax.jit
    def sample_innner(key):
        samples = jax.vmap(produce_output)(jax.random.split(key, INNER_SAMPLES))

        results = jax.pure_callback(
            batched_n_islands,
            jnp.empty((100, 1,)),  # example output with the expected shape and dtype that's required for jit
            samples,
            # pass all elements generated by nested vmaps (i.e. in the current call-stack) to the callback
            # otherwise the function is vmapped over all elments individually.
            vectorized=True
        )

        return results

    @jax.jit
    def sample_outer(key):
        results = jax.vmap(sample_innner)(jax.random.split(key, OUTER_SAMPLES))
        return results

    key = jax.random.PRNGKey(0)

    with SnippetTimer():
        results = sample_outer(key)

My question is, how much thread contention there is between jax and the multiprocessing library and how can I mitigate it. The options above can in principle limit jax's ability to interfere with the allocation of resources in the callback (for this example I get 3s vs 1.3s) but it doesn't seem ideal, or more precisely I am not sure what the trade-offs are (I know one is slower compile time, but that's fine). Is there a way to specify what resources each library can use? Are there any best practices at all?

Thanks in advance!