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

tuft/config.py

@@ -2,23 +2,20 @@
 
 from __future__ import annotations
 
-from dataclasses import dataclass, field
 from pathlib import Path
-from typing import Dict, Iterable, List
-
-from .persistence import PersistenceConfig
+from typing import Any, Iterable
 
+from pydantic import BaseModel, Field, model_validator
 
-def _default_checkpoint_dir() -> Path:
-    return Path.home() / ".cache" / "tuft" / "checkpoints"
+from .persistence import PersistenceConfig
 
 
-def _default_persistence_config() -> PersistenceConfig:
-    return PersistenceConfig()
+def _default_checkpoint_dir() -> Path | None:
+    """Return None to let CLI set the default based on TUFT_HOME."""
+    return None
 
 
-@dataclass
-class TelemetryConfig:
+class TelemetryConfig(BaseModel):
     """Configuration for OpenTelemetry integration.
 
     Attributes:
@@ -31,29 +28,65 @@ class TelemetryConfig:
     enabled: bool = False
     service_name: str = "tuft"
     otlp_endpoint: str | None = None
-    resource_attributes: Dict[str, str] = field(default_factory=dict)
+    resource_attributes: dict[str, str] = Field(default_factory=dict)
+
+
+class ModelConfig(BaseModel):
+    model_config = {"arbitrary_types_allowed": True}
+
+    model_name: str  # name used in APIs
+    model_path: Path  # path to model checkpoint
+    max_model_len: int  # maximum context length supported by the model
+    tensor_parallel_size: int = 1  # tensor parallel size
+
+    # default sampling parameters for this model
+    temperature: float = 1.0
+    top_p: float = 1.0
+    top_k: int = -1
+    logprobs: int = 0
+    seed: int = 42
+    min_response_tokens: int = 0
+
+    # default lora setting
+    max_lora_rank: int = 16  # maximum rank for LoRA adapters
+    max_loras: int = 1  # maximum number of LoRA adapters that can be applied simultaneously
 
+    # default training setting
+    micro_batch_size: int = 1  # micro-batch size for training
 
-def _default_telemetry_config() -> TelemetryConfig:
-    return TelemetryConfig()
+    # whether to colocate sampling and training on the same device
+    # only for local testing purposes
+    colocate: bool = False
+    sampling_memory_fraction: float = 0.2  # fraction of GPU memory for sampling
+
+    @model_validator(mode="after")
+    def validate_colocate(self) -> "ModelConfig":
+        if self.colocate and self.tensor_parallel_size != 1:
+            raise ValueError("Colocate option is only supported for tensor_parallel_size=1.")
+        return self
 
 
-@dataclass
-class AppConfig:
-    """Runtime configuration for the TuFT server."""
+class AppConfig(BaseModel):
+    """Runtime configuration for the TuFT server.
 
