softs 0.1.0__py3-none-any.whl

softs/__init__.py

@@ -0,0 +1,25 @@
+"""Soft labels: async on-the-fly training data generation for PyTorch."""
+
+import logging
+
+__version__ = "0.6.0"
+
+
+def setup_logging(level: int | str = logging.INFO) -> None:
+    if isinstance(level, str):
+        level = getattr(logging, level.upper(), logging.INFO)
+    logging.basicConfig(
+        level=level,
+        format="%(asctime)s [%(name)s] %(levelname)s: %(message)s",
+        datefmt="%H:%M:%S",
+    )
+
+
+# User-facing
+from .configs import BatchConfig, TensorSpec, make_xy_config
+from .market import EndpointConfig
+from .dataset import SoftIterableDataset, SoftDataLoader, Batch, make_collate_fn
+
+# Market internals (for advanced use)
+from .market import Broker, Supplier, Client, BrokerStats
+from .market import ShmMedium, FilesystemMedium, TCPMedium

softs/configs.py

@@ -0,0 +1,196 @@
+"""Batch config, tensor specs, and encoding utilities."""
+
+import math
+from dataclasses import dataclass
+
+import torch
+import yaml
+
+_DTYPE_MAP: dict[str, torch.dtype] = {
+    "float64": torch.float64,
+    "float32": torch.float32,
+    "float16": torch.float16,
+    "bfloat16": torch.bfloat16,
+    "int64": torch.int64,
+    "int32": torch.int32,
+    "int16": torch.int16,
+    "int8": torch.int8,
+    "uint8": torch.uint8,
+    "bool": torch.bool,
+}
+_TORCH_DTYPE_TO_STR: dict[torch.dtype, str] = {
+    v: k for k, v in _DTYPE_MAP.items()
+}
+_DTYPE_SIZES: dict[torch.dtype, int] = {
+    d: torch.empty(1, dtype=d).element_size() for d in _DTYPE_MAP.values()
+}
+
+
+def dtype_to_torch(dtype_str: str) -> torch.dtype:
+    try:
+        return _DTYPE_MAP[dtype_str]
+    except KeyError:
+        raise ValueError(f"Unsupported dtype: {dtype_str}") from None
+
+
+def torch_dtype_to_str(dtype: torch.dtype) -> str:
+    try:
+        return _TORCH_DTYPE_TO_STR[dtype]
+    except KeyError:
+        raise ValueError(f"Unsupported torch dtype: {dtype}") from None
+
+
+@dataclass
+class TensorSpec:
+    name: str
+    shape: tuple[int, ...]
+    dtype: str
+
+    @property
+    def torch_dtype(self) -> torch.dtype:
+        return dtype_to_torch(self.dtype)
+
+    @property
+    def numel(self) -> int:
+        return math.prod(self.shape)
+
+    @property
+    def nbytes(self) -> int:
+        return self.numel * _DTYPE_SIZES[self.torch_dtype]
+
+    def encode(self, tensor: torch.Tensor) -> bytes:
+        if tuple(tensor.shape) != self.shape:
+            raise ValueError(
+                f"Shape mismatch: {tuple(tensor.shape)} != {self.shape}"
+            )
+        if tensor.dtype != self.torch_dtype:
+            raise ValueError(
+                f"Dtype mismatch: {tensor.dtype} != {self.torch_dtype}"
+            )
+        tensor = tensor.contiguous()
+        if tensor.dtype == torch.bfloat16:
+            return tensor.view(torch.uint16).numpy().tobytes()
+        return tensor.numpy().tobytes()
+
+    def decode(self, data: bytes) -> torch.Tensor:
+        if len(data) != self.nbytes:
+            raise ValueError(
+                f"Data length {len(data)} != expected {self.nbytes}"
+            )
+        if self.dtype == "bfloat16":
+            tensor = torch.frombuffer(bytearray(data), dtype=torch.uint16).view(
+                torch.bfloat16
+            )
+        else:
+            tensor = torch.frombuffer(bytearray(data), dtype=self.torch_dtype)
+        return tensor.reshape(self.shape).clone()
+
+
+@dataclass
+class BatchConfig:
+    specs: list[TensorSpec]
+
+    def __post_init__(self):
+        converted = []
+        for spec in self.specs:
+            if isinstance(spec, TensorSpec):
+                converted.append(spec)
+            elif hasattr(spec, "items"):
+                converted.append(
+                    TensorSpec(
+                        **{
+                            k: v
+                            for k, v in spec.items()
+                            if not str(k).startswith("_")
+                        }
+                    )
+                )
+            else:
+                raise TypeError(
+                    f"Expected TensorSpec or dict-like, got {type(spec)}"
+                )
+        object.__setattr__(self, "specs", converted)
+
+        names = [s.name for s in self.specs]
+        if len(names) != len(set(names)):
+            raise ValueError(f"Duplicate tensor names: {names}")
+
+        self._name_to_idx = {s.name: i for i, s in enumerate(self.specs)}
+        self._offsets = []
+        off = 0
+        for s in self.specs:
+            self._offsets.append(off)
+            off += s.nbytes
+
+    @property
+    def tensor_names(self) -> list[str]:
+        return [s.name for s in self.specs]
+
+    def nbytes(self) -> int:
+        return sum(s.nbytes for s in self.specs)
+
+    def get_spec(self, name: str) -> TensorSpec:
+        idx = self._name_to_idx.get(name)
+        if idx is None:
+            raise KeyError(
+                f"Unknown tensor '{name}'. Available: {self.tensor_names}"
+            )
+        return self.specs[idx]
+
+    def get_offset(self, name: str) -> int:
+        idx = self._name_to_idx.get(name)
+        if idx is None:
+            raise KeyError(f"Unknown tensor '{name}'")
+        return self._offsets[idx]
+
+    def encode(self, **tensors: torch.Tensor) -> bytes:
+        provided = set(tensors.keys())
+        expected = set(self.tensor_names)
+        if provided != expected:
+            raise ValueError(
+                f"Tensor mismatch: missing={expected - provided}, extra={provided - expected}"
+            )
+        return b"".join(spec.encode(tensors[spec.name]) for spec in self.specs)
+
+    def decode(self, data: bytes) -> dict[str, torch.Tensor]:
+        if len(data) != self.nbytes():
+            raise ValueError(
+                f"Data length {len(data)} != expected {self.nbytes()}"
+            )
+        return {
+            spec.name: spec.decode(data[off : off + spec.nbytes])
+            for spec, off in zip(self.specs, self._offsets)
+        }
+
+    def decode_single(self, data: bytes, name: str) -> torch.Tensor:
+        spec = self.get_spec(name)
+        off = self.get_offset(name)
+        return spec.decode(data[off : off + spec.nbytes])
+
+    @classmethod
+    def from_dict(cls, config: dict) -> "BatchConfig":
+        return cls(config.get("specs", []))
+
+    @classmethod
+    def from_yaml(cls, path: str) -> "BatchConfig":
+        with open(path) as f:
+            return cls.from_dict(yaml.safe_load(f))
+
+    def to_dict(self) -> dict:
+        return {
+            "specs": [
+                {"name": s.name, "shape": list(s.shape), "dtype": s.dtype}
+                for s in self.specs
+            ]
+        }
+
+
+def make_xy_config(
+    x_shape: tuple[int, ...],
+    x_dtype: str,
+    y_shape: tuple[int, ...],
+    y_dtype: str,
+) -> BatchConfig:
+    return BatchConfig(
+        [TensorSpec("x", x_shape, x_dtype), TensorSpec("y", y_shape, y_dtype)]
+    )

