How to pass the keyword arguments when customizing the serialization by `__reduce__` function?

As code below, I want to serialize the DBConnection by overriding the __reduce__ function.

But the DBConnection class needs multiple args (including keyword args) to initialize. How can I pass these args by serialized_data?

import ray
import sqlite3

ray.init(address='local')

class DBConnection:
    def __init__(self, path, **kwargs):
        self.path = path
        self.kwargs = kwargs
        self.conn = sqlite3.connect(path)
        
    def __repr__(self) -> str:
        return f"id: {id(self)}, path: {self.path}, kwargs: {self.kwargs}"

    # without '__reduce__', the instance is unserializable.
    def __reduce__(self):
        deserializer = DBConnection
        serialized_data = (self.path, self.kwargs) # not work in this way
        return deserializer, serialized_data

original = DBConnection("/tmp/db", first=1)
print(original)

copied = ray.get(ray.put(original))
print(copied)

import ray

ray.init()

class DBConnection:
    def __init__(self, path, **kwargs):
        self.path = path
        self.kwargs = kwargs
        self.conn = sqlite3.connect(path)

    def __repr__(self) -> str:
        return f"id: {id(self)}, path: {self.path}, kwargs: {self.kwargs}"

    # without '__reduce__', the instance is unserializable.
    def __reduce__(self):
        def deserializer(path, kwargs):
            return DBConnection(path, **kwargs)

        serialized_data = (self.path, self.kwargs) # not work in this way
        return deserializer, serialized_data

original = DBConnection("/tmp/db", first=1)
print(original)

copied = ray.get(ray.put(original))
print(copied)

Try this?

1 Like

Wow, it works for me. thanks a lot. :smiling_face_with_three_hearts: