gradientmail 0.0.1__py3-none-any.whl

gradientmail/lease.py

@@ -0,0 +1,44 @@
+"""Learner liveness: a lease per learner, renewed by beats and judged on the observer's clock.
+
+A learner's lease is the newest of the create-once beat objects under ``leases/{learner}/``
+(``protocol`` module). The learner renews it every ``lease_every`` seconds while it runs and
+deletes it when it closes. An observer lists the leases on every poll and remembers, per
+learner, the newest beat and the moment on its own clock when it first saw that beat: a
+learner is live while its beat has changed within ``lease_timeout`` seconds, and gone the
+moment its lease leaves the listing. No wall clock crosses machines, so a lease written by a
+node with a skewed clock reads like any other; the cost is that an observer trusts a lease it
+has just met for one timeout before doubting it.
+"""
+
+from collections.abc import Mapping
+from dataclasses import dataclass, field, replace
+
+__all__ = ["Liveness"]
+
+
+@dataclass(frozen=True)
+class Liveness:
+    """Which learners hold a live lease, as one observer has seen it.
+
+    ``seen`` maps a learner to its newest beat and when this observer first saw that beat.
+    """
+
+    timeout: float
+    seen: Mapping[str, tuple[int, float]] = field(default_factory=dict)
+
+    def observed(self, leases: Mapping[str, int], now: float) -> "Liveness":
+        """Fold one listing of the leases, taken at ``now``."""
+
+        def since(learner: str, beat: int) -> float:
+            known = self.seen.get(learner)
+            return known[1] if known is not None and known[0] == beat else now
+
+        return replace(
+            self, seen={learner: (beat, since(learner, beat)) for learner, beat in leases.items()}
+        )
+
+    def live(self, now: float) -> frozenset[str]:
+        """The learners whose beat has changed within ``timeout`` seconds of ``now``."""
+        return frozenset(
+            learner for learner, (_, since) in self.seen.items() if now - since < self.timeout
+        )

gradientmail/mailbox/__init__.py

@@ -0,0 +1,57 @@
+"""Transport: four operations over immutable, create-once objects.
+
+An object is written at most once (``put_if_absent``), read back by key (``get``), discovered
+by prefix in UTF-8 byte order (``list``) and removed idempotently (``delete``). Nothing here
+knows about rounds or snapshots; ``gradientmail.protocol`` builds the exchange on top.
+"""
+
+import logging
+from collections.abc import Awaitable, Callable, Iterable, Sequence
+from typing import Protocol, runtime_checkable
+
+from gradientmail.clock import Clock
+
+log = logging.getLogger(__name__)
+
+
+class MailboxUnavailable(Exception):
+    """Transient transport failure (timeouts, dropped connections); callers retry."""
+
+
+@runtime_checkable
+class Mailbox(Protocol):
+    async def put_if_absent(self, key: str, data: bytes) -> bool:
+        """Create ``key`` holding ``data``; True iff this call created it.
+
+        An existing object is left untouched and False is returned.
+        """
+        ...
+
+    async def get(self, key: str) -> bytes | None:
+        """The bytes under ``key``, or None when no such object exists."""
+        ...
+
+    async def list(self, prefix: str, *, after: str = "") -> Sequence[str]:
+        """Every key under ``prefix`` greater than ``after``, in UTF-8 byte order, all pages."""
+        ...
+
+    async def delete(self, keys: Iterable[str]) -> None:
+        """Remove ``keys`` in any number; missing keys are not an error."""
+        ...
+
+
+async def retrying[T](
+    op: Callable[[], Awaitable[T]], *, clock: Clock, max_delay: float = 30.0
+) -> T:
+    """Run ``op`` until it returns, sleeping through ``clock`` after each ``MailboxUnavailable``.
+
+    Delays start at 0.5 s and double up to ``max_delay``. Any other exception propagates at once.
+    """
+    delay = min(0.5, max_delay)
+    while True:
+        try:
+            return await op()
+        except MailboxUnavailable as exc:
+            log.warning("mailbox unavailable: %s; retrying in %.1f s", exc, delay)
+            await clock.sleep(delay)
+            delay = min(delay * 2, max_delay)

gradientmail/mailbox/memory.py

