Experiment With Ray Distributed for STUMPED

[seanlaw]

Ray has become more popular over the years but has had some difficulty supporting numba and Apple Silicon chips. However that seems to have changed. STUMPY purposely designed stumped (and other distributed functions) in a way that can be easily adopted for other distributed systems like Ray. The tricky part is figuring out how to not add Ray as a dependency (similar to how Dask is not a STUMPY dependency). This post is to demonstrate/document a minimum viable example of how one might use Ray with STUMPY. Below is an initial set of experimentation steps for running the basic stump (not stumped) function on Ray:

Install Ray on Apple M2

1. Install miniconda - conda create -n ray python=3.11
2. Install STUMPY dependencies (conda.sh)
3. Install STUMPY (setup.sh)
4. conda install -c conda-forge grpcio
5. pip install ray

Decorate STUMPY

Note that we don't need to import ray here!

# foo.py
import stumpy
import inspect


def stump(client, T, m):
    if inspect.ismodule(client) and str(client).startswith("<module 'ray'"):
        ray_stump_func = client.remote(stumpy.stump)
        future = ray_stump_func.remote(T, m)
        return client.get(future)
    elif str(client).startswith("<Client"):
        return stumpy.stumped(client, T, m)
    else:
        raise ValueError("`ray` is not a ray module")

Note that you can gather a set of futures from ray via:

results = client.get([ray_stump_func.remote(T, m) for i in range(4)])

Test Ray and Dask

# test.py
import numpy as np
import ray
from dask.distributed import Client
import foo


if __name__ == "__main__":
    T = np.random.rand(10_000)
    m = 50

    ray.init()  # Initialize the ray cluster
    mp = foo.stump(ray, T, m)
    print(mp)

    dask_client = Client()
    mp = foo.stump(dask_client, T, m)
    print(mp)

    if ray.is_initialized():
        ray.shutdown()  # Tear down the ray cluster

Distributed Process Pool

https://docs.ray.io/en/latest/ray-more-libs/multiprocessing.html#run-on-a-cluster