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

tuft/backends/training_backend.py

@@ -0,0 +1,327 @@
+import asyncio
+import logging
+from dataclasses import dataclass, field
+from typing import Sequence
+
+import numpy as np
+from tinker import types
+
+from tuft.backends.base_backend import BaseTrainingBackend
+from tuft.checkpoints import CheckpointRecord
+from tuft.config import ModelConfig
+from tuft.telemetry.tracing import get_tracer, inject_context
+
+
+_get_tracer = lambda: get_tracer("tuft.training_backend")  # noqa: E731
+
+logger = logging.getLogger(__name__)
+
+
+class HFTrainingBackend(BaseTrainingBackend):
+    """A training backend using Hugging Face transformers."""
+
+    def __init__(self, config: ModelConfig) -> None:
+        from .hf_training_model import HFTrainingModel
+
+        self.config = config
+        logger.info("Ray actor created: HFTrainingModel(%s)", config.model_name)
+        self.model = HFTrainingModel.get_actor(config)
+
+    async def async_init(self) -> None:
+        await self.model.async_init.remote()
+
+    async def create_adapter(self, lora_id: str, lora_config: types.LoraConfig) -> None:
+        """Create a LoRA adapter with the given ID and configuration."""
+        with _get_tracer().start_as_current_span("training_backend.create_adapter") as span:
+            span.set_attribute("tuft.lora_id", lora_id)
+            span.set_attribute("tuft.lora_rank", lora_config.rank)
+            # Inject trace context for Ray actor
+            trace_context: dict[str, str] = {}
+            inject_context(trace_context)
+            await self.model.create_adapter.remote(lora_id, lora_config, trace_context)
+
+    async def remove_adapter(self, lora_id: str) -> None:
+        with _get_tracer().start_as_current_span("training_backend.remove_adapter") as span:
+            span.set_attribute("tuft.lora_id", lora_id)
+            await self.model.remove_adapter.remote(lora_id)
+
+    async def forward(
+        self,
+        data: list[types.Datum],
+        lora_id: str,
+        loss_fn: types.LossFnType,
+        loss_fn_config: dict[str, float] | None,
+        backward: bool = False,
+    ) -> types.ForwardBackwardOutput:
+        """Forward pass (and backward if specified).
+
+        Args:
+            data: List of Datum objects containing input data.
+            lora_id: The LoRA adapter ID to use.
+            loss_fn: The loss function to apply.
+            loss_fn_config: Optional configuration for the loss function.
+            backward: Whether to perform backward pass.
+
+        Returns:
+            ForwardBackwardOutput: The output of the forward (and backward) pass.
+        """
+        span_name = "training_backend.forward_backward" if backward else "training_backend.forward"
+        with _get_tracer().start_as_current_span(span_name) as span:
+            span.set_attribute("tuft.lora_id", lora_id)
+            span.set_attribute("tuft.backward", backward)
+            span.set_attribute("tuft.data_count", len(data))
+            # Inject trace context for Ray actor
+            trace_context: dict[str, str] = {}
+            inject_context(trace_context)
+            return await self.model.forward.remote(
+                data=data,
+                lora_id=lora_id,
+                loss_fn=loss_fn,
+                loss_fn_config=loss_fn_config,
+                backward=backward,
+                trace_context=trace_context,
+            )
+
+    async def optim_step(
+        self, adam_params: types.AdamParams, lora_id: str
+    ) -> types.OptimStepResponse:
+        """Perform an optimization step using Adam optimizer.
+
+        Args:
+            adam_params: Parameters for the Adam optimizer.
+
+        Returns:
+            OptimStepResponse: The response containing optimization metrics.
+        """
+        with _get_tracer().start_as_current_span("training_backend.optim_step") as span:
+            span.set_attribute("tuft.lora_id", lora_id)
+            # Inject trace context for Ray actor
+            trace_context: dict[str, str] = {}
+            inject_context(trace_context)
+            return await self.model.optim_step.remote(adam_params, lora_id, trace_context)
+
+    async def save_state(
+        self, lora_id: str, checkpoint_record: "CheckpointRecord", optimizer: bool
+    ) -> None:
+        """Save the state of the specified LoRA adapter."""
+        with _get_tracer().start_as_current_span("training_backend.save_state") as span:
+            span.set_attribute("tuft.lora_id", lora_id)
+            span.set_attribute("tuft.optimizer", optimizer)
+            # Inject trace context for Ray actor
+            trace_context: dict[str, str] = {}
+            inject_context(trace_context)
+            await self.model.save_state.remote(
+                lora_id=lora_id,
+                checkpoint_record=checkpoint_record,
+                optimizer=optimizer,
+                trace_context=trace_context,
+            )
+
+    async def load_state(
+        self, lora_id: str, checkpoint_record: "CheckpointRecord", optimizer: bool
+    ) -> None:
+        """Load the state of the specified LoRA adapter from the given path."""
+        with _get_tracer().start_as_current_span("training_backend.load_state") as span:
+            span.set_attribute("tuft.lora_id", lora_id)
+            span.set_attribute("tuft.optimizer", optimizer)
+            # Inject trace context for Ray actor
+            trace_context: dict[str, str] = {}
+            inject_context(trace_context)
+            await self.model.load_state.remote(
+                lora_id=lora_id,
+                checkpoint_record=checkpoint_record,
+                optimizer=optimizer,
+                trace_context=trace_context,
+            )
+
+
+@dataclass
+class DummyTrainingBackend(BaseTrainingBackend):
+    """A dummy training backend for testing purposes."""
+
+    _lock: asyncio.Lock = field(init=False, repr=False)
+    _weights: np.ndarray = field(init=False, repr=False)
+    _adam_m: np.ndarray = field(init=False, repr=False)
+    _adam_v: np.ndarray = field(init=False, repr=False)
+    _beta1_power: float = field(init=False, default=1.0, repr=False)
+    _beta2_power: float = field(init=False, default=1.0, repr=False)
+    _pending_grad: np.ndarray | None = field(init=False, default=None, repr=False)
+    _pending_examples: int = field(init=False, default=0, repr=False)
+    _embedding_cache: dict[int, np.ndarray] = field(init=False, default_factory=dict, repr=False)
+    step: int = field(init=False, default=0)
+
+    def __init__(self, config: ModelConfig) -> None:
+        self.config = config
+        self.seed = config.seed
+        self.hidden_dim = 16
+        rng = np.random.default_rng(self.seed or 0)
+        self._lock = asyncio.Lock()
+        self._weights = rng.standard_normal(self.hidden_dim, dtype=np.float32)
+        self._adam_m = np.zeros_like(self._weights)
+        self._adam_v = np.zeros_like(self._weights)
+        self._embedding_cache = {}
+        self._adapters = dict()
+
+    async def async_init(self) -> None:
+        pass
+
+    # ------------------------------------------------------------------
+    # Forward / backward helpers
+    # ------------------------------------------------------------------
+    async def forward(
+        self,
+        data: list[types.Datum],
+        lora_id: str,
+        loss_fn: types.LossFnType,
+        loss_fn_config: dict[str, float] | None,
+        backward: bool = False,
+    ) -> types.ForwardBackwardOutput:
+        return await self._run_step(data, backward=backward)
+
+    async def _run_step(
+        self, data: list[types.Datum], *, backward: bool
+    ) -> types.ForwardBackwardOutput:
+        outputs: list[types.LossFnOutput] = []
+        total_loss = 0.0
+        grad_accum = np.zeros_like(self._weights)
+        for datum in data:
+            prompt_tokens = datum.model_input.to_ints()
+            target_tokens = self._target_tokens(datum)
+            prompt_vec = self._vectorize(prompt_tokens)
+            target_scalar = self._target_scalar(target_tokens)
+            prediction = float(np.dot(self._weights, prompt_vec))
+            loss = (prediction - target_scalar) ** 2
+            total_loss += loss
+            if backward:
+                grad = 2 * (prediction - target_scalar) * prompt_vec
+                grad_accum += grad
+            logprob_tensor = types.TensorData(
+                data=[float(-abs(prediction - target_scalar))] * max(len(target_tokens), 1),
+                dtype="float32",
+                shape=[max(len(target_tokens), 1)],
+            )
+            outputs.append({"logprobs": logprob_tensor})
+
+        metrics = {
+            "loss:sum": total_loss,
+            "step:max": float(self.step),
+        }
+        if backward:
+            grad_norm = float(np.linalg.norm(grad_accum) / max(len(data), 1))
+            metrics["grad_norm:mean"] = grad_norm
+            async with self._lock:
+                if self._pending_grad is None:
+                    self._pending_grad = grad_accum
+                else:
+                    self._pending_grad += grad_accum
+                self._pending_examples += len(data)
+
+        return types.ForwardBackwardOutput(
+            loss_fn_output_type="ToyLoss",
+            loss_fn_outputs=outputs,
+            metrics=metrics,
+        )
+
+    # ------------------------------------------------------------------
+    # Optimizer
+    # ------------------------------------------------------------------
+    async def optim_step(
+        self, adam_params: types.AdamParams, lora_id: str
+    ) -> types.OptimStepResponse:
+        async with self._lock:
+            grad = self._pending_grad
+            examples = self._pending_examples
+            self._pending_grad = None
+            self._pending_examples = 0
+
+        if grad is None or not np.any(grad):
+            return types.OptimStepResponse(
+                metrics={
+                    "learning_rate:mean": adam_params.learning_rate,
+                    "step:max": float(self.step),
+                }
+            )
+
+        grad = grad / max(examples, 1)
+        if adam_params.grad_clip_norm > 0:
+            norm = np.linalg.norm(grad)
+            if norm > adam_params.grad_clip_norm:
+                grad *= adam_params.grad_clip_norm / max(norm, 1e-12)
+
+        if adam_params.weight_decay:
+            grad += adam_params.weight_decay * self._weights
+
+        beta1 = adam_params.beta1
+        beta2 = adam_params.beta2
+
+        self._adam_m = beta1 * self._adam_m + (1 - beta1) * grad
+        self._adam_v = beta2 * self._adam_v + (1 - beta2) * (grad**2)
+        self._beta1_power *= beta1
+        self._beta2_power *= beta2
+        m_hat = self._adam_m / (1 - self._beta1_power + 1e-12)
+        v_hat = self._adam_v / (1 - self._beta2_power + 1e-12)
+
+        update = adam_params.learning_rate * m_hat / (np.sqrt(v_hat) + adam_params.eps)
+        self._weights -= update
+        self.step += 1
+
+        metrics = {
+            "learning_rate:mean": adam_params.learning_rate,
+            "step:max": float(self.step),
+            "update_norm:mean": float(np.linalg.norm(update)),
+        }
+        return types.OptimStepResponse(metrics=metrics)
+
+    async def save_state(
+        self, lora_id: str, checkpoint_record: "CheckpointRecord", optimizer: bool
+    ) -> None:
+        if lora_id not in self._adapters:
+            raise ValueError(f"Adapter {lora_id} does not exist.")
+        # dummy save
+
+    async def load_state(
+        self, lora_id: str, checkpoint_record: "CheckpointRecord", optimizer: bool
+    ) -> None:
+        # create a dummy adapter on load
+        self._adapters[lora_id] = types.LoraConfig(rank=4)
+
+    async def create_adapter(self, lora_id: str, lora_config: types.LoraConfig) -> None:
+        self._adapters[lora_id] = lora_config
+
+    async def remove_adapter(self, lora_id: str) -> None:
+        self._adapters.pop(lora_id, None)
+
+    # ------------------------------------------------------------------
+    # Internal helpers
+    # ------------------------------------------------------------------
+    def _target_tokens(self, datum: types.Datum) -> list[int]:
+        if not datum.loss_fn_inputs:
+            return datum.model_input.to_ints()
+        tensor = datum.loss_fn_inputs.get("target_tokens")
+        if tensor is None:
+            return datum.model_input.to_ints()
+        return [int(value) for value in tensor.data]
+
+    def _vectorize(self, tokens: Sequence[int]) -> np.ndarray:
+        if not tokens:
+            return np.zeros(self.hidden_dim, dtype=np.float32)
+        vecs = [self._token_embedding(token) for token in tokens]
+        return np.mean(vecs, axis=0)
+
+    def _token_embedding(self, token_id: int) -> np.ndarray:
+        cached = self._embedding_cache.get(token_id)
+        if cached is None:
+            rng = np.random.default_rng(self.seed + token_id)
+            cached = rng.standard_normal(self.hidden_dim, dtype=np.float32)
+            self._embedding_cache[token_id] = cached
+        return cached
+
+    def _safe_mean(self, values: Sequence[int]) -> float:
+        if not values:
+            return 0.0
+        return float(sum(values) / len(values))
+
+    def _target_scalar(self, tokens: Sequence[int]) -> float:
+        if not tokens:
+            return 0.0
+        return np.tanh(self._safe_mean(tokens) / 100.0)

