baseten-loops 0.1.0__tar.gz

baseten_loops-0.1.0/.gitignore

@@ -0,0 +1,10 @@
+.env
+__pycache__/
+*.pyc
+*.pyo
+.venv/
+dist/
+*.egg-info/
+.pytest_cache/
+.mypy_cache/
+.ruff_cache/

baseten_loops-0.1.0/PKG-INFO

@@ -0,0 +1,13 @@
+Metadata-Version: 2.4
+Name: baseten-loops
+Version: 0.1.0
+Summary: SDK for Baseten training workers
+Requires-Python: >=3.12
+Requires-Dist: datasets>=4.8.4
+Requires-Dist: httpx>=0.24.1
+Requires-Dist: pydantic>=2.10.0
+Requires-Dist: tenacity>=8.0.0
+Requires-Dist: transformers>=4.40.0
+Requires-Dist: truss>=0.16.0
+Provides-Extra: tinker
+Requires-Dist: baseten-loops-tinker>=0.1.0; extra == 'tinker'

baseten_loops-0.1.0/baseten/loops/__init__.py

@@ -0,0 +1,65 @@
+from baseten.loops.client import create_training_client
+from baseten.loops.promise_client import (
+    RemoteOpError,
+    ServerShutdownError,
+    UnknownRequestError,
+)
+from baseten.loops.training_client import TrainingClient, OperationFuture
+from baseten.loops.sampling_client import SamplingClient
+from baseten.loops.service_client import ServiceClient
+from baseten.loops.models import (
+    AdamParams,
+    Checkpoint,
+    CheckpointFile,
+    CheckpointFilesResponse,
+    Datum,
+    EncodedTextChunk,
+    ForwardBackwardOutput,
+    InitTrainerServerResponse,
+    ImageChunk,
+    LoadWeightsResponse,
+    ModelInput,
+    ModelInputChunk,
+    OptimStepResponse,
+    SampledSequence,
+    SampleResponse,
+    SampleResult,
+    SamplingParams,
+    SaveWeightsForSamplerResponse,
+    SaveWeightsResponse,
+    TensorData,
+)
+
+__all__ = [
+    # Client classes
+    "create_training_client",
+    "ServiceClient",
+    "TrainingClient",
+    "SamplingClient",
+    "OperationFuture",
+    # Promise-protocol exceptions
+    "RemoteOpError",
+    "UnknownRequestError",
+    "ServerShutdownError",
+    # Types
+    "AdamParams",
+    "Checkpoint",
+    "CheckpointFile",
+    "CheckpointFilesResponse",
+    "Datum",
+    "EncodedTextChunk",
+    "ForwardBackwardOutput",
+    "InitTrainerServerResponse",
+    "ImageChunk",
+    "LoadWeightsResponse",
+    "ModelInput",
+    "ModelInputChunk",
+    "OptimStepResponse",
+    "SampledSequence",
+    "SampleResponse",
+    "SampleResult",
+    "SamplingParams",
+    "SaveWeightsForSamplerResponse",
+    "SaveWeightsResponse",
+    "TensorData",
+]

baseten_loops-0.1.0/baseten/loops/client.py

