rlightning.buffer.utils.table

Episode-to-storage shard assignment table.

class rlightning.buffer.utils.table.EpisodeTable(num_storages: int, env_ids: Sequence[str] | None = None, num_train_workers: int | None = None, component_distribution: Dict[str, Dict[str, Dict[str, Any]]] | None = None, node_affinity_env: bool = False, node_affinity_train: bool = False)

Bases: object

Track episode -> storage shard mapping and storage -> train worker mapping.

The MVP version uses a simple, even distribution strategy for both: - env_ids are assigned to storage shards in round-robin order with load balance. - train workers are assigned to storage shards in contiguous, even blocks.

When node affinity is enabled, envs and train workers are bound to storages on the same node using the component distribution view.

get_env_to_storage() → Dict[str, int]

Return env -> storage mapping.

get_envs_for_storage(storage_idx: int) → List[str]

List env_ids currently mapped to the given storage shard.

get_storage_idx_for_env(env_id: str) → int

Return storage index for env_id, assigning it if unseen.

get_storage_to_train_workers() → Dict[int, List[int]]

Return storage -> train worker mapping.

get_train_workers(storage_idx: int) → List[int]

Return train workers assigned to the given storage shard.

register_envs(env_ids: Iterable[str]) → None

Register a batch of env_ids with the current assignment strategy.