softs/dataset.py

@@ -0,0 +1,201 @@
+"""PyTorch Dataset and DataLoader for soft label generation."""
+
+import ctypes
+import multiprocessing
+import time
+from typing import Callable, Iterator
+
+import torch
+from torch.utils.data import DataLoader, IterableDataset
+
+from .configs import BatchConfig
+from .market import Client
+
+
+class _SharedStr:
+    """Process-safe mutable string via shared memory."""
+
+    def __init__(self, value: str = "", max_len: int = 256):
+        self._buf = multiprocessing.RawArray(ctypes.c_char, max_len)
+        self._len = multiprocessing.RawValue(ctypes.c_int, 0)
+        self._max = max_len
+        if value:
+            self.set(value)
+
+    def set(self, value: str) -> None:
+        encoded = value.encode()[: self._max]
+        self._buf[: len(encoded)] = encoded
+        self._len.value = len(encoded)
+
+    def get(self) -> str:
+        return bytes(self._buf[: self._len.value]).decode()
+
+
+class SoftIterableDataset(IterableDataset[dict[str, torch.Tensor]]):
+    """Infinite dataset yielding decoded tensor dicts.
+
+    Call ``set_model(model_id)`` to switch models at any time.
+    DataLoader workers detect the change and discard pending work.
+    """
+
+    def __init__(
+        self,
+        model_id: str,
+        endpoint: str,
+        batch_config: BatchConfig,
+        medium_cls,
+        num_slots: int = 8,
+        max_retries: int = 10,
+        retry_delay: float = 0.01,
+    ):
+        self._model = _SharedStr(model_id)
+        self.endpoint = endpoint
+        self.num_slots = num_slots
+        self.batch_config = batch_config
+        self.medium_cls = medium_cls
+        self.max_retries = max_retries
+        self.retry_delay = retry_delay
+        self._client: Client | None = None
+
+    @property
+    def model_id(self) -> str:
+        return self._model.get()
+
+    def set_model(self, model_id: str) -> None:
+        self._model.set(model_id)
+
+    def _ensure_client(self) -> Client:
+        if self._client is None:
+            self._client = Client(
+                endpoint=self.endpoint,
+                slot_size=self.batch_config.nbytes(),
+                medium_cls=self.medium_cls,
+                num_slots=self.num_slots,
+            )
+            self._client.hello()
+        return self._client
+
+    def __iter__(self) -> Iterator[dict[str, torch.Tensor]]:
+        client = self._ensure_client()
+        retries = 0
+        current = self.model_id
+        while True:
+            wanted = self.model_id
+            if wanted != current:
+                client.discard()
+                current = wanted
+                retries = 0
+            slot_id = client.request_sample(
+                current, timeout_ms=int(self.retry_delay * 1000)
+            )
+            if slot_id is None:
+                retries += 1
+                if retries >= self.max_retries:
+                    time.sleep(self.retry_delay)
+                continue
+            retries = 0
+            tensors = self.batch_config.decode(client.medium.read(slot_id))
+            client.release_slot(slot_id)
+            yield tensors
+
+    def __del__(self):
+        if self._client is not None:
+            self._client.close()
+
+
+class SoftDataLoader(DataLoader):
+    """DataLoader with model switching support.
+
+    Usage::
+
+        loader = SoftDataLoader(model_id="teacher_v1", slot_count=8,
+                                batch_config=config, endpoints=ep,
+                                medium_cls=ShmMedium, batch_size=4)
+        for batch in loader:
+            train(batch)
+
+        loader.set_model("teacher_v2")  # all workers switch automatically
+    """
+
+    def __init__(
+        self,
+        model_id: str,
+        endpoint: str,
+        batch_config: BatchConfig,
+        medium_cls,
+        num_slots: int = 8,
+        max_retries: int = 10,
+        retry_delay: float = 0.01,
+        **dataloader_kwargs,
+    ):
+        self.dataset = SoftIterableDataset(
+            model_id=model_id,
+            endpoint=endpoint,
+            batch_config=batch_config,
+            medium_cls=medium_cls,
+            num_slots=num_slots,
+            max_retries=max_retries,
+            retry_delay=retry_delay,
+        )
+        super().__init__(self.dataset, **dataloader_kwargs)
+
+    def set_model(self, model_id: str) -> None:
+        self.dataset.set_model(model_id)
+
+    @property
+    def model_id(self) -> str:
+        return self.dataset.model_id
+
+
+class Batch:
+    def __init__(
+        self,
+        tensors: dict[str, torch.Tensor],
+        slot_ids: list[int] | None = None,
+        client: Client | None = None,
+    ):
+        self.tensors = tensors
+        self.slot_ids = slot_ids or []
+        self._client = client
+        self._released = False
+
+    def __getitem__(self, key: str) -> torch.Tensor:
+        return self.tensors[key]
+
+    def __contains__(self, key: str) -> bool:
+        return key in self.tensors
+
+    def keys(self):
+        return self.tensors.keys()
+
+    def release(self) -> None:
+        if self._released or self._client is None:
+            return
+        for slot_id in self.slot_ids:
+            self._client.release_slot(slot_id)
+        self._released = True
+
+    def __del__(self):
+        self.release()
+
+
+def make_collate_fn(
+    client: Client, batch_config: BatchConfig, auto_release: bool = True
+) -> Callable[[list[int]], Batch]:
+    def collate_fn(slot_ids: list[int]) -> Batch:
+        tensor_lists: dict[str, list[torch.Tensor]] = {
+            name: [] for name in batch_config.tensor_names
+        }
+        for slot_id in slot_ids:
+            for name, tensor in batch_config.decode(
+                client.medium.read(slot_id)
+            ).items():
+                tensor_lists[name].append(tensor)
+        batched = {name: torch.stack(ts) for name, ts in tensor_lists.items()}
+        if auto_release:
+            for slot_id in slot_ids:
+                client.release_slot(slot_id)
+            return Batch(batched, slot_ids, None)
+        return Batch(batched, slot_ids, client)
+
+    return collate_fn

softs/market/__init__.py

@@ -0,0 +1,7 @@
+"""Marketplace: broker, supplier, client, and transfer mediums."""
+
+from .broker import Broker
+from .supplier import Supplier
+from .client import Client
+from .protocol import EndpointConfig, BrokerStats
+from .mediums import Medium, ShmMedium, FilesystemMedium, TCPMedium