@@ -0,0 +1,150 @@
+"""Deploy a training worker and return a TrainingClient connected directly to it."""
+
+from __future__ import annotations
+
+import json
+import uuid
+from pathlib import Path
+from typing import Optional
+
+import httpx
+
+from baseten.loops.training_client import TrainingClient
+from truss.base import truss_config
+from truss_train import definitions
+from truss_train.public_api import push
+
+
+_DEFAULT_WORKSPACE_ROOT = Path(__file__).resolve().parent.parent / "server"
+
+
+def create_training_client(
+    base_model: str,
+    worker_url: str,
+    *,
+    api_key: str | None = None,
+    gpu_count: int = 1,
+    accelerator: str = "H100",
+    training_gpus: Optional[list[int]] = None,
+    inference_gpus: Optional[list[int]] = None,
+    max_seq_len: int = 4096,
+    worker_port: int = 8001,
+    namespace: str = "default",
+    remote: str = "baseten",
+    workspace_root: Optional[Path] = None,
+    deploy: bool = True,
+    timeout: float = 600.0,
+) -> TrainingClient:
+    """Deploy a training worker and return a TrainingClient connected to it.
+
+    Args:
+        base_model: HuggingFace model ID (e.g. "Qwen/Qwen3-8B").
+        worker_url: URL of the dp_worker to connect to. If deploy=True, this
+            is computed automatically from the job ID and namespace.
+        api_key: API key for authentication.
+        gpu_count: Total number of GPUs to request.
+        accelerator: GPU type (H100, H200, B200).
+        training_gpus: GPU indices for training. Defaults to [0].
+        inference_gpus: GPU indices for inference. Defaults to [0].
+        max_seq_len: Maximum sequence length for training.
+        worker_port: Port the dp_worker will listen on.
+        namespace: K8s namespace where the training job runs.
+        remote: Baseten remote name from .trussrc.
+        workspace_root: Path to thinker workspace root.
+        deploy: If True, deploy a training job. If False, just connect to worker_url.
+        timeout: HTTP timeout for training operations.
+    """
+    if not deploy:
+        return TrainingClient(worker_url, api_key=api_key, timeout=timeout)
+
+    suffix = uuid.uuid4().hex[:7]
+    project_name = f"trainer-{base_model.replace('/', '-')}-{suffix}"
+    print(f"Project: {project_name}")
+
+    if training_gpus is None:
+        training_gpus = [0]
+    if inference_gpus is None:
+        inference_gpus = [0]
+
+    rl_config = {
+        "model_id": base_model,
+        "training": {
+            "tensor_parallel_size": 1,
+            "pipeline_parallel_size": 1,
+            "max_length": max_seq_len,
+            "gpus": training_gpus,
+        },
+        "inference": {
+            "tensor_parallel_size": 1,
+            "gpus": inference_gpus,
+            "gpu_memory_utilization": 0.9,
+        },
+    }
+
+    ws_root = workspace_root or _DEFAULT_WORKSPACE_ROOT
+    if not ws_root.exists():
+        raise FileNotFoundError(
+            f"Workspace root not found: {ws_root}. "
+            "Pass workspace_root= pointing to the thinker repo."
+        )
+
+    rl_config_path = ws_root / "rl_config.json"
+    rl_config_path.write_text(json.dumps(rl_config))
+
+    accel_enum = getattr(
+        truss_config.Accelerator,
+        accelerator.upper(),
+        truss_config.Accelerator.H100,
+    )
+
+    project = definitions.TrainingProject(
+        name=project_name,
+        job=definitions.TrainingJob(
+            compute=definitions.Compute(
+                accelerator=truss_config.AcceleratorSpec(
+                    accelerator=accel_enum,
+                    count=gpu_count,
+                ),
+            ),
+            runtime=definitions.Runtime(
+                start_commands=[
+                    "apt-get update && apt-get install -y python3-dev curl",
+                    "curl -LsSf https://astral.sh/uv/install.sh | sh",
+                    ". $HOME/.local/bin/env && uv sync --extra worker",
+                    f".venv/bin/python -m trainers_server.dp_worker.main --config $RL_CONFIG_PATH --port {worker_port}",
+                ],
+                environment_variables={
+                    "RL_CONFIG_PATH": "rl_config.json",
+                    "BASETEN_API_KEY": definitions.SecretReference(name="baseten_api_key"),
+                },
+            ),
+            image=definitions.Image(
+                base_image="nvcr.io/nvidia/cuda:12.8.1-cudnn-devel-ubuntu24.04",
+            ),
+            workspace=definitions.Workspace(
+                workspace_root=str(ws_root),
+                exclude_dirs=[
+                    str(ws_root / ".venv"),
+                    str(ws_root / ".git"),
+                ],
+            ),
+        ),
+    )
+
+    try:
+        result = push(project, remote=remote, source_dir=ws_root)
+        job_id = result["id"]
+        print(f"Training Job ID: {job_id}")
+    finally:
+        rl_config_path.unlink(missing_ok=True)
+
+    # Build the deterministic pod DNS for the worker.
+    worker_host = (
+        f"baseten-training-job-{job_id}-multinode-0"
+        f".baseten-training-job-{job_id}-multinode"
+        f".{namespace}.svc.cluster.local"
+    )
+    resolved_url = f"http://{worker_host}:{worker_port}"
+    print(f"Worker URL: {resolved_url}")
+
+    return TrainingClient(resolved_url, api_key=api_key, timeout=timeout)

baseten_loops-0.1.0/baseten/loops/models.py

