ZKDV core¶

Shared ZKDV configuration and lazy backend frontends.

FloatTolerance ¶

SamplingPolicy ¶

Configure transcript-derived update sampling.

state_check_probability is kept for backwards compatibility and should not be configured. Its only valid value is the default sentinel, -1.0.

ZKDVConfig ¶

ZKDVError ¶

Bases: RuntimeError

Base error for the public ZKDV API.

ZKDVPoisonedError ¶

Bases: ZKDVError

The driver cannot proceed after an asynchronous failure.

ZKDVVerificationError ¶

Bases: ZKDVError

A sampled update failed verification.

compilation ¶

Backend-neutral mechanics for compiled training functions.

CompiledFunction ¶

CompiledFunction(driver: Any, function: Callable[..., Any], trees: PyTreeAdapter, api: str)

Submit one compiled backend update through the shared pipeline ABI.

Source code in zkdv/compilation.py

def __init__(
    self,
    driver: Any,
    function: Callable[..., Any],
    trees: PyTreeAdapter,
    api: str,
) -> None:
    functools.update_wrapper(self, function)
    self._driver = driver
    self._function = function
    self._trees = trees
    self._signature = FunctionSignature(function, api)
    self._schema = None
    self._result_schema: ResultSchema | None = None
    self._compiled = None

PyTreeAdapter `dataclass` ¶

PyTreeAdapter(flatten: Callable[[Any], tuple[list[Any], Any]], leaves: Callable[[Any], list[Any]], unflatten: Callable[[Any, list[Any]], Any])

The three tree operations used by backend-neutral result handling.

ResultSchema `dataclass` ¶

ResultSchema(tree: Any, hidden_tree: Any, sources: tuple[int, ...], hidden_leaves: int, delta_leaves: int, private_delta: bool, extra_leaves: int)

Reconstruct user results without returning transaction values twice.

control ¶

Small host decisions, independent of backend accelerator collective ordering.

HostControl ¶

HostControl(rank: int, size: int, *, address: tuple[str, int] | None = None, token: str | None = None, timeout: float = 600)

Broadcast bounded JSON decisions from the sole tape owner.

Each calling thread has its own connection: a training admission must not block delivery of a check decision being awaited by a completion thread. Named rounds may arrive in different orders on these independent threads. No model arrays or replay snapshots belong on this channel.

Source code in zkdv/control.py

def __init__(
    self, rank: int, size: int, *, address: tuple[str, int] | None = None,
    token: str | None = None, timeout: float = 600,
) -> None:
    if size < 2 or not 0 <= rank < size or timeout <= 0:
        raise ValueError("invalid ZKDV host-control configuration")
    self.rank, self.size, self.timeout = rank, size, timeout
    self.token = secrets.token_hex(32) if rank == 0 else token
    self._ready = Condition()
    self._rounds: dict[str, tuple[bytes, set[int]]] = {}
    self._closed = Event()
    self._connections: list[socket.socket] = []
    self._connection_lock = Lock()
    self._local = local()
    self._listener = None
    self._acceptor = None
    if rank == 0:
        self._listener = socket.socket()
        self._listener.bind(address or ("0.0.0.0", 0))
        self._listener.listen()
        self._listener.settimeout(0.25)
        self.address = self._listener.getsockname()
        self._acceptor = Thread(target=self._accept, daemon=True, name="zkdv-control")
        self._acceptor.start()
    else:
        if address is None or token is None:
            raise ValueError("a follower requires the owner's address and token")
        self.address = address

close ¶

close() -> None

Wake blocked decisions and release every transport connection.

Source code in zkdv/control.py

def close(self) -> None:
    """Wake blocked decisions and release every transport connection."""
    self._closed.set()
    with self._ready:
        self._ready.notify_all()
    if self._listener is not None:
        self._listener.close()
    with self._connection_lock:
        for connection in self._connections:
            try:
                connection.shutdown(socket.SHUT_RDWR)
            except OSError:
                pass
            connection.close()
        self._connections.clear()
    if self._acceptor is not None:
        self._acceptor.join(timeout=1)

exchange ¶

exchange(key: str, value: Any = None) -> Any

Deliver one owner decision to every follower, with bounded lifetime.

Source code in zkdv/control.py

