merge main

Author: pan-x-c

tests/buffer/file_test.py

@@ -0,0 +1,94 @@
+import os
+import unittest
+
+import ray
+
+from tests.tools import (
+    get_checkpoint_path,
+    get_template_config,
+    get_unittest_dataset_config,
+)
+from trinity.buffer.buffer import get_buffer_reader, get_buffer_writer
+from trinity.buffer.utils import default_storage_path
+from trinity.common.config import StorageConfig
+from trinity.common.constants import StorageType
+
+
+class TestFileBuffer(unittest.TestCase):
+    def test_file_reader(self):
+        """Test file reader."""
+        reader = get_buffer_reader(self.config.buffer.explorer_input.taskset, self.config.buffer)
+
+        tasks = []
+        while True:
+            try:
+                tasks.extend(reader.read())
+            except StopIteration:
+                break
+        self.assertEqual(len(tasks), 16)
+
+        # test epoch and offset
+        self.config.buffer.explorer_input.taskset.total_epochs = 2
+        self.config.buffer.explorer_input.taskset.index = 4
+        reader = get_buffer_reader(self.config.buffer.explorer_input.taskset, self.config.buffer)
+        tasks = []
+        while True:
+            try:
+                tasks.extend(reader.read())
+            except StopIteration:
+                break
+        self.assertEqual(len(tasks), 16 * 2 - 4)
+
+        # test offset > dataset_len
+        self.config.buffer.explorer_input.taskset.total_epochs = 3
+        self.config.buffer.explorer_input.taskset.index = 20
+        reader = get_buffer_reader(self.config.buffer.explorer_input.taskset, self.config.buffer)
+        tasks = []
+        while True:
+            try:
+                tasks.extend(reader.read())
+            except StopIteration:
+                break
+        self.assertEqual(len(tasks), 16 * 3 - 20)
+
+    def test_file_writer(self):
+        writer = get_buffer_writer(
+            self.config.buffer.trainer_input.experience_buffer, self.config.buffer
+        )
+        writer.write(
+            [
+                {"prompt": "hello world"},
+                {"prompt": "hi"},
+            ]
+        )
+        file_wrapper = ray.get_actor("json-test_buffer")
+        self.assertIsNotNone(file_wrapper)
+        file_path = default_storage_path(
+            self.config.buffer.trainer_input.experience_buffer, self.config.buffer
+        )
+        with open(file_path, "r") as f:
+            self.assertEqual(len(f.readlines()), 2)
+
+    def setUp(self):
+        self.config = get_template_config()
+        self.config.checkpoint_root_dir = get_checkpoint_path()
+        dataset_config = get_unittest_dataset_config("countdown", "train")
+        self.config.buffer.explorer_input.taskset = dataset_config
+        self.config.buffer.trainer_input.experience_buffer = StorageConfig(
+            name="test_buffer", storage_type=StorageType.FILE
+        )
+        self.config.buffer.trainer_input.experience_buffer.name = "test_buffer"
+        self.config.buffer.cache_dir = os.path.join(
+            self.config.checkpoint_root_dir, self.config.project, self.config.name, "buffer"
+        )
+        os.makedirs(self.config.buffer.cache_dir, exist_ok=True)
+        if os.path.exists(
+            default_storage_path(
+                self.config.buffer.trainer_input.experience_buffer, self.config.buffer
+            )
+        ):
+            os.remove(
+                default_storage_path(
+                    self.config.buffer.trainer_input.experience_buffer, self.config.buffer
+                )
+            )

tests/buffer/queue_test.py

@@ -1,3 +1,5 @@
+import os
+
 import torch
 
 from tests.tools import RayUnittestBase
@@ -7,6 +9,8 @@
 from trinity.common.constants import StorageType
 from trinity.common.experience import Experience
 
+BUFFER_FILE_PATH = os.path.join(os.path.dirname(__file__), "test_queue_buffer.jsonl")
+
 
 class TestQueueBuffer(RayUnittestBase):
     def test_queue_buffer(self):
@@ -17,6 +21,7 @@ def test_queue_buffer(self):
             name="test_buffer",
             algorithm_type="ppo",
             storage_type=StorageType.QUEUE,
