fedmesh 0.1.0__py3-none-any.whl

fedmesh/__init__.py

@@ -0,0 +1,90 @@
+"""fedmesh: privacy-preserving federated training across organisations.
+
+fedmesh coordinates training runs that span several organisations, each of
+which keeps its data local. It is built on iroh-tinker's public surface:
+``LocalTransport``/``IrohTransport`` for the wire, ``CompressedTransport``
+for bandwidth, and ``JobManager``/``LaunchSpec`` for the launch.
+
+Two things make it different from a plain multi-node job:
+
+* **No raw gradient leaves a party.** Contributions are masked with pairwise
+  secrets before they cross the wire, and the masks cancel in the aggregate.
+* **No single party holds a global secret.** Every pair of parties derives
+  its own secret, so compromising the coordinator does not open the other
+  pairs' channels.
+
+The guarantees, and their limits, are written down in ``docs/threat-model.md``.
+"""
+from __future__ import annotations
+
+__version__ = "0.1.0"
+
+from .dp import (
+    DPConfig,
+    PrivacyAccountant,
+    PrivacyBudgetError,
+    add_gaussian_noise,
+    clip_by_l2_norm,
+    epsilon_for_gaussian_mechanism,
+)
+from .masking import (
+    InsecureTransportError,
+    MaskedTransport,
+    MaskRegistry,
+    broadcast_within_party,
+    exchange_party_sums,
+    is_lossy_transport,
+)
+from .mesh import Mesh, MeshError, PartyManifest
+from .party import Party, PartyError, pairwise_secrets
+from .sampler import Shard, ShardAudit, ShardPolicy, audit_shards
+from .trust import (
+    PairwiseSecrets,
+    PartyKeyPair,
+    TrustError,
+    derive_pairwise_secret,
+    hkdf_expand,
+    hkdf_extract,
+    sign_ticket,
+    verify_ticket,
+)
+
+__all__ = [
+    "__version__",
+    # trust
+    "PartyKeyPair",
+    "PairwiseSecrets",
+    "TrustError",
+    "derive_pairwise_secret",
+    "hkdf_expand",
+    "hkdf_extract",
+    "sign_ticket",
+    "verify_ticket",
+    # mesh
+    "Mesh",
+    "MeshError",
+    "PartyManifest",
+    # party
+    "Party",
+    "PartyError",
+    "pairwise_secrets",
+    # masking
+    "MaskedTransport",
+    "MaskRegistry",
+    "InsecureTransportError",
+    "broadcast_within_party",
+    "exchange_party_sums",
+    "is_lossy_transport",
+    # differential privacy
+    "DPConfig",
+    "PrivacyAccountant",
+    "PrivacyBudgetError",
+    "add_gaussian_noise",
+    "clip_by_l2_norm",
+    "epsilon_for_gaussian_mechanism",
+    # sampling
+    "Shard",
+    "ShardAudit",
+    "ShardPolicy",
+    "audit_shards",
+]

fedmesh/cli.py

