buffer writer support async

Author: pan-x-c

tests/buffer/file_test.py

@@ -16,7 +16,7 @@
 from trinity.common.constants import StorageType
 
 
-class TestFileBuffer(unittest.TestCase):
+class TestFileBuffer(unittest.IsolatedAsyncioTestCase):
     temp_output_path = "tmp/test_file_buffer/"
 
     @classmethod
@@ -93,25 +93,31 @@ def test_file_reader(self):
                 break
         self.assertEqual(len(tasks), 16 * 3 - 20)
 
-    def test_file_writer(self):
+    async def test_file_writer(self):
         writer = get_buffer_writer(
             self.config.buffer.trainer_input.experience_buffer, self.config.buffer
         )
-        writer.acquire()
+        await writer.acquire()
         writer.write(
             [
                 {"prompt": "hello world"},
                 {"prompt": "hi"},
             ]
         )
-        writer.release()
+        await writer.write_async(
+            [
+                {"prompt": "My name is"},
+                {"prompt": "What is your name?"},
+            ]
+        )
+        await writer.release()
         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)
+            self.assertEqual(len(f.readlines()), 4)
 
     def setUp(self):
         self.config = get_template_config()

tests/buffer/queue_test.py

@@ -3,7 +3,7 @@
 
 import torch
 
-from tests.tools import RayUnittestBase
+from tests.tools import RayUnittestBaseAysnc
 from trinity.buffer.reader.queue_reader import QueueReader
 from trinity.buffer.writer.queue_writer import QueueWriter
 from trinity.common.config import BufferConfig, StorageConfig
@@ -13,8 +13,8 @@
 BUFFER_FILE_PATH = os.path.join(os.path.dirname(__file__), "test_queue_buffer.jsonl")
 
 
-class TestQueueBuffer(RayUnittestBase):
-    def test_queue_buffer(self):
+class TestQueueBuffer(RayUnittestBaseAysnc):
+    async def test_queue_buffer(self):
         total_num = 8
         put_batch_size = 2
         read_batch_size = 4
@@ -32,7 +32,7 @@ def test_queue_buffer(self):
         )
         writer = QueueWriter(meta, config)
         reader = QueueReader(meta, config)
-        self.assertEqual(writer.acquire(), 1)
+        self.assertEqual(await writer.acquire(), 1)
         exps = [
             Experience(
                 tokens=torch.tensor([float(j) for j in range(i + 1)]),
@@ -43,7 +43,7 @@ def test_queue_buffer(self):
             for i in range(1, put_batch_size + 1)
         ]
         for _ in range(total_num // put_batch_size):
-            writer.write(exps)
+            await writer.write_async(exps)
         for _ in range(total_num // read_batch_size):
             exps = reader.read()
             self.assertEqual(len(exps), read_batch_size)
@@ -62,7 +62,7 @@ def test_queue_buffer(self):
         )
         exps = reader.read(batch_size=put_batch_size * 2)
         self.assertEqual(len(exps), put_batch_size * 2)
-        self.assertEqual(writer.release(), 0)
+        self.assertEqual(await writer.release(), 0)
         self.assertRaises(StopIteration, reader.read)
         with open(BUFFER_FILE_PATH, "r") as f:
             self.assertEqual(len(f.readlines()), total_num + put_batch_size * 2)

tests/buffer/sql_test.py

@@ -1,9 +1,9 @@
 import os
-import unittest
 
 import ray
 import torch
 
+from tests.tools import RayUnittestBaseAysnc
 from trinity.buffer.reader.sql_reader import SQLReader
 from trinity.buffer.writer.sql_writer import SQLWriter
 from trinity.common.config import BufferConfig, StorageConfig
@@ -13,8 +13,8 @@
 db_path = os.path.join(os.path.dirname(__file__), "test.db")
 
 
-class TestSQLBuffer(unittest.TestCase):
-    def test_create_sql_buffer(self) -> None:
+class TestSQLBuffer(RayUnittestBaseAysnc):
+    async def test_create_sql_buffer(self) -> None:
         total_num = 8
         put_batch_size = 2
         read_batch_size = 4
@@ -42,9 +42,9 @@ def test_create_sql_buffer(self) -> None:
             )
             for i in range(1, put_batch_size + 1)
         ]
-        self.assertEqual(sql_writer.acquire(), 1)
+        self.assertEqual(await sql_writer.acquire(), 1)
         for _ in range(total_num // put_batch_size):
-            sql_writer.write(exps)
+            await sql_writer.write_async(exps)
         for _ in range(total_num // read_batch_size):
             exps = sql_reader.read()
             self.assertEqual(len(exps), read_batch_size)
@@ -66,5 +66,5 @@ def test_create_sql_buffer(self) -> None:
         self.assertEqual(len(exps), put_batch_size * 2)
         db_wrapper = ray.get_actor("sql-test_buffer")
         self.assertIsNotNone(db_wrapper)
-        self.assertEqual(sql_writer.release(), 0)
+        self.assertEqual(await sql_writer.release(), 0)
         self.assertRaises(StopIteration, sql_reader.read)

tests/tools.py

@@ -182,3 +182,13 @@ def setUpClass(cls):
     @classmethod
     def tearDownClass(cls):
         ray.shutdown(_exiting_interpreter=True)
+
+
+class RayUnittestBaseAysnc(unittest.IsolatedAsyncioTestCase):
+    @classmethod
+    def setUpClass(cls):
+        ray.init(ignore_reinit_error=True, namespace="trinity_unittest")
+
+    @classmethod
+    def tearDownClass(cls):
+        ray.shutdown(_exiting_interpreter=True)

trinity/buffer/buffer_writer.py

@@ -11,15 +11,19 @@ def write(self, data: List) -> None:
         """Write to buffer."""
 
     @abstractmethod
-    def acquire(self) -> int:
+    async def write_async(self, data: List) -> None:
+        """Write to buffer asynchronously."""
+
+    @abstractmethod
+    async def acquire(self) -> int:
         """Acquire the buffer writer.
 
         Returns:
             `int`: The reference count of the buffer after acquiring.
         """
 
     @abstractmethod
-    def release(self) -> int:
+    async def release(self) -> int:
         """Release the buffer writer. After release, the buffer writer can not be used again.
 
         Returns:

trinity/buffer/queue.py

@@ -57,7 +57,7 @@ async def release(self) -> int:
         self.ref_count -= 1
         if self.ref_count <= 0:
             await self.queue.put(self.FINISH_MESSAGE)
-            self.writer.release()
+            await self.writer.release()
         return self.ref_count
 
     def length(self) -> int:

trinity/buffer/writer/file_writer.py

@@ -20,15 +20,21 @@ def write(self, data: List) -> None:
         else:
             self.writer.write(data)
 
-    def acquire(self) -> int:
+    async def write_async(self, data):
         if self.wrap_in_ray:
-            return ray.get(self.writer.acquire.remote())
+            await self.writer.write.remote(data)
+        else:
+            self.writer.write(data)
+
+    async def acquire(self) -> int:
+        if self.wrap_in_ray:
+            return await self.writer.acquire.remote()
         else:
             return 0
 
-    def release(self) -> int:
+    async def release(self) -> int:
         if self.wrap_in_ray:
-            return ray.get(self.writer.release.remote())
+            return await self.writer.release.remote()
         else:
             self.writer.release()
             return 0

trinity/buffer/writer/queue_writer.py

@@ -23,8 +23,11 @@ def __init__(self, meta: StorageConfig, config: BufferConfig):
     def write(self, data: List) -> None:
         ray.get(self.queue.put_batch.remote(data))
 
-    def acquire(self) -> int:
-        return ray.get(self.queue.acquire.remote())
+    async def write_async(self, data):
+        return await self.queue.put_batch.remote(data)
 
-    def release(self) -> int:
-        return ray.get(self.queue.release.remote())
+    async def acquire(self) -> int:
+        return await self.queue.acquire.remote()
+
+    async def release(self) -> int:
+        return await self.queue.release.remote()

trinity/buffer/writer/sql_writer.py

@@ -23,15 +23,21 @@ def write(self, data: list) -> None:
         else:
             self.db_wrapper.write(data)
 
-    def acquire(self) -> int:
+    async def write_async(self, data):
         if self.wrap_in_ray:
-            return ray.get(self.db_wrapper.acquire.remote())
+            await self.db_wrapper.write.remote(data)
+        else:
+            self.db_wrapper.write(data)
+
+    async def acquire(self) -> int:
+        if self.wrap_in_ray:
+            return await self.db_wrapper.acquire.remote()
         else:
             return 0
 
-    def release(self) -> int:
+    async def release(self) -> int:
         if self.wrap_in_ray:
-            return ray.get(self.db_wrapper.release.remote())
+            return await self.db_wrapper.release.remote()
         else:
             self.db_wrapper.release()
             return 0

trinity/explorer/explorer.py

@@ -43,12 +43,12 @@ def __init__(self, config: Config):
         self.config = config
         self.algorithm_manager = AlgorithmManager(config)
         self.models, self.auxiliary_models = create_inference_models(config)
+        self.experience_buffer = None
         if self.config.mode != "bench":
             self.experience_buffer = get_buffer_writer(
                 self.config.buffer.explorer_output,  # type: ignore
                 self.config.buffer,
             )
-            self.experience_buffer.acquire()
         self.config.buffer.explorer_input.taskset.index = explorer_meta.get("latest_task_index", 0)
         self.taskset = get_buffer_reader(
             self.config.buffer.explorer_input.taskset, self.config.buffer
@@ -169,6 +169,8 @@ async def prepare(self) -> None:
                 asyncio.create_task(self.setup_weight_sync_group(master_address, master_port))
             )
         asyncio.gather(*futures, return_exceptions=True)
+        if self.experience_buffer:
+            await self.experience_buffer.acquire()
         if self.config.explorer.eval_on_startup and self.explore_step_num == 0:
             self.eval()
 
@@ -217,7 +219,7 @@ async def explore_step(self) -> bool:
             self.logger.warning("No more tasks to explore. Stop exploring.")
             await self.save_checkpoint(sync_weight=False)
             self.status = RunningStatus.STOPPED
-            self.experience_buffer.release()
+            await self.experience_buffer.release()
             return False
         self.scheduler.schedule(tasks, batch_id=self.explore_step_num + 1)
         self.explore_step_num += 1