Checkpoints¶

from kura.checkpoints import JSONLCheckpointManager

# Default JSONL checkpointing
checkpoint_manager = JSONLCheckpointManager("./checkpoints", enabled=True)

# Use in pipeline functions
summaries = await summarise_conversations(
    conversations,
    model=summary_model,
    checkpoint_manager=checkpoint_manager
)

from kura.checkpoints import ParquetCheckpointManager

# Parquet checkpointing (requires PyArrow)
checkpoint_manager = ParquetCheckpointManager("./checkpoints", enabled=True)

# Use in pipeline functions
summaries = await summarise_conversations(
    conversations,
    model=summary_model,
    checkpoint_manager=checkpoint_manager
)

from kura.checkpoints import HFDatasetCheckpointManager

# With HuggingFace Hub integration
checkpoint_manager = HFDatasetCheckpointManager(
    "./checkpoints",
    enabled=True,
    hub_repo="your-username/kura-checkpoints",
    hub_token="your_hf_token"
)

# Handle datasets larger than RAM
checkpoint_manager = HFDatasetCheckpointManager(
    "./checkpoints",
    streaming=True
)

# Filter checkpoints without loading everything
filtered_summaries = checkpoint_manager.filter_checkpoint(
    "summaries",
    lambda x: x["concerning_score"] > 3,
    ConversationSummary,
    "summaries"
)

from kura.checkpoints import HFDatasetCheckpointManager

# Basic HF dataset checkpointing
checkpoint_manager = HFDatasetCheckpointManager("./checkpoints", enabled=True)

# Save with explicit checkpoint type
checkpoint_manager.save_checkpoint("summaries", summaries, "summaries")

# Load with checkpoint type
loaded = checkpoint_manager.load_checkpoint(
    "summaries",
    ConversationSummary,
    checkpoint_type="summaries"
)

from kura.checkpoints import JSONLCheckpointManager

checkpoint_manager = JSONLCheckpointManager(
    checkpoint_dir="./checkpoints",
    enabled=True  # Set to False to disable checkpointing
)

from kura.checkpoints import ParquetCheckpointManager

checkpoint_manager = ParquetCheckpointManager(
    checkpoint_dir="./checkpoints",
    enabled=True,
    compression='snappy'  # Options: 'snappy', 'gzip', 'brotli', 'lz4', 'zstd'
)

from kura.checkpoints import HFDatasetCheckpointManager

checkpoint_manager = HFDatasetCheckpointManager(
    checkpoint_dir="./checkpoints",
    enabled=True,
    hub_repo="username/dataset-name",  # Optional: HF Hub repository
    hub_token="your_token",           # Optional: HF token for private repos
    streaming=False,                  # Enable streaming for large datasets
    compression="gzip"               # Compression: 'gzip', 'lz4', 'zstd', None
)

import os
from pathlib import Path

def compare_checkpoint_sizes(base_dir):
    """Compare checkpoint sizes across formats."""
    formats = ["jsonl", "parquet", "hf"]
    sizes = {}

    for fmt in formats:
        fmt_dir = Path(base_dir) / fmt
        if fmt_dir.exists():
            if fmt == "hf":
                # HF datasets are directories
                sizes[fmt] = sum(f.stat().st_size for f in fmt_dir.rglob("*") if f.is_file())
            else:
                # JSONL and Parquet are single files
                checkpoint_files = list(fmt_dir.glob(f"*.{fmt}"))
                sizes[fmt] = sum(f.stat().st_size for f in checkpoint_files)

    return sizes

# Usage example with tutorial_checkpoints
sizes = compare_checkpoint_sizes("./tutorial_checkpoints")
for fmt, size in sizes.items():
    print(f"{fmt.upper()}: {size / 1024:.0f}KB")

# Expected output:
# JSONL: 200KB
# PARQUET: 100KB
# HF: 186KB

try:
    summaries = checkpoint_manager.load_checkpoint("summaries", ConversationSummary)
    if summaries is None:
        print("No checkpoint found, will generate from scratch")