@@ -0,0 +1,39 @@
+"""In-memory Mailbox for development, tests and the Phase 2 example."""
+
+import threading
+from collections.abc import Iterable, Sequence
+
+from gradientmail.mailbox import Mailbox
+
+
+class MemoryMailbox(Mailbox):
+    """A dict behind a ``threading.Lock``.
+
+    The lock is not an asyncio primitive because one instance is shared between event loops
+    running on different threads (a ``Session``'s background loop and a ``Syncer``).
+    """
+
+    def __init__(self) -> None:
+        self._store: dict[str, bytes] = {}
+        self._lock = threading.Lock()
+
+    async def put_if_absent(self, key: str, data: bytes) -> bool:
+        with self._lock:
+            if key in self._store:
+                return False
+            self._store[key] = data
+            return True
+
+    async def get(self, key: str) -> bytes | None:
+        with self._lock:
+            return self._store.get(key)
+
+    async def list(self, prefix: str, *, after: str = "") -> Sequence[str]:
+        # Code point order is UTF-8 byte order, so str comparison matches S3's listing order.
+        with self._lock:
+            return sorted(k for k in self._store if k.startswith(prefix) and k > after)
+
+    async def delete(self, keys: Iterable[str]) -> None:
+        with self._lock:
+            for key in keys:
+                self._store.pop(key, None)

gradientmail/mailbox/s3.py