tuft/checkpoints.py

@@ -0,0 +1,193 @@
+"""Module for managing checkpoints on disk."""
+
+import contextlib
+import shutil
+from datetime import datetime, timezone
+from pathlib import Path
+
+from pydantic import BaseModel, ConfigDict, Field, ValidationError, field_serializer
+from tinker import types
+
+from .exceptions import CheckpointMetadataReadException
+
+
+class CheckpointMetadata(BaseModel):
+    """A representation of checkpoint metadata."""
+
+    model_id: str
+    name: str
+    base_model: str
+    checkpoint_type: types.CheckpointType
+    created_at: str
+    session_id: str
+    tinker_path: str
+    owner_name: str
+    size_bytes: int = 0
+    lora_rank: int | None = None
+    public: bool = False
+    future_id: int = 0
+    seq_id: int | None = None
+
+
+class CheckpointRecord(BaseModel):
+    """A record representing a checkpoint on disk."""
+
+    model_config = ConfigDict(arbitrary_types_allowed=True)
+
+    checkpoint_id: str
+    owner_name: str
+    checkpoint_type: types.CheckpointType
+    training_run_id: str
+    path: Path
+    size_bytes: int = 0
+    public: bool = False
+    created_at: datetime = Field(default_factory=lambda: datetime.now(timezone.utc))
+    future_id: int = 0
+    seq_id: int | None = None
+
+    @field_serializer("path")
+    def serialize_path(self, path: Path) -> str:
+        """Serialize Path to string for JSON."""
+        return str(path)
+
+    @property
+    def tinker_checkpoint(self) -> types.Checkpoint:
+        """Get a Tinker Checkpoint instance representing this record."""
+        return types.Checkpoint(
+            checkpoint_id=self.checkpoint_id,
+            checkpoint_type=self.checkpoint_type,
+            time=self.created_at,
+            tinker_path=self.tinker_path,
+            size_bytes=self.size_bytes,
+            public=self.public,
+        )
+
+    @property
+    def metadata(self) -> CheckpointMetadata:
+        """Get the checkpoint metadata.
+
+        Raises:
+            CheckpointMetadataReadException: If the metadata file does not
+                exist or is invalid.
+        """
+        try:
+            return CheckpointMetadata.model_validate_json(
+                self.metadata_path.read_text(encoding="utf-8")
+            )
+        except FileNotFoundError as exc:
+            raise CheckpointMetadataReadException(checkpoint_id=self.checkpoint_id) from exc
+        except ValidationError as exc:
+            raise CheckpointMetadataReadException(checkpoint_id=self.checkpoint_id) from exc
+
+    @property
+    def tinker_path(self) -> str:
+        """Get the tinker style path for this checkpoint."""
+        folder = "weights" if self.checkpoint_type == "training" else "sampler_weights"
+        return f"tinker://{self.training_run_id}/{folder}/{self.checkpoint_id}"
+
+    @property
+    def adapter_path(self) -> Path:
+        """Get the path to the adapter weights file."""
+        return self.path / "adapter"
+
+    @property
+    def optimizer_path(self) -> Path:
+        """Get the path to the optimizer state file."""
+        return self.path / "optimizer"
+
+    @property
+    def metadata_path(self) -> Path:
+        """Get the path to the metadata JSON file."""
+        return self.path / "metadata.json"
+
+    def set_visibility(self, public: bool) -> None:
+        """Set the visibility of the checkpoint."""
+        self.public = public
+        metadata = self.metadata
+        metadata.public = public
+        self.save_metadata(
+            base_model=metadata.base_model,
+            session_id=metadata.session_id,
+            lora_rank=metadata.lora_rank,
+        )
+
+    def save_metadata(self, base_model: str, session_id: str, lora_rank: int | None) -> None:
+        """Save the checkpoint metadata to disk."""
+        # check the format of metadata
+        try:
+            metadata = CheckpointMetadata(
+                model_id=self.training_run_id,
+                name=self.checkpoint_id,
+                base_model=base_model,
+                checkpoint_type=self.checkpoint_type,
+                created_at=self.created_at.isoformat(),
+                session_id=session_id,
+                tinker_path=self.tinker_path,
+                owner_name=self.owner_name,
+                lora_rank=lora_rank,
+                public=self.public,
+                size_bytes=self.size_bytes,
+                future_id=self.future_id,
+                seq_id=self.seq_id,
+            )
+        except Exception as e:
+            raise ValueError(f"Invalid checkpoint metadata: {e}") from e
+        self.metadata_path.write_text(metadata.model_dump_json(indent=2), encoding="utf-8")
+
+    @classmethod
+    def from_tinker_path(cls, path: str, checkpoint_root_dir: Path) -> "CheckpointRecord":
+        """Create a CheckpointRecord from a Tinker path.
+
+        Raises:
+            FileNotFoundError: If the checkpoint directory or metadata.json is missing.
+            json.JSONDecodeError: If metadata.json cannot be parsed as JSON.
+        """
+        parsed = types.ParsedCheckpointTinkerPath.from_tinker_path(path)
+        checkpoint_path = (
+            checkpoint_root_dir / parsed.training_run_id / parsed.checkpoint_id.split("/", 1)[-1]
+        )
+        record = cls(
+            checkpoint_id=parsed.checkpoint_id.split("/", 1)[-1],
+            checkpoint_type=parsed.checkpoint_type,
+            training_run_id=parsed.training_run_id,
+            path=checkpoint_path,
+            owner_name="",  # Will be filled from metadata later
+            size_bytes=0,  # Will be filled from metadata later
+        )
+        metadata = record.metadata  # This may raise FileNotFoundError or JSONDecodeError
+        record.owner_name = metadata.owner_name
+        record.size_bytes = metadata.size_bytes
+        record.public = metadata.public
+        record.created_at = datetime.fromisoformat(metadata.created_at)
+        record.future_id = metadata.future_id
+        record.seq_id = metadata.seq_id
+        return record
+
+    def delete(self) -> None:
+        """Delete the checkpoint from disk."""
+        with contextlib.suppress(FileNotFoundError):
+            shutil.rmtree(self.path)
+
+    @classmethod
+    def from_training_run(
+        cls,
+        training_run_id: str,
+        checkpoint_name: str,
+        owner_name: str,
+        checkpoint_type: types.CheckpointType,
+        checkpoint_root_dir: Path,
+        exist_ok: bool = True,
+    ) -> "CheckpointRecord":
+        """Create a CheckpointRecord from a training run."""
+        checkpoint_dir = checkpoint_root_dir / training_run_id / checkpoint_name
+        if not exist_ok and checkpoint_dir.exists():
+            raise FileExistsError(f"Checkpoint directory already exists: {checkpoint_dir}")
+        checkpoint_dir.mkdir(parents=True, exist_ok=exist_ok)
+        return cls(
+            checkpoint_id=checkpoint_name,
+            owner_name=owner_name,
+            checkpoint_type=checkpoint_type,
+            training_run_id=training_run_id,
+            path=checkpoint_dir,
+            size_bytes=0,
+        )

