flashruntime 0.3.0__py3-none-any.whl

flashml_workloads/fedavg_weights.py

@@ -0,0 +1,223 @@
+"""Weight/delta encoding and the federated-averaging reduce.
+
+Weights cross the wire as JSON so the driver never imports torch: it runs
+inside the cloud API (spec §5.4.5), which must stay a light service. Only
+`fedavg_worker` needs torch, and it converts at the boundary.
+
+    {"<param>": {"shape": [int, ...], "data": [float, ...]}}
+
+`data` is the flattened tensor in row-major order; `shape` restores it.
+Pure stdlib on purpose — same rule as kmeans_shard and sgd_trainer, so
+this module runs on any device, including inside a --network none
+container.
+"""
+
+from __future__ import annotations
+
+import math
+
+__all__ = [
+    "NonFiniteWeights",
+    "WeightShapeMismatch",
+    "apply_delta",
+    "decode",
+    "encode",
+    "reduce_deltas",
+    "require_finite",
+    "subtract",
+]
+
+
+class WeightShapeMismatch(ValueError):
+    """Two weight blobs do not describe the same parameter set.
+
+    Never coerce past this: averaging mismatched blobs would emit weights
+    that load fine and train to nonsense.
+    """
+
+
+class NonFiniteWeights(ValueError):
+    """A weight/delta blob contains NaN or +/-Infinity.
+
+    This is not a paranoid check, and it is not primarily about attackers:
+    Python's `json` module both EMITS and PARSES the non-standard `NaN`,
+    `Infinity` and `-Infinity` literals, so a single shard whose learning
+    rate diverged writes `NaN` into its delta, the reduce turns every
+    weight into `NaN`, every later round trains from `NaN` weights, and the
+    run still reports success. There is no recovering from it afterwards —
+    NaN is absorbing — so it has to fail at the boundary where it enters.
+    """
+
+
+def require_finite(blob: dict, where: str) -> dict:
+    """Raise `NonFiniteWeights` if any value in `blob` is NaN or Inf.
+
+    Named (not `_private`) because it is the containment boundary every
+    caller of this module is entitled to use; the message names the
+    parameter and the flat index so an operator can point at the shard that
+    diverged rather than bisecting a megabyte of JSON.
+    """
+    for name, param in blob.items():
+        for i, v in enumerate(param["data"]):
+            try:
+                finite = math.isfinite(v)
+            except TypeError:  # a node sent a string/null where a float belongs
+                finite = False
+            if not finite:
+                raise NonFiniteWeights(
+                    f"{where}: parameter {name!r} index {i} is not a finite "
+                    f"number ({v!r})"
+                )
+    return blob
+
+
+def encode(state: dict[str, tuple[list[int], list[float]]]) -> dict:
+    return {name: {"shape": list(shape), "data": list(data)}
+            for name, (shape, data) in state.items()}
+
+
+def decode(blob: dict) -> dict[str, tuple[list[int], list[float]]]:
+    return {name: (list(p["shape"]), list(p["data"])) for name, p in blob.items()}
+
+
+def _require_same_params(a: dict, b: dict) -> None:
+    if a.keys() != b.keys():
+        raise WeightShapeMismatch(
+            f"parameter names differ: {sorted(a.keys())} vs {sorted(b.keys())}"
+        )
+    for name in a:
+        shape_a = a[name]["shape"]
+        shape_b = b[name]["shape"]
+
+        if shape_a != shape_b:
+            raise WeightShapeMismatch(
+                f"parameter {name!r} shape {shape_a} vs {shape_b}"
+            )
+
+        # Both shapes are the same; compute expected data length.
+        # Empty shape (scalar) has product 1.
+        expected_len = math.prod(shape_a) if shape_a else 1
+
+        # Validate that data length matches declared shape in both blobs.
+        data_len_a = len(a[name]["data"])
+        if data_len_a != expected_len:
+            raise WeightShapeMismatch(
+                f"parameter {name!r} declared shape {shape_a} (product {expected_len}) but data has length {data_len_a}"
+            )
+
+        data_len_b = len(b[name]["data"])
+        if data_len_b != expected_len:
+            raise WeightShapeMismatch(
+                f"parameter {name!r} declared shape {shape_b} (product {expected_len}) but data has length {data_len_b}"
+            )
+
+
+def subtract(new: dict, base: dict) -> dict:
+    _require_same_params(new, base)
+    # Checked on the RESULT, not the inputs: it is the cheapest single place
+    # that catches a diverged local step (NaN weights) before the delta is
+    # uploaded, and NaN/Inf in either input propagates into the difference.
+    return require_finite({
+        name: {
+            "shape": list(new[name]["shape"]),
+            "data": [x - y for x, y in zip(new[name]["data"], base[name]["data"])],
+        }
+        for name in new
+    }, "subtract")
+
+
+def apply_delta(base: dict, delta: dict, scale: float = 1.0) -> dict:
+    _require_same_params(base, delta)
+    return require_finite({
+        name: {
+            "shape": list(base[name]["shape"]),
+            "data": [b + scale * d
+                     for b, d in zip(base[name]["data"], delta[name]["data"])],
+        }
+        for name in base
+    }, "apply_delta")
+
+
+def reduce_deltas(contributions: list[tuple[dict, int]]) -> dict:
+    """Sample-weighted mean of per-worker deltas (FedAvg).
+
+    Weighting by sample count, not by worker, is what keeps the result
+    equal to centralized training on the union of the shards when the
+    shards are unequal — which they always are once machines differ.
+    """
+    if not contributions:
+        raise ValueError("reduce_deltas: no contributions")
+    total = sum(n for _, n in contributions)
+    if total < 0:
+        raise ValueError("reduce_deltas: negative total samples")
+    if total == 0:
+        raise ValueError("reduce_deltas: zero total samples")
+
+    # Validating only the TOTAL is not enough, and the gap is a model-
+    # poisoning primitive rather than a hygiene nit. `samples` is chosen by
+    # an untrusted volunteer node. With contributions (delta=-999, n=-999)
+    # and (delta=1.0, n=1000) the total is a healthy 1, but the weights
+    # w = n/total are -999 and 1000, so the "average" of two updates of
+    # magnitude ~1 is 999001.0 — six orders of magnitude outside the convex
+    # hull the average is supposed to stay inside. A weight is only a convex
+    # combination when every count is positive, so reject non-positive
+    # counts outright. (Ordered AFTER the total guard so an all-zero
+    # contribution set still reports the clearer "zero total samples".)
+    for i, (_, n) in enumerate(contributions):
+        # A NaN/Inf sample count slips BOTH the total guard above and the
+        # `n <= 0` guard below, because every comparison against NaN is
+        # False: `total <= 0`, `total == 0` and `n <= 0` are all False for a
+        # NaN contribution. It happens to be caught downstream by the
+        # finiteness check on the reduced result, but only by accident, not
+        # by design, so it is rejected explicitly here, by index, before it
+        # can reach that accidental safety net.
+        if isinstance(n, float) and not math.isfinite(n):
+            raise ValueError(
+                f"reduce_deltas: contribution {i} has a non-finite sample "
+                f"count {n!r}; sample counts must be a finite positive "
+                "integer"
+            )
+        # `bool` is a subclass of `int` (`True == 1`, `False == 0`), so a
+        # bool sample count would otherwise slide through silently. A
+        # sample count is a count of training examples, not a flag, and
+        # this boundary receives untrusted JSON from a volunteer node where
+        # `true`/`false` is a plausible malformed value for a field that
+        # should be an integer — reject it rather than coerce it. Likewise
+        # a non-integer float (e.g. `2.5`) does not correspond to any real
+        # shard size; it happens not to break the convex-combination math
+        # below, but silently accepting it is the same kind of
+        # accidental-safety gap the NaN/Inf case above is about.
+        if isinstance(n, bool) or not isinstance(n, int):
+            raise ValueError(
+                f"reduce_deltas: contribution {i} has a non-integer sample "
+                f"count {n!r} (type {type(n).__name__}); sample counts must "
+                "be a plain positive int, not a bool or a fractional float"
+            )
+        if n <= 0:
+            raise ValueError(
+                f"reduce_deltas: contribution {i} has non-positive sample "
+                f"count {n!r}; sample counts must be > 0 (a negative or zero "
+                "count makes the sample weight fall outside [0, 1] and lets "
+                "one shard amplify its delta arbitrarily)"
+            )
+
+    first = contributions[0][0]
+    # Validate first blob's internal consistency and all subsequent blobs,
+    # and reject NaN/Inf on the way IN: one NaN anywhere in one shard's delta
+    # makes every output weight NaN, and every subsequent round then trains
+    # from NaN. Rejecting the round is recoverable; poisoning the model is not.
+    for i, (blob, _) in enumerate(contributions):
+        # i == 0 validates the first blob against itself (internal
+        # shape/data-length consistency), the rest against the first.
+        _require_same_params(first, blob)
+        require_finite(blob, f"reduce_deltas: contribution {i}")
+
+    out: dict = {}
+    for name in first:
+        acc = [0.0] * len(first[name]["data"])
+        for blob, n in contributions:
+            w = n / total
+            for i, v in enumerate(blob[name]["data"]):
+                acc[i] += w * v
+        out[name] = {"shape": list(first[name]["shape"]), "data": acc}
+    return out