@@ -0,0 +1,377 @@
+"""``fedmesh`` command line interface.
+
+Subcommands:
+
+``init``      create a party's key pair and namespace
+``manifest``  print this party's manifest
+``mesh``      combine manifests into a mesh and elect the coordinator
+``shards``    show which party owns which dataset rows
+``dp``        project a differential privacy budget
+``spike``     run the compression/masking spike and print the numbers
+
+Every command supports ``--json``.
+"""
+from __future__ import annotations
+
+import argparse
+import json
+import stat
+import sys
+from collections.abc import Sequence
+from pathlib import Path
+from typing import Any
+
+from .dp import DPConfig, PrivacyAccountant
+from .masking import MaskRegistry
+from .mesh import Mesh
+from .party import Party
+from .sampler import ShardPolicy
+from .trust import PartyKeyPair
+
+__all__ = ["build_parser", "main"]
+
+KEY_FILE = "party_key.json"
+
+
+def _write_private(
+    path: Path,
+    party_id: str,
+    keypair: PartyKeyPair,
+    ranks: int,
+    rank_offset: int = 0,
+) -> None:
+    """Write a party's private key with owner-only permissions."""
+    path.parent.mkdir(parents=True, exist_ok=True)
+    payload = {
+        "party_id": party_id,
+        "ranks": int(ranks),
+        "rank_offset": int(rank_offset),
+        "private_exponent": str(keypair.private_exponent),
+        "public_key": keypair.public_bytes().hex(),
+    }
+    path.write_text(json.dumps(payload, indent=2) + "\n")
+    path.chmod(stat.S_IRUSR | stat.S_IWUSR)
+
+
+def _read_key(path: Path) -> PartyKeyPair:
+    raw = json.loads(path.read_text())
+    return PartyKeyPair.from_private(
+        str(raw["party_id"]), int(str(raw["private_exponent"]))
+    )
+
+
+def _load_party(directory: Path, party_id: str | None = None) -> Party:
+    key_path = directory / KEY_FILE
+    if not key_path.exists():
+        raise FileNotFoundError(f"no party key at {key_path}; run 'fedmesh init' first")
+    keypair = _read_key(key_path)
+    raw_ranks = json.loads(key_path.read_text()).get("ranks", 1)
+    if party_id is not None and party_id != keypair.party_id:
+        raise ValueError(
+            f"key at {key_path} belongs to {keypair.party_id!r}, not {party_id!r}"
+        )
+    offset = int(json.loads(key_path.read_text()).get("rank_offset", 0))
+    ranks = tuple(range(offset, offset + int(raw_ranks)))
+    return Party(
+        party_id=keypair.party_id,
+        keypair=keypair,
+        ranks=ranks,
+        rendezvous_uri=str(directory),
+    )
+
+
+def build_parser() -> argparse.ArgumentParser:
+    parser = argparse.ArgumentParser(
+        prog="fedmesh",
+        description="Privacy-preserving federated training across organisations.",
+    )
+    parser.add_argument("--json", action="store_true", help="emit JSON on stdout")
+    sub = parser.add_subparsers(dest="command", required=True)
+
+    init = sub.add_parser("init", help="create a party key pair")
+    init.add_argument("--party-id", required=True, help="the party's identifier")
+    init.add_argument("--dir", required=True, help="party directory")
+    init.add_argument("--ranks", type=int, default=1, help="number of ranks")
+    init.add_argument(
+        "--rank-offset",
+        type=int,
+        default=0,
+        help="the party's first global rank in the mesh",
+    )
+    init.add_argument("--json", action="store_true", help="emit JSON on stdout")
+
+    manifest = sub.add_parser("manifest", help="print this party's manifest")
+    manifest.add_argument("--dir", required=True, help="party directory")
+    manifest.add_argument("--json", action="store_true", help="emit JSON on stdout")
+
+    mesh = sub.add_parser("mesh", help="combine manifests into a mesh")
+    mesh.add_argument("manifests", nargs="+", help="manifest JSON files")
+    mesh.add_argument("--json", action="store_true", help="emit JSON on stdout")
+
+    shards = sub.add_parser("shards", help="show the shard assignment")
+    shards.add_argument("--parties", required=True, help="comma separated party ids")
+    shards.add_argument("--rows", type=int, default=1000, help="number of dataset rows")
+    shards.add_argument("--json", action="store_true", help="emit JSON on stdout")
+
+    dp = sub.add_parser("dp", help="project a privacy budget")
+    dp.add_argument("--epsilon", type=float, default=8.0, help="target epsilon")
+    dp.add_argument("--delta", type=float, default=1e-5, help="target delta")
+    dp.add_argument("--noise", type=float, default=10.0, help="noise multiplier")
+    dp.add_argument("--clip", type=float, default=1.0, help="clipping norm")
+    dp.add_argument("--steps", type=int, default=10, help="number of steps to project")
+    dp.add_argument("--json", action="store_true", help="emit JSON on stdout")
+
+    spike = sub.add_parser("spike", help="run the compression/masking spike")
+    spike.add_argument("--json", action="store_true", help="emit JSON on stdout")
+
+    return parser
+
+
+def _cmd_init(args: argparse.Namespace) -> int:
+    directory = Path(args.dir)
+    key_path = directory / KEY_FILE
+    if key_path.exists():
+        raise FileExistsError(f"party key already exists at {key_path}")
+    keypair = PartyKeyPair.generate(args.party_id)
+    _write_private(key_path, args.party_id, keypair, args.ranks, args.rank_offset)
+    party = Party(
+        party_id=args.party_id,
+        keypair=keypair,
+        ranks=tuple(range(args.ranks)),
+        rendezvous_uri=str(directory),
+    )
+    payload = {
+        "party_id": party.party_id,
+        "ranks": list(party.ranks),
+        "public_key": party.public_key().hex(),
+        "key_file": str(key_path),
+    }
+    if args.json:
+        _emit_json(payload)
+        return 0
+    print(f"party {payload['party_id']} initialised with {len(party.ranks)} rank(s)")
+    print(f"key file: {key_path} (mode 600)")
+    return 0
+
+
+def _cmd_manifest(args: argparse.Namespace) -> int:
+    party = _load_party(Path(args.dir))
+    manifest = party.manifest().to_dict()
+    if args.json:
+        _emit_json(manifest)
+    else:
+        print(json.dumps(manifest, indent=2))
+    return 0
+
+
+def _cmd_mesh(args: argparse.Namespace) -> int:
+    mesh = Mesh()
+    for path in args.manifests:
+        raw = json.loads(Path(path).read_text())
+        for entry in raw.get("parties", [raw]):
+            from .mesh import PartyManifest
+
+            mesh.add(PartyManifest.from_dict(entry))
+    mesh.validate()
+    payload = mesh.to_dict()
+    if args.json:
+        _emit_json(payload)
+        return 0
+    print(f"parties: {len(mesh.parties)}  world size: {mesh.world_size()}")
+    print(f"coordinator: {mesh.coordinator()}")
+    print(f"election proof: {mesh.election_proof()}")
+    for party in mesh.party_ids():
+        ranks = mesh.ranks_of(party)
+        print(f"  {party:<12} ranks {list(ranks)}")
+    return 0
+
+
+def _cmd_shards(args: argparse.Namespace) -> int:
+    party_ids = [p.strip() for p in args.parties.split(",") if p.strip()]
+    policy = ShardPolicy(party_ids)
+    shards = policy.assign(args.rows)
+    payload = {
+        party: {
+            "rows": len(shard),
+            "indices": list(shard.indices[:20]),
+            "truncated": len(shard.indices) > 20,
+        }
+        for party, shard in sorted(shards.items())
+    }
+    if args.json:
+        _emit_json(payload)
+        return 0
+    for party, entry in payload.items():
+        print(f"  {party:<12} {entry['rows']:>6} rows")
+    print(f"total: {args.rows} rows, one owner per row")
+    return 0
+
+
+def _cmd_dp(args: argparse.Namespace) -> int:
+    config = DPConfig(
+        clip_norm=args.clip,
+        noise_multiplier=args.noise,
+        target_epsilon=args.epsilon,
+        target_delta=args.delta,
+    )
+    accountant = PrivacyAccountant(args.epsilon, args.delta)
+    steps = 0
+    try:
+        for _ in range(args.steps):
+            accountant.spend_step(config)
+            steps += 1
+    except Exception as exc:  # noqa: BLE001 - the budget error is the answer
+        payload = {
+            "config": {
+                "clip_norm": config.clip_norm,
+                "noise_multiplier": config.noise_multiplier,
+                "target_epsilon": config.target_epsilon,
+                "target_delta": config.target_delta,
+            },
+            "steps_completed": steps,
+            "budget_exhausted": True,
+            "reason": str(exc),
+            "accountant": accountant.to_dict(),
+        }
+        if args.json:
+            _emit_json(payload)
+        else:
+            print(f"budget exhausted after {steps} step(s): {exc}")
+        return 1
+    payload = {
+        "config": {
+            "clip_norm": config.clip_norm,
+            "noise_multiplier": config.noise_multiplier,
+            "target_epsilon": config.target_epsilon,
+            "target_delta": config.target_delta,
+        },
+        "steps_completed": steps,
+        "budget_exhausted": False,
+        "accountant": accountant.to_dict(),
+    }
+    if args.json:
+        _emit_json(payload)
+        return 0
+    print(
+        f"{steps} step(s) within budget: epsilon {accountant.epsilon:.4f} of "
+        f"{accountant.target_epsilon}, delta {accountant.delta:.3e}"
+    )
+    return 0
+
+
+def _cmd_spike(args: argparse.Namespace) -> int:
+    """Run the P2.1 spike: masks cancel exactly, and not through compression."""
+    import numpy as np
+
+    from .party import pairwise_secrets
+
+    size = 1024
+    party_ids = ("alpha", "beta", "gamma")
+    parties: dict[str, Party] = {}
+    for index, party_id in enumerate(party_ids):
+        parties[party_id] = Party(
+            party_id=party_id,
+            keypair=PartyKeyPair.generate(party_id, seed=b"fedmesh-spike"),
+            ranks=(index,),
+        )
+    for i, first in enumerate(party_ids):
+        for second in party_ids[i + 1 :]:
+            parties[first].establish(second, parties[second].public_key())
+            parties[second].establish(first, parties[first].public_key())
+    registry = MaskRegistry(
+        party_ids=party_ids,
+        secrets=pairwise_secrets(parties),
+        size=size,
+        scale=1.0,
+    )
+    masks = registry.all_masks()
+    rng = np.random.default_rng(17)
+    local_sums = {
+        index: rng.normal(0.0, 1.0, size=size).astype(np.float32)
+        for index in range(len(party_ids))
+    }
+    truth = sum(local_sums.values())
+
+    from iroh_tinker.comm.transport import (
+        CompressedTransport,
+        LocalNetwork,
+        Transport,
+    )
+
+    from .masking import MaskedTransport, exchange_party_sums
+
+    def run(compress: bool) -> np.ndarray:
+        network = LocalNetwork(len(party_ids))
+        transports: dict[int, MaskedTransport] = {}
+        for rank in range(len(party_ids)):
+            base: Transport = network.transport(rank)
+            if compress:
+                base = CompressedTransport(base, scheme="8bit", block_size=128)
+            transports[rank] = MaskedTransport(
+                base,
+                party_id=party_ids[rank],
+                party_of_rank=list(party_ids),
+                mask=masks[party_ids[rank]],
+                allow_lossy=True,
+            )
+        results = exchange_party_sums(
+            transports=transports,
+            party_of_rank=list(party_ids),
+            local_sums=local_sums,
+            masks=masks,
+        )
+        return results[0]
+
+    def relative(value: np.ndarray) -> float:
+        error = float(np.linalg.norm(value - truth))
+        return error / max(float(np.linalg.norm(truth)), 1e-9)
+
+    exact_error = relative(run(compress=False))
+    lossy_error = relative(run(compress=True))
+    payload = {
+        "exact_relative_error": exact_error,
+        "lossy_relative_error": lossy_error,
+        "lossy_is_worse_by": lossy_error / max(exact_error, 1e-18),
+        "decision": "the masked path must be lossless",
+    }
+    if args.json:
+        _emit_json(payload)
+        return 0
+    print(f"exact fp32 wire : relative error {exact_error:.3e}")
+    print(f"8-bit wire      : relative error {lossy_error:.3e}")
+    print(f"lossy is {payload['lossy_is_worse_by']:.0f}x worse -> {payload['decision']}")
+    return 0
+
+
+def _emit_json(payload: Any) -> None:
+    sys.stdout.write(json.dumps(payload, indent=2) + "\n")
+
+
+def main(argv: Sequence[str] | None = None) -> int:
+    parser = build_parser()
+    args = parser.parse_args(list(argv) if argv is not None else None)
+    try:
+        if args.command == "init":
+            return _cmd_init(args)
+        if args.command == "manifest":
+            return _cmd_manifest(args)
+        if args.command == "mesh":
+            return _cmd_mesh(args)
+        if args.command == "shards":
+            return _cmd_shards(args)
+        if args.command == "dp":
+            return _cmd_dp(args)
+        if args.command == "spike":
+            return _cmd_spike(args)
+    except (FileExistsError, FileNotFoundError, ValueError, RuntimeError) as exc:
+        if getattr(args, "json", False):
+            _emit_json({"error": str(exc)})
+        else:
+            print(f"error: {exc}", file=sys.stderr)
+        return 2
+    parser.error(f"unknown command {args.command!r}")
+    return 2
+
+
+if __name__ == "__main__":  # pragma: no cover
+    raise SystemExit(main())

