Upload 3 files

aworld/replay_buffer/storage/multi_proc_mem.py

@@ -0,0 +1,155 @@
+import multiprocessing
+import traceback
+import pickle
+from typing import Dict, List
+from aworld.replay_buffer.base import Storage, DataRow
+from aworld.replay_buffer.query_filter import QueryCondition, QueryFilter
+from aworld.logs.util import logger
+
+
+class MultiProcMemoryStorage(Storage):
+
+    """
+    Memory storage for multi-process.
+    """
+
+    def __init__(self,
+                 data_dict: Dict[str, str],
+                 fifo_queue: List[str],
+                 lock: multiprocessing.Lock,
+                 max_capacity: int = 10000):
+        self._data: Dict[str, str] = data_dict
+        self._fifo_queue = fifo_queue
+        self._max_capacity = max_capacity
+        self._lock = lock
+
+    def _save_to_shared_memory(self, data, task_id):
+        serialized_data = pickle.dumps(data)
+        try:
+            if task_id not in self._data or not self._data[task_id]:
+                shm = multiprocessing.shared_memory.SharedMemory(
+                    create=True, size=len(serialized_data))
+                shm.buf[:len(serialized_data)] = serialized_data
+                self._data[task_id] = shm.name
+                shm.close()
+                return
+            shm = multiprocessing.shared_memory.SharedMemory(
+                name=self._data[task_id], create=False)
+            if len(serialized_data) > shm.size:
+                shm.close()
+                shm.unlink()
+                shm = multiprocessing.shared_memory.SharedMemory(
+                    create=True, size=len(serialized_data))
+                shm.buf[:len(serialized_data)] = serialized_data
+                self._data[task_id] = shm.name
+            else:
+                shm.buf[:len(serialized_data)] = serialized_data
+        except FileNotFoundError:
+            shm = multiprocessing.shared_memory.SharedMemory(
+                create=True, size=len(serialized_data))
+            shm.buf[:len(serialized_data)] = serialized_data
+            self._data[task_id] = shm.name
+        shm.close()
+
+    def _load_from_shared_memory(self, task_id):
+        try:
+            if task_id not in self._data or not self._data[task_id]:
+                return []
+            try:
+                multiprocessing.shared_memory.SharedMemory(
+                    name=self._data[task_id], create=False)
+            except FileNotFoundError:
+                return []
+            shm = multiprocessing.shared_memory.SharedMemory(
+                name=self._data[task_id])
+            data = pickle.loads(shm.buf.tobytes())
+            shm.close()
+            return data
+        except Exception as e:
+            stack_trace = traceback.format_exc()
+            logger.error(
+                f"_load_from_shared_memory error: {e}\nStack trace:\n{stack_trace}")
+            return []
+
+    def _delete_from_shared_memory(self, task_id):
+        try:
+            if task_id not in self._data or not self._data[task_id]:
+                return
+            shm = multiprocessing.shared_memory.SharedMemory(
+                name=self._data[task_id])
+            shm.close()
+            shm.unlink()
+            del self._data[task_id]
+        except FileNotFoundError:
+            pass
+
+    def add(self, data: DataRow):
+        if not data:
+            raise ValueError("Data is required")
+        if not data.exp_meta:
+            raise ValueError("exp_meta is required")
+
+        with self._lock:
+            current_size = sum(len(self._load_from_shared_memory(task_id))
+                               for task_id in self._data.keys())
+            while current_size >= self._max_capacity and self._fifo_queue:
+                oldest_task_id = self._fifo_queue.pop(0)
+                if oldest_task_id in self._data.keys():
+                    current_size -= len(self._load_from_shared_memory(oldest_task_id))
+                    self._delete_from_shared_memory(oldest_task_id)
+
+            task_id = data.exp_meta.task_id
+            existing_data = self._load_from_shared_memory(task_id)
+            existing_data.append(data)
+            self._save_to_shared_memory(existing_data, task_id)
+            self._fifo_queue.append(task_id)
+
+    def add_batch(self, data_batch: List[DataRow]):
+        with self._lock:
+            for data in data_batch:
+                self.add(data)
+
+    def size(self, query_condition: QueryCondition = None) -> int:
+        with self._lock:
+            return len(self._get_all_without_lock(query_condition))
+
+    def get_paginated(self, page: int, page_size: int, query_condition: QueryCondition = None) -> List[DataRow]:
+        with self._lock:
+            if page < 1:
+                raise ValueError("Page must be greater than 0")
+            if page_size < 1:
+                raise ValueError("Page size must be greater than 0")
+            all_data = self._get_all_without_lock(query_condition)
+            start_index = (page - 1) * page_size
+            end_index = start_index + page_size
+            return all_data[start_index:end_index]
+
+    def _get_all_without_lock(self, query_condition: QueryCondition = None) -> List[DataRow]:
+        all_data = []
+        query_filter = None
+        if query_condition:
+            query_filter = QueryFilter(query_condition)
+        for task_id in self._data.keys():
+            local_data = self._load_from_shared_memory(task_id)
+            if query_filter:
+                all_data.extend(query_filter.filter(local_data))
+            else:
+                all_data.extend(local_data)
+        return all_data
+
+    def get_all(self, query_condition: QueryCondition = None) -> List[DataRow]:
+        with self._lock:
+            return self._get_all_without_lock(query_condition)
+
+    def get_by_task_id(self, task_id: str) -> List[DataRow]:
+        with self._lock:
+            if task_id in self._data.keys():
+                return self._load_from_shared_memory(task_id)
+
+    def get_bacth_by_task_ids(self, task_ids: List[str]) -> Dict[str, List[DataRow]]:
+        with self._lock:
+            result = {}
+            for task_id in task_ids:
+                if task_id in self._data.keys():
+                    result[task_id] = self._load_from_shared_memory(task_id)
+            return result

