fail safely when pickling (#4085)

* fail safely when pickling

* why did i do that
This commit is contained in:
Khaleel Al-Adhami 2024-10-07 11:59:33 -07:00 committed by Masen Furer
parent 08a493882d
commit 3810161873
No known key found for this signature in database
GPG Key ID: B0008AD22B3B3A95

View File

@ -1958,7 +1958,14 @@ class BaseState(Base, ABC, extra=pydantic.Extra.allow):
Returns:
The serialized state.
"""
try:
return pickle.dumps((state_to_schema(self), self))
except pickle.PicklingError:
console.warn(
f"Failed to serialize state {self.get_full_name()} due to unpicklable object. "
"This state will not be persisted."
)
return b""
@classmethod
def _deserialize(
@ -2812,6 +2819,7 @@ class StateManagerDisk(StateManager):
self.states[substate_token] = substate
state_dilled = substate._serialize()
if state_dilled:
if not self.states_directory.exists():
self.states_directory.mkdir(parents=True, exist_ok=True)
self.token_path(substate_token).write_bytes(state_dilled)
@ -3107,6 +3115,7 @@ class StateManagerRedis(StateManager):
if state._get_was_touched():
pickle_state = state._serialize()
self._warn_if_too_large(state, len(pickle_state))
if pickle_state:
await self.redis.set(
_substate_key(client_token, state),
pickle_state,