@@ -0,0 +1,291 @@
+"""Types for the loops SDK — wire-compatible with trainers-server."""
+
+from __future__ import annotations
+
+import base64
+from datetime import datetime
+from typing import Annotated, Literal, Self, Union
+
+from pydantic import BaseModel, Discriminator, Field, Tag
+
+
+# ── Tensor / model-input primitives ─────────────────────────────────
+
+
+class TensorData(BaseModel):
+    data: list
+    dtype: str
+    shape: list[int]
+
+    def tolist(self) -> list:
+        """Return the tensor data as a plain Python list."""
+        return list(self.data)
+
+    def to_torch(self):  # type: ignore[return]
+        """Return the tensor data as a torch.Tensor with the correct shape and dtype."""
+        import torch
+
+        _DTYPE_MAP = {
+            "float32": torch.float32,
+            "float16": torch.float16,
+            "bfloat16": torch.bfloat16,
+            "float64": torch.float64,
+            "int32": torch.int32,
+            "int64": torch.int64,
+            "int8": torch.int8,
+            "bool": torch.bool,
+        }
+        torch_dtype = _DTYPE_MAP.get(self.dtype, torch.float32)
+        # The wire format encodes masked / inactive positions as JSON null
+        # (since JSON has no NaN literal); the server emits them via
+        # masked_fill(..., float("nan")) and serialisation flips NaN → None.
+        # Convert back here so torch.tensor can ingest the list. Only floats
+        # carry the None-as-NaN convention — int / bool tensors never have
+        # masked positions.
+        data = self.data
+        if torch_dtype.is_floating_point and any(v is None for v in data):
+            data = [float("nan") if v is None else v for v in data]
+        return torch.tensor(data, dtype=torch_dtype).reshape(self.shape)
+
+    @classmethod
+    def from_torch(cls, tensor) -> Self:
+        """Construct TensorData from a torch.Tensor."""
+        dtype_name = str(tensor.dtype).replace("torch.", "")
+        return cls(
+            data=tensor.flatten().tolist(),
+            dtype=dtype_name,
+            shape=list(tensor.shape),
+        )
+
+    @classmethod
+    def from_list(cls, data: list, dtype: str = "float32") -> Self:
+        flat = data
+        shape = []
+        level = data
+        while isinstance(level, list):
+            shape.append(len(level))
+            level = level[0] if level else []
+        return cls(data=flat, dtype=dtype, shape=shape)
+
+
+class EncodedTextChunk(BaseModel):
+    type: Literal["encoded_text"] = "encoded_text"
+    tokens: list[int]
+
+    @property
+    def length(self) -> int:
+        return len(self.tokens)
+
+
+class ImageChunk(BaseModel):
+    type: Literal["image"] = "image"
+    data: str  # base64-encoded
+    format: str
+    expected_tokens: int = 0
+
+    @classmethod
+    def from_bytes(cls, raw: bytes, fmt: str, expected_tokens: int = 0) -> Self:
+        return cls(
+            data=base64.b64encode(raw).decode(),
+            format=fmt,
+            expected_tokens=expected_tokens,
+        )
+
+    def to_bytes(self) -> bytes:
+        return base64.b64decode(self.data)
+
+    @property
+    def length(self) -> int:
+        return self.expected_tokens
+
+
+ModelInputChunk = Annotated[
+    Union[
+        Annotated[EncodedTextChunk, Tag("encoded_text")],
+        Annotated[ImageChunk, Tag("image")],
+    ],
+    Discriminator("type"),
+]
+
+
+class ModelInput(BaseModel):
+    chunks: list[ModelInputChunk] = Field(default_factory=list)
+
+    @classmethod
+    def from_ints(cls, tokens: list[int]) -> Self:
+        return cls(chunks=[EncodedTextChunk(tokens=tokens)])
+
+    @classmethod
+    def empty(cls) -> Self:
+        return cls(chunks=[])
+
+    @property
+    def length(self) -> int:
+        return sum(c.length for c in self.chunks)
+
+    def to_ints(self) -> list[int]:
+        out: list[int] = []
+        for chunk in self.chunks:
+            if not isinstance(chunk, EncodedTextChunk):
+                raise TypeError(f"Cannot convert {type(chunk).__name__} to ints")
+            out.extend(chunk.tokens)
+        return out
+
+    def append(self, chunk: ModelInputChunk) -> Self:
+        return type(self)(chunks=[*self.chunks, chunk])
+
+
+class Datum(BaseModel):
+    model_input: ModelInput
+    loss_fn_inputs: dict[str, TensorData] = Field(default_factory=dict)
+
+
+# ── Sampling ─────────────────────────────────────────────────────────
+
+
+class SamplingParams(BaseModel):
+    max_tokens: int | None = None
+    seed: int | None = None
+    stop: str | list[str] | list[int] | None = None
+    temperature: float = 1.0
+    top_k: int = -1
+    top_p: float = 1.0
+
+
+class SampledSequence(BaseModel):
+    tokens: list[int] = Field(default_factory=list)
+    logprobs: list[float] | None = None
+    stop_reason: str = "length"
+
+
+class PromptLogprobs(BaseModel):
+    """Per-token log-probabilities for the prompt.
+
+    ``logprobs[i]`` is the logprob of the actual prompt token at position
+    ``i``. Index 0 is always ``None`` — the first token has no preceding
+    context to score against.
+
+    ``top_logprobs[i]`` (only populated when the caller passed
+    ``topk_prompt_logprobs > 0``) maps ``token_id -> logprob`` for the top-k
+    alternatives at position ``i``. The actual prompt token is also included
+    even if it falls outside the top-k.
+    """
+
+    logprobs: list[float | None] = Field(default_factory=list)
+    top_logprobs: list[dict[int, float] | None] | None = None
+
+
+class SampleResult(BaseModel):
+    sequences: list[SampledSequence] = Field(default_factory=list)
+    policy_version: int | None = None
+    prompt_logprobs: PromptLogprobs | None = None
+
+
+# Alias kept for API surface compatibility.
+SampleResponse = SampleResult
+
+
+# ── Optimizer ────────────────────────────────────────────────────────
+
+
+class AdamParams(BaseModel):
+    learning_rate: float = 1e-4
+    beta1: float = 0.9
+    beta2: float = 0.95
+    eps: float = 1e-12
+    weight_decay: float = 0.0
+    grad_clip_norm: float = 0.0  # 0 = disabled
+
+
+# ── Training response types ──────────────────────────────────────────
+
+
+class ForwardBackwardOutput(BaseModel):
+    loss: float = 0.0
+    loss_fn_output_type: str = ""
+    loss_fn_outputs: list[dict[str, TensorData]] = Field(default_factory=list)
+    metrics: dict[str, float] = Field(default_factory=dict)
+
+
+class OptimStepResponse(BaseModel):
+    metrics: dict[str, float] | None = None
+
+
+class SaveWeightsResponse(BaseModel):
+    mode: str = ""
+    path: str = ""
+
+
+class SaveWeightsForSamplerResponse(BaseModel):
+    """Response for ``/save_weights_for_sampler`` (the trainer-side
+    publish-LoRA endpoint).
+
+    ``version`` is the controller's ``step_count`` at the moment of the
+    call — clients don't pass it. Two syncs at the same version mean
+    "no new weights since last check"; a sampler watching the URI keys
+    its hot-swap on this number.
+
+    ``path`` is the checkpoint subdirectory the adapters were written
+    to (``{weight_sync_uri}/{name}/``).
+    """
+    version: int = 0
+    path: str = ""
+
+
+class LoadWeightsResponse(BaseModel):
+    step: int = 0
+
+
+class InitTrainerServerResponse(BaseModel):
+    """Response for ``/init_trainer_server``."""
+
+    step: int = 0
+    lora_rank: int
+
+
+# ── Checkpoints ──────────────────────────────────────────────────────
+
+
+class Checkpoint(BaseModel):
+    model_config = {"populate_by_name": True}
+
+    trainer_server_id: str = Field(alias="trainer_id")
+    checkpoint_id: str
+    checkpoint_type: str
+    created_at: datetime
+    base_model: str | None = None
+    lora_adapter_config: dict | None = None
+    size_bytes: int = 0
+    sync_status: str | None = None
+
+
+class CheckpointFile(BaseModel):
+    url: str
+    relative_file_name: str
+    size_bytes: int
+    last_modified: datetime
+
+
+class CheckpointFilesResponse(BaseModel):
+    presigned_urls: list[CheckpointFile]
+    next_page_token: int | None = None
+    total_count: int
+
+
+# ── Model info ───────────────────────────────────────────────────────
+
+
+class ModelData(BaseModel):
+    arch: str | None = None
+    model_name: str | None = None
+    tokenizer_id: str | None = None
+
+
+class GetInfoResponse(BaseModel):
+    model_data: ModelData = Field(default_factory=ModelData)
+    model_name: str | None = None
+    is_lora: bool | None = None
+    lora_rank: int | None = None
+    max_seq_len: int | None = None
+
+    model_config = {"protected_namespaces": ()}