except Exception as e:
    print(f"Checkpoint loading failed: {e}")
    # Fallback to regeneration

# Use HuggingFace with streaming
hf_manager = HFDatasetCheckpointManager(
    "./checkpoints",
    streaming=True,           # Enable streaming
    compression="zstd"        # Better compression
)

# Use Parquet with appropriate compression
parquet_manager = ParquetCheckpointManager(
    "./checkpoints",
    compression="snappy"      # Fast decompression
)

# Use JSONL for easy debugging
jsonl_manager = JSONLCheckpointManager("./checkpoints")
# Files are human-readable and can be inspected with text editors

# Handle datasets that don't fit in memory
manager = HFDatasetCheckpointManager("./checkpoints", streaming=True)

# Save normally
manager.save_checkpoint("large_summaries", huge_summaries, "summaries")

# Load with streaming - processes chunks automatically
loaded = manager.load_checkpoint("large_summaries", ConversationSummary,
                                checkpoint_type="summaries", streaming=True)

# Sync with HuggingFace Hub
manager = HFDatasetCheckpointManager(
    "./checkpoints",
    hub_repo="your-org/kura-analysis-2024",
    hub_token="hf_your_token_here"
)

# Automatically pushes to Hub after saving
manager.save_checkpoint("results", summaries, "summaries")

# Get detailed information about checkpoints
info = manager.get_checkpoint_info("summaries")
print(f"Rows: {info['num_rows']}")
print(f"Size: {info['size_bytes']} bytes")
print(f"Features: {info['features']}")

# Complex filtering without loading full dataset
recent_concerns = manager.filter_checkpoint(
    "summaries",
    lambda x: x["concerning_score"] >= 4 and x["user_frustration"] >= 3,
    ConversationSummary,
    "summaries"
)

# Topic-based filtering
python_questions = manager.filter_checkpoint(
    "summaries",
    lambda x: "python" in x["topic"].lower(),
    ConversationSummary,
    "summaries"
)

# Use streaming for large datasets
manager = HFDatasetCheckpointManager("./checkpoints", streaming=True)

# Or increase batch processing for Parquet
# Process in smaller chunks before saving

from kura.checkpoints import CheckpointManager, MultiCheckpointManager, ParquetCheckpointManager

# Create individual managers
jsonl_mgr = CheckpointManager("./checkpoints/jsonl")
parquet_mgr = ParquetCheckpointManager("./checkpoints/parquet")

# Coordinate them with MultiCheckpointManager
multi_mgr = MultiCheckpointManager(
    [jsonl_mgr, parquet_mgr],
    save_strategy="all_enabled",  # Save to both
    load_strategy="first_found"    # Load from first available
)

# Use in pipeline
summaries = await summarise_conversations(
    conversations,
    model=summary_model,
    checkpoint_manager=multi_mgr
)

multi_mgr = MultiCheckpointManager(
    managers,
    save_strategy="all_enabled"
)
# Saves to: JSONL ✓, Parquet ✓, HF ✓

multi_mgr = MultiCheckpointManager(
    managers,
    save_strategy="primary_only"
)
# Saves to: JSONL ✓, Parquet ✗, HF ✗

multi_mgr = MultiCheckpointManager(
    managers,
    load_strategy="first_found"
)
# Tries: JSONL → Parquet → HF (stops at first success)

multi_mgr = MultiCheckpointManager(
    managers,
    load_strategy="priority"
)
# Tries: JSONL only (fails if not found)

# Fast local storage for development
local_mgr = CheckpointManager("./checkpoints/local")

# Network storage for backup
network_mgr = CheckpointManager("/mnt/shared/checkpoints")

multi_mgr = MultiCheckpointManager(
    [local_mgr, network_mgr],
    save_strategy="all_enabled",    # Save to both
    load_strategy="first_found"     # Load from local first
)

# Migrate from JSONL to Parquet
old_mgr = CheckpointManager("./old_jsonl")
new_mgr = ParquetCheckpointManager("./new_parquet")

