tuft 0.1.1__py3-none-any.whl → 0.1.3__py3-none-any.whl

tuft/persistence/redis_store.py

@@ -8,9 +8,14 @@ Key Design:
 - Nested records: {namespace}::{type}::{parent_id}::{nested_type}::{nested_id}
 
 Persistence Modes:
-- disabled: No persistence, all data is in-memory only
-- redis_url: Use external Redis server via URL
-- file_redis: Use file-backed storage for tests and demos
+- DISABLE: No persistence, all data is in-memory only
+- REDIS: Use external Redis server via URL
+- FILE: Use file-backed storage for tests and demos
+
+Config Validation:
+- On startup, the current config signature is compared with the stored signature
+- If mismatch is detected, server stops with an error message
+- Use `tuft clear persistence` to override and clear existing data
 """
 
 from __future__ import annotations
@@ -19,12 +24,12 @@ import logging
 import os
 import threading
 import time
-from dataclasses import dataclass
+from datetime import datetime, timezone
 from enum import Enum
 from pathlib import Path
 from typing import Any, TypeVar
 
-from pydantic import BaseModel
+from pydantic import BaseModel, Field
 
 
 logger = logging.getLogger(__name__)
@@ -50,71 +55,104 @@ T = TypeVar("T", bound=BaseModel)
 class PersistenceMode(str, Enum):
     """Persistence mode options."""
 
-    DISABLED = "disabled"  # No persistence
-    REDIS_URL = "redis_url"  # Use external Redis server
-    FILE_REDIS = "file_redis"  # Use file-backed storage for tests/demos
+    DISABLE = "DISABLE"  # No persistence
+    REDIS = "REDIS"  # Use external Redis server
+    FILE = "FILE"  # Use file-backed storage for tests/demos
 
 
 # Default TTL values in seconds
 DEFAULT_FUTURE_TTL_SECONDS = 24 * 3600  # 1 day for future records (short-lived)
 
 
-@dataclass
-class PersistenceConfig:
+class ConfigCheckField:
+    """Available fields that can be checked for configuration validation.
+
+    Field names correspond directly to AppConfig attribute names.
+    SUPPORTED_MODELS is always required (mandatory) for restore safety.
+    """
+
+    SUPPORTED_MODELS = "SUPPORTED_MODELS"
+    CHECKPOINT_DIR = "CHECKPOINT_DIR"
+    MODEL_OWNER = "MODEL_OWNER"
+    TOY_BACKEND_SEED = "TOY_BACKEND_SEED"
+    AUTHORIZED_USERS = "AUTHORIZED_USERS"
+    TELEMETRY = "TELEMETRY"
+
+
+# Default fields to check (supported_models is mandatory)
+DEFAULT_CHECK_FIELDS: list[str] = [ConfigCheckField.SUPPORTED_MODELS]
+
+
+class PersistenceConfig(BaseModel):
     """Configuration for Redis persistence.
 
     Attributes:
-        mode: Persistence mode - disabled, redis_url, or file_redis
-        redis_url: Redis server URL (only used when mode=redis_url)
-        file_path: JSON file path (only used when mode=file_redis)
-        namespace: Key namespace prefix
+        mode: Persistence mode - DISABLE, REDIS, or FILE
+        redis_url: Redis server URL (only used when mode=REDIS)
+        file_path: JSON file path (only used when mode=FILE)
+        namespace: Key namespace prefix for Redis keys. Defaults to "persistence-tuft-server".
         future_ttl_seconds: TTL for future records in seconds. None means no expiry.