fedmesh/dp.py

@@ -0,0 +1,163 @@
+"""Differential privacy for the federated update.
+
+fedmesh does not implement DP in the first version of the training loop; this
+module provides the pieces a training run needs and, importantly, a budget
+accountant that *refuses* to go over budget.
+
+The privacy guarantee is the standard Gaussian mechanism: per-example
+gradients are clipped to a bounded L2 norm and Gaussian noise proportional to
+that norm is added before the update leaves the party. The accountant uses
+the basic (not the tighter Rényi) composition bound, so the reported epsilon
+is an upper bound.
+"""
+from __future__ import annotations
+
+import math
+from dataclasses import dataclass, field
+from typing import Any
+
+import numpy as np
+
+__all__ = [
+    "DPConfig",
+    "PrivacyAccountant",
+    "PrivacyBudgetError",
+    "add_gaussian_noise",
+    "clip_by_l2_norm",
+    "epsilon_for_gaussian_mechanism",
+]
+
+
+class PrivacyBudgetError(RuntimeError):
+    """Raised when a step would exceed the privacy budget."""
+
+
+@dataclass(frozen=True)
+class DPConfig:
+    """Parameters of the differential privacy mechanism."""
+
+    clip_norm: float = 1.0
+    noise_multiplier: float = 10.0
+    target_epsilon: float = 8.0
+    target_delta: float = 1e-5
+
+    def __post_init__(self) -> None:
+        if self.clip_norm <= 0:
+            raise ValueError("clip_norm must be positive")
+        if self.noise_multiplier <= 0:
+            raise ValueError("noise_multiplier must be positive")
+        if not 0 < self.target_delta < 1:
+            raise ValueError("target_delta must be in (0, 1)")
+        if self.target_epsilon <= 0:
+            raise ValueError("target_epsilon must be positive")
+
+
+def clip_by_l2_norm(vector: np.ndarray, max_norm: float) -> np.ndarray:
+    """Clip a vector to ``max_norm`` in L2, preserving direction."""
+    flat = np.asarray(vector, dtype=np.float64).reshape(-1)
+    norm = float(np.linalg.norm(flat))
+    if norm <= max_norm or norm == 0.0:
+        return flat.astype(np.float32)
+    return (flat * (max_norm / norm)).astype(np.float32)
+
+
+def add_gaussian_noise(
+    vector: np.ndarray,
+    noise_multiplier: float,
+    rng: np.random.Generator,
+) -> np.ndarray:
+    """Add Gaussian noise with the given multiplier of the sensitivity."""
+    flat = np.asarray(vector, dtype=np.float64).reshape(-1)
+    noise = rng.normal(0.0, noise_multiplier, size=flat.shape)
+    return (flat + noise).astype(np.float32)
+
+
+def epsilon_for_gaussian_mechanism(
+    noise_multiplier: float,
+    clip_norm: float,
+    delta: float,
+) -> float:
+    """The basic (upper bound) epsilon of the Gaussian mechanism.
+
+    Uses the standard bound ``eps = C * sqrt(2 * ln(1.25 / delta)) / sigma``
+    for the Gaussian mechanism with sensitivity ``C`` and noise ``sigma *
+    C``. It is not the tightest bound available, but it is the one that is
+    easy to audit and it never understates the privacy loss.
+    """
+    if not 0 < delta < 1:
+        raise ValueError("delta must be in (0, 1)")
+    if noise_multiplier <= 0:
+        raise ValueError("noise_multiplier must be positive")
+    return clip_norm * math.sqrt(2.0 * math.log(1.25 / delta)) / noise_multiplier
+
+
+@dataclass
+class PrivacyAccountant:
+    """Tracks the privacy budget spent by a run.
+
+    Basic composition: epsilon adds up, delta combines as
+    ``1 - prod(1 - delta_i)``.
+    """
+
+    target_epsilon: float
+    target_delta: float
+    _epsilon: float = 0.0
+    _delta_product: float = 1.0
+    steps: int = 0
+    history: list[dict[str, Any]] = field(default_factory=list)
+
+    def __post_init__(self) -> None:
+        if self.target_epsilon <= 0:
+            raise ValueError("target_epsilon must be positive")
+        if not 0 < self.target_delta < 1:
+            raise ValueError("target_delta must be in (0, 1)")
+
+    @property
+    def epsilon(self) -> float:
+        return self._epsilon
+
+    @property
+    def delta(self) -> float:
+        return 1.0 - self._delta_product
+
+    @property
+    def remaining(self) -> float:
+        return max(0.0, self.target_epsilon - self._epsilon)
+
+    def spend(self, epsilon: float, delta: float) -> None:
+        """Record a step, refusing to exceed the budget."""
+        if epsilon < 0 or delta < 0:
+            raise ValueError("epsilon and delta must be non-negative")
+        if epsilon + self._epsilon > self.target_epsilon + 1e-12:
+            raise PrivacyBudgetError(
+                f"step would spend epsilon {epsilon:.4f} on top of "
+                f"{self._epsilon:.4f}, exceeding the budget {self.target_epsilon:.4f}"
+            )
+        if self.delta + delta > self.target_delta + 1e-15:
+            raise PrivacyBudgetError(
+                f"step would spend delta {delta:.3e} on top of {self.delta:.3e}, "
+                f"exceeding the budget {self.target_delta:.3e}"
+            )
+        self._epsilon += epsilon
+        self._delta_product *= 1.0 - delta
+        self.steps += 1
+        self.history.append({"step": self.steps, "epsilon": epsilon, "delta": delta})
+
+    def spend_step(self, config: DPConfig) -> float:
+        """Spend one step of the Gaussian mechanism described by ``config``."""
+        epsilon = epsilon_for_gaussian_mechanism(
+            config.noise_multiplier, config.clip_norm, config.target_delta
+        )
+        self.spend(epsilon, config.target_delta)
+        return epsilon
+
+    def to_dict(self) -> dict[str, Any]:
+        return {
+            "target_epsilon": self.target_epsilon,
+            "target_delta": self.target_delta,
+            "epsilon": self.epsilon,
+            "delta": self.delta,
+            "steps": self.steps,
+            "remaining_epsilon": self.remaining,
+            "history": list(self.history),
+        }