tuft/cli.py

@@ -0,0 +1,91 @@
+"""Command line utilities for the local TuFT server."""
+
+from __future__ import annotations
+
+import logging
+from pathlib import Path
+
+import typer
+import uvicorn
+
+from .config import AppConfig, load_yaml_config
+from .server import create_root_app
+from .telemetry import init_telemetry
+from .telemetry.metrics import ResourceMetricsCollector
+
+
+app = typer.Typer(help="TuFT - Tenant-unified Fine-Tuning Server.")
+
+_HOST_OPTION = typer.Option("127.0.0.1", "--host", help="Interface to bind", envvar="TUFT_HOST")
+_PORT_OPTION = typer.Option(10610, "--port", "-p", help="Port to bind", envvar="TUFT_PORT")
+_LOG_LEVEL_OPTION = typer.Option("info", "--log-level", help="Uvicorn log level")
+_RELOAD_OPTION = typer.Option(False, "--reload", help="Enable auto-reload (development only)")
+_CONFIG_OPTION = typer.Option(
+    None,
+    "--config",
+    "-c",
+    help="Path to a TuFT configuration file (YAML)",
+)
+_CHECKPOINT_DIR_OPTION = typer.Option(
+    None,
+    "--checkpoint-dir",
+    help="Override checkpoint_dir from config file. Defaults to ~/.cache/tuft/checkpoints.",
+)
+
+
+def _build_config(
+    config_path: Path | None,
+    checkpoint_dir: Path | None,
+) -> AppConfig:
+    if config_path is None:
+        raise typer.BadParameter("Configuration file must be provided via --config")
+    config = load_yaml_config(config_path)
+    if checkpoint_dir is not None:
+        config.checkpoint_dir = checkpoint_dir.expanduser()
+    config.ensure_directories()
+    return config
+
+
+def _init_telemetry(config: AppConfig, log_level: str) -> None:
+    """Initialize OpenTelemetry if enabled."""
+    # Configure root logger level to ensure logs flow to OTel
+    numeric_level = getattr(logging, log_level.upper(), logging.INFO)
+
+    if not config.telemetry.enabled:
+        logging.basicConfig(level=numeric_level)
+        return
+
+    init_telemetry(config.telemetry)
+    # Start resource metrics collection
+    ResourceMetricsCollector.start(str(config.checkpoint_dir))
+
+
+@app.command()
+def launch(
+    host: str = _HOST_OPTION,
+    port: int = _PORT_OPTION,
+    log_level: str = _LOG_LEVEL_OPTION,
+    reload: bool = _RELOAD_OPTION,
+    config_path: Path | None = _CONFIG_OPTION,
+    checkpoint_dir: Path | None = _CHECKPOINT_DIR_OPTION,
+) -> None:
+    """Launch the TuFT server."""
+    app_config = _build_config(config_path, checkpoint_dir)
+    # Initialize telemetry before starting the server
+    _init_telemetry(app_config, log_level)
+    logging.getLogger("tuft").info("Server starting on %s:%s", host, port)
+    uvicorn.run(
+        create_root_app(app_config),
+        host=host,
+        port=port,
+        log_level=log_level,
+        reload=reload,
+    )
+
+
+def main() -> None:
+    app()
+
+
+if __name__ == "__main__":
+    main()