Skip to content

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 
 89
 90
 91
 92
 93
 94
 95
 96
 97
 98
 99
100
101
102
103
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.

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 
15
16
17
18
19
20
21
22
23
24
25
26
27
28
29
30
31
32
33
34
35
36
37
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 
28
29
30
31
32
33
34
35
36
37
38
39
40
41
42
43
44
45
46
47
48
49
50
51
52
53
54
55
56
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 
65
66
67
68
69
70
71
72
73
74
75
76
77
78
79
80
81
82
83
84
85
86
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 
210
211
212
213
214
215
216
217
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 
14
15
16
17
18
19
20
21
22
23
24
25
26
27
28
29
30
31
32
33
34
35
36
37
38
39
40
41
42
43
44
45
46
47
48
49
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 
204
205
206
207
208
209
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

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 
11
12
13
14
15
16
17
18
19
20
21
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 
23
24
25
26
27
28
29
30
31
32
33
34
35
36
37
38
39
40
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)