-
Notifications
You must be signed in to change notification settings - Fork 1.2k
New issue
Have a question about this project? Sign up for a free GitHub account to open an issue and contact its maintainers and the community.
By clicking “Sign up for GitHub”, you agree to our terms of service and privacy statement. We’ll occasionally send you account related emails.
Already on GitHub? Sign in to your account
cannot pickle '_thread.RLock' object when save Agent Object to Redis #1170
Comments
To save an
Here's an example: import pickle
import redis
import uuid
class Agent:
def __init__(self, dfs, config):
self.dfs = dfs
self.config = config
self.conversation_id = uuid.uuid4()
self.logger = self.initialize_logger()
self.context = self.initialize_context()
self.pipeline = self.initialize_pipeline()
def initialize_logger(self):
return Logger(save_logs=self.config.get('save_logs', False), verbose=self.config.get('verbose', False))
def initialize_context(self):
return PipelineContext(dfs=self.dfs, config=self.config)
def initialize_pipeline(self):
return GenerateChatPipeline(self.context, self.logger)
def __getstate__(self):
state = self.__dict__.copy()
del state['logger']
del state['context']
del state['pipeline']
return state
def __setstate__(self, state):
self.__dict__.update(state)
self.logger = self.initialize_logger()
self.context = self.initialize_context()
self.pipeline = self.initialize_pipeline()
def save_to_redis(self, redis_client, key):
serialized_agent = pickle.dumps(self)
redis_client.set(key, serialized_agent)
@classmethod
def load_from_redis(cls, redis_client, key):
serialized_agent = redis_client.get(key)
if serialized_agent is None:
return None
return pickle.loads(serialized_agent)
def chat(self, message):
pass
def get_redis_connection():
return redis.StrictRedis(host='localhost', port=6379, db=0)
df = Agent([df], config={"llm": llm})
df.chat('Which are the 5 happiest countries?')
def save_agent_to_redis(agent, conversation_id):
redis_conn = get_redis_connection()
agent.save_to_redis(redis_conn, conversation_id)
def load_agent_from_redis(conversation_id):
redis_conn = get_redis_connection()
return Agent.load_from_redis(redis_conn, conversation_id)
save_agent_to_redis(df, str(df.conversation_id))
loaded_agent = load_agent_from_redis(str(df.conversation_id))
loaded_agent.chat('Which are the 5 happiest countries?') This approach ensures that the
|
the Agent class is part of a third-party library (pandasai), I cannot directly modify its code to implement custom serialization and deserialization methods. import dill try: and detection But I checked in the Agent class, and in the Pandas AI REPO, but couldn't find any specific properties of Agent. Can you help me find it correctly? @dosu |
The issue you're encountering with the
Here’s an example implementation: import pickle
from pandasai.agent.base import Agent
class AgentWrapper:
def __init__(self, agent: Agent):
self.agent = agent
def __getstate__(self):
state = self.__dict__.copy()
# Remove non-picklable attributes
if hasattr(self.agent, 'context') and hasattr(self.agent.context, 'dfs'):
dfs = self.agent.context.dfs
if isinstance(dfs, list):
for df in dfs:
if hasattr(df, 'connection'):
del df.connection
elif hasattr(dfs, 'connection'):
del dfs.connection
return state
def __setstate__(self, state):
self.__dict__.update(state)
# Re-instantiate the non-picklable attributes if needed
if hasattr(self.agent, 'context') and hasattr(self.agent.context, 'dfs'):
dfs = self.agent.context.dfs
if isinstance(dfs, list):
for df in dfs:
if hasattr(df, 'reconnect'):
df.connection = df.reconnect()
elif hasattr(dfs, 'reconnect'):
dfs.connection = dfs.reconnect()
# Example usage
agent = Agent(dfs=[...], config=...)
wrapper = AgentWrapper(agent)
# Serialize the wrapper object
serialized_wrapper = pickle.dumps(wrapper)
# Deserialize the wrapper object
deserialized_wrapper = pickle.loads(serialized_wrapper) In this solution:
Replace the placeholder
|
I checked but I don't see the source code mentioning the "connection" attribute belong to of Can you provide me with the relevant source codes to verify the problem @dosu |
To handle the non-picklable import pickle
from pandasai.agent.base import Agent
class AgentWrapper:
def __init__(self, agent: Agent):
self.agent = agent
def __getstate__(self):
state = self.__dict__.copy()
# Remove the non-picklable attributes
if hasattr(self.agent.context.dfs, 'connection'):
del self.agent.context.dfs.connection
return state
def __setstate__(self, state):
self.__dict__.update(state)
# Reinitialize the non-picklable attributes if necessary
# self.agent.context.dfs.connection = <reinitialize_connection_here>
def save_to_redis(self, redis_client, key):
serialized_agent = pickle.dumps(self)
redis_client.set(key, serialized_agent)
@classmethod
def load_from_redis(cls, redis_client, key):
serialized_agent = redis_client.get(key)
return pickle.loads(serialized_agent)
# Usage example
# agent = Agent(...)
# wrapper = AgentWrapper(agent)
# wrapper.save_to_redis(redis_client, 'agent_key')
# loaded_wrapper = AgentWrapper.load_from_redis(redis_client, 'agent_key')
# loaded_agent = loaded_wrapper.agent This code defines an Additionally, there are other non-picklable attributes within the import pickle
from pandasai.agent.base import Agent
class AgentWrapper:
def __init__(self, agent: Agent):
self.agent = agent
def __getstate__(self):
state = self.__dict__.copy()
# Remove the non-picklable attributes
if hasattr(self.agent.context.dfs, 'connection'):
del self.agent.context.dfs.connection
if hasattr(self.agent, 'logger'):
del self.agent.logger
if hasattr(self.agent, 'pipeline'):
del self.agent.pipeline
if hasattr(self.agent.context.config, 'llm'):
del self.agent.context.config.llm
return state
def __setstate__(self, state):
self.__dict__.update(state)
# Reinitialize the non-picklable attributes if necessary
# self.agent.context.dfs.connection = <reinitialize_connection_here>
# self.agent.logger = <reinitialize_logger_here>
# self.agent.pipeline = <reinitialize_pipeline_here>
# self.agent.context.config.llm = <reinitialize_llm_here>
def save_to_redis(self, redis_client, key):
serialized_agent = pickle.dumps(self)
redis_client.set(key, serialized_agent)
@classmethod
def load_from_redis(cls, redis_client, key):
serialized_agent = redis_client.get(key)
return pickle.loads(serialized_agent)
# Usage example
# agent = Agent(...)
# wrapper = AgentWrapper(agent)
# wrapper.save_to_redis(redis_client, 'agent_key')
# loaded_wrapper = AgentWrapper.load_from_redis(redis_client, 'agent_key')
# loaded_agent = loaded_wrapper.agent This extended
|
@dobahoang - I managed to change this a bit and got it working for my use case. Maybe this will help: import pickle
from pandasai import Agent
from pandasai.helpers.cache import Cache
class AgentWrapper:
def __init__(self, agent: Agent):
self.agent = agent
def __getstate__(self):
state = self.__dict__.copy()
self.remove_unpicklable()
return state
def __setstate__(self, state):
self.__dict__.update(state)
def remove_unpicklable(self):
# Remove the non-picklable attributes
if hasattr(self.agent.context, 'cache'):
del self.agent.context.cache
if hasattr(self.agent, '_vectorstore'):
del self.agent._vectorstore
if hasattr(self.agent.context, 'vectorstore'):
del self.agent.context.vectorstore
if hasattr(self.agent.context.config, 'llm'):
del self.agent.context.config.llm
@classmethod
def restore_unpicklable(cls, agent, llm, vector_store=None):
# Reinitialize the non-picklable attributes if necessary
if agent.context.config.enable_cache:
agent.context.cache = Cache()
else:
agent.context.cache = None
agent._vectorstore = vector_store
agent.context.vectorstore = vector_store
agent.context.config.llm = llm
return agent
def save_to_redis(self, key, redis_client):
serialized_agent = pickle.dumps(self)
redis_client.set(key, serialized_agent)
@classmethod
def load_from_redis(cls, key, redis_client, llm, vector_store=None):
serialized_agent = redis_client.get(key)
wrapper = pickle.loads(serialized_agent)
wrapper.agent = cls.restore_unpicklable(wrapper.agent, llm, vector_store)
return wrapper.agent
def save_to_pkl(self, key):
self.remove_unpicklable()
with open(key, 'wb') as f:
pickle.dump(self.agent, f)
@classmethod
def load_from_pkl(cls, key, llm, vector_store=None):
with open(key, 'rb') as f:
agent = pickle.load(f)
agent = cls.restore_unpicklable(agent, llm, vector_store)
return agent p.s. I havent tried writing to Redis with these functions but writing and restoring from a dict worked. So I think it shouldnt be a problem. |
System Info
I use last version pandas AI,
OS: WINDOWS
🐛 Describe the bug
How can I save Object Agent to Redis successfully?
My goal is to serve multiple users, each user has its own context. So the idea is that when the user sends the conversation_id, IT WILL load in Redis, to get the correct Agent object, then I can call the function chat(). for example :
df = Agent([df], config={"llm": llm})
df.chat('Which are the 5 happiest countries?')
def save_agent_to_redis(agent, conversation_id):
"""
Serialize and save the Agent object to Redis.
"""
redis_conn = get_redis_connection()
agent_bytes = pickle.dumps(agent)
redis_conn.set(conversation_id, agent_bytes)
but when I try to save to Redis, the error I get is:
The error TypeError: cannot pickle '_thread.RLock' object occurs when you try to pickle (serialize) an object that contains a thread lock or other non-picklable objects. Pickling is the process of converting a Python object into a byte stream that can be saved to a file or sent over a network.( in my case, i need to save to Redis)
I have used some libraries specialized for serialization (pickle, dill) but without success
The Agent class is also quite special, it will automatically create a Conversation_id every time a new Agent is created, we do not control the Conversation_id, so if we choose to serialize some properties and ignore the remaining properties, then recreate them, will generate a different conversation_id than before
Can anyone give me suggestions, please
The text was updated successfully, but these errors were encountered: