Add default batch size for SQLReader

Author: chenyushuo

tests/buffer/sql_test.py

@@ -1,4 +1,5 @@
 import os
+from copy import deepcopy
 
 import ray
 import torch
@@ -40,7 +41,9 @@ async def test_sql_exp_buffer_read_write(self, enable_replay: bool) -> None:
         )
         if enable_replay:
             config.replay_buffer = ReplayBufferConfig(enable=True)
-        sql_writer = SQLWriter(config.to_storage_config())
+        writer_config = deepcopy(config)
+        writer_config.batch_size = put_batch_size
+        sql_writer = SQLWriter(writer_config.to_storage_config())
         sql_reader = SQLReader(config.to_storage_config())
         exps = [
             Experience(

trinity/buffer/reader/sql_reader.py

@@ -16,15 +16,18 @@ class SQLReader(BufferReader):
     def __init__(self, config: StorageConfig) -> None:
         assert config.storage_type == StorageType.SQL.value
         self.wrap_in_ray = config.wrap_in_ray
+        self.read_batch_size = config.batch_size
         self.storage = SQLStorage.get_wrapper(config)
 
     def read(self, batch_size: Optional[int] = None, **kwargs) -> List:
+        batch_size = batch_size or self.read_batch_size
         if self.wrap_in_ray:
             return ray.get(self.storage.read.remote(batch_size, **kwargs))
         else:
             return self.storage.read(batch_size, **kwargs)
 
     async def read_async(self, batch_size: Optional[int] = None, **kwargs) -> List:
+        batch_size = batch_size or self.read_batch_size
         if self.wrap_in_ray:
             try:
                 return await self.storage.read.remote(batch_size, **kwargs)