diff --git a/src/orion/core/io/database/pickleddb.py b/src/orion/core/io/database/pickleddb.py index 6650daf84a..ca8efd0a17 100644 --- a/src/orion/core/io/database/pickleddb.py +++ b/src/orion/core/io/database/pickleddb.py @@ -101,6 +101,9 @@ def __init__(self, host="", timeout=60, *args, **kwargs): host = DEFAULT_HOST super().__init__(host) + # NOTE: Save the original value of `host`, so this object can be pickled and unpickled on + # different machines more easily if it's a relative path. + self.original_host = host self.host = os.path.abspath(host) self.timeout = timeout @@ -255,6 +258,17 @@ def get_defaults(cls): """ return {"host": DEFAULT_HOST} + def __getstate__(self): + """Return state to be pickled.""" + return self.__dict__.copy() + + def __setstate__(self, state: dict) -> None: + """Restore state from pickled object.""" + self.__dict__.update(state) + # NOTE: `original_host` might not be present when unpickling old databases. + self.original_host = state.setdefault("original_host", self.host) + self.host = os.path.abspath(self.original_host) + local_file_systems = ["ext2", "ext3", "ext4", "ntfs"]