Primary async SQL storage for experiences. Used directly as Ray actor.
| 33 | |
| 34 | |
| 35 | class SQLExperienceStorage: |
| 36 | """Primary async SQL storage for experiences. Used directly as Ray actor.""" |
| 37 | |
| 38 | def __init__(self, config: StorageConfig) -> None: |
| 39 | self.logger = get_logger(f"sql_{config.name}") |
| 40 | self.config = config |
| 41 | self.max_timeout = config.max_read_timeout |
| 42 | self.batch_size = config.batch_size |
| 43 | self.enable_replay = config.replay_buffer is not None and config.replay_buffer.enable |
| 44 | self.max_experience_bytes = int(os.getenv(MAX_EXP_BYTES_ENV_VAR, 1024 * 1024 * 32)) |
| 45 | self.max_retry_times = config.max_retry_times |
| 46 | self.max_retry_interval = config.max_retry_interval |
| 47 | self.ref_count = 0 |
| 48 | self.stopped = False |
| 49 | self.offset = config.index |
| 50 | self._initialized = False |
| 51 | |
| 52 | if config.schema_type == "experience": |
| 53 | self._read_method = self._read_priority |
| 54 | else: |
| 55 | self._read_method = self._read_fifo |
| 56 | |
| 57 | async def prepare(self) -> None: |
| 58 | """Initialize async engine and create tables.""" |
| 59 | if self._initialized: |
| 60 | return |
| 61 | result = await init_async_engine( |
| 62 | self.config.path, self.config.name, self.config.schema_type # type: ignore |
| 63 | ) |
| 64 | self.engine, self.table_model_cls, self.blob_model_cls = result |
| 65 | self.session = async_sessionmaker(self.engine, expire_on_commit=False) |
| 66 | self._initialized = True |
| 67 | self.logger.info(f"SQL storage initialized at {self.config.path}") |
| 68 | |
| 69 | async def write(self, data: List[Experience]) -> None: |
| 70 | await self.prepare() |
| 71 | |
| 72 | async def operation(session: AsyncSession): |
| 73 | for exp in data: |
| 74 | exp_bytes = exp.serialize() |
| 75 | if ( |
| 76 | self.max_experience_bytes > 0 |
| 77 | and exp_bytes is not None |
| 78 | and len(exp_bytes) > self.max_experience_bytes |
| 79 | ): |
| 80 | self.logger.warning( |
| 81 | f"Experience size {len(exp_bytes)} bytes exceeds " |
| 82 | f"max_experience_bytes {self.max_experience_bytes}, skipping." |
| 83 | ) |
| 84 | continue |
| 85 | meta_row = self.table_model_cls.from_experience(exp) |
| 86 | session.add(meta_row) |
| 87 | await session.flush() |
| 88 | blob_row = self.blob_model_cls(id=meta_row.id, experience_bytes=exp_bytes) |
| 89 | session.add(blob_row) |
| 90 | |
| 91 | await async_run_with_retry_session( |
| 92 | self.session, operation, self.max_retry_times, self.max_retry_interval |
no outgoing calls