fail safely when pickling (#4085)
* fail safely when pickling * why did i do that
This commit is contained in:
parent
59dd54c049
commit
c7c830de80
@ -1999,7 +1999,14 @@ class BaseState(Base, ABC, extra=pydantic.Extra.allow):
|
||||
Returns:
|
||||
The serialized state.
|
||||
"""
|
||||
return pickle.dumps((self._to_schema(), self))
|
||||
try:
|
||||
return pickle.dumps((self._to_schema(), 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(
|
||||
@ -2826,9 +2833,10 @@ class StateManagerDisk(StateManager):
|
||||
if substate._get_was_touched():
|
||||
substate._was_touched = False # Reset the touched flag after serializing.
|
||||
pickle_state = substate._serialize()
|
||||
if not self.states_directory.exists():
|
||||
self.states_directory.mkdir(parents=True, exist_ok=True)
|
||||
self.token_path(substate_token).write_bytes(pickle_state)
|
||||
if pickle_state:
|
||||
if not self.states_directory.exists():
|
||||
self.states_directory.mkdir(parents=True, exist_ok=True)
|
||||
self.token_path(substate_token).write_bytes(pickle_state)
|
||||
|
||||
for substate_substate in substate.substates.values():
|
||||
await self.set_state_for_substate(client_token, substate_substate)
|
||||
@ -3121,11 +3129,12 @@ class StateManagerRedis(StateManager):
|
||||
if state._get_was_touched():
|
||||
pickle_state = state._serialize()
|
||||
self._warn_if_too_large(state, len(pickle_state))
|
||||
await self.redis.set(
|
||||
_substate_key(client_token, state),
|
||||
pickle_state,
|
||||
ex=self.token_expiration,
|
||||
)
|
||||
if pickle_state:
|
||||
await self.redis.set(
|
||||
_substate_key(client_token, state),
|
||||
pickle_state,
|
||||
ex=self.token_expiration,
|
||||
)
|
||||
|
||||
# Wait for substates to be persisted.
|
||||
for t in tasks:
|
||||
|
Loading…
Reference in New Issue
Block a user