File size: 1,167 Bytes
8c9ba62 | 1 2 3 4 5 6 7 8 9 10 11 12 13 14 15 16 17 18 19 20 21 22 23 24 25 26 27 28 29 30 31 32 33 34 35 36 37 38 39 40 41 | from typing import List
import ray
from trinity.buffer.buffer_writer import BufferWriter
from trinity.buffer.storage.file import FileStorage
from trinity.common.config import StorageConfig
from trinity.common.constants import StorageType
class JSONWriter(BufferWriter):
def __init__(self, config: StorageConfig):
assert config.storage_type == StorageType.FILE.value
self.writer = FileStorage.get_wrapper(config)
self.wrap_in_ray = config.wrap_in_ray
def write(self, data: List) -> None:
if self.wrap_in_ray:
ray.get(self.writer.write.remote(data))
else:
self.writer.write(data)
async def write_async(self, data):
if self.wrap_in_ray:
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
async def release(self) -> int:
if self.wrap_in_ray:
return await self.writer.release.remote()
else:
self.writer.release()
return 0
|