+            path=BUFFER_FILE_PATH,
         )
         config = BufferConfig(
             max_retry_times=3,
@@ -36,9 +41,29 @@ def test_queue_buffer(self):
         ]
         for _ in range(total_num // put_batch_size):
             writer.write(exps)
-        writer.finish()
         for _ in range(total_num // read_batch_size):
             exps = reader.read()
             self.assertEqual(len(exps), read_batch_size)
             print(f"finish read {read_batch_size} experience")
+        writer.write(
+            [
+                Experience(
+                    tokens=torch.tensor([float(j) for j in range(i + 1)]),
+                    prompt_length=i,
+                    reward=float(i),
+                    logprobs=torch.tensor([0.1]),
+                    action_mask=torch.tensor([j % 2 for j in range(i + 1)]),
+                )
+                for i in range(1, put_batch_size * 2 + 1)
+            ]
+        )
+        exps = reader.read(batch_size=put_batch_size * 2)
+        self.assertEqual(len(exps), put_batch_size * 2)
+        writer.finish()
         self.assertRaises(StopIteration, reader.read)
+        with open(BUFFER_FILE_PATH, "r") as f:
+            self.assertEqual(len(f.readlines()), total_num + put_batch_size * 2)
+
+    def setUp(self):
+        if os.path.exists(BUFFER_FILE_PATH):
+            os.remove(BUFFER_FILE_PATH)

tests/buffer/sql_test.py

@@ -47,5 +47,21 @@ def test_create_sql_buffer(self) -> None:
         for _ in range(total_num // read_batch_size):
             exps = sql_reader.read()
             self.assertEqual(len(exps), read_batch_size)
+
+        # dynamic read/write
+        sql_writer.write(
+            [
+                Experience(
+                    tokens=torch.tensor([float(j) for j in range(i + 1)]),
+                    prompt_length=i,
+                    reward=float(i),
+                    logprobs=torch.tensor([0.1]),
+                    action_mask=torch.tensor([j % 2 for j in range(i + 1)]),
+                )
+                for i in range(1, put_batch_size * 2 + 1)
+            ]
+        )
+        exps = sql_reader.read(batch_size=put_batch_size * 2)
+        self.assertEqual(len(exps), put_batch_size * 2)
         db_wrapper = ray.get_actor("sql-test_buffer")
         self.assertIsNotNone(db_wrapper)

tests/explorer/runner_pool_test.py

@@ -106,6 +106,7 @@ def setUp(self):
             name="test",
             storage_type=StorageType.QUEUE,
             algorithm_type="ppo",
+            path="",
         )
         self.queue = QueueReader(
             self.config.buffer.trainer_input.experience_buffer, self.config.buffer

trinity/buffer/buffer.py

@@ -61,5 +61,9 @@ def get_buffer_writer(storage_config: StorageConfig, buffer_config: BufferConfig
         from trinity.buffer.writer.queue_writer import QueueWriter
 
         return QueueWriter(storage_config, buffer_config)
+    elif storage_config.storage_type == StorageType.FILE:
+        from trinity.buffer.writer.file_writer import JSONWriter
+
+        return JSONWriter(storage_config, buffer_config)
     else:
         raise ValueError(f"{storage_config.storage_type} not supported.")

trinity/buffer/buffer_reader.py

@@ -9,5 +9,7 @@ class BufferReader(ABC):
     """Interface of the buffer reader."""
 
     @abstractmethod
-    def read(self, strategy: Optional[ReadStrategy] = None) -> List:
+    def read(
+        self, batch_size: Optional[int] = None, strategy: Optional[ReadStrategy] = None
+    ) -> List:
         """Read from buffer."""

trinity/buffer/queue.py

@@ -5,9 +5,19 @@
 
 import ray
 
+from trinity.buffer.writer.file_writer import JSONWriter
 from trinity.buffer.writer.sql_writer import SQLWriter
 from trinity.common.config import BufferConfig, StorageConfig
 from trinity.common.constants import StorageType
+from trinity.utils.log import get_logger
+
+
+def is_database_url(path: str) -> bool:
+    return any(path.startswith(prefix) for prefix in ["sqlite:///", "postgresql://", "mysql://"])
+
+
+def is_json_file(path: str) -> bool:
+    return path.endswith(".json") or path.endswith(".jsonl")
 
 
 @ray.remote
@@ -17,16 +27,26 @@ class QueueActor:
     FINISH_MESSAGE = "$FINISH$"
 
     def __init__(self, storage_config: StorageConfig, config: BufferConfig) -> None:
+        self.logger = get_logger(__name__)
         self.config = config
         self.capacity = getattr(config, "capacity", 10000)
         self.queue = asyncio.Queue(self.capacity)
-        if storage_config.path is not None and len(storage_config.path) > 0:
-            sql_config = deepcopy(storage_config)
-            sql_config.storage_type = StorageType.SQL
-            sql_config.wrap_in_ray = False
-            self.sql_writer = SQLWriter(sql_config, self.config)
+        st_config = deepcopy(storage_config)
+        st_config.wrap_in_ray = False
+        if st_config.path is not None:
+            if is_database_url(st_config.path):
+                st_config.storage_type = StorageType.SQL
+                self.writer = SQLWriter(st_config, self.config)
+            elif is_json_file(st_config.path):
+                st_config.storage_type = StorageType.FILE
+                self.writer = JSONWriter(st_config, self.config)
+            else:
+                self.logger.warning("Unknown supported storage path: %s", st_config.path)
+                self.writer = None
         else:
-            self.sql_writer = None
+            st_config.storage_type = StorageType.FILE
+            self.writer = JSONWriter(st_config, self.config)
+        self.logger.warning(f"Save experiences in {st_config.path}.")
 
     def length(self) -> int:
         """The length of the queue."""
@@ -35,8 +55,8 @@ def length(self) -> int:
     async def put_batch(self, exp_list: List) -> None:
         """Put batch of experience."""
         await self.queue.put(exp_list)
-        if self.sql_writer is not None:
-            self.sql_writer.write(exp_list)
+        if self.writer is not None:
+            self.writer.write(exp_list)
 
     async def finish(self) -> None:
         """Stop the queue."""

trinity/buffer/db_wrapper.py trinity/buffer/ray_wrapper.pytrinity/buffer/db_wrapper.py renamed to trinity/buffer/ray_wrapper.py

@@ -1,3 +1,5 @@
+import json
+import os
 import time
 from typing import List, Optional
 
@@ -8,9 +10,11 @@
 from sqlalchemy.pool import NullPool
 
 from trinity.buffer.schema import Base, create_dynamic_table
-from trinity.buffer.utils import retry_session
+from trinity.buffer.utils import default_storage_path, retry_session
 from trinity.common.config import BufferConfig, StorageConfig
 from trinity.common.constants import ReadStrategy
+from trinity.common.experience import Experience
+from trinity.common.workflows import Task
 from trinity.utils.log import get_logger
 
 
@@ -27,6 +31,8 @@ class DBWrapper:
 
     def __init__(self, storage_config: StorageConfig, config: BufferConfig) -> None:
         self.logger = get_logger(__name__)
+        if storage_config.path is None:
+            storage_config.path = default_storage_path(storage_config, config)
         self.engine = create_engine(storage_config.path, poolclass=NullPool)
         self.table_model_cls = create_dynamic_table(
             storage_config.algorithm_type, storage_config.name
@@ -61,7 +67,9 @@ def write(self, data: list) -> None:
             experience_models = [self.table_model_cls.from_experience(exp) for exp in data]
             session.add_all(experience_models)
 
-    def read(self, strategy: Optional[ReadStrategy] = None) -> List:
+    def read(
+        self, batch_size: Optional[int] = None, strategy: Optional[ReadStrategy] = None
+    ) -> List:
         if strategy is None:
             strategy = ReadStrategy.LFU
 
@@ -78,7 +86,8 @@ def read(self, strategy: Optional[ReadStrategy] = None) -> List:
             raise NotImplementedError(f"Unsupported strategy {strategy} by SQLStorage")
 
         exp_list = []
-        while len(exp_list) < self.batch_size:
+        batch_size = batch_size or self.batch_size
+        while len(exp_list) < batch_size:
             if len(exp_list):
                 self.logger.info("waiting for experiences...")
                 time.sleep(1)
@@ -90,7 +99,7 @@ def read(self, strategy: Optional[ReadStrategy] = None) -> List:
                     session.query(self.table_model_cls)
                     .filter(self.table_model_cls.reward.isnot(None))
                     .order_by(*sortOrder)  # TODO: very slow
-                    .limit(self.batch_size - len(exp_list))
+                    .limit(batch_size - len(exp_list))
                     .with_for_update()
                     .all()
                 )
@@ -103,3 +112,63 @@ def read(self, strategy: Optional[ReadStrategy] = None) -> List:
         self.logger.info(f"first prompt_text = {exp_list[0].prompt_text}")
         self.logger.info(f"first response_text = {exp_list[0].response_text}")
         return exp_list
+
+
+class _Encoder(json.JSONEncoder):
+    def default(self, o):
+        if isinstance(o, Experience):
+            return o.to_dict()
+        if isinstance(o, Task):
+            return o.to_dict()
+        return super().default(o)
+
+
+class FileWrapper:
+    """
+    A wrapper of a local jsonl file.
+
+    If `wrap_in_ray` in `StorageConfig` is `True`, this class will be run as
+    a Ray Actor, and provide a remote interface to the local file.
+
+    This wrapper is only for writing, if you want to read from the file, use
+    StorageType.QUEUE instead.
+    """
+
+    def __init__(self, storage_config: StorageConfig, config: BufferConfig) -> None:
+        if storage_config.path is None:
+            storage_config.path = default_storage_path(storage_config, config)
+        ext = os.path.splitext(storage_config.path)[-1]
+        if ext != ".jsonl" and ext != ".json":
+            raise ValueError(
+                f"File path must end with '.json' or '.jsonl', got {storage_config.path}"
+            )
+        self.file = open(storage_config.path, "a", encoding="utf-8")
+        self.encoder = _Encoder(ensure_ascii=False)
+
+    @classmethod
+    def get_wrapper(cls, storage_config: StorageConfig, config: BufferConfig):
+        if storage_config.wrap_in_ray:
+            return (
+                ray.remote(cls)
+                .options(
+                    name=f"json-{storage_config.name}",
+                    get_if_exists=True,
+                )
+                .remote(storage_config, config)
+            )
+        else:
+            return cls(storage_config, config)
+
+    def write(self, data: List) -> None:
+        for item in data:
+            json_str = self.encoder.encode(item)
+            self.file.write(json_str + "\n")
+        self.file.flush()
+
+    def read(self) -> List:
+        raise NotImplementedError(
+            "read() is not implemented for FileWrapper, please use QUEUE instead"
+        )
+
+    def finish(self) -> None:
+        self.file.close()