fix redis pool

modelscope · Sep 9, 2024 · 8c1a668 · 8c1a668
1 parent acd4b77
commit 8c1a668
Showing 1 changed file with 3 additions and 4 deletions.
diff --git a/src/agentscope/server/async_result_pool.py b/src/agentscope/server/async_result_pool.py
@@ -118,10 +118,8 @@ def prepare(self) -> int:
         return self._get_object_id()
 
     def set(self, key: int, value: bytes) -> None:
-        pipe = self.pool.pipeline()
-        pipe.set(key, value, ex=self.max_timeout)
-        pipe.rpush(RedisPool.TASK_QUEUE_PREFIX + str(key), key)
-        pipe.execute()
+        self.pool.set(key, value, ex=self.max_timeout)
+        self.pool.rpush(RedisPool.TASK_QUEUE_PREFIX + str(key), key)
 
     def get(self, key: int) -> bytes:
         result = self.pool.get(key)
@@ -132,6 +130,7 @@ def get(self, key: int) -> bytes:
                 keys=RedisPool.TASK_QUEUE_PREFIX + str(key),
                 timeout=self.max_timeout,
             )
+            self.pool.rpush(RedisPool.TASK_QUEUE_PREFIX + str(key), key)
             if keys is None:
                 raise ValueError(
                     f"Waiting timeout for async result of task[{key}]",