flashml_workloads/fedavg_worker.py

@@ -0,0 +1,166 @@
+"""One federated-averaging round on one shard, runnable as a leased task.
+
+Executor contract (same as kmeans_shard / sgd_trainer):
+
+    python -m flashml_workloads.fedavg_worker --spec spec.json --out OUTDIR
+
+`spec.json`:
+    params: round, shard, num_shards, local_steps, lr, batch_size, seed,
+            in_dim, hidden, out_dim, dataset_size
+    inputs: weights (optional path to the round's weights JSON; absent on
+            round 0, where the seed determines the starting point)
+
+Outputs `OUTDIR/delta.json` (this worker's weight change) and
+`OUTDIR/metrics.json` (the commit artifact — only a root-level
+metrics.json sets the commit hash).
+
+Why a delta and not the new weights: the driver averages contributions,
+and averaging deltas keeps the arithmetic correct when a worker joins a
+round late with stale weights — its delta is still a valid direction from
+the weights it actually saw.
+
+torch is imported inside functions so this module can be inspected (and
+the rest of flashml_workloads used) without torch installed.
+"""
+
+from __future__ import annotations
+
+import argparse
+import json
+from pathlib import Path
+
+from flashml_workloads.fedavg_weights import decode, encode, require_finite, subtract
+
+
+def build_model(seed: int, in_dim: int, hidden: int, out_dim: int):
+    import torch
+    from torch import nn
+
+    torch.manual_seed(seed)
+    return nn.Sequential(
+        nn.Linear(in_dim, hidden), nn.ReLU(), nn.Linear(hidden, out_dim)
+    )
+
+
+def state_to_blob(model) -> dict:
+    state = {}
+    for name, t in model.state_dict().items():
+        state[name] = (list(t.shape), [float(x) for x in t.flatten().tolist()])
+    return encode(state)
+
+
+def blob_to_state(model, blob: dict) -> None:
+    """Load a weight blob into `model`, failing loudly on any mismatch."""
+    import torch
+
+    from flashml_workloads.fedavg_weights import WeightShapeMismatch
+
+    # This is the worker's entry point for weights that arrived over the
+    # network as a staged input file — the last line of defence on the node
+    # side before a NaN/Inf is loaded straight into the torch model. Gate it
+    # before anything else so a corrupted or attacker-written weights.json
+    # never reaches `load_state_dict`.
+    require_finite(blob, "blob_to_state")
+
+    current = state_to_blob(model)
+    if current.keys() != blob.keys():
+        raise WeightShapeMismatch(
+            f"weights name mismatch: expected {sorted(current)}, got {sorted(blob)}"
+        )
+    new_state = {}
+    for name, (shape, data) in decode(blob).items():
+        if shape != list(current[name]["shape"]):
+            raise WeightShapeMismatch(
+                f"parameter {name!r} shape {shape} != {current[name]['shape']}"
+            )
+        new_state[name] = torch.tensor(data, dtype=torch.float32).reshape(shape)
+    model.load_state_dict(new_state)
+
+
+def _make_shard(params: dict):
+    """Deterministic synthetic data, sliced by shard.
+
+    Strided slicing (`x[shard::num_shards]`) rather than contiguous blocks
+    so every shard sees the same label distribution — a contiguous split of
+    sorted data would give workers disjoint classes and FedAvg would
+    diverge for reasons unrelated to the runtime.
+    """
+    import torch
+
+    g = torch.Generator().manual_seed(params["seed"])
+    n, d = params["dataset_size"], params["in_dim"]
+    x = torch.randn(n, d, generator=g)
+    w = torch.randn(d, 1, generator=g)
+    y = ((x @ w).squeeze(1) > 0).long()
+    shard, num = params["shard"], params["num_shards"]
+    return x[shard::num], y[shard::num]
+
+
+def run_worker(spec: dict, outdir: Path) -> dict:
+    import torch
+    from torch import nn
+
+    p = spec["params"]
+    outdir = Path(outdir)
+
+    model = build_model(p["seed"], p["in_dim"], p["hidden"], p["out_dim"])
+    weights_path = (spec.get("inputs") or {}).get("weights")
+    if weights_path:
+        blob_to_state(model, json.loads(Path(weights_path).read_text()))
+    base = state_to_blob(model)
+
+    x, y = _make_shard(p)
+    samples = int(x.shape[0])
+    if samples == 0:
+        raise ValueError(
+            f"shard {p['shard']} of num_shards={p['num_shards']} is empty "
+            f"for dataset_size={p['dataset_size']}: num_shards must not "
+            f"exceed dataset_size"
+        )
+    opt = torch.optim.SGD(model.parameters(), lr=p["lr"])
+    loss_fn = nn.CrossEntropyLoss()
+    batch = p["batch_size"]
+
+    # Batches are indexed by step via a wrapped, step-indexed gather (each
+    # index individually wrapped mod `samples`, not a single truncating
+    # slice) so every step sees exactly `batch` rows even when `batch` does
+    # not divide `samples` evenly — same rule as sgd_trainer. Pure function
+    # of `step`, no RNG state, so a retried attempt reproduces the same delta.
+    last_loss = 0.0
+    for step in range(p["local_steps"]):
+        idx = [(step * batch + i) % samples for i in range(batch)]
+        xb, yb = x[idx], y[idx]
+        opt.zero_grad()
+        loss = loss_fn(model(xb), yb)
+        loss.backward()
+        opt.step()
+        last_loss = float(loss.item())
+
+    delta = subtract(state_to_blob(model), base)
+    (outdir / "delta.json").write_text(json.dumps(delta))
+
+    metrics = {
+        "round": p["round"],
+        "shard": p["shard"],
+        "samples": samples,
+        "loss": last_loss,
+        "local_steps": p["local_steps"],
+        "delta_file": "delta.json",
+    }
+    (outdir / "metrics.json").write_text(json.dumps(metrics, sort_keys=True))
+    return metrics
+
+
+def main() -> int:
+    ap = argparse.ArgumentParser()
+    ap.add_argument("--spec", required=True)
+    ap.add_argument("--out", required=True)
+    args = ap.parse_args()
+    outdir = Path(args.out)
+    outdir.mkdir(parents=True, exist_ok=True)
+    run_worker(json.loads(Path(args.spec).read_text()), outdir)
+    return 0
+
+
+if __name__ == "__main__":
+    raise SystemExit(main())