+        check_fields: List of AppConfig fields to validate on restart.
+                     Defaults to ["SUPPORTED_MODELS"]. SUPPORTED_MODELS is always
+                     checked regardless of this setting for restore safety.
+                     Available fields: SUPPORTED_MODELS, CHECKPOINT_DIR, MODEL_OWNER,
+                     TOY_BACKEND_SEED, AUTHORIZED_USERS, TELEMETRY.
     """
 
-    mode: PersistenceMode = PersistenceMode.DISABLED
+    # Allow Path type
+    model_config = {"arbitrary_types_allowed": True}
+
+    mode: PersistenceMode = PersistenceMode.DISABLE
     redis_url: str = "redis://localhost:6379/0"
     file_path: Path | None = None
-    namespace: str = "tuft"
-    future_ttl_seconds: int | None = DEFAULT_FUTURE_TTL_SECONDS  # Futures expire after 1 day
+    namespace: str = "persistence-tuft-server"  # Default namespace for Redis keys
+    future_ttl_seconds: int | None = DEFAULT_FUTURE_TTL_SECONDS
+    check_fields: list[str] = Field(default_factory=lambda: DEFAULT_CHECK_FIELDS.copy())
 
     @property
     def enabled(self) -> bool:
         """Check if persistence is enabled."""
-        return self.mode != PersistenceMode.DISABLED
+        return self.mode != PersistenceMode.DISABLE
 
-    @classmethod
-    def disabled(cls, namespace: str = "tuft") -> "PersistenceConfig":
-        """Create a disabled persistence config."""
-        return cls(mode=PersistenceMode.DISABLED, namespace=namespace)
+    def get_check_fields(self) -> list[str]:
+        """Get the fields to check, ensuring SUPPORTED_MODELS is always included."""
+        fields = list(self.check_fields)
+        if ConfigCheckField.SUPPORTED_MODELS not in fields:
+            fields.insert(0, ConfigCheckField.SUPPORTED_MODELS)
+        return fields
 
     @classmethod
     def from_redis_url(
         cls,
         redis_url: str,
-        namespace: str = "tuft",
+        namespace: str = "persistence-tuft-server",
         future_ttl_seconds: int | None = DEFAULT_FUTURE_TTL_SECONDS,
+        check_fields: list[str] | None = None,
     ) -> "PersistenceConfig":
         """Create a config using external Redis server."""
         return cls(
-            mode=PersistenceMode.REDIS_URL,
+            mode=PersistenceMode.REDIS,
             redis_url=redis_url,
             namespace=namespace,
             future_ttl_seconds=future_ttl_seconds,
+            check_fields=check_fields or DEFAULT_CHECK_FIELDS.copy(),
         )
 
     @classmethod
-    def from_file_redis(
+    def from_file(
         cls,
         file_path: Path | None = None,
-        namespace: str = "tuft",
+        namespace: str = "persistence-tuft-server",
         future_ttl_seconds: int | None = DEFAULT_FUTURE_TTL_SECONDS,
+        check_fields: list[str] | None = None,
     ) -> "PersistenceConfig":
         """Create a config using file-backed storage."""
         return cls(
-            mode=PersistenceMode.FILE_REDIS,
+            mode=PersistenceMode.FILE,
             file_path=file_path,
             namespace=namespace,
             future_ttl_seconds=future_ttl_seconds,
+            check_fields=check_fields or DEFAULT_CHECK_FIELDS.copy(),
         )
 
 
@@ -123,7 +161,7 @@ class RedisStore:
 
     Supports two modes:
     - External Redis server (via redis-py)
-    - No persistence (disabled mode)
+    - No persistence (DISABLE mode)
     """
 
     _instance: "RedisStore | None" = None
@@ -166,7 +204,7 @@ class RedisStore:
                 if self._redis is None or self._pid != current_pid:
                     self._close_connections()
 
-                    if self._config.mode in (PersistenceMode.REDIS_URL, PersistenceMode.FILE_REDIS):
+                    if self._config.mode in (PersistenceMode.REDIS, PersistenceMode.FILE):
                         logger.info("Redis connection begin")
                         self._redis = self._create_redis_client()
 
@@ -181,7 +219,7 @@ class RedisStore:
         if self._config is None:
             return None
         try:
-            if self._config.mode == PersistenceMode.FILE_REDIS:
+            if self._config.mode == PersistenceMode.FILE:
                 from .file_redis import FileRedis
 
                 file_path = self._config.file_path or (
@@ -201,7 +239,7 @@ class RedisStore:
 
     @property
     def namespace(self) -> str:
-        return self._config.namespace if self._config else "tuft"
+        return self._config.namespace if self._config else "persistence-tuft-server"
 
     @property
     def future_ttl(self) -> int | None:
@@ -486,3 +524,286 @@ def is_persistence_enabled() -> bool:
 def get_redis_store() -> RedisStore:
     """Get the global Redis store instance."""
     return RedisStore.get_instance()
+
+
+def _now() -> datetime:
+    return datetime.now(timezone.utc)
+
+
+class ConfigSignature(BaseModel):
+    """Stores a complete snapshot of AppConfig for validation on restart.
+
+    Since AppConfig is now a Pydantic model, we directly store its serialized
+    form (excluding the persistence field which is runtime-only).
+    """
+
+    # Serialized AppConfig data (excludes persistence)
+    config_data: dict[str, Any] = Field(default_factory=dict)
+
+    # Metadata
+    created_at: datetime = Field(default_factory=_now)
+    namespace: str = "persistence-tuft-server"
+
+    @classmethod
+    def from_app_config(cls, config: Any) -> "ConfigSignature":
+        """Create a signature by serializing the AppConfig."""
+        # Use the method on AppConfig to get persistence-safe data
+        config_data = config.get_config_for_persistence()
+        namespace = (
+            config.persistence.namespace if config.persistence else "persistence-tuft-server"
+        )
+        return cls(config_data=config_data, namespace=namespace)
+
+    def _get_field_value(self, field_name: str) -> Any:
+        """Get the value of a field by name."""
+        lowercase_field = field_name.lower()
+        return self.config_data.get(lowercase_field)
+
+    def _normalize_for_comparison(self, value: Any) -> Any:
+        if isinstance(value, list):
+            normalized_items = []
+            for item in value:
+                if isinstance(item, dict):
+                    normalized_items.append(tuple(sorted(item.items())))
+                else:
+                    normalized_items.append(item)
+            # Sort for order-independent comparison
+            return sorted(normalized_items, key=lambda x: str(x))
+        return value
+
+    def _compare_field(self, other: "ConfigSignature", field_name: str) -> bool:
+        """Compare a single field between two signatures."""
+        current_value = self._get_field_value(field_name)
+        other_value = other._get_field_value(field_name)
+        current_normalized = self._normalize_for_comparison(current_value)
+        other_normalized = self._normalize_for_comparison(other_value)
+
+        return current_normalized == other_normalized
+
+    def _get_field_diff(self, other: "ConfigSignature", field_name: str) -> dict[str, Any] | None:
+        """Get the difference for a single field.
+
+        Returns:
+            {"current": value, "stored": value} if different, None otherwise.
+        """
+        current_value = self._get_field_value(field_name)
+        other_value = other._get_field_value(field_name)
+
+        current_normalized = self._normalize_for_comparison(current_value)
+        other_normalized = self._normalize_for_comparison(other_value)
+
+        if current_normalized != other_normalized:
+            return {"current": current_value, "stored": other_value}
+        return None
+
+    def matches(
+        self,
+        other: "ConfigSignature",
+        check_fields: list[str] | None = None,
+    ) -> bool:
+        """Check if this signature matches another signature.
+
+        Args:
+            other: The other signature to compare against.
+            check_fields: List of field names to check. If None, uses DEFAULT_CHECK_FIELDS.
+                         SUPPORTED_MODELS is always included (mandatory).
+
+        Returns:
+            True if all specified fields match, False otherwise.
+        """
+        fields_to_check = self._get_fields_to_check(check_fields)
+
+        for field_name in fields_to_check:
+            if not self._compare_field(other, field_name):
+                return False
+        return True
+
+    def get_diff(
+        self,
+        other: "ConfigSignature",
+        check_fields: list[str] | None = None,
+    ) -> dict[str, dict[str, Any]]:
+        """Get the differences between this signature and another.
+
+        Args:
+            other: The other signature to compare against.
+            check_fields: List of field names to check. If None, uses DEFAULT_CHECK_FIELDS.
+                         SUPPORTED_MODELS is always included (mandatory).
+
+        Returns:
+            Dict mapping field names to their differences.
+        """
+        fields_to_check = self._get_fields_to_check(check_fields)
+        diff: dict[str, dict[str, Any]] = {}
+
+        for field_name in fields_to_check:
+            field_diff = self._get_field_diff(other, field_name)
+            if field_diff is not None:
+                diff[field_name] = field_diff
+
+        return diff
+
+    def _get_fields_to_check(self, check_fields: list[str] | None) -> list[str]:
+        """Get the list of fields to check, ensuring mandatory fields are included."""
+        if check_fields is None:
+            return DEFAULT_CHECK_FIELDS.copy()
+
+        # Ensure SUPPORTED_MODELS is always included (mandatory)
+        fields = list(check_fields)
+        if ConfigCheckField.SUPPORTED_MODELS not in fields:
+            fields.insert(0, ConfigCheckField.SUPPORTED_MODELS)
+        return fields
+
+
+CONFIG_SIGNATURE_KEY = "config_signature"
+
+
+def save_config_signature(config: Any) -> bool:
+    """Save the config signature to Redis.
+
+    Args:
+        config: The AppConfig to create a signature from.
+
+    Returns:
+        True if saved successfully, False otherwise.
+    """
+    store = RedisStore.get_instance()
+    if not store.is_enabled:
+        return False
+
+    signature = ConfigSignature.from_app_config(config)
+    key = store.build_key(CONFIG_SIGNATURE_KEY)
+
+    try:
+        json_str = signature.model_dump_json()
+        return store.set(key, json_str)
+    except Exception:
+        logger.exception("Failed to save config signature to Redis")
+        return False
+
+
+def load_config_signature() -> ConfigSignature | None:
+    """Load the config signature from Redis.
+
+    Returns:
+        The stored ConfigSignature, or None if not found.
+    """
+    store = RedisStore.get_instance()
+    if not store.is_enabled:
+        return None
+
+    key = store.build_key(CONFIG_SIGNATURE_KEY)
+
+    try:
+        json_str = store.get(key)
+        if json_str is None:
+            return None
+        return ConfigSignature.model_validate_json(json_str)
+    except Exception:
+        logger.exception("Failed to load config signature from Redis")
+        return None
+
+
+def has_existing_data() -> bool:
+    """Check if there is any existing data in the current namespace.
+
+    Returns:
+        True if any keys exist in the namespace, False otherwise.
+    """
+    store = RedisStore.get_instance()
+    if not store.is_enabled:
+        return False
+
+    pattern = f"{store.namespace}::*"
+    keys = store.keys(pattern)
+    return len(keys) > 0
+
+
+def validate_config_signature(config: Any) -> bool:
+    """Validate that the current config matches the stored config signature.
+
+    This function ONLY reads from Redis, it does NOT write.
+    The signature should be saved after successful restore using
+    save_config_signature().
+
+    The fields to check are read from config.persistence.check_fields.
+    SUPPORTED_MODELS is always checked regardless of this setting.
+
+    This function handles several cases:
+    1. No signature AND no other data in namespace -> fresh start (return True)
+    2. No signature BUT other data exists -> corrupted/incompatible state, raise error
+    3. Signature exists and matches -> OK (return False, not fresh)
+    4. Signature exists but doesn't match -> raise error
+
+    Args:
+        config: The current AppConfig to validate.
+
+    Returns:
+        True if this is a fresh start (no existing data), False otherwise.
+
+    Raises:
+        ConfigMismatchError: If the configs don't match or state is corrupted.
+    """
+    from tuft.exceptions import ConfigMismatchError
+
+    stored = load_config_signature()
+
+    if stored is None:
+        # Check if there's any other data in the namespace
+        if has_existing_data():
+            # Data exists but no signature -> corrupted/incompatible state
+            logger.warning(
+                "Redis namespace has data but no config signature. "
+                "This indicates a corrupted or incompatible persistence state."
+            )
+            raise ConfigMismatchError(
+                diff={
+                    "_state": {
+                        "current": "valid configuration",
+                        "stored": "missing signature (corrupted or legacy data)",
+                    }
+                }
+            )
+        else:
+            # No data at all -> fresh start
+            logger.info("No stored config signature found - fresh start")
+            return True
+
+    # Get check_fields from persistence config
+    check_fields = config.persistence.get_check_fields() if config.persistence else None
+
+    current = ConfigSignature.from_app_config(config)
+    if not current.matches(stored, check_fields=check_fields):
+        diff = current.get_diff(stored, check_fields=check_fields)
+        raise ConfigMismatchError(diff)
+
+    logger.debug("Config signature validated successfully")
+    return False
+
+
+def get_current_namespace() -> str:
+    """Get the current Redis namespace.
+
+    Returns:
+        The namespace string, or 'tuft' if not configured.
+    """
+    store = RedisStore.get_instance()
+    return store.namespace
+
+
+def flush_all_data() -> tuple[int, str]:
+    """Clear all data from the current Redis namespace.
+
+    This removes all keys with the current namespace prefix.
+    Use with caution - this is destructive!
+
+    Returns:
+        A tuple of (number of keys deleted, namespace that was cleared).
+    """
+    store = RedisStore.get_instance()
+    if not store.is_enabled:
+        return 0, store.namespace
+
+    pattern = f"{store.namespace}::*"
+    deleted_count = store.delete_pattern(pattern)
+    return deleted_count, store.namespace

tuft/sampling_controller.py

@@ -20,12 +20,17 @@ from .exceptions import (
     CheckpointAccessDeniedException,
     CheckpointNotFoundException,
     MissingSequenceIDException,
-    SequenceConflictException,
     SessionNotFoundException,
     UnknownModelException,
     UserMismatchException,
 )
-from .persistence import get_redis_store, is_persistence_enabled, load_record, save_record
+from .persistence import (
+    get_redis_store,
+    is_persistence_enabled,
+    load_record,
+    save_record,
+)
+from .sequence_executor import SequenceExecutor
 from .telemetry.metrics import get_metrics
 from .telemetry.tracing import get_tracer
 
@@ -64,7 +69,12 @@ class SamplingSessionRecord(BaseModel):
     model_path: str | None = None
     session_seq_id: int
     last_seq_id: int = -1
+    max_submitted_seq_id: int = -1
     history: list[SamplingHistoryEntry] = Field(default_factory=list)
+    executor: SequenceExecutor = Field(default_factory=SequenceExecutor, exclude=True)
+
+    class Config:
+        arbitrary_types_allowed = True
 
 
 class SamplingController:
@@ -100,6 +110,9 @@ class SamplingController:
             if record.base_model and record.base_model not in self._base_backends:
                 invalid_sessions.append(record.sampling_session_id)
                 continue
+            # Initialize executor with next_sequence_id based on max_submitted_seq_id
+            # to avoid hanging on new requests after restore
+            record.executor = SequenceExecutor(next_sequence_id=record.max_submitted_seq_id + 1)
             self.sampling_sessions[record.sampling_session_id] = record
         for session_id in invalid_sessions:
             store.delete(self._build_key(session_id))
@@ -181,8 +194,10 @@ class SamplingController:
                 if model_path:
                     # model_path should have higher priority than base_model
                     try:
+                        assert self.config.checkpoint_dir is not None
                         parsed_checkpoint = CheckpointRecord.from_tinker_path(
-                            model_path, self.config.checkpoint_dir
+                            model_path,
+                            self.config.checkpoint_dir,
                         )
                     except FileNotFoundError as exc:
                         raise CheckpointNotFoundException(checkpoint_id=model_path) from exc
@@ -221,7 +236,8 @@ class SamplingController:
                     model_path=str(adapter_path) if adapter_path else None,
                     session_seq_id=session_seq_id,
                 )
