torch-vault

Zero-copy PyTorch checkpoint management and SQL-driven dataset streaming built with Rust, PyO3, SQLite, and DuckDB.

Features

• Safe Checkpointing: Replaces pickle-based torch.save with safetensors memory-mapping backed by an atomic SQLite WAL index.
• Selective Weight Saving: vault.save(model, metrics=..., save_if_better="val_loss") continuously records metric history into SQLite while writing .safetensors weight files to disk only when metrics improve.
• SQL Dataset Streaming: Segment("data.parquet").where("quality > 0.8") runs DuckDB predicate pushdown on disk before streaming records into PyTorch DataLoaders.
• Zero-Copy Memory Interchange: Uses the Arrow C Stream Interface (__arrow_c_stream__ PyCapsule) to pass record batches directly into memory.
• CLI & WebApp Inspector: Includes torch-vault list and torch-vault ui for inspecting checkpoints and charting metrics in a local browser interface.

Installation

pip install torch-vault

For development build:

git clone https://github.com/ethqnol/torch-vault.git
cd torch-vault
pip install maturin
maturin develop --release

Quickstart

Checkpoint Management (Vault)

import torch
import torch.nn as nn
from torch_vault import Vault

model = nn.Linear(10, 2)
optimizer = torch.optim.Adam(model.parameters(), lr=1e-3)
vault = Vault("./checkpoints")

for epoch in range(1, 11):
    train_loss = 0.5 # compute training loss
    val_loss = 0.4   # compute validation loss

    # Log metrics to SQLite; save weight file on disk only if val_loss improves
    vault.save(
        model=model,
        optimizer=optimizer,
        metrics={"train_loss": train_loss, "val_loss": val_loss},
        epoch=epoch,
        save_if_better="val_loss",
        mode="min",
    )

# Load best model weights automatically by metric
best_weights = vault.load(metric="val_loss", mode="min")
model.load_state_dict(best_weights["model"])

SQL Dataset Streaming & Train/Val/Test Split (Segment)

from torch_vault import Segment
from torch_vault.transforms import Compose, SelectKeys, RenameKey, StandardScaler

pipeline = Compose([
    RenameKey(src="score", dst="quality"),
    StandardScaler(keys=["quality"]),
    SelectKeys("quality", "label"),
])

# Stream filtered Parquet data with a 80%/10%/10% split
train_loader, val_loader, test_loader = (
    Segment("data/dataset.parquet")
    .where("quality > 0.7 AND split = 'train'")
    .transform(pipeline)
    .to_dataloader(batch_size=32, split=(0.8, 0.1, 0.1), seed=42)
)

CLI & WebApp Inspector

# Print terminal ASCII table of logged checkpoints
torch-vault list ./checkpoints

# Launch browser WebApp inspector at http://127.0.0.1:8000/
torch-vault ui ./checkpoints --port 8000

Documentation

Full interactive API reference and architecture specifications are available at ethqnol.github.io/torch-vault.

License

MIT License.