flashml_workloads/kmeans_driver.py

@@ -0,0 +1,134 @@
+"""K-means driver: composes lease jobs into Lloyd's algorithm.
+
+One iteration = one Mode A job (N independent shard tasks); the driver
+reduces the shard partials into new centroids and submits the next
+iteration. This is the stage-composition pattern: pipelines are jobs
+chained by a driver, not a new execution mode — a dead worker inside an
+iteration costs one shard retry, and a dead *driver* can resume from the
+last completed iteration's artifacts.
+
+The same broadcast → partial-sums → reduce shape is, conceptually, what
+gradient synchronization does with gradients instead of cluster sums.
+
+Pure stdlib (urllib); usable as a library (`run_kmeans`) or CLI.
+"""
+
+from __future__ import annotations
+
+import json
+import random
+import time
+import urllib.request
+
+
+def init_centroids(rows: list[list[float]], k: int, seed: int = 0) -> list[list[float]]:
+    """Deterministic sample of k distinct starting points."""
+    rng = random.Random(seed)
+    picks = rng.sample(range(len(rows)), k)
+    return [list(rows[i]) for i in sorted(picks)]
+
+
+def reduce_partials(
+    partials: list[dict], prev_centroids: list[list[float]]
+) -> tuple[list[list[float]], float]:
+    """Sum the shard partials; new centroid = sums/counts. An empty cluster
+    keeps its previous centroid (the standard degenerate-cluster rule).
+    Returns (new_centroids, total_inertia)."""
+    k = len(prev_centroids)
+    dims = len(prev_centroids[0])
+    sums = [[0.0] * dims for _ in range(k)]
+    counts = [0] * k
+    inertia = 0.0
+    for p in partials:
+        for j in range(k):
+            counts[j] += p["counts"][j]
+            for d in range(dims):
+                sums[j][d] += p["sums"][j][d]
+        inertia += p["inertia"]
+    centroids = [
+        [s / counts[j] for s in sums[j]] if counts[j] else list(prev_centroids[j])
+        for j in range(k)
+    ]
+    return centroids, inertia
+
+
+# ---------------------------------------------------------------------------
+# Coordinator plumbing (stdlib HTTP)
+# ---------------------------------------------------------------------------
+
+
+def _http(method: str, url: str, data: bytes | None = None, ctype="application/json"):
+    req = urllib.request.Request(
+        url, data=data, method=method, headers={"Content-Type": ctype} if data else {}
+    )
+    with urllib.request.urlopen(req, timeout=15) as r:
+        body = r.read()
+        return json.loads(body) if body else None
+
+
+def shard_and_upload(
+    base_url: str, rows: list[list[float]], n_shards: int, prefix: str
+) -> list[str]:
+    """Round-robin the rows into shards, upload each as a CSV artifact to
+    the coordinator; returns the artifact:// URIs."""
+    uris = []
+    for s in range(n_shards):
+        shard_rows = rows[s::n_shards]
+        csv_bytes = ("\n".join(",".join(f"{v:.6f}" for v in r) for r in shard_rows) + "\n").encode()
+        key = f"{prefix}/shard-{s:03d}.csv"
+        _http("PUT", f"{base_url}/v1alpha1/artifacts/{key}", csv_bytes,
+              ctype="application/octet-stream")
+        uris.append(f"artifact://{key}")
+    return uris
+
+
+def run_kmeans(
+    base_url: str,
+    shard_uris: list[str],
+    centroids: list[list[float]],
+    iterations: int = 5,
+    lease_seconds: float = 30.0,
+    poll_seconds: float = 0.5,
+    timeout_per_iteration_s: float = 300.0,
+) -> dict:
+    """Run Lloyd's algorithm as a sequence of lease jobs. Returns
+    {centroids, inertia_history, job_ids}."""
+    history: list[float] = []
+    job_ids: list[str] = []
+    for it in range(iterations):
+        job = _http("POST", f"{base_url}/v1alpha1/jobs", json.dumps({
+            "apiVersion": "flashml.dev/v1alpha1", "kind": "Job",
+            "metadata": {"name": f"kmeans-it{it:02d}"},
+            "spec": {
+                "execution": {"backend": "leases"},
+                "image": {"repository": "local/tier1", "tag": "dev"},
+                "workload": {"type": "sharded_kmeans", "parameters": {
+                    "shards": shard_uris, "centroids": centroids,
+                    "iteration": it, "lease_seconds": lease_seconds}},
+            }}).encode())
+        job_id = job["job_id"]
+        job_ids.append(job_id)
+
+        deadline = time.monotonic() + timeout_per_iteration_s
+        while True:
+            state = _http("GET", f"{base_url}/v1alpha1/jobs/{job_id}")["state"]
+            if state == "SUCCEEDED":
+                break
+            if state in ("FAILED", "CANCELLED"):
+                raise RuntimeError(f"kmeans iteration {it} ended {state} (job {job_id})")
+            if time.monotonic() > deadline:
+                raise TimeoutError(f"kmeans iteration {it} timed out (job {job_id})")
+            time.sleep(poll_seconds)
+
+        partials = []
+        for a in _http("GET", f"{base_url}/v1alpha1/jobs/{job_id}/artifacts"):
+            if a["key"].endswith("metrics.json"):
+                partials.append(_http("GET", f"{base_url}/v1alpha1/artifacts/{a['key']}"))
+        if len(partials) != len(shard_uris):
+            raise RuntimeError(
+                f"iteration {it}: expected {len(shard_uris)} partials, got {len(partials)}"
+            )
+        centroids, inertia = reduce_partials(partials, centroids)
+        history.append(inertia)
+
+    return {"centroids": centroids, "inertia_history": history, "job_ids": job_ids}

