tuft 0.1.0__py3-none-any.whl → 0.1.2__py3-none-any.whl

tuft/state.py

@@ -0,0 +1,352 @@
+"""In-memory state containers backing the FastAPI endpoints."""
+
+from __future__ import annotations
+
+import uuid
+from datetime import datetime, timezone
+from typing import Dict, TypeVar
+
+from pydantic import BaseModel, Field
+from tinker import types
+
+from .auth import AuthenticationDB, User
+from .checkpoints import CheckpointRecord
+from .config import AppConfig
+from .exceptions import SessionNotFoundException, UserMismatchException
+from .futures import FutureStore
+from .persistence import get_redis_store, is_persistence_enabled, load_record, save_record
+from .sampling_controller import SamplingController
+from .training_controller import TrainingController, TrainingRunRecord
+
+
+T = TypeVar("T")
+
+
+def _now() -> datetime:
+    return datetime.now(timezone.utc)
+
+
+class SessionRecord(BaseModel):
+    """Session record with persistence support.
+
+    Sessions are permanent records (no TTL) as they represent user sessions
+    that may need to be accessed at any time.
+    """
+
+    session_id: str
+    tags: list[str]
+    user_metadata: dict[str, str] | None = None
+    user_id: str
+    sdk_version: str
+    created_at: datetime = Field(default_factory=_now)
+    last_heartbeat: datetime = Field(default_factory=_now)
+
+
+class SessionManager:
+    """Maintains session metadata and heartbeats so other controllers can enforce ownership."""
+
+    REDIS_KEY_PREFIX = "session"
+
+    def __init__(self) -> None:
+        self._sessions: Dict[str, SessionRecord] = {}
+        self._restore_from_redis()
+
+    def _build_key(self, session_id: str) -> str:
+        return get_redis_store().build_key(self.REDIS_KEY_PREFIX, session_id)
+
+    def _restore_from_redis(self) -> None:
+        if not is_persistence_enabled():
+            return
+        store = get_redis_store()
+        pattern = store.build_key(self.REDIS_KEY_PREFIX, "*")
+        for key in store.keys(pattern):
+            record = load_record(key, SessionRecord)
+            if record is not None:
+                self._sessions[record.session_id] = record
+
+    def _save_session(self, session_id: str) -> None:
+        """Save session to Redis (no TTL - permanent record)."""
+        if not is_persistence_enabled():
+            return
+        record = self._sessions.get(session_id)
+        if record is not None:
+            save_record(self._build_key(session_id), record)
+
+    def _delete_session(self, session_id: str) -> None:
+        if not is_persistence_enabled():
+            return
+        get_redis_store().delete(self._build_key(session_id))
+
+    def create_session(self, request: types.CreateSessionRequest, user: User) -> SessionRecord:
+        """Create a new session for the given user and request."""
+        session_id = str(uuid.uuid4())
+        record = SessionRecord(
+            session_id=session_id,
+            tags=request.tags,
+            user_id=user.user_id,
+            user_metadata=request.user_metadata,
+            sdk_version=request.sdk_version,
+        )
+        self._sessions[session_id] = record
+        self._save_session(session_id)
+        return record
+
+    def require(self, session_id: str) -> SessionRecord:
+        record = self._sessions.get(session_id)
+        if record is None:
+            raise SessionNotFoundException(session_id)
+        return record
+
+    def heartbeat(self, session_id: str, user_id: str) -> None:
+        record = self.require(session_id)
+        if record.user_id != user_id:
+            raise UserMismatchException()
+        record.last_heartbeat = _now()
+        self._save_session(session_id)
+
+    def list_sessions(self, user_id: str) -> list[str]:
+        return [k for k, v in self._sessions.items() if v.user_id == user_id]
+
+
+class ServerState:
+    """Application-wide container that wires controllers together
+    and exposes a simple façade to FastAPI.
+    """
+
+    def __init__(self, config: AppConfig | None = None) -> None:
+        self.config = config or AppConfig()
+        self.config.ensure_directories()
+        self.config.check_validity()
+        self.sessions = SessionManager()
+        self.training = TrainingController(self.config)
+        self.sampling = SamplingController(self.config)
+        self.auth_db = AuthenticationDB(self.config.authorized_users)
+        self.future_store = FutureStore()
+
+    async def async_init(self) -> None:
+        """Put any async initialization logic here"""
+        await self.sampling.async_init()
+        await self._restore_from_checkpoints()
+
+    async def _restore_from_checkpoints(self) -> None:
+        """Restore server state from checkpoints after Redis restore.
+
+        This method handles checkpoint-based recovery:
+        1. For each training run restored from Redis, create adapter and load latest checkpoint
+        2. Mark ALL futures created after checkpoint's future_id as failed
+        3. For training runs without checkpoints, mark all futures as failed
+        """
+        # Restore training runs (adapter + checkpoint)
+        for model_id, record in self.training.training_runs.items():
+            if record.backend is None or record.corrupted:
+                continue
+            latest_ckpt = await self.training.restore_from_checkpoint(model_id)
+
+            if latest_ckpt is None:
+                self.future_store.mark_futures_failed_after_checkpoint(
+                    model_id=model_id,
+                    checkpoint_future_id=None,
+                    error_message=f"No checkpoint found for model {model_id}. Please retry.",
+                )
+            else:
+                self.future_store.mark_futures_failed_after_checkpoint(
+                    model_id=model_id,
+                    checkpoint_future_id=latest_ckpt.future_id,
+                    error_message=(
+                        f"Server restored from checkpoint {latest_ckpt.checkpoint_id}. "
+                        "Operations after this checkpoint need to be retried."
+                    ),
+                )
+
+    def create_session(self, request: types.CreateSessionRequest, user: User) -> SessionRecord:
+        return self.sessions.create_session(request, user)
+
+    def heartbeat(self, session_id: str, user_id: str) -> None:
+        self.sessions.heartbeat(session_id, user_id)
+
+    async def create_model(
+        self,
+        session_id: str,
+        base_model: str,
+        lora_config: types.LoraConfig,
+        model_owner: str,
+        user_metadata: dict[str, str] | None,
+    ) -> TrainingRunRecord:
+        self.sessions.require(session_id)
+        return await self.training.create_model(
+            session_id=session_id,
+            base_model=base_model,
+            lora_config=lora_config,
+            model_owner=model_owner,
+            user_metadata=user_metadata,
+        )
+
+    def build_supported_models(self) -> list[types.SupportedModel]:
+        return self.training.build_supported_models()
+
+    def get_user(self, api_key: str) -> User | None:
+        return self.auth_db.authenticate(api_key)
+
+    async def run_forward(
+        self,
+        model_id: str,
+        user_id: str,
+        data: list[types.Datum],
+        loss_fn: types.LossFnType,
+        loss_fn_config: dict[str, float] | None,
+        seq_id: int | None,
+        *,
+        backward: bool,
+    ) -> types.ForwardBackwardOutput:
+        return await self.training.run_forward(
+            model_id=model_id,
+            user_id=user_id,
+            data=data,
+            loss_fn=loss_fn,
+            loss_fn_config=loss_fn_config,
+            seq_id=seq_id,
+            backward=backward,
+        )
+
+    async def run_optim_step(
+        self, model_id: str, user_id: str, params: types.AdamParams, seq_id: int | None
+    ) -> types.OptimStepResponse:
+        return await self.training.run_optim_step(
+            model_id=model_id, user_id=user_id, params=params, seq_id=seq_id
+        )
+
+    async def create_sampling_session(
+        self,
+        session_id: str,
+        base_model: str | None,
+        model_path: str | None,
+        user_id: str,
+        *,
+        session_seq_id: int,
+    ) -> str:
+        self.sessions.require(session_id)
+        return await self.sampling.create_sampling_session(
+            session_id=session_id,
+            user_id=user_id,
+            base_model=base_model,
+            model_path=model_path,
+            session_seq_id=session_seq_id,
+        )
+
+    async def run_sample(self, request: types.SampleRequest, user_id: str) -> types.SampleResponse:
+        return await self.sampling.run_sample(request, user_id=user_id)
+
+    async def save_checkpoint(
+        self,
+        model_id: str,
+        user_id: str,
+        name: str | None,
+        checkpoint_type: types.CheckpointType,
+        seq_id: int | None = None,
+    ) -> CheckpointRecord:
+        current_future_id = self.future_store.get_current_future_id()
+        return await self.training.save_checkpoint(
+            model_id=model_id,
+            user_id=user_id,
+            name=name,
+            checkpoint_type=checkpoint_type,
+            future_id=current_future_id,
+            seq_id=seq_id,
+        )
+
+    async def load_checkpoint(
+        self, model_id: str, user_id: str, path: str, optimizer: bool, seq_id: int | None = None
+    ) -> None:
+        return await self.training.load_checkpoint(
+            model_id=model_id,
+            user_id=user_id,
+            path=path,
+            optimizer=optimizer,
+            seq_id=seq_id,
+        )
+
+    def delete_checkpoint(self, model_id: str, user_id: str, checkpoint_id: str) -> None:
+        self.training.delete_checkpoint(model_id, user_id, checkpoint_id)
+
+    def list_checkpoints(self, model_id: str, user_id: str) -> list[types.Checkpoint]:
+        return self.training.list_checkpoints(model_id, user_id)
+
+    def list_user_checkpoints(self, user_id: str) -> list[types.Checkpoint]:
+        return self.training.list_user_checkpoints(user_id)
+
+    def set_checkpoint_visibility(
+        self,
+        model_id: str,
+        user_id: str,
+        checkpoint_id: str,
+        *,
+        public: bool,
+    ) -> None:
+        self.training.set_visibility(
+            model_id=model_id,
+            user_id=user_id,
+            checkpoint_id=checkpoint_id,
+            public=public,
+        )
+
+    def get_weights_info(self, tinker_path: str, user_id: str) -> types.WeightsInfoResponse:
+        parsed = types.ParsedCheckpointTinkerPath.from_tinker_path(tinker_path)
+        return self.training.get_weights_info(parsed.training_run_id, user_id)
+
+    def build_archive_url(
+        self,
+        model_id: str,
+        user_id: str,
+        checkpoint_id: str,
+    ) -> types.CheckpointArchiveUrlResponse:
+        return self.training.build_archive_url(model_id, user_id, checkpoint_id)
+
+    def list_training_runs(
+        self, *, user_id: str, limit: int | None = None, offset: int = 0
+    ) -> types.TrainingRunsResponse:
+        return self.training.list_training_runs(user_id=user_id, limit=limit, offset=offset)
+
+    def get_training_run_view(self, model_id: str, user_id: str) -> types.TrainingRun:
+        return self.training.get_training_run_view(model_id, user_id)
+
+    def get_model_info(self, model_id: str, user_id: str) -> types.GetInfoResponse:
+        return self.training.get_model_info(model_id, user_id=user_id)
+
+    async def unload_model(self, model_id: str, user_id: str) -> None:
+        await self.training.unload_model(model_id, user_id=user_id)
+        await self.sampling.evict_model(model_id, user_id=user_id)
+
+    def get_session_overview(self, session_id: str, user_id: str) -> types.GetSessionResponse:
+        record = self.sessions.require(session_id)
+        if record.user_id != user_id:
+            raise UserMismatchException()
+        training_run_ids = [
+            run_id
+            for run_id, run in self.training.training_runs.items()
+            if run.session_id == session_id
+        ]
+        sampler_ids = [
+            sid
+            for sid, record in self.sampling.sampling_sessions.items()
+            if record.session_id == session_id
+        ]
+        return types.GetSessionResponse(training_run_ids=training_run_ids, sampler_ids=sampler_ids)
+
+    def list_sessions(
+        self, user_id: str, *, limit: int | None = None, offset: int = 0
+    ) -> types.ListSessionsResponse:
+        sessions = self.sessions.list_sessions(user_id=user_id)
+        total = len(sessions)
+        start = min(offset, total)
+        if limit is None:
+            subset = sessions[start:]
+        else:
+            subset = sessions[start : min(start + limit, total)]
+        return types.ListSessionsResponse(sessions=subset)
+
+    def get_sampler_info(self, sampler_id: str, user_id: str) -> types.GetSamplerResponse:
+        return self.sampling.get_sampler_info(
+            sampler_id=sampler_id,
+            user_id=user_id,
+            default_base_model=self.config.supported_models[0].model_name,
+        )

