Why not making ReplayBuffers serializabe?

[jamartinh]

It would be very good to make replay Buffers serializable.

For instance, how can I pass and share buffer using ray library?
I tried to use the MPreplaybuffer but ray says it is not serializabe.

Perhaps the serialization/deserialization of the MPrepayBuffer can be done easy as just passing the neccesary info to recreate a Buffer pointing to the shared memory.

[ymd-h]

@jamartinh
Thank you for your feedback!

I am not sure about usage of ray, however, MPReplayBuffer and MPPrioritizedReplayBuffer can be pickled.
We can pass them to standard multiprocessing.Process etc. (Both fork and spawn methods are supported.)

If you give us a sample code, OS, and versions of libraries, we can investigate it.

---

FYI: Here are the internal implementations for pickling.

cpprb/cpprb/PyReplayBuffer.pyx

Lines 463 to 464 in d513587

	def __reduce__(self):
	return (SharedBuffer,(self.view.shape,self.dtype,self.data))

cpprb/cpprb/PyReplayBuffer.pyx

Lines 2138 to 2142 in d513587

	def __reduce__(self):
	return (ThreadSafePrioritizedSampler,
	(self.size,self.alpha,self.eps,self.max_p,
	self.sum,self.sum_a,
	self.min,self.min_a))

[jamartinh]

Hi, many thanks @ymd-h

I was reading of an apparent very usefull feature of multiprocesing shared memory:

https://docs.python.org/3/library/multiprocessing.shared_memory.html#multiprocessing.shared_memory.SharedMemory

In particular:

"Creates a new shared memory block or attaches to an existing shared memory block. Each shared memory block is assigned a unique name. In this way, one process can create a shared memory block with a particular name and a different process can attach to that same shared memory block using that same name."

[ymd-h]

@jamartinh
Great thanks!

The SharedMemory seems to be promising.
I will investigate more.

The weakness is that the module was introduced at Python 3.8, however, we still support older version.

At least Python 3.7 is still active.
https://www.python.org/downloads/

[jamartinh]

I will also try to make some experiments

[ymd-h]

I tested and found that SharedMemory can be passed to functions wrapped by@ray.remote directly.

In the following test code, the both cases works fine.

import multiprocessing
from multiprocessing.shared_memory import SharedMemory

import numpy as np

import ray


def ray_test():
    ray.init()

    shm = SharedMemory(create=True, size=32 * 3)
    a = np.ndarray(shape=(3,), dtype=np.int32, buffer=shm.buf)
    print(a)

    @ray.remote
    def add(name, shape, dtype):
        m = SharedMemory(name=name)
        b = np.ndarray(shape=shape, dtype=dtype, buffer=m.buf)
        print(b)
        b += 2
        print(b)

    @ray.remote
    def add_shm(shm, shape, dtype):
        b = np.ndarray(shape=shape, dtype=dtype, buffer=shm.buf)
        print(b)
        b += 2
        print(b)

    ray.get(add.remote(shm.name, a.shape, a.dtype))
    print(a)

    ray.get(add_shm.remote(shm, a.shape, a.dtype))
    print(a)

    shm.close()
    shm.unlink()

if __name__ == "__main__":
    ray_test()

Its __reduce__ implementation pass the name.
https://github.com/python/cpython/blob/1d81fdc4c004511c25f74db0e04ddbbb8a04ce6d/Lib/multiprocessing/shared_memory.py#L188-L196

[ymd-h]

I realized we need Lock as well as shared memory to manage critical section.

Unfortunately, Lock is not supported in Ray, ether.
Maybe we can implement custom lock object on shared memory in principle...

[ymd-h]

@jamartinh
As we discussed at another thread, I'm investigating and try to solve the challenges, however, the goal is still far away.
Meanwhile, I propose a workaround to use with Ray.

from multiprocessing.shared_memory import SharedMemory
import time

from cpprb import PrioritizedReplayBuffer
import numpy as np
import ray


@ray.remote
def worker(rb, v, shm):
    """
    Worker

    Notes
    -----
    We must pass SharedMemory directly. When we pass np.ndarray,
    the array is copied and doesn't point to shared memory any more.
    """
    done = np.ndarray(shape=tuple(), dtype=np.int32, buffer=shm.buf)
    while not done:
        print(done)
        rb.add.remote(a=v)
        time.sleep(1)

    return None