aworld/replay_buffer/storage/odps.py

@@ -0,0 +1,223 @@
+import json
+from pydantic import parse_obj_as
+from typing import Any, List, Dict
+from aworld.replay_buffer.base import Storage, DataRow, ExpMeta, Experience
+from aworld.replay_buffer.query_filter import QueryCondition, QueryBuilder
+from aworld.core.common import Observation, ActionModel
+from aworld.logs.util import logger
+from aworld.utils.import_package import import_package
+import_package("odps")  # noqa
+from odps import ODPS  # noqa
+from odps.models.record import Record  # noqa
+
+
+class OdpsSQLBuilder:
+    ''' Example:
+            query_condition = QueryBuilder().eq("field1", "value1").and_().eq("field2", "value2")
+            sql_builder = OdpsSQLBuilder(query_condition)
+            sql = sql_builder.build_sql()
+            print(sql)  # 输出: "field1 = 'value1' AND field2 = 'value2'"
+    '''
+
+    def __init__(self, query_condition: QueryCondition):
+        self.query_condition = query_condition
+
+    def _build_condition(self, condition: QueryCondition) -> str:
+        if condition is None:
+            return ""
+
+        if "field" in condition and "op" in condition:
+            field = condition["field"].split('.')[-1]
+            op = condition["op"]
+            value = condition.get("value")
+
+            if op == "eq":
+                return f"{field} = {self._format_value(value)}"
+            elif op == "ne":
+                return f"{field} != {self._format_value(value)}"
+            elif op == "gt":
+                return f"{field} > {self._format_value(value)}"
+            elif op == "gte":
+                return f"{field} >= {self._format_value(value)}"
+            elif op == "lt":
+                return f"{field} < {self._format_value(value)}"
+            elif op == "lte":
+                return f"{field} <= {self._format_value(value)}"
+            elif op == "in":
+                return f"{field} IN ({self._format_value(value)})"
+            elif op == "not_in":
+                return f"{field} NOT IN ({self._format_value(value)})"
+            elif op == "like":
+                return f"{field} LIKE '{value}'"
+            elif op == "not_like":
+                return f"{field} NOT LIKE '{value}'"
+            elif op == "is_null":
+                return f"{field} IS NULL"
+            elif op == "is_not_null":
+                return f"{field} IS NOT NULL"
+
+        elif "and_" in condition:
+            return f"({' AND '.join(self._build_condition(c) for c in condition['and_'])})"
+        elif "or_" in condition:
+            return f"({' OR '.join(self._build_condition(c) for c in condition['or_'])})"
+
+        return ""
+
+    def _format_value(self, value: Any) -> str:
+        if isinstance(value, str):
+            return f"'{value}'"
+        elif isinstance(value, (list, tuple)):
+            return ", ".join(self._format_value(v) for v in value)
+        return str(value)
+
+    def build_sql(self) -> str:
+        if not self.query_condition:
+            return ""
+        return self._build_condition(self.query_condition)
+
+
+class OdpsStorage(Storage):
+    '''
+        Aliyun ODPS storage.
+        Table schema:
+            id: int
+            task_id: string
+            task_name: string
+            agent_id: string
+            step: int
+            execute_time: string
+            state: string
+            actions: string
+            reward_t: string
+            adv_t: string
+            v_t: string
+    '''
+
+    def __init__(self, table_name: str, project: str, endpoint: str, access_id: str, access_key: str, **kwargs):
+        self.table_name = table_name
+        self.project = project
+        self.endpoint = endpoint
+        self.access_id = access_id
+        self.access_key = access_key
+        self.kwargs = kwargs
+        self._init_odps()
+
+    def _init_odps(self):
+        self.odps = ODPS(self.access_id, self.access_key,
+                         self.project, self.endpoint)
+
+    def _get_table(self):
+        return self.odps.get_table(self.table_name)
+
+    def _convert_row_to_record(self, row: DataRow) -> Record:
+        table = self._get_table()
+        record = table.new_record()
+        record["id"] = row.id
+        record["task_id"] = row.exp_meta.task_id
+        record["task_name"] = row.exp_meta.task_name
+        record["agent_id"] = row.exp_meta.agent_id
+        record["step"] = row.exp_meta.step
+        record["execute_time"] = row.exp_meta.execute_time
+        if row.exp_data.state:
+            record["state"] = row.exp_data.state.model_dump_json()
+        if row.exp_data.actions:
+            record["actions"] = "[" + ", ".join(action.model_dump_json()
+                                                for action in row.exp_data.actions) + "]"
+        if row.exp_data.reward_t:
+            record["reward_t"] = row.exp_data.reward_t
+        if row.exp_data.adv_t:
+            record["adv_t"] = row.exp_data.adv_t
+        if row.exp_data.v_t:
+            record["v_t"] = row.exp_data.v_t
+        return record
+
+    def _convert_record_to_row(self, record: Record) -> DataRow:
+        return DataRow(
+            id=record.id,
+            exp_meta=ExpMeta(
+                task_id=record['task_id'],
+                task_name=record['task_name'],
+                agent_id=record['agent_id'],
+                step=record['step'],
+                execute_time=record['execute_time'],
+                pre_agent=record['pre_agent'] if 'pre_agent' in record else None
+            ),
+            exp_data=Experience(
+                state=parse_obj_as(Observation, json.loads(record['state'])),
+                actions=[parse_obj_as(ActionModel, item)
+                         for item in json.loads(record['actions'])],
+                reward_t=record['reward_t'] if 'reward_t' in record else None,
+                adv_t=record['adv_t'] if 'adv_t' in record else None,
+                v_t=record['v_t'] if 'v_t' in record else None,
+            )
+        )
+
+    def _build_paginated_sql(self, page: int = None, page_size: int = None):
+        if page and page_size:
+            offset = (page - 1) * page_size
+            limit = page_size
+            return f" LIMIT {offset}, {limit}"
+        return ""
+
+    def _build_sql(self, query_condition: QueryCondition, page: int = None, page_size: int = None):
+        if not query_condition:
+            return f"SELECT * FROM {self.table_name}" + self._build_paginated_sql(page, page_size)
+        where_builder = OdpsSQLBuilder(query_condition)
+        sql = f"SELECT * FROM {self.table_name} WHERE {where_builder.build_sql()}" + self._build_paginated_sql(page,
+                                                                                                               page_size)
+        return sql
+
+    def _build_count_sql(self, query_condition: QueryCondition):
+        if not query_condition:
+            return f"SELECT count(1) as count FROM {self.table_name}"
+        where_builder = OdpsSQLBuilder(query_condition)
+        sql = f"SELECT count(1) as count FROM {self.table_name} WHERE {where_builder.build_sql()}"
+        return sql
+
+    def add(self, row: DataRow):
+        record = self._convert_row_to_record(row)
+        self.odps.write_table(self.table_name, [record])
+
+    def add_batch(self, rows: list[DataRow]):
+        records = [self._convert_row_to_record(row) for row in rows]
+        self.odps.write_table(self.table_name, records)
+
+    def size(self, query_condition: QueryCondition = None) -> int:
+        sql = self._build_count_sql(query_condition)
+        with self.odps.execute_sql(sql).open_reader() as reader:
+            return reader[0]["count"]
+
+    def get_all(self, query_condition: QueryCondition = None) -> list[DataRow]:
+        sql = self._build_sql(query_condition)
+        logger.info(f"get_all sql: {sql}")
+        with self.odps.execute_sql(sql).open_reader(tunnel=True) as reader:
+            rows = []
+            for record in reader:
+                rows.append(self._convert_record_to_row(record))
+            return rows
+
+    def get_paginated(self, page: int, page_size: int, query_condition: QueryCondition = None) -> List[DataRow]:
+        sql = self._build_sql(query_condition, page, page_size)
+        logger.info(f"get_paginated sql: {sql}")
+        with self.odps.execute_sql(sql).open_reader(tunnel=True) as reader:
+            rows = []
+            for record in reader:
+                rows.append(self._convert_record_to_row(record))
+            return rows
+
+    def get_by_task_id(self, task_id: str) -> List[DataRow]:
+        query_condition = QueryBuilder().eq("task_id", task_id).build()
+        return self.get_all(query_condition)
+
+    def get_bacth_by_task_ids(self, task_ids: List[str]) -> Dict[str, List[DataRow]]:
+        query_condition = QueryBuilder().in_("task_id", task_ids).build()
+        sql = self._build_sql(query_condition)
+        logger.info(f"get_bacth_by_task_ids sql: {sql}")
+        result = {}
+        with self.odps.execute_sql(sql).open_reader(tunnel=True) as reader:
+            for record in reader:
+                row = self._convert_record_to_row(record)
+                if row.exp_meta.task_id not in result:
+                    result[row.exp_meta.task_id] = []
+                result[row.exp_meta.task_id].append(row)
+        return result