tuft/telemetry/__init__.py

@@ -0,0 +1,17 @@
+"""OpenTelemetry integration for TuFT.
+
+This module provides observability through Traces, Metrics, and Logs.
+"""
+
+from __future__ import annotations
+
+from tuft.config import TelemetryConfig
+
+from .provider import init_telemetry, shutdown_telemetry
+
+
+__all__ = [
+    "TelemetryConfig",
+    "init_telemetry",
+    "shutdown_telemetry",
+]

tuft/telemetry/metrics.py

@@ -0,0 +1,335 @@
+"""Metrics utilities for TuFT.
+
+Provides meter access and predefined metric instruments.
+"""
+
+from __future__ import annotations
+
+import logging
+import threading
+from collections.abc import Iterable
+from typing import Any
+
+import psutil
+import pynvml
+from opentelemetry import metrics
+from opentelemetry.metrics import CallbackOptions, Observation
+
+
+logger = logging.getLogger(__name__)
+
+# Module-level meter cache
+_meters: dict[str, Any] = {}
+
+
+def get_meter(name: str = "tuft"):
+    """Get a meter instance by name.
+
+    Args:
+        name: Name for the meter (typically module name).
+
+    Returns:
+        A Meter instance. When no MeterProvider is configured,
+        OpenTelemetry automatically returns a NoOpMeter.
+    """
+    if name in _meters:
+        return _meters[name]
+
+    meter = metrics.get_meter(name)
+    _meters[name] = meter
+    return meter
+
+
+def clear_meters() -> None:
+    """Clear the meter cache. Used during shutdown."""
+    _meters.clear()
+
+
+class TuftMetrics:
+    """Centralized metrics registry for TuFT.
+
+    Provides access to all predefined metrics with proper typing.
+    """
+
+    _instance: "TuftMetrics | None" = None
+    _lock = threading.Lock()
+
+    def __init__(self):
+        meter = get_meter("tuft")
+
+        # Training metrics
+        self.training_models_active = meter.create_up_down_counter(
+            "tuft.training.models.active",
+            description="Number of active training models",
+        )
+
+        self.training_tokens_per_second = meter.create_histogram(
+            "tuft.training.tokens_per_second",
+            description="Training tokens per second",
+            unit="tokens/s",
+        )
+
+        self.training_checkpoints_saved = meter.create_counter(
+            "tuft.training.checkpoints.saved",
+            description="Number of checkpoints saved",
+        )
+
+        self.training_checkpoint_size = meter.create_histogram(
+            "tuft.training.checkpoint.size_bytes",
+            description="Checkpoint size in bytes",
+            unit="bytes",
+        )
+
+        # Sampling metrics
+        self.sampling_sessions_active = meter.create_up_down_counter(
+            "tuft.sampling.sessions.active",
+            description="Number of active sampling sessions",
+        )
+
+        self.sampling_requests = meter.create_counter(
+            "tuft.sampling.requests",
+            description="Number of sampling requests",
+        )
+
+        self.sampling_duration = meter.create_histogram(
+            "tuft.sampling.duration",
+            description="Sampling request duration",
+            unit="s",
+        )
+
+        self.sampling_tokens_per_second = meter.create_histogram(
+            "tuft.sampling.tokens_per_second",
+            description="Sampling tokens per second",
+            unit="tokens/s",
+        )
+
+        self.sampling_output_tokens = meter.create_histogram(
+            "tuft.sampling.output_tokens",
+            description="Number of output tokens per sample",
+            unit="tokens",
+        )
+
+        # Future queue metrics
+        self.futures_queue_length = meter.create_up_down_counter(
+            "tuft.futures.queue_length",
+            description="Number of futures in queue",
+        )
+
+        self.futures_created = meter.create_counter(
+            "tuft.futures.created",
+            description="Number of futures created",
+        )
+
+        self.futures_completed = meter.create_counter(
+            "tuft.futures.completed",
+            description="Number of futures completed",
+        )
+
+        self.futures_wait_time = meter.create_histogram(
+            "tuft.futures.wait_time",
+            description="Time waiting for future completion",
+            unit="s",
+        )
+
+        self.futures_execution_time = meter.create_histogram(
+            "tuft.futures.execution_time",
+            description="Future execution time",
+            unit="s",
+        )
+
+        # Redis metrics
+        self.redis_operation_duration = meter.create_histogram(
+            "tuft.redis.operation.duration",
+            description="Redis operation duration",
+            unit="s",
+        )
+
+    @classmethod
+    def get_instance(cls) -> "TuftMetrics":
+        """Get the singleton metrics instance."""
+        if cls._instance is None:
+            with cls._lock:
+                if cls._instance is None:
+                    cls._instance = cls()
+        return cls._instance
+
+
+def get_metrics() -> TuftMetrics:
+    """Get the TuFT metrics instance."""
+    return TuftMetrics.get_instance()
+
+
+class ResourceMetricsCollector:
+    """Collects system resource metrics periodically.
+
+    Collects CPU, memory, disk, GPU, and network metrics.
+    """
+
+    _instance: "ResourceMetricsCollector | None" = None
+    _lock = threading.Lock()
+
+    def __init__(self, checkpoint_dir: str | None = None):
+        self._checkpoint_dir = checkpoint_dir
+        self._running = False
+        self._thread: threading.Thread | None = None
+        self._nvml_initialized = False
+        self._gpu_available = self._check_and_init_gpu()
+        self._setup_metrics()
+
+    def _check_and_init_gpu(self) -> bool:
+        """Check if GPU monitoring is available and initialize NVML once."""
+        try:
+            pynvml.nvmlInit()
+            self._nvml_initialized = True
+            return True
+        except pynvml.NVMLError as e:
+            logger.debug("GPU monitoring not available: %s", e)
+            return False
+
+    def _shutdown_gpu(self) -> None:
+        """Shutdown NVML if it was initialized."""
+        if self._nvml_initialized:
+            try:
+                pynvml.nvmlShutdown()
+                self._nvml_initialized = False
+            except pynvml.NVMLError as e:
+                logger.warning("Failed to shutdown NVML: %s", e)
+
+    def _setup_metrics(self) -> None:
+        """Set up observable gauges for resource metrics."""
+        meter = metrics.get_meter("tuft.resources")
+
+        # CPU metrics
+        meter.create_observable_gauge(
+            "tuft.resource.cpu.utilization_percent",
+            callbacks=[self._cpu_utilization_callback],
+            description="CPU utilization percentage",
+            unit="%",
+        )
+
+        # Memory metrics
+        meter.create_observable_gauge(
+            "tuft.resource.memory.used_bytes",
+            callbacks=[self._memory_used_callback],
+            description="Memory used in bytes",
+            unit="bytes",
+        )
+        meter.create_observable_gauge(
+            "tuft.resource.memory.total_bytes",
+            callbacks=[self._memory_total_callback],
+            description="Total memory in bytes",
+            unit="bytes",
+        )
+        meter.create_observable_gauge(
+            "tuft.resource.memory.utilization_percent",
+            callbacks=[self._memory_utilization_callback],
+            description="Memory utilization percentage",
+            unit="%",
+        )
+
+        # GPU metrics (if available)
+        if self._gpu_available:
+            meter.create_observable_gauge(
+                "tuft.resource.gpu.utilization_percent",
+                callbacks=[self._gpu_utilization_callback],
+                description="GPU utilization percentage",
+                unit="%",
+            )
+            meter.create_observable_gauge(
+                "tuft.resource.gpu.memory_used_bytes",
+                callbacks=[self._gpu_memory_used_callback],
+                description="GPU memory used in bytes",
+                unit="bytes",
+            )
+            meter.create_observable_gauge(
+                "tuft.resource.gpu.memory_total_bytes",
+                callbacks=[self._gpu_memory_total_callback],
+                description="GPU total memory in bytes",
+                unit="bytes",
+            )
+
+        # Process metrics
+        meter.create_observable_gauge(
+            "tuft.resource.process.memory_used_bytes",
+            callbacks=[self._process_memory_callback],
+            description="Process memory usage in bytes",
+            unit="bytes",
+        )
+
+    def _cpu_utilization_callback(self, options: CallbackOptions) -> Iterable[Observation]:
+        """Callback for CPU utilization metric."""
+        yield Observation(psutil.cpu_percent(interval=None))
+
+    def _memory_used_callback(self, options: CallbackOptions) -> Iterable[Observation]:
+        """Callback for memory used metric."""
+        yield Observation(psutil.virtual_memory().used)
+
+    def _memory_total_callback(self, options: CallbackOptions) -> Iterable[Observation]:
+        """Callback for total memory metric."""
+        yield Observation(psutil.virtual_memory().total)
+
+    def _memory_utilization_callback(self, options: CallbackOptions) -> Iterable[Observation]:
+        """Callback for memory utilization metric."""
+        yield Observation(psutil.virtual_memory().percent)
+
+    def _gpu_utilization_callback(self, options: CallbackOptions) -> Iterable[Observation]:
+        """Callback for GPU utilization metric."""
+        if not self._nvml_initialized:
+            return
+        try:
+            device_count = pynvml.nvmlDeviceGetCount()
+            for i in range(device_count):
+                handle = pynvml.nvmlDeviceGetHandleByIndex(i)
+                util = pynvml.nvmlDeviceGetUtilizationRates(handle)
+                yield Observation(int(util.gpu), {"gpu_id": str(i)})
+        except pynvml.NVMLError as e:
+            logger.debug("Failed to get GPU utilization: %s", e)
+
+    def _gpu_memory_used_callback(self, options: CallbackOptions) -> Iterable[Observation]:
+        """Callback for GPU memory used metric."""
+        if not self._nvml_initialized:
+            return
+        try:
+            device_count = pynvml.nvmlDeviceGetCount()
+            for i in range(device_count):
+                handle = pynvml.nvmlDeviceGetHandleByIndex(i)
+                mem_info = pynvml.nvmlDeviceGetMemoryInfo(handle)
+                yield Observation(int(mem_info.used), {"gpu_id": str(i)})
+        except pynvml.NVMLError as e:
+            logger.debug("Failed to get GPU memory used: %s", e)
+
+    def _gpu_memory_total_callback(self, options: CallbackOptions) -> Iterable[Observation]:
+        """Callback for GPU total memory metric."""
+        if not self._nvml_initialized:
+            return
+        try:
+            device_count = pynvml.nvmlDeviceGetCount()
+            for i in range(device_count):
+                handle = pynvml.nvmlDeviceGetHandleByIndex(i)
+                mem_info = pynvml.nvmlDeviceGetMemoryInfo(handle)
+                yield Observation(int(mem_info.total), {"gpu_id": str(i)})
+        except pynvml.NVMLError as e:
+            logger.debug("Failed to get GPU memory total: %s", e)
+
+    def _process_memory_callback(self, options: CallbackOptions) -> Iterable[Observation]:
+        """Callback for process memory metric."""
+        process = psutil.Process()
+        yield Observation(process.memory_info().rss)
+
+    @classmethod
+    def start(cls, checkpoint_dir: str | None = None) -> "ResourceMetricsCollector":
+        """Start the resource metrics collector."""
+        if cls._instance is None:
+            with cls._lock:
+                if cls._instance is None:
+                    cls._instance = cls(checkpoint_dir)
+        return cls._instance
+
+    @classmethod
+    def shutdown(cls) -> None:
+        """Shutdown the resource metrics collector."""
+        if cls._instance is not None:
+            with cls._lock:
+                if cls._instance is not None:
+                    cls._instance._shutdown_gpu()
+                    cls._instance = None