def exchange(self, key: str, value: Any = None) -> Any:
    """Deliver one owner decision to every follower, with bounded lifetime."""
    if self._closed.is_set():
        raise RuntimeError("ZKDV host control is closed")
    if self.rank == 0:
        encoded = json.dumps(value, allow_nan=False).encode() + b"\n"
        if len(encoded) > self.limit:
            raise ValueError("ZKDV host decisions must be small")
        with self._ready:
            if key in self._rounds:
                raise RuntimeError(f"duplicate ZKDV control round: {key}")
            readers: set[int] = set()
            self._rounds[key] = encoded, readers
            self._ready.notify_all()
            completed = self._ready.wait_for(
                lambda: len(readers) == self.size - 1 or self._closed.is_set(),
                timeout=self.timeout,
            )
            self._rounds.pop(key, None)
            if not completed or self._closed.is_set():
                raise TimeoutError(f"ZKDV followers did not receive {key}")
        return value
    if not hasattr(self._local, "stream"):
        connection = socket.create_connection(self.address, timeout=self.timeout)
        with self._connection_lock:
            self._connections.append(connection)
        self._local.stream = connection.makefile("rwb")
    stream = self._local.stream
    request = json.dumps([self.rank, self.token, key]).encode() + b"\n"
    if len(request) > self.limit:
        raise ValueError("ZKDV control key is too large")
    stream.write(request)
    stream.flush()
    response = stream.readline(self.limit + 1)
    if not response.endswith(b"\n") or len(response) > self.limit:
        raise RuntimeError(f"ZKDV owner disconnected during {key}")
    return json.loads(response)

for_backend `classmethod` ¶

for_backend(backend: str) -> HostControl | None

Bootstrap once through the user's initialized process group.

Source code in zkdv/control.py

@classmethod
def for_backend(cls, backend: str) -> "HostControl | None":
    """Bootstrap once through the user's initialized process group."""
    if backend == "jax":
        import jax
        from jax.experimental import multihost_utils
        import numpy as np

        rank, size = jax.process_index(), jax.process_count()
    elif backend == "torch":
        import torch.distributed as distributed

        if not distributed.is_initialized():
            return None
        rank, size = distributed.get_rank(), distributed.get_world_size()
    else:
        raise ValueError(f"unsupported ZKDV control backend: {backend}")
    if size == 1:
        return None
    owner = cls(rank, size) if rank == 0 else None
    descriptor = (
        [socket.getfqdn(), owner.address[1], owner.token]
        if owner is not None else None
    )
    try:
        if backend == "jax":
            encoded = json.dumps(descriptor).encode()
            if len(encoded) >= 1024:
                raise ValueError("ZKDV control address is too long")
            payload = np.zeros(1024, dtype=np.uint8)
            payload[:len(encoded)] = np.frombuffer(encoded, dtype=np.uint8)
            payload = multihost_utils.broadcast_one_to_all(payload, is_source=rank == 0)
            descriptor = json.loads(payload.tobytes().rstrip(b"\0"))
        else:
            payload = [descriptor]
            distributed.broadcast_object_list(payload, src=0)
            descriptor = payload[0]
        if owner is not None:
            return owner
        hostname, port, token = descriptor
        return cls(rank, size, address=(hostname, port), token=token)
    except BaseException:
        if owner is not None:
            owner.close()
        raise

driver ¶

Shared lifecycle for backend-specific ZKDV drivers.

Driver ¶

Driver(path: str | Path, config: ZKDVConfig | None, *, backend: str, annotation: Any, max_in_flight: int, replay_snapshot_interval: int)

Own one backend transcript and its asynchronous pipeline lifecycle.

Source code in zkdv/driver/__init__.py

def __init__(
    self,
    path: str | Path,
    config: ZKDVConfig | None,
    *,
    backend: str,
    annotation: Any,
    max_in_flight: int,
    replay_snapshot_interval: int,
) -> None:
    if replay_snapshot_interval < 1:
        raise ValueError("ZKDV replay_snapshot_interval must be positive")
    self.path = Path(path)
    self.config = config if config is not None else ZKDVConfig()
    self._backend = backend
    self._core: ZKDVCore | None = None
    self._program = None
    self._pipeline = None
    self._queue = PipelineQueue(max_in_flight, annotation)
    self._replay_snapshot_interval = replay_snapshot_interval
    self._next_index = 0
    self._poison: BaseException | None = None
    self._closed = False

errors ¶

Semantic errors raised by the public ZKDV API.

ZKDVError ¶

Bases: RuntimeError

Base error for the public ZKDV API.

ZKDVPoisonedError ¶

Bases: ZKDVError

The driver cannot proceed after an asynchronous failure.

ZKDVVerificationError ¶

Bases: ZKDVError

A sampled update failed verification.

flowcontrol ¶

Bounded completion queue for asynchronous pipeline transactions.

PipelineQueue ¶

PipelineQueue(capacity: int, annotation: Any = unannotated)