@ray.remote
class RemoteReplayBuffer:
    """
    Wrapper for Replay Buffer

    Notes
    -----
    1. All method calls are executed serially.

    2. We cannot pass the buffer class directly.

    >>> rb = ray.remote(PrioritizedReplayBuffer).remote(buffer_size, env_dict)
    TypeError: __cinit__() takes at least 1 positional argument (0 given)
    """
    def __init__(self, *args, **kwargs):
        self.rb = PrioritizedReplayBuffer(*args, **kwargs)

    def add(self, **kwargs):
        return self.rb.add(**kwargs)

    def sample(self, *args, **kwargs):
        return self.rb.sample(*args, **kwargs)

    def update_priorities(self, *args, **kwargs):
        return self.rb.update_priorities(*args, **kwargs)

    def get_stored_size(self):
        return self.rb.get_stored_size()

    def get_all_transitions(self):
        return self.rb.get_all_transitions()

def run():
    buffer_size = 32
    env_dict = {"a": {}}
    alpha = 0.5

    shm = SharedMemory(create=True, size=32)
    try:
        done = np.ndarray(shape=tuple(), dtype=np.int32, buffer=shm.buf)
        done[...] = 0

        ray.init()

        rb = RemoteReplayBuffer.remote(buffer_size, env_dict, alpha=alpha)

        w1 = worker.remote(rb, 1, shm)
        w2 = worker.remote(rb, np.asarray([2, 3]), shm)

        while True:
            stored_size = ray.get(rb.get_stored_size.remote())
            print(stored_size)
            if stored_size < 20:
                time.sleep(1)
            else:
                break


        done[...] = 1
        ray.get([w1, w2])

        print(ray.get(rb.get_stored_size.remote()))
        print(ray.get(rb.get_all_transitions.remote()))

    finally:
        # To avoid (shared) memory leak, we must close() and unlink().
        # On Linux, you might find shared memory file at /dev/shm
        shm.close()
        shm.unlink()

if __name__ == "__main__":
    run()

[jamartinh]

Hi @ymd-h tanks! I think I am close!
The idea is just to specify the serialization protocol for the cython classes droping all non-picackle objects and reconstructing them in set_state

Also, changing all syncrhonization types to memoryManager ones. that should work.

[jamartinh]

It is working for me now, for simple MPReplayBuffer with Ray.

I have pushed to a draft pool request the code: #19

For this to work, you either have to put in every actor __init__:

authkey = b'abc'
multiprocessing.current_process().authkey = authkey

and pass the global_replay_bufffer during a set_method(global_replay_buffer)

or do:

def auth_fn(*args):
    authkey = b"abc"
    multiprocessing.current_process().authkey = authkey

And then after that:

ray.worker.global_worker.cached_functions_to_run.append(auth_fn)

Before calling ray.init

All the history resumes to make the unpickle (deserialization) after setting the authkey in each python process that does not inherit from fork or spawn

[ymd-h]

@jamartinh
Thank you for waiting.

Based on your work, we finally released cpprb v10.6.0.
MPReplayBuffer and MPPrioritizedReplayBuffer can be used with Ray now.

After some design modification, these buffers take new construction parameters for context (ctx) and backend (backend).

import base64
import multiprocessing as mp

from cpprb import MPReplayBuffer
import ray

ray.init()

encoded = base64.b64encode(mp.current_process().authkey)
def auth_fn(*args):
    mp.current_process().authkey = base64.b64decode(encoded)
ray.worker.global_worker.run_function_on_all_workers(auth_fn)

buffer_size = 1e+6
m = mp.get_context().Manager()

# Use `SyncManager` as context, "SharedMemory" as backend
rb = MPReplayBuffer(buffer_size, {"done": {}}, ctx=m, backend="SharedMemory")

Please see doc and example.

Thank you for your great contribution.

[jamartinh]

Thanks @ymd-h !!!

[ymd-h]

I found Ray's (private) method run_function_on_all_workers() was deprecated and will be deleted at version 2.4.
ray-project/ray#31528

We re-considered the usage.
Now, we recommend to use Ray Actors;
https://docs.ray.io/en/latest/ray-core/actors.html

@ray.remote
class RemoteWorker:
    # Encode base64 to avoid following error:
    #   TypeError: Pickling an AuthenticationString object is disallowed for security reasons
    encoded = base64.b64encode(mp.current_process().authkey)

    def __init__(self):
        # Set up 'authkey' to communicate with `SyncManager`.
        # Important: Do not pass `MPReplayBuffer` here, because it is not ready.
        mp.current_process().authkey = base64.b64decode(self.encoded)

    def run(self, rb):
        pass

w = RemoteWorker.remote()
w.run.remote(rb)

We updated the example, too.
https://ymd_h.gitlab.io/cpprb/examples/mp_with_ray/