@@ -0,0 +1,210 @@
+"""S3-compatible object storage as a ``Mailbox``: AWS S3, Cloudflare R2, Hugging Face Storage
+Buckets, MinIO, moto.
+
+Consistency assumptions
+-----------------------
+Three properties of the store carry the whole protocol: strong read-after-write and
+list-after-write (AWS general-purpose buckets since 2020-12, R2, MinIO on local xfs, zfs or
+btrfs, moto), an atomic ``If-None-Match: *`` on PUT (AWS since 2024-08, R2, MinIO since
+2024-04, moto since 5.0.15) and ``ListObjectsV2`` in UTF-8 byte order. ``StartAfter`` is sent
+but not relied on: Hugging Face's gateway accepts and ignores it, so ``list`` cuts at ``after``
+itself and a poll there lists the whole prefix. Nothing else is needed because every object is
+written once, payload keys carry the sha256 of their bytes and a record is written after the
+payloads it points at.
+
+Hugging Face Storage Buckets are reached at ``https://s3.hf.co/<namespace>`` with
+``request_checksums="when_required"`` and S3 credentials generated from a Hub token; a PUT there
+takes about half a second. Verified with the contract and training suites in October 2026.
+
+Not assumed: FIFO delivery, exactly-once, atomicity across keys, server clocks or ETag-as-MD5.
+Without multi-key atomicity a listing may show a payload before its record, so readers follow
+records and never list ``blobs/``.
+
+Unsupported: S3 Express directory buckets (unordered listing, no ``StartAfter``), R2 behind a
+caching custom domain (a cached read breaks read-after-write), MinIO on NFS (not a supported
+MinIO backend).
+
+Failure modes
+-------------
+A refused connection, a connect or read timeout and a connection closed mid-request become
+``MailboxUnavailable`` once botocore's standard retry mode has made its five attempts; the
+protocol layer then retries with backoff for as long as the outage lasts. A 412
+``PreconditionFailed`` is the create-once contract (``put_if_absent`` returns False). A 409
+``ConditionalRequestConflict`` means another conditional write on the same key was in flight
+and the outcome is unknown; the put is retried three times, then raised as
+``MailboxUnavailable``. Every other ``ClientError`` (a missing bucket, access denied,
+throttling that outlasts botocore's retries) propagates unchanged: on a learner it is raised
+from the ``Session`` constructor while joining and from the next ``step()`` afterwards, on the
+syncer from ``run()``. A payload above 5 GiB raises
+``ValueError`` instead of switching to multipart uploads; choose more fragments. A bulk delete
+reports per-key failures in a 200 response; the first one is raised as a ``ClientError``. On
+AWS and R2 a LIST costs 12.5 times a GET, and every learner and the syncer issue one LIST per
+poll, so ``poll_interval`` is the cost lever.
+
+Every call runs in ``asyncio.to_thread``; a ``threading.BoundedSemaphore`` acquired inside the
+thread keeps the connection pool from being oversubscribed. No asyncio primitive lives on the
+instance, so one mailbox can be driven from several event loops in turn.
+"""
+
+import asyncio
+import threading
+from collections.abc import Callable, Iterable, Sequence
+from functools import partial
+from itertools import batched
+from typing import TYPE_CHECKING, Literal
+
+import boto3.session
+from botocore.config import Config
+from botocore.exceptions import (
+    ClientError,
+    ConnectionClosedError,
+    ConnectTimeoutError,
+    EndpointConnectionError,
+    ReadTimeoutError,
+)
+
+from gradientmail.mailbox import Mailbox, MailboxUnavailable
+
+if TYPE_CHECKING:
+    from mypy_boto3_s3 import S3Client
+
+MAX_OBJECT_BYTES = 5 * 2**30
+"""Largest single PUT S3 accepts; larger payloads need more fragments, not multipart uploads."""
+
+
+class S3Mailbox(Mailbox):
+    """Keys live under ``prefix`` in ``bucket``; credentials come from the boto3 default chain.
+
+    Parameters
+    ----------
+    endpoint_url
+        None for AWS; the service URL for R2, MinIO or any other S3-compatible store.
+    region
+        Signing region; None lets boto3 resolve it from the environment.
+    max_concurrency
+        Requests in flight at once, which is also the size of the connection pool.
+    request_checksums
+        ``"when_supported"`` asks the store to verify a SHA-256 of every uploaded object;
+        ``"when_required"`` sends only the checksums the protocol mandates, for stores that
+        reject flexible checksums (Hugging Face Storage Buckets).
+    """
+
+    def __init__(
+        self,
+        bucket: str,
+        prefix: str,
+        *,
+        endpoint_url: str | None = None,
+        region: str | None = None,
+        max_concurrency: int = 8,
+        request_checksums: Literal["when_supported", "when_required"] = "when_supported",
+    ) -> None:
+        self._bucket = bucket
+        self._prefix = f"{prefix.strip('/')}/" if prefix.strip("/") else ""
+        self._slots = threading.BoundedSemaphore(max_concurrency)
+        self._client: S3Client = boto3.session.Session().client(
+            "s3",
+            endpoint_url=endpoint_url,
+            region_name=region,
+            config=Config(
+                retries={"mode": "standard", "total_max_attempts": 5},
+                max_pool_connections=max_concurrency,
+                connect_timeout=10,
+                read_timeout=120,
+                s3={"addressing_style": "path"},
+                request_checksum_calculation=request_checksums,
+                response_checksum_validation=request_checksums,
+            ),
+        )
+        put_object = partial(self._client.put_object, Bucket=bucket, IfNoneMatch="*")
+        self._put_object = (
+            partial(put_object, ChecksumAlgorithm="SHA256")
+            if request_checksums == "when_supported"
+            else put_object
+        )
+
+    async def put_if_absent(self, key: str, data: bytes) -> bool:
+        if len(data) > MAX_OBJECT_BYTES:
+            raise ValueError(
+                f"{key}: {len(data)} bytes exceeds the 5 GiB single-object limit; "
+                "use more fragments"
+            )
+
+        def put() -> bool:
+            # A 409 ConditionalRequestConflict means another conditional write on the same
+            # key was in progress; the outcome is unknown until one of them has landed.
+            for _ in range(3):
+                try:
+                    self._put_object(Key=self._prefix + key, Body=data)
+                except ClientError as exc:
+                    code = exc.response.get("Error", {}).get("Code")
+                    if code == "PreconditionFailed":
+                        return False
+                    if code != "ConditionalRequestConflict":
+                        raise
+                else:
+                    return True
+            raise MailboxUnavailable(f"{key}: conditional write conflicted three times")
+
+        return await self._call(put)
+
+    async def get(self, key: str) -> bytes | None:
+        def get() -> bytes | None:
+            try:
+                response = self._client.get_object(Bucket=self._bucket, Key=self._prefix + key)
+            except self._client.exceptions.NoSuchKey:
+                return None
+            return response["Body"].read()
+
+        return await self._call(get)
+
+    async def list(self, prefix: str, *, after: str = "") -> Sequence[str]:
+        def list_keys() -> Sequence[str]:
+            paginator = self._client.get_paginator("list_objects_v2")
+            full_prefix = self._prefix + prefix
+            pages = (
+                paginator.paginate(
+                    Bucket=self._bucket, Prefix=full_prefix, StartAfter=self._prefix + after
+                )
+                if after
+                else paginator.paginate(Bucket=self._bucket, Prefix=full_prefix)
+            )
+            # Hugging Face's gateway accepts StartAfter and ignores it, so the cut is made here
+            # as well; on stores that honour it the comparison never drops anything.
+            return tuple(
+                stripped
+                for page in pages
+                for obj in page.get("Contents", ())
+                if (key := obj.get("Key")) is not None
+                and (stripped := key[len(self._prefix) :]) > after
+            )
+
+        return await self._call(list_keys)
+
+    async def delete(self, keys: Iterable[str]) -> None:
+        def delete_chunk(chunk: Sequence[str]) -> None:
+            response = self._client.delete_objects(
+                Bucket=self._bucket,
+                Delete={"Objects": [{"Key": self._prefix + key} for key in chunk], "Quiet": True},
+            )
+            # Per-key failures arrive in a 200 response, not as an exception.
+            if errors := response.get("Errors"):
+                raise ClientError({"Error": errors[0]}, "DeleteObjects")
+
+        for chunk in batched(keys, 1000):
+            await self._call(partial(delete_chunk, chunk))
+
+    async def _call[T](self, op: Callable[[], T]) -> T:
+        def guarded() -> T:
+            with self._slots:
+                try:
+                    return op()
+                except (
+                    EndpointConnectionError,
+                    ConnectTimeoutError,
+                    ReadTimeoutError,
+                    ConnectionClosedError,
+                ) as exc:
+                    raise MailboxUnavailable(str(exc)) from exc
+
+        return await asyncio.to_thread(guarded)

