ViperEkura/AstrAI
1
0
Fork 0
Code Issues Pull Requests Packages Projects Releases Wiki Activity

refactor: remove H5 storage backend in favor of mmap bin

Browse Source 
- Remove H5Store, H5Writer, save_h5/load_h5 and h5py dependency
- MmapStore (bin) is the sole pre-tokenized storage backend
- Move setup_logging after imports to fix E402 in __init__.py
- Clean up unused imports across test files
- Move inline test imports to file top
This commit is contained in:
ViperEkura
2026-07-29 12:50:27 +08:00
parent c2b04d8458
commit 115192c67c
17 changed files with 91 additions and 209 deletions
Download Patch File Download Diff File Expand all files Collapse all files
astrai/__init__.py
+24 -24
View File
@@ -4,30 +4,6 @@ __author__ = "ViperEkura"
import logging
import os
def setup_logging(level: str = "INFO"):
"""Attach a handler to the ``astrai`` logger (only, not root).
Call once per process, e.g. at the top of CLI scripts.
Set ``ASTR_LOG_LEVEL`` to override the default ``INFO``.
"""
_logger = logging.getLogger("astrai")
if _logger.handlers:
return
_level = getattr(
logging, os.environ.get("ASTR_LOG_LEVEL", level).upper(), logging.INFO
)
_logger.setLevel(_level)
_handler = logging.StreamHandler()
_handler.setFormatter(
logging.Formatter(
"%(asctime)s | %(levelname)-7s | %(name)s | %(message)s",
datefmt="%Y-%m-%d %H:%M:%S",
)
)
_logger.addHandler(_handler)
from astrai.config import (
AutoRegressiveLMConfig,
BaseModelConfig,
@@ -80,6 +56,30 @@ from astrai.trainer import (
Trainer,
)
def setup_logging(level: str = "INFO"):
"""Attach a handler to the ``astrai`` logger (only, not root).
Call once per process, e.g. at the top of CLI scripts.
Set ``ASTR_LOG_LEVEL`` to override the default ``INFO``.
"""
_logger = logging.getLogger("astrai")
if _logger.handlers:
return
_level = getattr(
logging, os.environ.get("ASTR_LOG_LEVEL", level).upper(), logging.INFO
)
_logger.setLevel(_level)
_handler = logging.StreamHandler()
_handler.setFormatter(
logging.Formatter(
"%(asctime)s | %(levelname)-7s | %(name)s | %(message)s",
datefmt="%Y-%m-%d %H:%M:%S",
)
)
_logger.addHandler(_handler)
__all__ = [
"AutoRegressiveLM",
"AutoRegressiveLMConfig",
astrai/dataset/__init__.py
-6
View File
@@ -6,7 +6,6 @@ from astrai.dataset.dataset import (
)
from astrai.dataset.sampler import RDSampler
from astrai.dataset.storage import (
H5Store,
JsonlStore,
MmapStore,
Recordable,
@@ -17,9 +16,7 @@ from astrai.dataset.storage import (
)
from astrai.serialization import (
load_bin,
load_h5,
save_bin,
save_h5,
)
__all__ = [
@@ -31,12 +28,9 @@ __all__ = [
"Streamable",
"Recordable",
"StoreFactory",
"H5Store",
"MmapStore",
"JsonlStore",
"detect_format",
"save_h5",
"load_h5",
"save_bin",
"load_bin",
"RDSampler",
astrai/dataset/dataset.py
+3 -3
View File
@@ -314,7 +314,7 @@ class DatasetFactory(BaseFactory["BaseDataset"]):
stream datasets (SEQ/SFT). Record datasets ignore it.
stride: Stride between consecutive stream samples
(default: same as *window_size*).
storage_type: Storage backend ("h5", "bin", "jsonl") or
storage_type: Storage backend ("bin", "jsonl") or
None for auto-detection.
tokenizer_path: Path to tokenizer for lazy JSONL
tokenisation (record datasets only).
@@ -384,7 +384,7 @@ class DatasetFactory(BaseFactory["BaseDataset"]):
"""Build an on-the-fly tokenisation processor if applicable.
Only raw JSONL + record datasets (DPO/GRPO) need a processor;
pre-tokenised backends (H5/bin) and stream datasets (SEQ/SFT)
pre-tokenised backends (bin) and stream datasets (SEQ/SFT)
return ``None`` so no tokenizer is loaded.
"""
if tokenizer_path is None or storage_type != "jsonl":
@@ -451,7 +451,7 @@ class DPODataset(BaseDataset):
Two loading paths (handled by :class:`DatasetFactory`):
- **Pre-tokenized** (H5/bin): ``store.load(path)`` reads per-record
- **Pre-tokenized** (bin): ``store.load(path)`` reads per-record
tensors; ``__getitem__`` returns them directly.
- **Raw JSONL** (``tokenizer_path=...``): builds a lazy processor
via :func:`dpo_processor` that tokenises on the fly — no packing,
astrai/dataset/storage.py
+7 -45
View File
@@ -10,7 +10,6 @@ Architecture (composition over inheritance):
Streamable (mixin) — raw token slice fetch(begin, end, keys)
Recordable (mixin) — raw record slice fetch_record(idx, keys)
H5Store(Store, Streamable, Recordable)
MmapStore(Store, Streamable, Recordable)
JsonlStore(Store, Streamable, Recordable)
@@ -36,9 +35,9 @@ control. ``store.token_count`` is the total stream token count (what
``len(store)`` used to mean in the legacy stream-only API).
``segments_are_records`` (class attribute on each Store subclass)
tells ``_normalize`` whether segments are inherently per-record (H5/
JSONL) or opaque shards (bin). Record access for bin relies on
``_offsets`` instead.
tells ``_normalize`` whether segments are inherently per-record (JSONL)
or opaque shards (bin). Record access for bin relies on ``_offsets``
instead.
:class:`JsonlStore` supports a lazy mode (``processor=fn``) that keeps
raw records and defers tokenisation to ``fetch_record`` — used by DPO
@@ -62,7 +61,6 @@ from astrai.preprocessing.transform import TokenizeTransform
from astrai.serialization import (
load_bin,
load_bin_offsets,
load_h5,
)
logger = logging.getLogger(__name__)
@@ -83,19 +81,10 @@ def detect_format(load_path: str) -> str:
root = Path(load_path)
if root.is_file():
suffix = root.suffix.lower()
if suffix in (".h5", ".hdf5"):
return "h5"
if suffix == ".jsonl":
return "jsonl"
raise ValueError(f"Unsupported file format: {suffix}")
h5_files = [
Path(p)
for pattern in ("*.h5", "*.hdf5")
for p in glob.glob(str(root / "**" / pattern), recursive=True)
]
if h5_files:
return "h5"
bin_files = [Path(p) for p in glob.glob(str(root / "**" / "*.bin"), recursive=True)]
if bin_files:
has_meta = (root / "meta.json").exists() or len(
@@ -185,7 +174,7 @@ class Store(ABC):
"""Number of records available via :meth:`fetch_record`.
Non-zero only when the backing layout provides per-record
indexing (H5/JSONL segments or bin ``_offsets``).
indexing (JSONL segments or bin ``_offsets``).
"""
return self._num_records
@@ -269,7 +258,7 @@ class Store(ABC):
Record mode: if *offsets* is provided (bin layout),
``_offsets[key]`` stores cumulative per-record offsets into the
single concatenated segment. Otherwise, when
``segments_are_records`` is True (H5/JSONL), ``_data[key]`` is
``segments_are_records`` is True (JSONL), ``_data[key]`` is
a per-record list and ``fetch_record`` indexes it directly.
Nested keys (GRPO ``responses``/``masks`` as
@@ -305,7 +294,7 @@ class Store(ABC):
logger.warning(
"Key '%s' has %d segments with offsets — record mode "
"disabled for this key (multi-shard bin+offsets not "
"supported). Merge shards or use H5/JSONL.",
"supported). Merge shards or use JSONL.",
key,
len(segs),
)
@@ -330,7 +319,7 @@ class Streamable:
Stateless trait relying on ``self._data``, ``self._cum``,
``self._length`` maintained by :class:`Store`. Stream mode is
active when the owning store has ``window_size > 0``; for stores
that can also serve record access (H5/JSONL/bin+offsets), the
that can also serve record access (JSONL/bin+offsets), the
``fetch_record`` API from :class:`Recordable` is used instead.
"""
@@ -415,33 +404,6 @@ class StoreFactory(BaseFactory["Store"]):
"""Factory for creating Store instances by type name."""
@StoreFactory.register("h5")
class H5Store(Store, Streamable, Recordable):
"""HDF5-based storage backend (pre-tokenized data).
Each key is stored as a group of per-record datasets (``data_0``,
``data_1``, …). Supports both access modes:
- **Stream**: ``fetch(begin, end, key)`` and ``store[i]`` slice
across concatenated records via ``_cum`` — used by SEQ/SFT.
- **Record**: ``fetch_record(i, key)`` and ``store[i]`` (when
``window_size == 0``) index ``_data[key]`` directly — used by
DPO/GRPO.
"""
segments_are_records = True
def __init__(
self,
window_size: int = 0,
stride: Optional[int] = None,
):
super().__init__(window_size=window_size, stride=stride)
def load(self, path: str, **kwargs):
self._normalize(load_h5(path))
@StoreFactory.register("bin")
class MmapStore(Store, Streamable, Recordable):
"""Memory-mapped binary storage backend.
astrai/preprocessing/pipeline.py
+1 -1
View File
@@ -1,7 +1,7 @@
"""Config-driven JSONL preprocessing pipeline.
Composes a :class:`BaseMaskBuilder` (selected by ``input.type``) with
sharding and flush to ``.h5`` / ``.bin`` storage. Packing, position-id
sharding and flush to ``.bin`` storage. Packing, position-id
generation and storage writing are each delegated to pluggable strategies,
dispatched by configuration keys.
astrai/preprocessing/writer.py
+2 -21
View File
@@ -1,7 +1,7 @@
"""Storage writer strategies for pipeline output.
The :class:`StoreWriter` abstraction decouples the pipeline from the
concrete storage format (bin / h5). The pipeline builds a ``{key:
concrete storage format (bin). The pipeline builds a ``{key:
List[Tensor]}`` dict and delegates the write to the writer selected
by ``output.storage_format``.
"""
@@ -15,7 +15,7 @@ from typing import Dict, List
import torch
from astrai.factory import BaseFactory
from astrai.serialization import save_bin, save_h5
from astrai.serialization import save_bin
logger = logging.getLogger(__name__)
@@ -54,22 +54,3 @@ class BinWriter(StoreWriter):
exc_info=True,
)
raise
@StoreWriterFactory.register("h5")
class H5Writer(StoreWriter):
def save(self, output_dir, domain, shard_idx, tensors):
chunk_dir = os.path.join(output_dir, domain)
file_path = os.path.join(chunk_dir, f"data_{shard_idx:04d}.h5")
try:
save_h5(chunk_dir, f"data_{shard_idx:04d}", tensors)
except Exception:
if os.path.exists(file_path):
os.remove(file_path)
logger.error(
"Failed to write shard %s/data_%04d.h5, cleaned up partial output",
domain,
shard_idx,
exc_info=True,
)
raise
astrai/serialization/__init__.py
-4
View File
@@ -20,9 +20,7 @@ from astrai.serialization.checkpoint import (
from astrai.serialization.dataset import (
load_bin,
load_bin_offsets,
load_h5,
save_bin,
save_h5,
)
__all__ = [
@@ -39,7 +37,5 @@ __all__ = [
"save_torch",
"load_bin",
"load_bin_offsets",
"load_h5",
"save_bin",
"save_h5",
]
astrai/serialization/dataset.py
+4 -45
View File
@@ -1,55 +1,14 @@
"""Dataset storage serialization helpers (HDF5 / memory-mapped binary)."""
"""Dataset storage serialization helpers (memory-mapped binary)."""
import json
import os
from pathlib import Path
from typing import Any, Dict, List, Optional
import h5py
import numpy as np
import torch
from torch import Tensor
def save_h5(file_path: str, file_name: str, tensor_group: Dict[str, List[Tensor]]):
os.makedirs(file_path, exist_ok=True)
full_file_path = os.path.join(file_path, f"{file_name}.h5")
with h5py.File(full_file_path, "w") as f:
for key, tensors in tensor_group.items():
grp = f.create_group(key)
for idx, tensor in enumerate(tensors):
arr = tensor.cpu().numpy()
grp.create_dataset(f"data_{idx}", data=arr)
def load_h5(file_path: str, share_memory=True) -> Dict[str, List[Tensor]]:
tensor_group: Dict[str, List[Tensor]] = {}
root_path = Path(file_path)
if root_path.is_file() and root_path.suffix in (".h5", ".hdf5"):
h5_files = [root_path]
else:
h5_files = list(root_path.rglob("*.h5")) + list(root_path.rglob("*.hdf5"))
for h5_file in h5_files:
with h5py.File(h5_file, "r") as f:
for key in f.keys():
grp = f[key]
dsets = []
for dset_name in grp.keys():
dset = grp[dset_name]
tensor = torch.from_numpy(dset[:])
if share_memory:
tensor = tensor.share_memory_()
dsets.append(tensor)
if tensor_group.get(key) is None:
tensor_group[key] = []
tensor_group[key].extend(dsets)
return tensor_group
def save_bin(
file_path: str,
tensor_group: Dict[str, List[Tensor]],
@@ -65,7 +24,7 @@ def save_bin(
offsets, preserving backward compatibility.
Nested keys (``List[List[Tensor]]`` such as GRPO ``responses``) are
not supported in bin format — use H5 for those.
not supported in bin format — use JSONL for those.
"""
os.makedirs(file_path, exist_ok=True)
record_keys = set(record_keys or [])
@@ -74,7 +33,7 @@ def save_bin(
if tensors and isinstance(tensors[0], list):
raise ValueError(
f"Nested key '{key}' (List[List[Tensor]]) is not supported "
f"in bin format. Use H5 or JSONL storage instead."
f"in bin format. Use JSONL storage instead."
)
cat = torch.cat(tensors, dim=0)
entry: Dict[str, Any] = {
@@ -112,7 +71,7 @@ def load_bin_offsets(file_path: str) -> Dict[str, List[int]]:
Returns an empty dict when no key has offsets (legacy bin files),
in which case record-mode access falls back to per-record segment
indexing (H5/JSONL layout).
indexing (JSONL layout).
"""
with open(os.path.join(file_path, "meta.json"), "r") as f:
meta = json.load(f)