Primary async SQL storage for tasks. Used directly as Ray actor.
| 313 | |
| 314 | |
| 315 | class SQLTaskStorage: |
| 316 | """Primary async SQL storage for tasks. Used directly as Ray actor.""" |
| 317 | |
| 318 | def __init__(self, config: StorageConfig) -> None: |
| 319 | self.logger = get_logger(f"sql_{config.name}") |
| 320 | self.config = config |
| 321 | self.batch_size = config.batch_size |
| 322 | self.is_eval = config.is_eval |
| 323 | self.max_retry_times = config.max_retry_times |
| 324 | self.max_retry_interval = config.max_retry_interval |
| 325 | self.ref_count = 0 |
| 326 | self.stopped = False |
| 327 | self.offset = config.index |
| 328 | self._initialized = False |
| 329 | |
| 330 | if config.total_steps: |
| 331 | self.total_samples = self.batch_size * config.total_steps |
| 332 | else: |
| 333 | self.total_samples = float("inf") |
| 334 | |
| 335 | async def prepare(self) -> None: |
| 336 | """Initialize async engine and create tables.""" |
| 337 | if self._initialized: |
| 338 | return |
| 339 | from trinity.buffer.schema.formatter import TaskFormatter |
| 340 | |
| 341 | result = await init_async_engine( |
| 342 | self.config.path, self.config.name, self.config.schema_type # type: ignore |
| 343 | ) |
| 344 | self.engine, self.table_model_cls = result |
| 345 | self.session = async_sessionmaker(self.engine, expire_on_commit=False) |
| 346 | self.default_workflow_cls = WORKFLOWS.get(self.config.default_workflow_type) |
| 347 | self.default_reward_fn_cls = REWARD_FUNCTIONS.get(self.config.default_reward_fn_type) |
| 348 | self.formatter = TaskFormatter(self.config) |
| 349 | self._initialized = True |
| 350 | self.logger.info(f"SQL task storage initialized at {self.config.path}") |
| 351 | |
| 352 | async def write(self, data: List[Dict]) -> None: |
| 353 | await self.prepare() |
| 354 | |
| 355 | async def operation(session: AsyncSession): |
| 356 | tasks = [self.table_model_cls.from_dict(item) for item in data] |
| 357 | session.add_all(tasks) |
| 358 | |
| 359 | await async_run_with_retry_session( |
| 360 | self.session, operation, self.max_retry_times, self.max_retry_interval |
| 361 | ) |
| 362 | |
| 363 | async def read(self, batch_size: Optional[int] = None) -> List[Task]: |
| 364 | await self.prepare() |
| 365 | if self.stopped: |
| 366 | raise StopAsyncIteration() |
| 367 | if self.offset > self.total_samples: |
| 368 | raise StopAsyncIteration() |
| 369 | batch_size = self.batch_size if batch_size is None else batch_size |
| 370 | |
| 371 | table_cls = self.table_model_cls |
| 372 |
no outgoing calls