gradientmail/patterns/__init__.py

@@ -0,0 +1,4 @@
+from gradientmail.patterns.barrier import Barrier
+from gradientmail.patterns.shared_queue import shared_queue
+
+__all__ = ["Barrier", "shared_queue"]

gradientmail/patterns/barrier.py

@@ -0,0 +1,67 @@
+"""Hold nodes at a named point until a number of them have arrived."""
+
+import asyncio
+from functools import partial
+
+from gradientmail.clock import Clock, MonotonicClock
+from gradientmail.mailbox import Mailbox, retrying
+
+__all__ = ["Barrier"]
+
+
+class Barrier:
+    """A rendezvous through the mailbox: ``wait(name, parties=n)`` returns once ``n`` nodes arrived.
+
+    Nodes share a run-scoped mailbox, with a distinct learner ID per node. An arrival is the
+    create-once object ``barriers/{name}/{learner_id}``; ``wait`` creates this node's and lists
+    the prefix every ``poll_interval`` seconds until it holds ``parties`` keys or more. Arrivals
+    are never removed: a name serves once per run, a loop names each barrier (``f"epoch-{n}"``)
+    and a node arriving after the barrier opened passes at once. Transient mailbox failures are
+    retried. ``clock`` is the time source for polling and the timeout.
+    """
+
+    def __init__(
+        self,
+        *,
+        mailbox: Mailbox,
+        learner_id: str,
+        poll_interval: float = 1.0,
+        clock: Clock | None = None,
+    ) -> None:
+        self._mailbox = mailbox
+        self._learner_id = learner_id
+        self._poll_interval = poll_interval
+        self._clock = clock if clock is not None else MonotonicClock()
+
+    def wait(self, name: str, *, parties: int, timeout: float | None = None) -> frozenset[str]:
+        """Record this node's arrival at ``name`` and block until ``parties`` nodes have arrived.
+
+        ``parties`` is a count the caller fixes; it need not match the algorithm's quorum.
+        Returns the learner IDs that had arrived when the barrier opened. Raises ``TimeoutError``
+        after ``timeout`` seconds without opening; the arrival stays recorded, so calling again
+        resumes the wait. Calling again once the barrier is open returns at once.
+        """
+        if parties < 1:
+            raise ValueError("parties must be positive")
+        return asyncio.run(self._wait(name, parties, timeout))
+
+    async def _wait(self, name: str, parties: int, timeout: float | None) -> frozenset[str]:
+        prefix = f"barriers/{name}/"
+        deadline = None if timeout is None else self._clock.now() + timeout
+        await retrying(
+            partial(self._mailbox.put_if_absent, f"{prefix}{self._learner_id}", b""),
+            clock=self._clock,
+        )
+        while True:
+            arrived = await retrying(partial(self._mailbox.list, prefix), clock=self._clock)
+            if len(arrived) >= parties:
+                return frozenset(key.removeprefix(prefix) for key in arrived)
+            delay = self._poll_interval
+            if deadline is not None:
+                remaining = deadline - self._clock.now()
+                if remaining <= 0:
+                    raise TimeoutError(
+                        f"{len(arrived)} of {parties} nodes arrived at {name!r} within {timeout} s"
+                    )
+                delay = min(delay, remaining)
+            await self._clock.sleep(delay)