-                self._save_session(sampling_session_id)
+                loop = asyncio.get_event_loop()
+                await loop.run_in_executor(None, self._save_session, sampling_session_id)
 
                 # Update metrics
                 get_metrics().sampling_sessions_active.add(1, {"base_model": base_model_ref or ""})
@@ -236,11 +252,9 @@ class SamplingController:
         tokens = ",".join(str(token) for token in prompt.to_ints())
         return hashlib.sha1(tokens.encode("utf-8")).hexdigest()[:16]
 
-    def _record_sequence(
+    async def _record_sequence(
         self, record: SamplingSessionRecord, seq_id: int, prompt: types.ModelInput
     ) -> None:
-        if seq_id <= record.last_seq_id:
-            raise SequenceConflictException(expected=record.last_seq_id + 1, got=seq_id)
         record.last_seq_id = seq_id
         entry = SamplingHistoryEntry(
             seq_id=seq_id,
@@ -248,9 +262,10 @@ class SamplingController:
             prompt_hash=self._hash_prompt(prompt),
         )
         record.history.append(entry)
-        self._save_session(record.sampling_session_id)
+        loop = asyncio.get_event_loop()
+        await loop.run_in_executor(None, self._save_session, record.sampling_session_id)
 
-    def _resolve_backend(
+    async def _resolve_backend(
         self, request: types.SampleRequest, user_id: str
     ) -> Tuple[BaseSamplingBackend, str | None]:
         """Resolve the appropriate backend for the sampling request.
@@ -269,7 +284,18 @@ class SamplingController:
                 raise UserMismatchException()
             if request.seq_id is None:
                 raise MissingSequenceIDException()
-            self._record_sequence(record, request.seq_id, request.prompt)
+            # Track the maximum submitted seq_id for recovery purposes
+            if request.seq_id > record.max_submitted_seq_id:
+                record.max_submitted_seq_id = request.seq_id
+                loop = asyncio.get_event_loop()
+                await loop.run_in_executor(None, self._save_session, record.sampling_session_id)
+            await record.executor.submit(
+                sequence_id=request.seq_id,
+                func=self._record_sequence,
+                record=record,
+                seq_id=request.seq_id,
+                prompt=request.prompt,
+            )
             if record.base_model not in self._base_backends:
                 raise UnknownModelException(model_name=record.base_model)
             if record.model_path is None:
@@ -297,7 +323,7 @@ class SamplingController:
             logger.info("Sampling begin for %s", sampling_session_id)
             start_time = time.perf_counter()
 
-            backend, lora_id = self._resolve_backend(request, user_id=user_id)
+            backend, lora_id = await self._resolve_backend(request, user_id=user_id)
             prompt = request.prompt
             sampling_params = request.sampling_params
             num_samples = request.num_samples

tuft/sequence_executor.py

@@ -0,0 +1,72 @@
+import asyncio
+import heapq
+from typing import Any, Callable
+
+from .exceptions import SequenceConflictException, SequenceTimeoutException
+
+
+class SequenceExecutor:
+    """An executor that processes tasks strictly in the order of their `sequence_id`.
+    Out-of-order tasks are buffered until all previous sequence ids have been processed.
+    """
+
+    def __init__(self, timeout: float = 900, next_sequence_id: int = 0) -> None:
+        self.pending_heap = []  # (sequence_id, func, kwargs, future)
+        self.heap_lock = asyncio.Lock()
+        self.next_sequence_id = next_sequence_id
+        self._processing = False
+        self.timeout = timeout
+
+    async def submit(self, sequence_id: int, func: Callable, **kwargs) -> Any:
+        """Submit a task with a specific sequence_id.
+
+        Args:
+            sequence_id (int): The sequence ID of the task.
+            func (Callable): The async function to execute.
+            **kwargs: Keyword arguments to pass to the function.
+
+        Returns:
+            Any: The result of the function execution.
+
+        Raises:
+            SequenceTimeoutException: If the task times out waiting for its turn.
+            SequenceConflictException: If a task with a lower sequence_id has already been
+                processed.
+        """
+        if sequence_id < self.next_sequence_id:
+            raise SequenceConflictException(expected=self.next_sequence_id, got=sequence_id)
+        future = asyncio.Future()
+        async with self.heap_lock:
+            heapq.heappush(self.pending_heap, (sequence_id, func, kwargs, future))
+            # Start processing if not already running
+            if not self._processing:
+                self._processing = True
+                asyncio.create_task(self._process_tasks())
+        try:
+            result = await asyncio.wait_for(future, timeout=self.timeout)
+        except (asyncio.TimeoutError, asyncio.CancelledError) as e:
+            raise SequenceTimeoutException(sequence_id) from e
+        return result
+
+    async def _process_tasks(self):
+        while True:
+            async with self.heap_lock:
+                if not self.pending_heap:
+                    self._processing = False
+                    break
+                # Peek at the smallest sequence_id
+                sequence_id, func, kwargs, future = self.pending_heap[0]
+                if sequence_id != self.next_sequence_id:
+                    # wait next sequence_id
+                    self._processing = False
+                    break
+
+                heapq.heappop(self.pending_heap)
+                self.next_sequence_id += 1
+            try:
+                result = await func(**kwargs)
+                if not future.done():
+                    future.set_result(result)
+            except Exception as e:
+                if not future.done():
+                    future.set_exception(e)