Correct way to parallelize a large program using Ray

Viewed 2365

I have a decently large Python program (~800 lines) which has the following structure:

  • Setup instructions, where I process an input file provided by the user and define variables/objects which will be global to the program execution.
  • Main function, which makes use of the previous setup phase and calls the primary additional functions of the program.
  • Additional functions, which can be primary, in the sense that they are directly called by the main function, or secondary, in the sense that they are only called by primary additional functions.
  • Some final lines of code where I process the result of the main function.

The program is massively parallel as each execution of the main function is independent of the previous and next ones. Therefore, I am using Ray to execute the main function in parallel, over multiple worker nodes in a cluster. The operating system is CentOS Linux release 8.2.2004 (Core) and the cluster executes PBS Pro 19.2.4.20190830141245. I am using Python 3.7.4, Ray 0.8.7 and Redis 3.4.1.

I have the following in the Python script, where foo is the main function:

@ray.remote(memory=2.5 * 1024 * 1024 * 1024)
def foo(locInd):
    # Main function

if __name__ == '__main__':
    ray.init(address='auto', redis_password=args.pw,
             driver_object_store_memory=10 * 1024 * 1024 * 1024)
    futures = [foo.remote(i) for i in zip(*np.asarray(indArr == 0).nonzero())]
    waitingIds = list(futures)
    while len(waitingIds) > 0:
        readyIds, waitingIds = ray.wait(
            waitingIds, num_returns=min([checkpoint, len(waitingIds)]))
        for r0, r1, r2, r3, r4, r5, r6, r7 in ray.get(readyIds):
            # Process results
            indArr[r0[::-1]] = 1
            nodesComplete += 1
    ray.shutdown()

Below are the instructions I use to start Ray

# Head node
/path/to/ray start --head --port=6379 \
--redis-password=$redis_password \
--memory $((120 * 1024 * 1024 * 1024)) \
--object-store-memory $((20 * 1024 * 1024 * 1024)) \
--redis-max-memory $((10 * 1024 * 1024 * 1024)) \
--num-cpus 48 --num-gpus 0

# Worker nodes
/path/to/ray start --block --address=$1 \
--redis-password=$2 --memory $((120 * 1024 * 1024 * 1024)) \
--object-store-memory $((20 * 1024 * 1024 * 1024)) \
--redis-max-memory $((10 * 1024 * 1024 * 1024)) \
--num-cpus 48 --num-gpus 0

Everything runs as expected provided that I work on a sufficiently small dataset. Nevertheless, the execution produces the following warnings

  • 2020-08-17 17:16:44,289 WARNING worker.py:1134 -- Warning: The remote function __main__.foo has size 220019409 when pickled. It will be stored in Redis, which could cause memory issues. This may mean that its definition uses a large array or other object.
  • 2020-08-17 17:17:10,281 WARNING worker.py:1134 -- This worker was asked to execute a function that it does not have registered. You may have to restart Ray.

If I try to run the code on a larger dataset, I get the following error:

Traceback (most recent call last):
  File "/home/157/td5646/.local/lib/python3.7/site-packages/redis/connection.py", line 700, in send_packed_command
    sendall(self._sock, item)
  File "/home/157/td5646/.local/lib/python3.7/site-packages/redis/_compat.py", line 8, in sendall
2020-08-21 14:22:34,226 WARNING worker.py:1134 -- Warning: The remote function __main__.foo has size 898527351 when pickled. It will be stored in Redis, which could cause memory issues. This may mean that its definition uses a large array or other object.
    return sock.sendall(*args, **kwargs)
ConnectionResetError: [Errno 104] Connection reset by peer

During handling of the above exception, another exception occurred:

Traceback (most recent call last):
  File "./Program.py", line 1030, in <module>
    for i in zip(*np.asarray(indArr == 0).nonzero())]
  File "./Program.py", line 1030, in <listcomp>
    for i in zip(*np.asarray(indArr == 0).nonzero())]
  File "/home/157/td5646/.local/lib/python3.7/site-packages/ray/remote_function.py", line 95, in _remote_proxy
    return self._remote(args=args, kwargs=kwargs)
  File "/home/157/td5646/.local/lib/python3.7/site-packages/ray/remote_function.py", line 176, in _remote
    worker.function_actor_manager.export(self)
  File "/home/157/td5646/.local/lib/python3.7/site-packages/ray/function_manager.py", line 152, in export
    "max_calls": remote_function._max_calls
  File "/home/157/td5646/.local/lib/python3.7/site-packages/redis/client.py", line 3023, in hmset
    return self.execute_command('HMSET', name, *items)
  File "/home/157/td5646/.local/lib/python3.7/site-packages/redis/client.py", line 877, in execute_command
    conn.send_command(*args)
  File "/home/157/td5646/.local/lib/python3.7/site-packages/redis/connection.py", line 721, in send_command
    check_health=kwargs.get('check_health', True))
  File "/home/157/td5646/.local/lib/python3.7/site-packages/redis/connection.py", line 713, in send_packed_command
    (errno, errmsg))
redis.exceptions.ConnectionError: Error 104 while writing to socket. Connection reset by peer.

I am obviously doing something wrong regarding how I describe the program to Ray. I have Scipy Interpolator objects which I consider to be global, but, as pointed out already in this GitHub thread, I should call ray.put on them. The problem there is that I run into these ValueError: buffer source array is read-only which I have no idea how to diagnostic. Also, I am unsure if I should decorate all functions with @ray.remote or only the main function. I guess I could do @ray.remote(num_cpus=1) for all additional functions, as it really only should be the main function that is executed in parallel, but I do not know if that makes sense.

Any help is greatly appreciated, and I am happy to provide more information if required.

1 Answers

I have potentially solved my problems, but I would not mind having someone else's opinion, as my Ray knowledge is indeed limited. Also, I guess this can potentially help others who are experiencing similar troubles (hopefully, I am not alone!).

As I mentioned in the question, the program ran decently fine for sufficiently small datasets (even though it seemed to bypass a few aspects of Ray's logic), but it eventually collapsed on large datasets. Using only Ray tasks, I did not manage to call Scipy Interpolator objects stored in the Object Store (ValueError: buffer source array is read-only), and decorating all functions did not make sense as there is really only the main one that should execute concurrently (while calling other functions).

As a result, I have decided to change the structure of the program to use Ray Actors. Setup instructions are now part of the __init__ method. In particular, Scipy Interpolator objects are defined within this method and set as attributes of self, as global variables are. Most functions (including the main fucntion) have become class methods, with the exceptions of functions that are compiled via Numba. For the latter, they are still independent functions decorated with @jit, but each of them now has an equivalent wrapper method within the class which calls the jitted function.

To have my program execute my now main method in parallel, I rely on an ActorPool. I create as many actors as I have available CPUs, and each actor executes the main method, successfully calling both methods and Numba-compiled functions while also managing to access the Interpolator objects. I only apply @ray.remote to the defined Python class. All of this translates to the following structure:

@ray.remote
class FooClass(object):
    def __init__(self, initArgs):
        # Initialisation

    @staticmethod
    def exampleStaticMethod(args):
        # Processing
        return

    def exampleMethod(self, args):
        # Processing
        return

    def exampleWrapperMethod(self, args):
        return numbaCompiledFunction(args)

    def mainMethod(self, poolMapArgs):
        # Processing
        return


@jit
def numbaCompiledFunction(args):
    # Processing
    return


ray.init(address='auto', redis_password=redPass)
actors = []
for actor in range(int(ray.cluster_resources()['CPU'])):
    actors.append(FooClass.remote(initArgs))
pool = ActorPool(actors)
for unpackedTuple in pool.map_unordered(
        lambda a, v: a.mainMethod.remote(v),
        poolMapArgs):
    # Processing
ray.shutdown()

This successfully runs on 192 CPUs distributed on 4 nodes, without any Warning or Error.

Related