gradientmail/patterns/shared_queue.py

@@ -0,0 +1,104 @@
+"""Distribute work through mailbox claims and prefetch its loaded items."""
+
+import asyncio
+from collections import deque
+from collections.abc import Callable, Generator, Iterator, Mapping
+from contextlib import closing, contextmanager
+from threading import Condition, Thread
+
+from gradientmail.mailbox import Mailbox
+
+__all__ = ["shared_queue"]
+
+
+@contextmanager
+def shared_queue[Work, Batch](
+    units: Mapping[str, Work],
+    *,
+    mailbox: Mailbox,
+    learner_id: str,
+    load: Callable[[Work], Iterator[Batch]],
+    prefetch: int = 8,
+) -> Iterator[Iterator[Batch]]:
+    """Reserve units in mapping order and load their items on a background thread.
+
+    Nodes share unit identifiers and a run-scoped mailbox, with a distinct learner ID
+    per node. Values stay local; reservations use ``units/{identifier}`` keys. Exclusive
+    ownership requires atomic ``put_if_absent``. Claims remain after completion or failure.
+
+    At most ``prefetch`` loaded items wait for consumption. Loader and mailbox failures
+    propagate after buffered items. Exhaustion means no unclaimed work remains, not that
+    other nodes have finished. Exiting stops prefetch and joins the thread, waiting for
+    any current mailbox or loader operation to return.
+    """
+    if prefetch < 1:
+        raise ValueError("prefetch must be positive")
+
+    owner = learner_id.encode()
+    condition = Condition()
+    pending = deque[Batch]()
+    stopped = False
+    finished = False
+    failure: BaseException | None = None
+
+    async def claim(key: str) -> bool:
+        return await mailbox.put_if_absent(key, owner) and await mailbox.get(key) == owner
+
+    def loaded() -> Generator[Batch, None, None]:
+        with asyncio.Runner() as runner:
+            while True:
+                taken = set(runner.run(mailbox.list("units/")))
+                for identifier, work in units.items():
+                    key = f"units/{identifier}"
+                    if key not in taken and runner.run(claim(key)):
+                        yield from load(work)
+                        break
+                else:
+                    return
+
+    def produce() -> None:
+        nonlocal finished, failure
+        try:
+            with closing(loaded()) as source:
+                while True:
+                    with condition:
+                        condition.wait_for(lambda: stopped or len(pending) < prefetch)
+                        if stopped:
+                            return
+                    batch = next(source)
+                    with condition:
+                        pending.append(batch)
+                        condition.notify_all()
+        except StopIteration:
+            pass
+        except BaseException as error:
+            failure = error
+        finally:
+            with condition:
+                finished = True
+                condition.notify_all()
+
+    def consume() -> Iterator[Batch]:
+        while True:
+            with condition:
+                condition.wait_for(lambda: stopped or pending or finished)
+                if stopped:
+                    return
+                if not pending:
+                    if failure is not None:
+                        raise failure
+                    return
+                batch = pending.popleft()
+                condition.notify_all()
+            yield batch
+
+    thread = Thread(target=produce, name=f"shared-queue-{learner_id}")
+    thread.start()
+    try:
+        yield consume()
+    finally:
+        with condition:
+            stopped = True
+            condition.notify_all()
+        thread.join()
+        pending.clear()