migration_mgr = MultiCheckpointManager(
    [old_mgr, new_mgr],
    save_strategy="all_enabled",  # Save to both during migration
    load_strategy="priority"      # Load from old format
)

# Run pipeline - will save to both formats
summaries = await summarise_conversations(
    conversations,
    model=summary_model,
    checkpoint_manager=migration_mgr
)

# Mix formats for optimal performance
jsonl_mgr = CheckpointManager("./jsonl")      # Human-readable
parquet_mgr = ParquetCheckpointManager("./parquet")  # Compressed
hf_mgr = HFDatasetCheckpointManager("./hf")  # Advanced features

multi_mgr = MultiCheckpointManager(
    [jsonl_mgr, parquet_mgr, hf_mgr],
    save_strategy="all_enabled",
    load_strategy="first_found"
)

# Benefits:
# - JSONL for debugging
# - Parquet for space efficiency
# - HF for advanced queries

# Handle failures gracefully
primary_mgr = CheckpointManager("/fast/ssd", enabled=True)
backup_mgr = CheckpointManager("/slow/hdd", enabled=True)
archive_mgr = CheckpointManager("/glacier", enabled=False)  # Disabled

multi_mgr = MultiCheckpointManager(
    [primary_mgr, backup_mgr, archive_mgr]
)
# Will use only enabled managers automatically

# Get information about the multi-checkpoint setup
stats = multi_mgr.get_stats()
print(f"Enabled managers: {stats['enabled_managers']}/{stats['num_managers']}")
print(f"Save strategy: {stats['save_strategy']}")
print(f"Load strategy: {stats['load_strategy']}")

for mgr_stat in stats["managers"]:
    print(f"- {mgr_stat['type']}: {mgr_stat['checkpoint_count']} files")

# List all unique checkpoints across managers
all_checkpoints = multi_mgr.list_checkpoints()
print(f"Available checkpoints: {all_checkpoints}")

# Delete from all managers
multi_mgr.delete_checkpoint("old_summaries.jsonl")

import asyncio
from kura.checkpoints import CheckpointManager, MultiCheckpointManager, ParquetCheckpointManager
from kura.summarisation import summarise_conversations, SummaryModel
from kura.cluster import generate_base_clusters_from_conversation_summaries, ClusterDescriptionModel
from kura.meta_cluster import reduce_clusters_from_base_clusters, MetaClusterModel

async def run_analysis_with_multi_checkpoints(conversations):
    # Set up multi-backend checkpointing
    jsonl_mgr = CheckpointManager("./checkpoints/jsonl")
    parquet_mgr = ParquetCheckpointManager("./checkpoints/parquet")

    checkpoint_mgr = MultiCheckpointManager(
        [jsonl_mgr, parquet_mgr],
        save_strategy="all_enabled",  # Redundancy
        load_strategy="first_found"   # Performance
    )

    # Run pipeline with automatic multi-backend checkpointing
    summaries = await summarise_conversations(
        conversations,
        model=SummaryModel(model="gpt-4o-mini"),
        checkpoint_manager=checkpoint_mgr
    )

    clusters = await generate_base_clusters_from_conversation_summaries(
        summaries,
        model=ClusterDescriptionModel(),
        checkpoint_manager=checkpoint_mgr
    )

    meta_clusters = await reduce_clusters_from_base_clusters(
        clusters,
        model=MetaClusterModel(max_clusters=10),
        checkpoint_manager=checkpoint_mgr
    )

    # Check storage usage
    stats = checkpoint_mgr.get_stats()
    print(f"\nCheckpoint storage summary:")
    for mgr_stat in stats["managers"]:
        print(f"- {mgr_stat['type']}: {mgr_stat['checkpoint_count']} files")

    return meta_clusters

# Run the analysis
conversations = load_conversations()  # Your data loading logic
results = asyncio.run(run_analysis_with_multi_checkpoints(conversations))