aworld/replay_buffer/storage/redis.py

@@ -0,0 +1,262 @@
+import json
+from typing import Dict, List
+from aworld.replay_buffer.base import Storage, DataRow, ExpMeta, Experience
+from aworld.logs.util import logger
+from aworld.utils.import_package import import_package
+from aworld.replay_buffer.query_filter import QueryCondition, QueryBuilder
+from aworld.core.common import Observation, ActionModel
+import_package("redis")  # noqa
+from redis import Redis  # noqa
+from redis.commands.json.path import Path  # noqa
+import redis.commands.search.aggregation as aggregations  # noqa
+import redis.commands.search.reducers as reducers  # noqa
+from redis.commands.search.field import TextField, NumericField, TagField  # noqa
+from redis.commands.search.index_definition import IndexDefinition, IndexType  # noqa
+from redis.commands.search.query import Query  # noqa
+import redis.exceptions  # noqa
+
+
+class RedisSearchQueryBuilder:
+    """
+    Build redis search query from query condition
+    """
+
+    def __init__(self, query_condition: QueryCondition):
+        self.query_condition = query_condition
+
+    def _build_condition(self, condition: QueryCondition) -> str:
+        if condition is None:
+            return ""
+
+        if "field" in condition and "op" in condition:
+            field = condition["field"].split('.')[-1]
+            op = condition["op"]
+            value = condition.get("value")
+
+            if op == "eq":
+                return f"@{field}:{{{value}}}"
+            elif op == "ne":
+                return f"-@{field}:{{{value}}}"
+            elif op == "gt":
+                return f"@{field}:[{value} +inf]"
+            elif op == "gte":
+                return f"@{field}:[{value} +inf]"
+            elif op == "lt":
+                return f"@{field}:[-inf {value}]"
+            elif op == "lte":
+                return f"@{field}:[-inf {value}]"
+            elif op == "in":
+                return f"@{field}:{{{'|'.join(str(v) for v in value)}}}"
+            elif op == "not_in":
+                return f"-@{field}:{{{'|'.join(str(v) for v in value)}}}"
+            elif op == "like":
+                return f"@{field}:*{value}*"
+            elif op == "not_like":
+                return f"-@{field}:*{value}*"
+            elif op == "is_null":
+                return f"-@{field}:*"
+            elif op == "is_not_null":
+                return f"@{field}:*"
+
+        elif "and_" in condition:
+            conditions = [self._build_condition(c) for c in condition["and_"]]
+            return " ".join(conditions)
+        elif "or_" in condition:
+            conditions = [self._build_condition(c) for c in condition["or_"]]
+            return f"({'|'.join(conditions)})"
+
+        return ""
+
+    def build(self) -> Query:
+        query_str = self._build_condition(self.query_condition)
+        logger.info(f"redis search query: {query_str}")
+        return Query(query_str)
+
+
+class RedisStorage(Storage):
+    def __init__(self,
+                 host: str = 'localhost',
+                 port: int = 6379,
+                 db: int = 0,
+                 password: str = None,
+                 key_prefix: str = 'AWORLD:RB:',
+                 index_name: str = 'idx:AWORLD:RB',
+                 recreate_idx_if_exists=False):
+        self._redis = Redis(host=host, port=port, db=db, password=password)
+        self._key_prefix = key_prefix
+        self._index_name = index_name
+        self._recreate_idx_if_exists = recreate_idx_if_exists
+        self._create_index()
+
+    def _create_index(self):
+        try:
+            existing_indices = self._redis.execute_command('FT._LIST')
+            if self._index_name.encode('utf-8') in existing_indices:
+                logger.info(f"Index {self._index_name} already exists")
+                if self._recreate_idx_if_exists:
+                    self._redis.ft(self._index_name).dropindex()
+                    logger.info(f"Index {self._index_name} dropped")
+                else:
+                    return
+            self._redis.ft(self._index_name).create_index(
+                (
+                    TagField("id"),
+                    TagField("task_id"),
+                    TextField("task_name"),
+                    TagField("agent_id"),
+                    NumericField("step"),
+                    NumericField("execute_time"),
+                    TagField("pre_agent")
+                ),
+                definition=IndexDefinition(
+                    prefix=[self._key_prefix], index_type=IndexType.HASH)
+            )
+        except redis.exceptions.ResponseError as e:
+            logger.error(f"Create index {self._index_name} failed. {e}")
+
+    def _get_object_key(self, key: str) -> str:
+        return f"{self._key_prefix}{key}"
+
+    def _serialize_to_str(self, value) -> str:
+        if str is None:
+            return ""
+        if isinstance(value, (int, float)):
+            return str(value)
+        return str(value) if value is not None else ""
+
+    def _serialize(self, data: DataRow) -> Dict[str, str]:
+        dict_data = {
+            'id': data.id,
+            'task_id': data.exp_meta.task_id,
+            'task_name': data.exp_meta.task_name,
+            'agent_id': data.exp_meta.agent_id,
+            'step': data.exp_meta.step,
+            'execute_time': data.exp_meta.execute_time,
+            'pre_agent': data.exp_meta.pre_agent,
+            'state': data.exp_data.state.model_dump_json(),
+            'actions': "[" + ", ".join(action.model_dump_json()
+                                       for action in data.exp_data.actions) + "]",
+            'reward_t': data.exp_data.reward_t,
+            'adv_t': data.exp_data.adv_t,
+            'v_t': data.exp_data.v_t
+        }
+        return {k: self._serialize_to_str(v) for k, v in dict_data.items()}
+
+    def _deserialize(self, data: Dict) -> DataRow:
+        if not data:
+            return None
+        return DataRow(
+            id=data.get('id'),
+            exp_meta=ExpMeta(
+                task_id=data.get('task_id'),
+                task_name=data.get('task_name'),
+                agent_id=data.get('agent_id'),
+                step=int(data.get('step', 0)),
+                execute_time=float(data.get('execute_time', 0)),
+                pre_agent=data.get('pre_agent')
+            ),
+            exp_data=Experience(
+                state=Observation.model_validate_json(data.get('state', '{}')),
+                actions=[ActionModel.model_validate_json(json.dumps(action))
+                         for action in json.loads(data.get('actions', '[]'))],
+                reward_t=float(data.get('reward_t', 0)) if data.get(
+                    'reward_t') is not '' else None,
+                adv_t=float(data.get('adv_t', 0)) if data.get(
+                    'adv_t') is not '' else None,
+                v_t=float(data.get('v_t', 0)) if data.get(
+                    'v_t') is not '' else None
+            )
+        )
+
+    def add(self, data: DataRow):
+        key = self._get_object_key(data.id)
+        self._redis.hset(key, mapping=self._serialize(data))
+
+    def add_batch(self, data_batch: List[DataRow]):
+        pipeline = self._redis.pipeline()
+        for data in data_batch:
+            if not data or not data.exp_meta:
+                continue
+            key = self._get_object_key(data.id)
+            pipeline.hset(key, mapping=self._serialize(data))
+        pipeline.execute()
+
+    def search(self, key: str, value: str) -> DataRow:
+        result = self._redis.ft(self._index_name).search(
+            Query(f"@{key}:{{{value}}}"))
+        logger.info(f"Search result: {result}")
+
+    def size(self, query_condition: QueryCondition = None) -> int:
+        '''
+        Get the size of the storage.
+        Returns:
+            int: Size of the storage.
+        '''
+        if not query_condition:
+            return self._redis.ft(self._index_name).info()['num_docs']
+        query_builder = RedisSearchQueryBuilder(query_condition)
+        query = query_builder.build()
+        return self._redis.ft(self._index_name).search(query).total
+
+    def get_paginated(self, page: int, page_size: int, query_condition: QueryCondition = None) -> List[DataRow]:
+        '''
+        Get paginated data from the storage.
+        Args:
+            page (int): Page number.
+            page_size (int): Number of data per page.
+        Returns:
+            List[DataRow]: List of data.
+        '''
+        if not query_condition:
+            result = self._redis.ft(self._index_name).search(
+                Query("*").paging(page, page_size))
+        else:
+            query_builder = RedisSearchQueryBuilder(query_condition)
+            query = query_builder.build().paging(page, page_size)
+            result = self._redis.ft(self._index_name).search(query)
+        return [self._deserialize(doc.__dict__) for doc in result.docs]
+
+    def get_all(self, query_condition: QueryCondition = None) -> List[DataRow]:
+        '''
+        Get all data from the storage.
+        Returns:
+            List[DataRow]: List of data.
+        '''
+        if not query_condition:
+            result = self._redis.ft(self._index_name).search(Query("*"))
+        else:
+            query_builder = RedisSearchQueryBuilder(query_condition)
+            query = query_builder.build()
+            result = self._redis.ft(self._index_name).search(query)
+        return [self._deserialize(doc.__dict__) for doc in result.docs]
+
+    def get_by_task_id(self, task_id: str) -> List[DataRow]:
+        '''
+        Get data by task_id from the storage.
+        Args:
+            task_id (str): Task id.
+        Returns:
+            List[DataRow]: List of data.
+        '''
+        query_condition = QueryBuilder().eq("task_id", task_id).build()
+        return self.get_all(query_condition)
+
+    def get_bacth_by_task_ids(self, task_ids: List[str]) -> Dict[str, List[DataRow]]:
+        '''
+        Get data by task_ids from the storage.
+        Args:
+            task_ids (List[str]): List of task ids.
+        Returns:
+            Dict[str, List[DataRow]]: Dict of task id and list of data.
+        '''
+        query_condition = QueryBuilder().in_("task_id", task_ids).build()
+        result = self.get_all(query_condition)
+        return {task_id: [data for data in result if data.exp_meta.task_id == task_id] for task_id in task_ids}
+
+    def clear(self):
+        '''
+        Clear the storage.
+        '''
+        keys = self._redis.keys(f"{self._key_prefix}*")
+        if keys:
+            self._redis.delete(*keys)