flashml_workloads/kmeans_shard.py

@@ -0,0 +1,69 @@
+"""One K-means map step over one data shard, runnable as a leased task.
+
+Executor contract (same as sklearn_trial):
+
+    python -m flashml_workloads.kmeans_shard --spec spec.json --out OUTDIR
+
+`spec.json`: {"params": {"centroids": [[…], …]}, "inputs": {"shard": path}}
+The shard is a headerless CSV of points. Output `metrics.json` carries the
+*partial* statistics — per-centroid coordinate sums, counts, and inertia —
+which the driver reduces into the next round's centroids. Pure stdlib:
+assignment is just nearest-squared-distance; no numpy required on devices.
+"""
+
+from __future__ import annotations
+
+import argparse
+import csv
+import json
+from pathlib import Path
+
+
+def run_shard(spec: dict) -> dict:
+    centroids: list[list[float]] = spec["params"]["centroids"]
+    k = len(centroids)
+    dims = len(centroids[0])
+    sums = [[0.0] * dims for _ in range(k)]
+    counts = [0] * k
+    inertia = 0.0
+    n = 0
+
+    with open(spec["inputs"]["shard"], newline="") as f:
+        for row in csv.reader(f):
+            if not row:
+                continue
+            point = [float(x) for x in row]
+            best, best_d = 0, float("inf")
+            for j, c in enumerate(centroids):
+                d = sum((p - q) ** 2 for p, q in zip(point, c))
+                if d < best_d:
+                    best, best_d = j, d
+            for dim in range(dims):
+                sums[best][dim] += point[dim]
+            counts[best] += 1
+            inertia += best_d
+            n += 1
+
+    return {
+        "task_id": spec.get("task_id", ""),
+        "sums": [[round(v, 10) for v in s] for s in sums],
+        "counts": counts,
+        "inertia": inertia,
+        "n": n,
+    }
+
+
+def main(argv: list[str] | None = None) -> int:
+    parser = argparse.ArgumentParser(prog="kmeans_shard")
+    parser.add_argument("--spec", required=True)
+    parser.add_argument("--out", required=True)
+    args = parser.parse_args(argv)
+    spec = json.loads(Path(args.spec).read_text())
+    out = Path(args.out)
+    out.mkdir(parents=True, exist_ok=True)
+    (out / "metrics.json").write_text(json.dumps(run_shard(spec), sort_keys=True))
+    return 0
+
+
+if __name__ == "__main__":
+    raise SystemExit(main())