Bound unresolved transactions and complete their native submissions.

Source code in zkdv/flowcontrol.py

def __init__(self, capacity: int, annotation: Any = unannotated) -> None:
    if capacity < 1:
        raise ValueError("ZKDV max_in_flight must be positive")
    self._capacity = capacity
    self._annotation = annotation
    self._retirements = deque()
    self._fences = []
    self._flow_lock = Lock()
    self._start_executor = ThreadPoolExecutor(
        max_workers=1, thread_name_prefix="zkdv-start"
    )
    self._prepare_executor = ThreadPoolExecutor(
        max_workers=capacity, thread_name_prefix="zkdv-prepare"
    )
    self._commit_executor = ThreadPoolExecutor(
        max_workers=1, thread_name_prefix="zkdv-commit"
    )
    self._check_executor = ThreadPoolExecutor(
        max_workers=1, thread_name_prefix="zkdv-check"
    )
    self._snapshot_executor = ThreadPoolExecutor(
        max_workers=1, thread_name_prefix="zkdv-snapshot"
    )
    self._start_tails: list[Future[None]] = []
    self._commit_tails: list[Future[bool]] = []
    self._check_tails: list[Future[bool]] = []
    self._snapshot_tail: Future[None] | None = None
    self._check_lock = Lock()
    self._overlapped_check: Future[bool] | None = None

acquire ¶

acquire(policy: Policy) -> None

Admit one transaction under every active transaction's policy.

Source code in zkdv/flowcontrol.py

def acquire(self, policy: Policy) -> None:
    """Admit one transaction under every active transaction's policy."""

    self.check()
    while True:
        with self._flow_lock:
            self._discard_retired()
            active = tuple(self._retirements)
            fences = tuple(
                retirement for retirement in self._fences if not retirement.is_set()
            )
            limits = [self._limit(policy)]
            limits.extend(self._limit(active_policy) for _, active_policy in active)
            admitted = not fences and len(active) < min(limits)
            wait_for = fences[0] if fences else (active[0][0] if active else None)
        if admitted:
            return
        if wait_for is None:
            raise RuntimeError("ZKDV flow control could not select a retirement")
        with self.annotate("zkdv.flow_control.wait"):
            wait_for.wait()
        self.check()

close ¶

close() -> None

Release executors after the owning driver has joined the pipeline.

Source code in zkdv/flowcontrol.py

def close(self) -> None:
    """Release executors after the owning driver has joined the pipeline."""

    self._start_executor.shutdown(wait=True, cancel_futures=False)
    self._prepare_executor.shutdown(wait=True, cancel_futures=False)
    self._commit_executor.shutdown(wait=True, cancel_futures=False)
    self._check_executor.shutdown(wait=True, cancel_futures=False)
    self._snapshot_executor.shutdown(wait=True, cancel_futures=False)

overlap ¶

Per-transaction training and verification overlap policies.

Policy ¶

Bases: IntEnum

Control how a transaction shares device-buffer residency with its neighbors.

pipeline ¶

Shared state machine for one pending backend update.

PendingUpdate ¶

PendingUpdate(pipeline: Any, index: int, submission: Any, batch: Any, opt_state: Any, batch_evidence: Any, parameter_evidence: Any, pre_projection: Any, pre_optimizer_hash: Any, overlap: Policy)

One native submission awaiting a fused update's device evidence.

Source code in zkdv/pipeline.py

def __init__(
    self,
    pipeline: Any,
    index: int,
    submission: Any,
    batch: Any,
    opt_state: Any,
    batch_evidence: Any,
    parameter_evidence: Any,
    pre_projection: Any,
    pre_optimizer_hash: Any,
    overlap: Policy,
) -> None:
    self._pipeline = pipeline
    self.index = index
    self._submission = submission
    self._retirement = submission.retirement()
    self._batch = batch
    self._opt_state = opt_state
    self._pre_optimizer_hash = pre_optimizer_hash
    self._post_optimizer_hash = Future()
    self._batch_evidence = batch_evidence
    self._parameter_evidence = parameter_evidence
    self._pre_projection = pre_projection
    self._post_projection = Future()
    self.overlap = overlap
    self._deltas = None
    self._params_after = None
    self._opt_state_after = None
    self._evidence = None
    self._complete = False
    self._predecessor = None
    self._successor = Future()
    self._successor_lock = Lock()
    self._successor_linked = False
    self._checked_successor_sealed = False

Pipeline ¶

Pipeline(core: Any, queue: Any, checkpoint: Any)

Bases: ABC

Backend-neutral lifecycle for the native transaction pipeline.

Source code in zkdv/pipeline.py

