torch-vault

Embedded Rust engine for zero-copy PyTorch checkpoint management and SQL-driven dataset streaming using safetensors, SQLite WAL mode, and DuckDB OLAP predicate pushdown.

Installation

Installing the Python Package & CLI

pip install torch-vault

Building from Source (Rust & PyO3)

Building from source requires Python >= 3.10 and a Rust toolchain (cargo and rustc):

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

Quickstart

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

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

# Save weight checkpoint ONLY IF val_loss improves, while continuously logging metrics to SQLite
ckpt_id = vault.save(
    model,
    optimizer,
    metrics={"val_loss": 0.342, "accuracy": 0.912},
    save_if_better="val_loss",
    mode="min"
)

# Export SQLite metrics history directly to a Pandas DataFrame
df = vault.to_dataframe()

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

CLI & Interactive WebApp UI Inspector

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

# Print terminal table of all checkpoints and logged metrics
torch-vault list ./checkpoints

About & Architecture Deep-Dive

torch-vault replaces un-indexed, pickle-based PyTorch state dict saving with a compiled Rust engine linked to Python via PyO3 bindings.

Checkpoint Subsystem (safetensors + SQLite WAL)

Model weight tensors are serialized into binary safetensors format using zero-copy memory mapping (memmap2). Checkpoint metadata (epoch, step, metrics, tags) is stored in an embedded SQLite index database (vault_index.db) configured with Write-Ahead Logging (WAL) for concurrent read access during training.

Dataset & Streaming Subsystem (DuckDB + Arrow C Stream)

Dataset filtering executes directly on NVMe disk inside an embedded DuckDB instance before crossing into Python memory. Filtered record batches cross the PyO3 boundary via zero-copy C-ABI Arrow memory pointers (__arrow_c_stream__ PyCapsule) into PyTorch DataLoaders.

Architectural Comparison Matrix

End-to-End Production Showcase (PyTorch LSTM Sentiment Classifier)

import torch
import torch.nn as nn
from torch_vault import Vault, Segment
from torch_vault.transforms import (
    Compose, LabelEncoder, Lambda, MinMaxScaler, RenameKey, SelectKeys
)

class SentimentLSTM(nn.Module):
    def __init__(self, vocab_size: int, embed_dim: int = 64, hidden_dim: int = 128):
        super().__init__()
        self.embedding = nn.Embedding(vocab_size, embed_dim, padding_idx=0)
        self.lstm = nn.LSTM(embed_dim, hidden_dim, batch_first=True)
        self.fc = nn.Linear(hidden_dim, 2)

    def forward(self, x):
        out, (hn, cn) = self.lstm(self.embedding(x))
        return self.fc(hn[-1])

vocab = {"<pad>": 0, "<unk>": 1, "good": 2, "bad": 3, "great": 4}

def preprocess_sample(row: dict) -> dict:
    tokens = [vocab.get(w, 1) for w in str(row["text"]).lower().split()[:50]]
    tokens += [0] * (50 - len(tokens))
    return {
        "input_ids": torch.tensor(tokens, dtype=torch.long),
        "target": torch.tensor(row["sentiment_score"], dtype=torch.long),
    }

label_encoder = LabelEncoder(key="raw_label", classes=["negative", "positive"])
pipeline = Compose([
    RenameKey(src="raw_label", dst="sentiment_score"),
    label_encoder,
    Lambda(preprocess_sample),
    SelectKeys("input_ids", "target"),
])

train_loader, val_loader, test_loader = (
    Segment("data/reviews.parquet")
    .where("length(text) > 10")
    .transform(pipeline)
    .to_dataloader(batch_size=32, split=(0.8, 0.1, 0.1), seed=42)
)

model = SentimentLSTM(vocab_size=len(vocab) + 10)
optimizer = torch.optim.Adam(model.parameters(), lr=1e-3)
criterion = nn.CrossEntropyLoss()
vault = Vault("./sentiment_vault")

for epoch in range(1, 6):
    model.train()
    for batch in train_loader:
        optimizer.zero_grad()
        loss = criterion(model(batch["input_ids"]), batch["target"])
        loss.backward()
        optimizer.step()

    # Evaluate validation loss
    model.eval()
    val_loss = sum(criterion(model(b["input_ids"]), b["target"]).item() for b in val_loader) / len(val_loader)

    vault.save(
        model=model,
        optimizer=optimizer,
        metrics={"val_loss": val_loss},
        epoch=epoch,
        save_if_better="val_loss",
        mode="min"
    )

best_weights = vault.load(metric="val_loss", mode="min")
model.load_state_dict(best_weights["model"])

Vault API Specification

Segment API Specification

Transforms & Preprocessing Reference