-    checkpoint_dir: Path = field(default_factory=_default_checkpoint_dir)
-    supported_models: List[ModelConfig] = field(default_factory=list)
+    This is a Pydantic model that can be serialized/deserialized for persistence.
+    """
+
+    model_config = {"arbitrary_types_allowed": True}
+
+    checkpoint_dir: Path | None = Field(default_factory=_default_checkpoint_dir)
+    supported_models: list[ModelConfig] = Field(default_factory=list)
     model_owner: str = "local-user"
     toy_backend_seed: int = 0
     # TODO: Temporary implementation for user authorization,
     # replace with proper auth system later
-    authorized_users: Dict[str, str] = field(default_factory=dict)
-    persistence: PersistenceConfig = field(default_factory=_default_persistence_config)
-    telemetry: TelemetryConfig = field(default_factory=_default_telemetry_config)
+    authorized_users: dict[str, str] = Field(default_factory=dict)
+    persistence: PersistenceConfig = Field(default_factory=PersistenceConfig)
+    telemetry: TelemetryConfig = Field(default_factory=TelemetryConfig)
 
     def ensure_directories(self) -> None:
-        self.checkpoint_dir.mkdir(parents=True, exist_ok=True)
+        if self.checkpoint_dir is not None:
+            self.checkpoint_dir.mkdir(parents=True, exist_ok=True)
 
     def check_validity(self) -> None:
         if not self.supported_models:
@@ -72,50 +105,21 @@ class AppConfig:
             self.supported_models = updated
         return self
 
-
-@dataclass
-class ModelConfig:
-    """Configuration for a specific model."""
-
-    model_name: str  # name used in APIs
-    model_path: Path  # path to model checkpoint
-    max_model_len: int  # maximum context length supported by the model
-    tensor_parallel_size: int = 1  # tensor parallel size
-
-    # default sampling parameters for this model
-    temperature: float = 1.0
-    top_p: float = 1.0
-    top_k: int = -1
-    logprobs: int = 0
-    seed: int = 42
-    min_response_tokens: int = 0
-
-    # default lora setting
-    max_lora_rank: int = 16  # maximum rank for LoRA adapters
-    max_loras: int = 1  # maximum number of LoRA adapters that can be applied simultaneously
-
-    # whether to colocate sampling and training on the same device
-    # only for local testing purposes
-    colocate: bool = False
-    sampling_memory_fraction: float = 0.2  # fraction of GPU memory for sampling
-
-    def __post_init__(self) -> None:
-        if self.colocate and self.tensor_parallel_size != 1:
-            raise ValueError("Colocate option is only supported for tensor_parallel_size=1.")
+    def get_config_for_persistence(self) -> dict[str, Any]:
+        """Get config fields for persistence signature (excludes persistence config itself)."""
+        return self.model_dump(mode="json", exclude={"persistence"})
 
 
 def load_yaml_config(config_path: Path) -> AppConfig:
     """Loads an AppConfig from a YAML file."""
     from omegaconf import OmegaConf
 
-    schema = OmegaConf.structured(AppConfig)
     loaded = OmegaConf.load(config_path)
     try:
-        config = OmegaConf.merge(schema, loaded)
-        app_config = OmegaConf.to_object(config)
-        assert isinstance(app_config, AppConfig), (
-            "Loaded config is not of type AppConfig, which should not happen."
-        )
-        return app_config
+        # Convert OmegaConf to plain dict for Pydantic
+        config_dict = OmegaConf.to_container(loaded, resolve=True)
+        if not isinstance(config_dict, dict):
+            raise ValueError("Config file must contain a dictionary at root level")
+        return AppConfig.model_validate(config_dict)
     except Exception as e:
         raise ValueError(f"Failed to load config from {config_path}: {e}") from e

tuft/exceptions.py

@@ -1,5 +1,7 @@
 """Some custom exceptions."""
 
+from typing import Any
+
 
 class TuFTException(Exception):
     """Base exception for TuFT errors."""
@@ -79,6 +81,15 @@ class SequenceConflictException(FutureException):
         self.got = got
 
 
+class SequenceTimeoutException(FutureException):
+    """Timeout waiting for the expected sequence ID."""
+
+    def __init__(self, expected_sequence_id: int):
+        detail = f"Timeout when waiting for sequence ID {expected_sequence_id}."
+        super().__init__(detail)
+        self.sequence_id = expected_sequence_id
+
+
 class MissingSequenceIDException(FutureException):
     """Missing sequence ID in the request."""
 
@@ -136,3 +147,58 @@ class LossFunctionInputShapeMismatchException(LossFunctionException):
         detail = f"Input tensors must have the same shape. Got shapes: {shapes}"
         super().__init__(detail)
         self.shapes = shapes
+
+
+class LossFunctionUnknownMetricReductionException(LossFunctionException):
+    def __init__(self, reduction_type: str):
+        detail = f"Unknown metric reduction type: {reduction_type}"
+        super().__init__(detail)
+        self.reduction_type = reduction_type
+
+
+class PersistenceException(TuFTException):
+    """Base exception for Persistence related errors."""
+
+
+class ConfigMismatchError(PersistenceException):
+    """Raised when current config doesn't match the stored config in Redis.
+
+    This error occurs during server startup when persistence is enabled and
+    the configuration has changed since the last run. This can cause data
+    corruption when restoring persisted state.
+    """
+
+    def __init__(
+        self,
+        diff: dict[str, dict[str, Any]],
+    ):
+        self.diff = diff
+
+        # Build detailed diff message
+        diff_parts = []
+        for field_name, field_diff in diff.items():
+            # Handle scalar fields (current/stored)
+            current = field_diff.get("current")
+            stored = field_diff.get("stored")
+
+            parts = []
+            if current is not None or stored is not None:
+                parts.append(f"current: {current}, stored: {stored}")
+
+            if parts:
+                diff_parts.append(f"{field_name} ({', '.join(parts)})")
+
+        diff_str = "; ".join(diff_parts) if diff_parts else "unknown difference"
+
+        message = (
+            f"Configuration mismatch detected: {diff_str}.\n"
+            "The current configuration does not match the stored configuration in Redis.\n"
+            "This can cause data corruption when restoring persisted state.\n\n"
+            "Options:\n"
+            "  1. Use a different Redis database (change redis_url in config)\n"
+            "  2. Run `tuft clear persistence -c <config_path>` to clear existing data\n"
+            "     Use `--force` or `-f` to skip confirmation prompt.\n"
+            "     (WARNING: This will delete all persisted sessions, training runs, etc.)\n"
+            "  3. Restore the original configuration that matches the stored data"
+        )
+        super().__init__(message)

tuft/futures.py

@@ -189,6 +189,24 @@ class FutureStore:
                 count += 1
         return count
 
+    def mark_pending_sample_futures_failed(
+        self,
+        error_message: str = "Server restarted while sample request was pending. Please retry.",
+    ) -> int:
+        """Mark all pending sample futures as failed."""
+        count = 0
+        for record in self._records.values():
+            if record.status == "pending" and record.operation_type == "sample":
+                record.status = "failed"
+                record.error = types.RequestFailedResponse(
+                    error=error_message,
+                    category=types.RequestErrorCategory.Server,
+                )
+                record.event.set()
+                self._save_future(record.request_id)
+                count += 1
+        return count
+
     def _store_record(self, record: FutureRecord) -> None:
         self._records[record.request_id] = record
         self._save_future(record.request_id)
@@ -327,8 +345,9 @@ class FutureStore:
             record.payload = payload
             record.status = "ready"
             record.error = None
+            loop = asyncio.get_event_loop()
+            await loop.run_in_executor(None, self._save_future, request_id)
             record.event.set()
-            self._save_future(request_id)
 
             # Update metrics
             get_metrics().futures_completed.add(
@@ -352,8 +371,9 @@ class FutureStore:
                 return
             record.status = "failed"
             record.error = failure
+            loop = asyncio.get_event_loop()
+            await loop.run_in_executor(None, self._save_future, request_id)
             record.event.set()
-            self._save_future(request_id)
 
             # Update metrics
             get_metrics().futures_completed.add(

tuft/loss_fn/__init__.py

@@ -1,5 +1,6 @@
 from typing import Callable, Dict, Tuple
 
+from tinker.lib.chunked_fwdbwd_helpers import REDUCE_MAP
 from torch import Tensor
 from typing_extensions import TypeAlias
 
@@ -7,6 +8,7 @@ from ..exceptions import (
     LossFunctionInputShapeMismatchException,
     LossFunctionMissingInputException,
     LossFunctionNotFoundException,
+    LossFunctionUnknownMetricReductionException,
 )
 
 
@@ -46,3 +48,34 @@ def _check_loss_fn_inputs(
         shapes = [loss_fn_inputs[key].shape for key in required_keys]
         if not all(shape == shapes[0] for shape in shapes):
             raise LossFunctionInputShapeMismatchException(shapes)
+
+
+def metrics_reduction(
+    metric_list: list[dict[str, float]],
+    weights: list[float],
+) -> dict[str, float]:
+    """Aggregate metrics from multiple batches.
+
+    Modified from tinker.lib.chunked_fwdbwd_helpers._metrics_reduction
+    """
+    if not metric_list:
+        return {}
+    keys = metric_list[0].keys()
+    result = {}
+    for key in keys:
+        _, reduction = key.split(":")
+        if reduction not in REDUCE_MAP:
+            raise LossFunctionUnknownMetricReductionException(reduction)
+        if not all(key in m for m in metric_list):
+            continue
+        reduce_fn = REDUCE_MAP[reduction]
+        values = [m[key] for m in metric_list]
+
+        if reduction in ["mean", "slack"]:
+            result[key] = reduce_fn(values, weights)
+        elif reduction in ["unique"]:
+            result[key] = values[0]
+            result.update({f"{key}_{i + 1}": v for i, v in enumerate(values[1:])})
+        else:
+            result[key] = reduce_fn(values)
+    return result

tuft/persistence/__init__.py

@@ -3,30 +3,38 @@
 from __future__ import annotations
 
 from .redis_store import (
-    DEFAULT_FUTURE_TTL_SECONDS,
+    ConfigCheckField,
     PersistenceConfig,
     PersistenceMode,
     RedisPipeline,
     RedisStore,
     delete_record,
+    flush_all_data,
+    get_current_namespace,
     get_redis_store,
     is_persistence_enabled,
     load_record,
+    save_config_signature,
     save_record,
     save_records_atomic,
+    validate_config_signature,
 )
 
 
 __all__ = [
-    "DEFAULT_FUTURE_TTL_SECONDS",
+    "ConfigCheckField",
     "PersistenceConfig",
     "PersistenceMode",
     "RedisPipeline",
     "RedisStore",
     "delete_record",
+    "flush_all_data",
+    "get_current_namespace",
     "get_redis_store",
     "is_persistence_enabled",
     "load_record",
+    "save_config_signature",
     "save_record",
     "save_records_atomic",
+    "validate_config_signature",
 ]