gradientmail 0.0.1__py3-none-any.whl

gradientmail/learner.py

@@ -0,0 +1,891 @@
+"""The learner: counters and publishing decisions as pure state, plus the two-thread runtime.
+
+The training thread owns the model and every counter: ``step()`` advances the state, applies
+the fragment versions the receive loop has queued, and hands one fragment snapshot, with the
+anchor it was trained from, to the publish loop when a round is due; the publish loop
+computes their difference in the backend and publishes that delta. A daemon thread runs the
+receive and publish loops on its own event loop; the two sides meet only through two
+``queue.SimpleQueue``s, so ``step()`` never awaits and never does I/O. It blocks only when
+the schedule bounds ``overlap_steps``, waiting in ``step()`` for the model to align with a
+round it published that many steps ago, and ``close()`` waits for the last ones.
+
+A fetched batch is queued only once every payload in it is in hand, so ``rounds_seen`` never
+runs ahead of a fragment that is not yet applied: that is what keeps a live learner's
+snapshots at the syncer's current version.
+
+The receive loop also renews this learner's lease and lists everyone's (``lease`` module), so
+``step()`` knows how many learners are live when it judges whether a round is about to close
+and ``batch_size`` can split a global batch among them; ``close()`` releases the lease once the
+last snapshot is out.
+"""
+
+import asyncio
+import logging
+import queue
+import threading
+import time
+from collections.abc import Callable, Iterator, Mapping, Sequence
+from dataclasses import dataclass, replace
+
+from gradientmail.adapters import Adapter, Backend, default_backend
+from gradientmail.algorithms import (
+    Algorithm,
+    Estimates,
+    Participant,
+    Precision,
+    RoundSchedule,
+    decoded,
+    encoded,
+)
+from gradientmail.clock import Clock, MonotonicClock
+from gradientmail.compression import Compression
+from gradientmail.events import (
+    AwaitingCreator,
+    Callback,
+    Closed,
+    Closing,
+    CreatingRun,
+    FragmentAligned,
+    GaveUp,
+    Joined,
+    LeaseReleased,
+    PollFailed,
+    PublishCancelled,
+    PublishFailed,
+    ReleaseFailed,
+    RoundRecorded,
+    SnapshotPublished,
+    TookOverCreation,
+    log_events,
+)
+from gradientmail.lease import Liveness
+from gradientmail.mailbox import Mailbox
+from gradientmail.protocol import (
+    Exchange,
+    Manifest,
+    ManifestDraft,
+    PayloadRef,
+    ProtocolError,
+    RoundRecord,
+    RunConfigMismatch,
+    SnapshotDraft,
+    theta_digest,
+)
+from gradientmail.serialization import DType, Schema, SchemaMismatch, TensorBytes, validate
+
+__all__ = [
+    "Aligned",
+    "Learner",
+    "LearnerState",
+    "Publishers",
+    "advanced",
+    "align",
+    "applied",
+    "due_round",
+    "overdue",
+    "publishable",
+]
+
+log = logging.getLogger(__name__)
+
+
+@dataclass(frozen=True)
+class LearnerState:
+    """Everything the training thread tracks; one value per fragment where it says so.
+
+    ``rounds_seen`` is the newest round whose fragment the model is aligned with,
+    ``last_published`` the newest round this learner has a snapshot queued for and
+    ``published`` the ``(round, local_step)`` of its last ``in_flight`` publishes, oldest first.
+    """
+
+    local_step: int
+    steps: tuple[int, ...]
+    tokens: tuple[int, ...]
+    applied_version: tuple[int, ...]
+    rounds_seen: int
+    last_published: int
+    published: tuple[tuple[int, int], ...]
+    step_seconds: float
+
+
+@dataclass(frozen=True)
+class Publishers:
+    """The other learners with a snapshot in one round, as the receive loop last listed it."""
+
+    round: int
+    others: frozenset[str]
+
+
+def advanced(state: LearnerState, samples: int, elapsed: float) -> LearnerState:
+    """One training step of ``samples`` tokens that took ``elapsed`` seconds."""
+    return replace(
+        state,
+        local_step=state.local_step + 1,
+        steps=tuple(count + 1 for count in state.steps),
+        tokens=tuple(count + samples for count in state.tokens),
+        step_seconds=(
+            elapsed
+            if state.local_step == 0
+            else state.step_seconds + _EMA_DECAY * (elapsed - state.step_seconds)
+        ),
+    )
+
+
+def applied(
+    state: LearnerState,
+    round: int,
+    fragment: int,
+    version: int,
+    *,
+    participant: Participant | None = None,
+    mixing: float = 0.0,
+) -> LearnerState:
+    """``fragment`` now holds ``version`` committed by ``round``; its counters start over.
+
+    They keep ``mixing`` of the steps and tokens the commit did not incorporate, rounded down:
+    those since this learner's snapshot in the round (``participant``), or all of them when it
+    took no part. That share of the local progress survives in the blended parameters.
+    """
+
+    def with_fragment(values: tuple[int, ...], value: int) -> tuple[int, ...]:
+        return (*values[:fragment], value, *values[fragment + 1 :])
+
+    def kept(counted: int, incorporated: int) -> int:
+        return int(mixing * (counted - incorporated))
+
+    steps, tokens = (0, 0) if participant is None else (participant.steps, participant.tokens)
+    return replace(
+        state,
+        steps=with_fragment(state.steps, kept(state.steps[fragment], steps)),
+        tokens=with_fragment(state.tokens, kept(state.tokens[fragment], tokens)),
+        applied_version=with_fragment(state.applied_version, version),
+        rounds_seen=max(state.rounds_seen, round),
+    )
+
+
+def publishable(state: LearnerState, schedule: RoundSchedule) -> int | None:
+    """The round this learner publishes next, if it may publish it now.
+
+    The candidate is the round after both the last published and the last seen. It may be
+    published while it is within ``in_flight`` rounds of the last seen and its fragment has
+    been trained since it was last overwritten.
+    """
+    round = max(state.last_published + 1, state.rounds_seen + 1)
+    allowed = (
+        round <= state.rounds_seen + schedule.in_flight
+        and state.steps[schedule.fragment_of(round)] >= 1
+    )
+    return round if allowed else None
+
+
+def due_round(
+    state: LearnerState, schedule: RoundSchedule, *, quorum: int, publishers: Publishers
+) -> int | None:
+    """The round to publish now, if any.
+
+    The publishable round is due on schedule, once the gap of steps has passed since the last
+    publish, or by opportunity, once the others' snapshots in it are one short of ``quorum``:
+    the round is about to close, and the steps trained so far go in rather than being
+    overwritten by its commit. ``quorum`` is the effective one, the configured size or the
+    number of live learners, whichever is smaller. Below quorum two, one other snapshot is that
+    signal.
+    """
+    round = publishable(state, schedule)
+    if round is None:
+        return None
+    last = state.published[-1][1] if state.published else 0
+    on_schedule = state.local_step - last >= schedule.gap(round)
+    closing = publishers.round == round and len(publishers.others) >= max(quorum - 1, 1)
+    return round if on_schedule or closing else None
+
+
+def overdue(state: LearnerState, overlap_steps: int) -> bool:
+    """Whether a round this learner published, and is not yet aligned with, went out
+    ``overlap_steps`` or more steps ago: training may not run further ahead of it."""
+    return any(
+        round > state.rounds_seen and state.local_step - at >= overlap_steps
+        for round, at in state.published
+    )
+
+
+class Learner:
+    """Wraps a training loop: ``start()`` once, ``step()`` after every optimizer step, ``close()``.
+
+    ``start`` blocks until the model holds the newest committed version of every fragment,
+    then starts the communication thread. An operation that fails on that thread is reported
+    (``PollFailed``, ``PublishFailed``) and its loop goes on (``communicate``); only the thread's
+    own death is raised from the next ``step()``. When the schedule bounds ``overlap_steps``,
+    ``step()`` waits for the model to align with a round published that many steps ago before
+    training on, as standard DiLoCo does at zero; otherwise training goes on and the commit
+    overwrites the steps taken meanwhile. The lease is renewed every ``lease_every`` seconds and
+    the other learners' leases are presumed lost after ``lease_timeout`` seconds without a
+    change.
+
+    What goes out is ``anchor - θ``, computed in ``backend`` (the installed one by default)
+    against the version of the fragment this learner aligned with last, of which it keeps one host
+    copy. ``compression`` is the codec its payloads are stored with. ``on_event`` receives every
+    event the learner reports (``events`` module), on the thread that produced it.
+    """
+
+    def __init__(
+        self,
+        *,
+        adapter: Adapter,
+        algorithm: Algorithm,
+        mailbox: Mailbox,
+        learner_id: str,
+        backend: Backend | None = None,
+        clock: Clock | None = None,
+        poll_interval: float = 1.0,
+        lease_every: float = 60.0,
+        lease_timeout: float = 120.0,
+        compression: Compression = "off",
+        on_event: Callback = log_events,
+    ) -> None:
+        if lease_timeout <= lease_every:
+            raise ValueError(
+                f"lease_timeout must exceed lease_every, got {lease_timeout} and {lease_every}"
+            )
+        self._adapter = adapter
+        self._algorithm = algorithm
+        self._learner_id = learner_id
+        self._backend = backend if backend is not None else default_backend()
+        self._clock = clock if clock is not None else MonotonicClock()
+        self._poll_interval = poll_interval
+        self._lease_every = lease_every
+        self._lease_timeout = lease_timeout
+        self._exchange = Exchange(mailbox, clock=self._clock, compression=compression)
+        self._on_event = on_event
+        self._inbox: queue.SimpleQueue[tuple[_Update, ...]] = queue.SimpleQueue()
+        self._outbox: queue.SimpleQueue[_Outgoing | None] = queue.SimpleQueue()
+        self._run: _Run | None = None
+        self._anchors: list[Mapping[str, TensorBytes]] = []
+        self._newest: dict[int, tuple[int, Mapping[str, TensorBytes]]] = {}
+        self._state = LearnerState(0, (), (), (), -1, -1, (), 0.0)
+        self._cursor = -1
+        self._publishers = Publishers(-1, frozenset())
+        self._liveness = Liveness(lease_timeout)
+        self._live: frozenset[str] = frozenset()
+        self._estimates = Estimates()
+        self._last_step_at = 0.0
+        self._thread: threading.Thread | None = None
+        self._main: asyncio.Task[None] | None = None
+        self._failure: BaseException | None = None
+
+    @property
+    def rounds_seen(self) -> int:
+        return self._state.rounds_seen
+
+    @property
+    def batch_size(self) -> int | None:
+        """This learner's share of the global batch, split among the learners live as of the
+        last poll and rounded up; ``None`` when the algorithm sets no global batch."""
+        global_batch = self._algorithm.global_batch()
+        if global_batch is None:
+            return None
+        live = len(self._live | {self._learner_id})
+        return (global_batch + live - 1) // live
+
+    def start(self) -> None:
+        """Join the run on the calling thread, then start the communication thread."""
+
+        def communicate() -> None:
+            def failure(exc: BaseException) -> BaseException:
+                # TaskGroups wrap a lone failure in one group per nesting level; step()
+                # should raise the failure itself.
+                match exc:
+                    case BaseExceptionGroup(exceptions=(BaseException() as only,)):
+                        return failure(only)
+                    case _:
+                        return exc
+
+            try:
+                asyncio.run(self.communicate())
+            except BaseException as exc:
+                self._failure = failure(exc)
+
+        asyncio.run(self.join())
+        self._last_step_at = self._clock.now()
+        self._thread = threading.Thread(
+            target=communicate, name=f"gradientmail-{self._learner_id}", daemon=True
+        )
+        self._thread.start()
+
+    def step(self, samples: int) -> None:
+        """Account for one training step of ``samples`` tokens; apply and publish as due.
+
+        Queued fragment versions are validated against the schema and applied here, while the
+        model is quiescent, and so is the snapshot of a fragment whose round is due. Only bytes
+        are copied; the delta, packing, hashing and I/O happen on the communication thread. When
+        the schedule bounds ``overlap_steps``, it then waits until no round published that many
+        steps ago is still outstanding; the wait does not count as step time.
+        """
+        if self._failure is not None:
+            raise self._failure
+        if self._run is None:
+            raise RuntimeError("Learner.step() before start()")
+        run = self._run
+        now = self._clock.now()
+        state = advanced(self._state, samples, now - self._last_step_at)
+        self._last_step_at = now
+        state = self._drained(run, state)
+        round = due_round(
+            state,
+            run.schedule,
+            quorum=min(run.quorum, len(self._live)),
+            publishers=self._publishers,
+        )
+        if round is not None:
+            state = self._publish(run, state, round, final=False)
+        if (limit := run.schedule.overlap_steps) is not None:
+            state = self._settled(run, state, lambda s: not overdue(s, limit), patience=None)
+            self._last_step_at = self._clock.now()
+        self._state = state
+
+    def close(self) -> None:
+        """Wait for the commits of this learner's rounds in flight and align with them, then stop.
+
+        The steps trained since a fragment's last round are dropped: the model is set back to
+        the newest version of every fragment it aligned with. With the algorithm's
+        ``flush_on_close`` they are published instead, as final snapshots, round after round,
+        each awaited and aligned with, so the run's parameters end with them and the model holds
+        them. The wait gives up once nothing has landed for the lease timeout plus the takeover
+        timeout the last record predicts (at least thirty polls): time for crashed learners to
+        be presumed lost, then for a round without them to be presumed not to commit, as when
+        nobody is left to close one. Then the lease is released, the outbox is drained, bounded
+        by ``_CLOSE_TIMEOUT`` seconds, and the communication thread is joined. Closing again
+        does nothing.
+        """
+        if self._thread is None:
+            return
+        try:
+            if self._run is not None and self._failure is None:
+                self._state = self._concluded(self._run, self._state)
+                if not self._algorithm.flush_on_close:
+                    for anchor in self._anchors:
+                        self._adapter.apply(anchor)
+        finally:
+            self._outbox.put(None)
+            self._thread.join(_CLOSE_TIMEOUT)
+            if self._thread.is_alive() and self._main is not None:
+                self._on_event(PublishCancelled(self._learner_id, _CLOSE_TIMEOUT))
+                self._main.get_loop().call_soon_threadsafe(self._main.cancel)
+                self._thread.join()
+            self._thread = None
+
+    async def join(self) -> None:
+        """Create the run or adopt it, then load the newest committed version of every fragment.
+
+        Adoption requires the model's schema and the algorithm's configuration to equal the
+        manifest's (``SchemaMismatch`` naming the differing parameters, ``RunConfigMismatch``).
+        Counters start at zero; a restart takes this same path.
+        """
+        schema = self._adapter.schema()
+        algorithm = self._algorithm.to_json()
+        manifest = await self._exchange.manifest()
+        if manifest is None:
+            manifest = await self.create(schema)
+        if manifest.algorithm != algorithm:
+            raise RunConfigMismatch(expected=algorithm, found=manifest.algorithm)
+
+        fragments = range(len(manifest.plan))
+        aligned = await align(self._adapter, self._exchange, manifest)
+        records = aligned.records
+        newest = _newest_per_fragment(records)
+        self._anchors = list(aligned.fragments)
+        self._newest = {
+            p: (newest[p].version if p in newest else 0, tensors)
+            for p, tensors in enumerate(aligned.fragments)
+        }
+        self._run = _Run(
+            schema=schema,
+            plan=manifest.plan,
+            schedule=self._algorithm.schedule(len(manifest.plan)),
+            quorum=self._algorithm.quorum(),
+            mixing=self._algorithm.mixing,
+            commit_dtype=self._algorithm.commit_dtype,
+            digest=schema.digest(),
+            dtypes={spec.name: spec.dtype for spec in schema.parameters},
+            shapes={spec.name: spec.shape for spec in schema.parameters},
+        )
+        self._state = LearnerState(
+            local_step=0,
+            steps=(0,) * len(manifest.plan),
+            tokens=(0,) * len(manifest.plan),
+            applied_version=tuple(newest[p].version if p in newest else 0 for p in fragments),
+            rounds_seen=records[-1].round if records else -1,
+            last_published=-1,
+            published=(),
+            step_seconds=0.0,
+        )
+        self._cursor = self._state.rounds_seen
+        if records:
+            self._estimates = records[-1].estimates
+        self._on_event(
+            Joined(
+                self._learner_id, manifest.created_by, len(manifest.plan), self._state.rounds_seen
+            )
+        )
+
+    async def create(self, schema: Schema[DType]) -> Manifest:
+        """Create the run from this model, or wait for the node that is creating it.
+
+        The first claim on the manifest elects the creator, as the first claim on a round elects
+        its closer: one node uploads Θ⁰ and writes the manifest while the others poll for it,
+        instead of each uploading a Θ⁰ of its own only to lose the race. A waiter counts the
+        init payloads as they land; when none lands for as long as the largest fragment would
+        take at ``_MIN_UPLOAD_RATE``, and never less than ``_CREATION_FLOOR_POLLS`` polls, it
+        takes the next claim and creates the run itself. A creator that was slow rather than
+        dead adopts the manifest that won, as ``create_run`` does. A restarted creator finds
+        its own claim and resumes.
+        """
+        plan = self._algorithm.plan(schema)
+
+        async def upload() -> Manifest:
+            draft = ManifestDraft(
+                schema=schema,
+                plan=plan,
+                algorithm=self._algorithm.to_json(),
+                created_by=self._learner_id,
+            )
+            return await self._exchange.create_run(
+                draft, [self._adapter.snapshot(names) for names in plan]
+            )
+
+        largest = max(sum(schema.nbytes(name) for name in names) for names in plan)
+        stall = max(_CREATION_FLOOR_POLLS * self._poll_interval, largest / _MIN_UPLOAD_RATE)
+        claimants = await self._exchange.creation_claims()
+        if self._learner_id in claimants or (
+            not claimants and await self._exchange.claim_creation(1, self._learner_id)
+        ):
+            self._on_event(CreatingRun(self._learner_id, len(plan)))
+            return await upload()
+        claimants = await self._exchange.creation_claims()
+        attempts = len(claimants)
+        self._on_event(AwaitingCreator(self._learner_id, tuple(claimants)))
+        uploaded = await self._exchange.init_uploaded()
+        seen_at = self._clock.now()
+        while True:
+            await self._clock.sleep(self._poll_interval)
+            manifest = await self._exchange.manifest()
+            if manifest is not None:
+                return manifest
+            count = await self._exchange.init_uploaded()
+            if count != uploaded:
+                uploaded, seen_at = count, self._clock.now()
+            elif self._clock.now() - seen_at >= stall:
+                if await self._exchange.claim_creation(attempts + 1, self._learner_id):
+                    self._on_event(TookOverCreation(self._learner_id, stall))
+                    return await upload()
+                attempts = len(await self._exchange.creation_claims())
+                seen_at = self._clock.now()
+
+    async def communicate(self) -> None:
+        """Receive, publish and renew the lease until ``close()``, then release the lease; what
+        the communication thread runs.
+
+        Neither loop dies of a failed operation: the failure is reported and the loop goes on,
+        the receive loop with its next poll (so the lease keeps being renewed), the publish loop
+        with the next snapshot (the failed round's steps stay counted and go out with the
+        fragment's next round). Only cancellation ends them, and the lease is released either
+        way.
+        """
+
+        async def receive() -> None:
+            renew_at = self._clock.now()
+            while True:
+                try:
+                    if self._clock.now() >= renew_at:
+                        await self._exchange.beat(self._learner_id)
+                        renew_at = self._clock.now() + self._lease_every
+                        log.debug("%s: lease renewed", self._learner_id)
+                    await self.poll()
+                except Exception as exc:
+                    self._on_event(PollFailed(self._learner_id, exc))
+                await self._clock.sleep(self._poll_interval)
+
+        async def publish() -> None:
+            while (outgoing := await asyncio.to_thread(self._outbox.get)) is not None:
+                draft = outgoing.draft
+                try:
+                    delta = await asyncio.to_thread(self._delta, run, outgoing)
+                    started = time.perf_counter()
+                    published = await self._exchange.publish_snapshot(draft, delta)
+                    self._on_event(
+                        SnapshotPublished(
+                            self._learner_id,
+                            draft.round,
+                            run.schedule.fragment_of(draft.round),
+                            draft.steps,
+                            draft.tokens,
+                            published.payload.size,
+                            time.perf_counter() - started,
+                        )
+                    )
+                except Exception as exc:
+                    self._on_event(PublishFailed(self._learner_id, draft.round, exc))
+
+        if self._run is None:
+            raise RuntimeError("Learner.communicate() before join()")
+        run = self._run
+        self._main = asyncio.current_task()
+        try:
+            async with asyncio.TaskGroup() as group:
+                receiving = group.create_task(receive())
+                await publish()
+                receiving.cancel()
+        finally:
+            # A cancelled publish leaves a worker thread blocked in outbox.get(), and asyncio.run
+            # waits for the executor's threads on shutdown; the sentinel releases it.
+            self._outbox.put(None)
+            try:
+                await self._exchange.release(self._learner_id)
+                self._on_event(LeaseReleased(self._learner_id))
+            except Exception as exc:
+                self._on_event(ReleaseFailed(self._learner_id, exc))
+
+    async def poll(self) -> None:
+        """One receive iteration: queue the newest version of each fragment committed meanwhile,
+        then list the leases and who has published the round this learner publishes next.
+
+        The batch is queued only once every payload in it has been fetched. A payload already
+        swept is skipped; the next poll finds that fragment's newer version.
+        """
+        if self._run is None:
+            raise RuntimeError("Learner.poll() before join()")
+        run = self._run
+        records = await self._exchange.rounds_after(self._cursor)
+        if records:
+            newest = tuple(_newest_per_fragment(records).values())
+            mine = {
+                record.fragment: participant
+                for record in records
+                for participant in record.participants
+                if participant.learner == self._learner_id
+            }
+            async with asyncio.TaskGroup() as group:
+                fetches = [group.create_task(self._fetched(run, record)) for record in newest]
+            updates = tuple(
+                _Update(
+                    record.round,
+                    record.fragment,
+                    record.version,
+                    tensors,
+                    mine.get(record.fragment),
+                )
+                for record, fetch in zip(newest, fetches, strict=True)
+                if (tensors := fetch.result()) is not None
+            )
+            if updates:
+                self._inbox.put(updates)
+            self._cursor = records[-1].round
+            self._estimates = records[-1].estimates
+            for record in records:
+                self._on_event(RoundRecorded(self._learner_id, record))
+        now = self._clock.now()
+        self._liveness = self._liveness.observed(await self._exchange.leases(), now)
+        self._live = self._liveness.live(now)
+        state = self._state
+        round = max(state.last_published + 1, state.rounds_seen + 1)
+        others = await self._exchange.publishers(round)
+        self._publishers = Publishers(round, others - {self._learner_id})
+
+    def _publish(
+        self, run: "_Run", state: LearnerState, round: int, *, final: bool
+    ) -> LearnerState:
+        """Snapshot ``round``'s fragment and its anchor for the publish loop; the bytes are
+        copied here. ``final`` when closing."""
+        fragment = run.schedule.fragment_of(round)
+        draft = SnapshotDraft(
+            learner=self._learner_id,
+            round=round,
+            local_step=state.local_step,
+            steps=state.steps[fragment],
+            tokens=state.tokens[fragment],
+            base_version=state.applied_version[fragment],
+            rounds_seen=state.rounds_seen,
+            step_seconds=state.step_seconds,
+            schema_digest=run.digest,
+            final=final,
+        )
+        self._outbox.put(
+            _Outgoing(draft, self._anchors[fragment], self._adapter.snapshot(run.plan[fragment]))
+        )
+        log.debug(
+            "%s: round %d queued: fragment %d, %d steps at local step %d",
+            self._learner_id,
+            round,
+            fragment,
+            draft.steps,
+            state.local_step,
+        )
+        published = (*state.published, (round, state.local_step))[-run.schedule.in_flight :]
+        return replace(state, last_published=round, published=published)
+
+    async def _fetched(self, run: "_Run", record: RoundRecord) -> Mapping[str, TensorBytes] | None:
+        """The committed bytes of ``record``'s fragment: rebuilt from the commit's update when
+        this learner holds the version before it and the result hashes to the record's theta,
+        otherwise fetched; None when swept."""
+        held = self._newest.get(record.fragment)
+        if record.update is not None and held is not None and held[0] == record.version - 1:
+            update = await self._exchange.fetch(record.update.payload)
+            if update is not None:
+                rebuilt = await asyncio.to_thread(
+                    self._rebuilt, run, record.fragment, held[1], update
+                )
+                if theta_digest(rebuilt) == record.update.theta_digest:
+                    self._newest[record.fragment] = (record.version, rebuilt)
+                    return rebuilt
+                log.debug(
+                    "%s: round %d: the update does not rebuild theta; fetching it",
+                    self._learner_id,
+                    record.round,
+                )
+        tensors = await self._exchange.fetch(record.theta)
+        if tensors is not None:
+            self._newest[record.fragment] = (record.version, tensors)
+        return tensors
+
+    def _rebuilt(
+        self,
+        run: "_Run",
+        fragment: int,
+        base: Mapping[str, TensorBytes],
+        update: Mapping[str, TensorBytes],
+    ) -> Mapping[str, TensorBytes]:
+        """``base`` plus the commit's update, rounded to the schema's dtypes."""
+        load = self._backend.load
+        return {
+            name: (
+                load(base[name])
+                + decoded(self._backend, update, name, run.shapes[name], run.commit_dtype)
+            ).to_bytes(run.dtypes[name])
+            for name in run.plan[fragment]
+        }
+
+    def _delta(self, run: "_Run", outgoing: "_Outgoing") -> Mapping[str, TensorBytes]:
+        """``anchor - θ`` of every tensor in the snapshot, at the algorithm's delta precision."""
+        load = self._backend.load
+        return {
+            key: tensor
+            for name, local in outgoing.snapshot.items()
+            for key, tensor in encoded(
+                name,
+                load(outgoing.anchor[name]) - load(local),
+                run.dtypes[name],
+                self._algorithm.delta_dtype,
+            ).items()
+        }
+
+    def _align(self, run: "_Run", state: LearnerState, update: "_Update") -> LearnerState:
+        """Set the fragment to the committed version, blended with the local parameters when
+        the algorithm mixes; the committed bytes become the anchor either way."""
+        names = run.plan[update.fragment]
+        validate(run.schema, update.tensors, names)
+        self._adapter.apply(
+            self._blended(run, names, update.tensors) if run.mixing else update.tensors
+        )
+        self._anchors[update.fragment] = update.tensors
+        mine = update.participant
+        self._on_event(
+            FragmentAligned(
+                self._learner_id,
+                update.round,
+                update.fragment,
+                update.version,
+                state.local_step,
+                None if mine is None else state.local_step - mine.local_step,
+            )
+        )
+        return applied(
+            state,
+            update.round,
+            update.fragment,
+            update.version,
+            participant=mine,
+            mixing=run.mixing,
+        )
+
+    def _blended(
+        self, run: "_Run", names: Sequence[str], theta: Mapping[str, TensorBytes]
+    ) -> Mapping[str, TensorBytes]:
+        """``mixing·θ + (1 - mixing)·Θ`` over the fragment, rounded to the schema's dtypes."""
+        load = self._backend.load
+        local = self._adapter.snapshot(names)
+        return {
+            name: (load(local[name]) * run.mixing + load(theta[name]) * (1 - run.mixing)).to_bytes(
+                run.dtypes[name]
+            )
+            for name in names
+        }
+
+    def _drained(self, run: "_Run", state: LearnerState) -> LearnerState:
+        """Align with every fragment version the receive loop has queued, without waiting."""
+
+        def queued() -> Iterator[_Update]:
+            while True:
+                try:
+                    batch = self._inbox.get_nowait()
+                except queue.Empty:
+                    return
+                yield from batch
+
+        for update in queued():
+            state = self._align(run, state, update)
+        return state
+
+    def _settled(
+        self,
+        run: "_Run",
+        state: LearnerState,
+        done: Callable[[LearnerState], bool],
+        *,
+        patience: float | None,
+    ) -> LearnerState:
+        """Align with fragment versions as they arrive until ``done(state)`` holds.
+
+        Blocks the calling thread. With ``patience``, returns once nothing has happened for that
+        long: no version aligned with and no new snapshot in the round this learner publishes
+        next. Raises what the communication thread died of.
+        """
+        progress = (state.rounds_seen, self._publishers)
+        since = self._clock.now()
+        while not done(state):
+            try:
+                batch = self._inbox.get(timeout=self._poll_interval)
+            except queue.Empty:
+                batch = ()
+            if self._failure is not None:
+                raise self._failure
+            for update in batch:
+                state = self._align(run, state, update)
+            now = self._clock.now()
+            if (state.rounds_seen, self._publishers) != progress:
+                progress, since = (state.rounds_seen, self._publishers), now
+            elif patience is not None and now - since >= patience:
+                break
+        return state
+
+    def _concluded(self, run: "_Run", state: LearnerState) -> LearnerState:
+        """Wait for the commits of the rounds in flight; with ``flush_on_close``, first publish
+        every fragment trained since its last overwrite as a final snapshot.
+
+        Flushed rounds go out as the schedule allows, ``in_flight`` at a time, each next one
+        once the commits make it publishable; done when nothing is publishable and nothing is
+        outstanding, or when the run stalls for the lease timeout plus the takeover timeout.
+        """
+
+        def flushable(state: LearnerState) -> int | None:
+            return publishable(state, run.schedule) if self._algorithm.flush_on_close else None
+
+        def settled(state: LearnerState) -> bool:
+            return flushable(state) is not None or state.rounds_seen >= state.last_published
+
+        self._on_event(Closing(self._learner_id, self._algorithm.flush_on_close))
+        patience = self._lease_timeout + self._estimates.timeout(self._poll_interval)
+        state = self._drained(run, state)
+        while True:
+            while (round := flushable(state)) is not None:
+                state = self._publish(run, state, round, final=True)
+            if state.rounds_seen >= state.last_published:
+                self._on_event(Closed(self._learner_id, state.last_published))
+                return state
+            state = self._settled(run, state, settled, patience=patience)
+            if not settled(state):
+                self._on_event(
+                    GaveUp(self._learner_id, state.last_published, patience, state.rounds_seen)
+                )
+                return state
+
+
+@dataclass(frozen=True)
+class Aligned:
+    """What ``align`` loaded: the round records read, in round order, and the tensors the
+    adapter now holds, one mapping per fragment."""
+
+    records: Sequence[RoundRecord]
+    fragments: tuple[Mapping[str, TensorBytes], ...]
+
+
+async def align(adapter: Adapter, exchange: Exchange, manifest: Manifest) -> Aligned:
+    """Load the newest committed version of every fragment of ``manifest``'s run into ``adapter``.
+
+    A fragment no round has committed yet is loaded from the manifest's initial parameters.
+    ``SchemaMismatch`` names the parameters on which the model and the run differ.
+    """
+    schema = adapter.schema()
+    if manifest.schema != schema:
+        differing = set(manifest.schema.parameters) ^ set(schema.parameters)
+        raise SchemaMismatch(
+            sorted({spec.name for spec in differing}), "model does not match the run's schema"
+        )
+    fragments = range(len(manifest.plan))
+    records = await exchange.rounds_after(-1)
+    newest = _newest_per_fragment(records)
+    refs = [newest[p].theta if p in newest else manifest.init[p] for p in fragments]
+    loaded: list[Mapping[str, TensorBytes]] = []
+    for fragment, (ref, tensors) in enumerate(zip(refs, await _fetch(exchange, refs), strict=True)):
+        if tensors is None:
+            raise ProtocolError(f"{ref.key} was swept while loading; start again")
+        validate(schema, tensors, manifest.plan[fragment])
+        adapter.apply(tensors)
+        loaded.append(tensors)
+    return Aligned(records, tuple(loaded))
+
+
+@dataclass(frozen=True)
+class _Run:
+    """The run as joined: fixed for the learner's lifetime."""
+
+    schema: Schema[DType]
+    plan: tuple[tuple[str, ...], ...]
+    schedule: RoundSchedule
+    quorum: int
+    mixing: float
+    commit_dtype: Precision
+    digest: str
+    dtypes: Mapping[str, DType]
+    shapes: Mapping[str, tuple[int, ...]]
+
+
+@dataclass(frozen=True)
+class _Update:
+    """A fragment version fetched by the receive loop, waiting for ``step()`` to apply it;
+    ``participant`` is this learner's snapshot in the fragment's rounds of the same listing,
+    when it took part in one."""
+
+    round: int
+    fragment: int
+    version: int
+    tensors: Mapping[str, TensorBytes]
+    participant: Participant | None
+
+
+@dataclass(frozen=True)
+class _Outgoing:
+    """A snapshot taken by ``step()`` with the anchor it was trained from, waiting for the
+    publish loop to subtract them."""
+
+    draft: SnapshotDraft
+    anchor: Mapping[str, TensorBytes]
+    snapshot: Mapping[str, TensorBytes]
+
+
+def _newest_per_fragment(records: Sequence[RoundRecord]) -> Mapping[int, RoundRecord]:
+    # Records come in round order, so the last one seen per fragment is its newest version.
+    return {record.fragment: record for record in records}
+
+
+async def _fetch(
+    exchange: Exchange, refs: Sequence[PayloadRef]
+) -> Sequence[Mapping[str, TensorBytes] | None]:
+    async with asyncio.TaskGroup() as group:
+        fetches = [group.create_task(exchange.fetch(ref)) for ref in refs]
+    return [fetch.result() for fetch in fetches]
+
+
+_EMA_DECAY = 0.1
+_CLOSE_TIMEOUT = 60.0
+_CREATION_FLOOR_POLLS = 30
+# A live creator is assumed to upload at least this fast; slower, and a waiter takes over.
+_MIN_UPLOAD_RATE = 2**20