def __init__(self, core: Any, queue: Any, checkpoint: Any) -> None:
    self._core = core
    self._queue = queue
    self._checkpoint = checkpoint
    self._latest: PendingUpdate | None = None
    self._last_projection = None

session ¶

One native prover owned by the coordinator of a distributed session.

OwnerChallenge ¶

OwnerChallenge(control, index, native, descriptor)

Carry the owner's indices into every rank's checked program.

Source code in zkdv/session.py

def __init__(self, control, index, native, descriptor):
    self.control, self.index = control, index
    self._native = native
    self.param_idx = descriptor["param_idx"]
    self.state_idx = descriptor["state_idx"]
    self.tolerance_abs = descriptor["tolerance_abs"]
    self.tolerance_rel = descriptor["tolerance_rel"]

OwnerSession ¶

OwnerSession(path, config, control, *, backend)

Keep native tape state on rank zero and device programs on every rank.

Source code in zkdv/session.py

def __init__(self, path, config, control, *, backend):
    self.control = control
    self._backend = backend
    self._native = None
    result = None
    if control.rank == 0:
        try:
            self._native = ZKDVCore(path, config, backend=backend)
        except BaseException as error:
            control.exchange("create", str(error))
            raise
    result = control.exchange("create", result)
    if result is not None:
        raise RuntimeError(f"ZKDV owner creation failed: {result}")
    self._context = None
    self._distributed_shards = 1

OwnerSubmission ¶

OwnerSubmission(control, index, native)

Mirror one native slot's decision without allocating a follower prover.

Source code in zkdv/session.py

def __init__(self, control, index, native):
    self.control, self.index, self._native = control, index, native
    self._retired = Event()

signature ¶

Framework-independent callable signature normalization.

FunctionSignature ¶

FunctionSignature(function: Callable[..., Any], api='ZKDV.jit')

Normalize ordinary named arguments into a stable positional ABI.

Source code in zkdv/signature.py

def __init__(self, function: Callable[..., Any], api="ZKDV.jit") -> None:
    self.signature = inspect.signature(function)
    unsupported = {
        inspect.Parameter.VAR_POSITIONAL,
        inspect.Parameter.VAR_KEYWORD,
    }
    if any(
        value.kind in unsupported for value in self.signature.parameters.values()
    ):
        raise TypeError(f"{api} does not support variadic training functions")
    self.names = tuple(self.signature.parameters)

transaction ¶

Backend-neutral bindings for symbolic training transactions.

InputBinding `dataclass` ¶

InputBinding(tree: Any, leaves: tuple[int, ...])

Reconstruct one registered pytree from a training invocation.

bind_input ¶

bind_input(inputs: Any, value: Any, name: str, trees: PyTreeAdapter) -> InputBinding

Bind a registered subtree to leaf positions in a function ABI.

Source code in zkdv/transaction.py

def bind_input(
    inputs: Any,
    value: Any,
    name: str,
    trees: PyTreeAdapter,
) -> InputBinding:
    """Bind a registered subtree to leaf positions in a function ABI."""

    input_leaves = trees.leaves(inputs)
    positions = {id(leaf): index for index, leaf in enumerate(input_leaves)}
    leaves, tree = trees.flatten(value)
    try:
        indices = tuple(positions[id(leaf)] for leaf in leaves)
    except KeyError as error:
        raise ValueError(
            f"ZKDV transaction {name} must be a subtree of the function inputs"
        ) from error
    return InputBinding(tree, indices)

ZKDV core¶

FloatTolerance ¶

SamplingPolicy ¶

ZKDVConfig ¶

ZKDVError ¶

ZKDVPoisonedError ¶

ZKDVVerificationError ¶

compilation ¶

CompiledFunction ¶

PyTreeAdapter dataclass ¶

ResultSchema dataclass ¶

control ¶

HostControl ¶

close ¶

exchange ¶

for_backend classmethod ¶

driver ¶

Driver ¶

errors ¶

ZKDVError ¶

ZKDVPoisonedError ¶

ZKDVVerificationError ¶

flowcontrol ¶

PipelineQueue ¶

acquire ¶

close ¶

overlap ¶

Policy ¶

pipeline ¶

PendingUpdate ¶

Pipeline ¶

session ¶

OwnerChallenge ¶

OwnerSession ¶

OwnerSubmission ¶

signature ¶

FunctionSignature ¶

transaction ¶

InputBinding dataclass ¶

bind_input ¶

PyTreeAdapter `dataclass` ¶

ResultSchema `dataclass` ¶

for_backend `classmethod` ¶

InputBinding `dataclass` ¶