diff --git a/benchmarks/wiflow-std/.gitignore b/benchmarks/wiflow-std/.gitignore index 56133dbc..5244b54e 100644 --- a/benchmarks/wiflow-std/.gitignore +++ b/benchmarks/wiflow-std/.gitignore @@ -17,3 +17,10 @@ downloads/ results/parity_fixture.json __pycache__/ *.onnx + +# Committed ground truth: corruption masks for the pristine Kaggle download. +# remote/clean_v2.py zeroes the corrupted source windows IN PLACE, so these +# masks CANNOT be regenerated from a cleaned copy (generate_corruption_masks.py +# documents the criteria and reproduces them only from a fresh download). +!results/nan_windows_mask.npy +!results/big_windows_mask.npy diff --git a/benchmarks/wiflow-std/RESULTS.md b/benchmarks/wiflow-std/RESULTS.md index 85820340..52306f22 100644 --- a/benchmarks/wiflow-std/RESULTS.md +++ b/benchmarks/wiflow-std/RESULTS.md @@ -66,6 +66,19 @@ mostly would-be *training* data — so this is not a split mismatch): Kaggle upload. Window masks: `results/nan_windows_mask.npy`, `results/big_windows_mask.npy`. +### Reproducing the corruption masks + +The two mask files (9,070 NaN/Inf windows, 9,072 with |amplitude| > 1.5; +union 9,072, all in dataset files 487–499) are **committed ground truth** +(gitignore-negated, ~352 KB each). They can only be regenerated from a +**pristine** Kaggle download: `remote/clean_v2.py` repairs the dataset by +zeroing the corrupted windows in place, after which the corruption evidence +is gone and a rescan returns all-False. `generate_corruption_masks.py` +re-derives them (chunked scan, criteria: any non-finite value OR +max |finite| > 1.5 per 540×20 window) and refuses to write all-False masks, +which indicate a cleaned copy. Verified 2026-06-11: a regeneration from the +local pristine download is bit-identical to the committed masks. + ### Retraining result (MEASURED, 2026-06-10): claims APPROXIMATELY REPRODUCED Since the shipped checkpoint is unusable, measurement (a) fell back to retraining diff --git a/benchmarks/wiflow-std/_bench_common.py b/benchmarks/wiflow-std/_bench_common.py new file mode 100644 index 00000000..f6a67bf0 --- /dev/null +++ b/benchmarks/wiflow-std/_bench_common.py @@ -0,0 +1,200 @@ +"""Shared infrastructure for the LOCAL wiflow-std benchmark scripts (ADR-152). + +This module is the single canonical implementation of the helpers that were +previously copy-pasted across eval_repro.py / quantize_bench.py / +onnx_bench.py / eval_ort_accuracy.py / export_to_safetensors.py: + + - ``import_upstream()`` -- sys.path setup + the models-package stub that + works around the upstream import bug, plus the >1GB np.load mmap patch + - ``install_np_load_mmap_patch()`` -- the mmap patch on its own + - ``remap_legacy_keys()`` / ``load_remapped_state()`` -- checkpoint + key remap for the pre-rename released checkpoint + - ``load_wiflow_model()`` -- WiFlowPoseModel from a checkpoint, eval mode + - ``set_seed()`` -- mirrors upstream run.py seeding exactly + - ``evaluate()`` -- THE canonical batch-weighted PCK/MPJPE evaluation loop + (thresholds 0.1-0.5, upstream utils/metrics.py math); accepts either a + torch nn.Module or an onnxruntime InferenceSession + +The scripts under remote/ deploy to ruvultra as standalone single files and +therefore intentionally inline private copies of these helpers; when editing +them, treat this module as the reference implementation and keep the copies +in sync. +""" + +import os +import random +import sys +import time +import types + +import numpy as np +import torch + +HERE = os.path.dirname(os.path.abspath(__file__)) +UPSTREAM = os.path.join(HERE, "upstream") +RESULTS = os.path.join(HERE, "results") + +DEFAULT_THRESHOLDS = (0.1, 0.2, 0.3, 0.4, 0.5) + +# --------------------------------------------------------------------------- +# >1GB np.load mmap patch +# --------------------------------------------------------------------------- + +# csi_windows.npy is ~13 GB; mmap large arrays instead of loading into RAM +# (loading it eagerly needs ~15 GB). +_np_load = np.load + + +def _np_load_mmap(path, *a, **kw): + if (isinstance(path, str) and path.endswith(".npy") + and os.path.getsize(path) > 1 << 30 and "mmap_mode" not in kw): + kw["mmap_mode"] = "r" + return _np_load(path, *a, **kw) + + +def install_np_load_mmap_patch(): + """Globally patch np.load so .npy files >1GB are mmap'd read-only. + + Idempotent. Patching the numpy module attribute is equivalent to the + historical ``upstream_dataset.np.load = _np_load_mmap`` (dataset.np IS + the numpy module), but works regardless of import order. + """ + np.load = _np_load_mmap + + +# --------------------------------------------------------------------------- +# upstream import shim +# --------------------------------------------------------------------------- + +def import_upstream(mmap_patch=True): + """Make the upstream WiFlow-STD clone importable; returns its path. + + Upstream bug: models/__init__.py imports TemporalConvNet, which + models/tcn.py does not define -- the package fails to import as + published. Register a stub package so the broken __init__ never + executes; submodules (models.pose_model etc.) still resolve via + __path__. Idempotent. + """ + if UPSTREAM not in sys.path: + sys.path.insert(0, UPSTREAM) + if "models" not in sys.modules: + _models_pkg = types.ModuleType("models") + _models_pkg.__path__ = [os.path.join(UPSTREAM, "models")] + sys.modules["models"] = _models_pkg + if mmap_patch: + install_np_load_mmap_patch() + return UPSTREAM + + +# --------------------------------------------------------------------------- +# checkpoint loading +# --------------------------------------------------------------------------- + +# The released checkpoint predates the published code: modules were renamed +# att -> attention, final_conv -> decoder (param count identical, 2.23M). +LEGACY_RENAMES = {"att.": "attention.", "final_conv.": "decoder."} + + +def remap_legacy_keys(state): + """Remap pre-rename state_dict keys; no-op for already-new-style keys.""" + return {next((new + k[len(old):] for old, new in LEGACY_RENAMES.items() + if k.startswith(old)), k): v + for k, v in state.items()} + + +def load_remapped_state(path, map_location="cpu"): + """torch.load (weights_only) + legacy key remap.""" + state = torch.load(path, map_location=map_location, weights_only=True) + return remap_legacy_keys(state) + + +def load_wiflow_model(checkpoint, map_location="cpu", dropout=0.5): + """Full-size WiFlowPoseModel from a checkpoint, strict load, eval mode.""" + import_upstream() + from models.pose_model import WiFlowPoseModel + model = WiFlowPoseModel(dropout=dropout) + model.load_state_dict(load_remapped_state(checkpoint, map_location), + strict=True) + model.eval() + return model + + +# --------------------------------------------------------------------------- +# seeding +# --------------------------------------------------------------------------- + +def set_seed(seed=42): + # mirror upstream run.py exactly + random.seed(seed) + np.random.seed(seed) + torch.manual_seed(seed) + if torch.cuda.is_available(): + torch.cuda.manual_seed(seed) + torch.cuda.manual_seed_all(seed) + torch.backends.cudnn.deterministic = True + torch.backends.cudnn.benchmark = False + + +# --------------------------------------------------------------------------- +# THE canonical evaluation loop +# --------------------------------------------------------------------------- + +def evaluate(model, loader, device=None, dtype=None, label="", + thresholds=DEFAULT_THRESHOLDS, progress_every=50): + """Batch-weighted PCK/MPJPE over a DataLoader (upstream metrics math). + + ``model`` may be a torch nn.Module (optionally evaluated on ``device`` + with inputs cast to ``dtype``) or an onnxruntime InferenceSession. + Per-threshold PCK values are independent in upstream calculate_pck, so + evaluating a superset of thresholds never changes any individual value. + + Returns {"samples", "mpjpe", "pck@10".."pck@50", "wall_seconds"}. + """ + import_upstream() + from utils.metrics import calculate_mpjpe, calculate_pck + + is_ort = hasattr(model, "get_inputs") # onnxruntime InferenceSession + if is_ort: + inp = model.get_inputs()[0].name + + def forward(bx): + return torch.from_numpy(model.run(None, {inp: bx.numpy()})[0]) + else: + model.eval() + + def forward(bx): + if device is not None: + bx = bx.to(device) + if dtype is not None: + bx = bx.to(dtype) + return model(bx).float() + + thresholds = list(thresholds) + totals = {t: 0.0 for t in thresholds} + total_mpe, n = 0.0, 0 + t0 = time.time() + with torch.no_grad(): + for batch_idx, (bx, by) in enumerate(loader): + out = forward(bx) + if device is not None and not is_ort: + by = by.to(device) + mpe = calculate_mpjpe(out, by) + pck = calculate_pck(out, by, thresholds=thresholds) + bs = by.size(0) + total_mpe += mpe * bs + for t in totals: + totals[t] += pck[t] * bs + n += bs + if batch_idx % progress_every == 0: + tag = f"[{label}] " if label else "" + pck20 = totals.get(0.2) + pck20_str = f"pck20={pck20 / n:.4f} " if pck20 is not None else "" + print(f" {tag}batch {batch_idx}: n={n} {pck20_str}" + f"mpjpe={total_mpe / n:.4f} ({time.time() - t0:.0f}s)", + flush=True) + return { + "samples": n, + "mpjpe": total_mpe / n, + **{f"pck@{int(t * 100)}": totals[t] / n for t in thresholds}, + "wall_seconds": time.time() - t0, + } diff --git a/benchmarks/wiflow-std/eval_ort_accuracy.py b/benchmarks/wiflow-std/eval_ort_accuracy.py index 280e627f..94f7016e 100644 --- a/benchmarks/wiflow-std/eval_ort_accuracy.py +++ b/benchmarks/wiflow-std/eval_ort_accuracy.py @@ -17,41 +17,17 @@ import argparse import json import os import sys -import time - -import numpy as np -import torch HERE = os.path.dirname(os.path.abspath(__file__)) -RESULTS = os.path.join(HERE, "results") sys.path.insert(0, HERE) +from _bench_common import RESULTS, evaluate # noqa: E402 from quantize_bench import build_test_subset # noqa: E402 (sets up upstream imports) -sys.path.insert(0, os.path.join(HERE, "upstream")) -from utils.metrics import calculate_mpjpe, calculate_pck # noqa: E402 - def evaluate_ort(sess, loader, label): - inp = sess.get_inputs()[0].name - totals = {0.2: 0.0, 0.5: 0.0} - total_mpe, n = 0.0, 0 - t0 = time.time() - for batch_idx, (bx, by) in enumerate(loader): - out = torch.from_numpy(sess.run(None, {inp: bx.numpy()})[0]) - pck = calculate_pck(out, by, thresholds=[0.2, 0.5]) - mpe = calculate_mpjpe(out, by) - bs = by.size(0) - total_mpe += mpe * bs - for t in totals: - totals[t] += pck[t] * bs - n += bs - if batch_idx % 50 == 0: - print(f" [{label}] batch {batch_idx}: n={n} " - f"pck20={totals[0.2]/n:.4f} mpjpe={total_mpe/n:.4f} " - f"({time.time()-t0:.0f}s)", flush=True) - return {"samples": n, "pck@20": totals[0.2] / n, "pck@50": totals[0.5] / n, - "mpjpe": total_mpe / n, "wall_seconds": time.time() - t0} + """ORT-session evaluation via the canonical _bench_common.evaluate loop.""" + return evaluate(sess, loader, label=label) def main(): diff --git a/benchmarks/wiflow-std/eval_repro.py b/benchmarks/wiflow-std/eval_repro.py index 7758c74a..6a35ea18 100644 --- a/benchmarks/wiflow-std/eval_repro.py +++ b/benchmarks/wiflow-std/eval_repro.py @@ -15,56 +15,18 @@ Usage: import argparse import json import os -import random import sys -import time -import numpy as np import torch from torch.utils.data import DataLoader -UPSTREAM = os.path.join(os.path.dirname(os.path.abspath(__file__)), "upstream") -sys.path.insert(0, UPSTREAM) +from _bench_common import (UPSTREAM, evaluate, import_upstream, + load_remapped_state, set_seed) -# Upstream bug: models/__init__.py imports TemporalConvNet, which models/tcn.py -# does not define (it defines TemporalBlock) — the package fails to import as -# published. Register a stub package so the broken __init__ never executes; -# submodules (models.pose_model etc.) still resolve via __path__. -import types # noqa: E402 +import_upstream() # sys.path + models stub + >1GB np.load mmap patch -_models_pkg = types.ModuleType("models") -_models_pkg.__path__ = [os.path.join(UPSTREAM, "models")] -sys.modules["models"] = _models_pkg - -import dataset as upstream_dataset # noqa: E402 from dataset import PreprocessedCSIKeypointsDataset, create_preprocessed_train_val_test_loaders # noqa: E402 from models.pose_model import WiFlowPoseModel # noqa: E402 -from utils.metrics import calculate_pck, calculate_mpjpe # noqa: E402 - -# csi_windows.npy is ~13 GB; mmap large arrays instead of loading into RAM. -_np_load = np.load - - -def _np_load_mmap(path, *a, **kw): - if (isinstance(path, str) and path.endswith(".npy") - and os.path.getsize(path) > 1 << 30 and "mmap_mode" not in kw): - kw["mmap_mode"] = "r" - return _np_load(path, *a, **kw) - - -upstream_dataset.np.load = _np_load_mmap - - -def set_seed(seed=42): - # mirror upstream run.py exactly - random.seed(seed) - np.random.seed(seed) - torch.manual_seed(seed) - if torch.cuda.is_available(): - torch.cuda.manual_seed(seed) - torch.cuda.manual_seed_all(seed) - torch.backends.cudnn.deterministic = True - torch.backends.cudnn.benchmark = False def find_data_dir(root): @@ -74,35 +36,6 @@ def find_data_dir(root): return None -def evaluate(model, loader, device): - model.eval() - totals = {t: 0.0 for t in (0.1, 0.2, 0.3, 0.4, 0.5)} - total_mpe = 0.0 - n = 0 - t0 = time.time() - with torch.no_grad(): - for batch_idx, (batch_x, batch_y) in enumerate(loader): - batch_x = batch_x.to(device) - batch_y = batch_y.to(device) - outputs = model(batch_x) - mpe = calculate_mpjpe(outputs, batch_y) - pck = calculate_pck(outputs, batch_y, thresholds=[0.1, 0.2, 0.3, 0.4, 0.5]) - bs = batch_y.size(0) - total_mpe += mpe * bs - for t in totals: - totals[t] += pck[t] * bs - n += bs - if batch_idx % 50 == 0: - print(f" batch {batch_idx}: n={n} pck20={totals[0.2]/n:.4f} " - f"mpjpe={total_mpe/n:.4f} ({time.time()-t0:.0f}s)", flush=True) - return { - "samples": n, - "mpjpe": total_mpe / n, - **{f"pck@{int(t*100)}": totals[t] / n for t in totals}, - "wall_seconds": time.time() - t0, - } - - def main(): parser = argparse.ArgumentParser() parser.add_argument("--data-dir", required=True, @@ -134,13 +67,9 @@ def main(): dataset=dataset, batch_size=args.batch_size, num_workers=0, random_seed=42) model = WiFlowPoseModel(dropout=0.5).to(device) - state = torch.load(args.checkpoint, map_location=device, weights_only=True) # released checkpoint predates the published code: modules were renamed # att -> attention, final_conv -> decoder (param count identical, 2.23M) - renames = {"att.": "attention.", "final_conv.": "decoder."} - state = {next((new + k[len(old):] for old, new in renames.items() - if k.startswith(old)), k): v - for k, v in state.items()} + state = load_remapped_state(args.checkpoint, map_location=device) model.load_state_dict(state, strict=True) n_params = sum(p.numel() for p in model.parameters()) print(f"checkpoint: {args.checkpoint} ({n_params/1e6:.2f}M params)") @@ -154,13 +83,13 @@ def main(): "device": str(device)} print("=== test set (full, drop_last=False) ===") - results["test_full"] = evaluate(model, test_loader, device) + results["test_full"] = evaluate(model, test_loader, device=device) print(json.dumps(results["test_full"], indent=2)) test_loader_dl = DataLoader(test_loader.dataset, batch_size=args.batch_size, shuffle=False, drop_last=True) print("=== test set (drop_last=True, as upstream train.py) ===") - results["test_drop_last"] = evaluate(model, test_loader_dl, device) + results["test_drop_last"] = evaluate(model, test_loader_dl, device=device) print(json.dumps(results["test_drop_last"], indent=2)) os.makedirs(os.path.dirname(args.out), exist_ok=True) diff --git a/benchmarks/wiflow-std/export_to_safetensors.py b/benchmarks/wiflow-std/export_to_safetensors.py index b5c9d4d5..75e3eded 100644 --- a/benchmarks/wiflow-std/export_to_safetensors.py +++ b/benchmarks/wiflow-std/export_to_safetensors.py @@ -41,24 +41,14 @@ Usage: import json import os import re -import sys import numpy as np import torch from safetensors.torch import save_file -HERE = os.path.dirname(os.path.abspath(__file__)) -UPSTREAM = os.path.join(HERE, "upstream") -RESULTS = os.path.join(HERE, "results") -sys.path.insert(0, UPSTREAM) +from _bench_common import RESULTS, import_upstream, remap_legacy_keys -# Upstream models/__init__.py is broken as published (imports a name tcn.py -# does not define); register a stub package so it never executes. -import types # noqa: E402 - -_models_pkg = types.ModuleType("models") -_models_pkg.__path__ = [os.path.join(UPSTREAM, "models")] -sys.modules["models"] = _models_pkg +import_upstream() # sys.path + models stub from models.pose_model import WiFlowPoseModel # noqa: E402 @@ -125,11 +115,8 @@ def main(): state = state[wrapper] break - # Legacy upstream names predate the published code (eval_repro.py). - renames = {"att.": "attention.", "final_conv.": "decoder."} - state = {next((new + k[len(old):] for old, new in renames.items() - if k.startswith(old)), k): v - for k, v in state.items()} + # Legacy upstream names predate the published code (_bench_common). + state = remap_legacy_keys(state) mapped = {} dropped = 0 diff --git a/benchmarks/wiflow-std/generate_corruption_masks.py b/benchmarks/wiflow-std/generate_corruption_masks.py new file mode 100644 index 00000000..2ab82c04 --- /dev/null +++ b/benchmarks/wiflow-std/generate_corruption_masks.py @@ -0,0 +1,148 @@ +"""Regenerate results/nan_windows_mask.npy + results/big_windows_mask.npy by +scanning a PRISTINE kagglehub download of the WiFlow-STD dataset +(kaka2434/wiflow-dataset v1, csi_windows.npy, 360,000 windows of 540x20). + +============================ READ THIS FIRST =============================== +This script MUST be run against an UNCLEANED copy of the dataset. + +remote/clean_v2.py (and its predecessor clean_nan.py) repair the dataset by +zeroing the corrupted windows IN PLACE, with no backup. A cleaned copy +contains no non-finite values and no out-of-range amplitudes, so on a cleaned +copy this scan produces ALL-FALSE masks -- silently wrong ground truth. The +script errors out loudly in that case (see the sanity check in main()). + +That irreversibility is exactly why the two committed mask files under +results/ (gitignore-negated) are the canonical ground truth: once a download +has been cleaned, the masks can NEVER be regenerated from it. Only run this +on a fresh `kagglehub.dataset_download("kaka2434/wiflow-dataset")`. +============================================================================ + +Criteria (per window; mirrors the original 2026-06-10 scan and the +remote/clean_v2.py repair criteria): + + nan mask: any non-finite value (NaN/Inf) anywhere in the 540x20 window + big mask: max |finite value| > 1.5 (the data is otherwise [0,1]-normalized; + the corrupted files contain garbage up to 3.4e38, float32 max) + +Expected result on the pristine Kaggle download (RESULTS.md defect 5): + nan: 9,070 True | big: 9,072 True | union: 9,072 -- all windows in dataset + files 487-499 (the final 13 files), window indices 350,922-359,999. + +Usage: + PYTHONUTF8=1 .venv/Scripts/python.exe generate_corruption_masks.py \ + [--data-dir ] [--out-dir results] +""" + +import argparse +import os +import sys + +import numpy as np + +HERE = os.path.dirname(os.path.abspath(__file__)) +RESULTS = os.path.join(HERE, "results") + +EXPECTED = {"nan": 9070, "big": 9072, "union": 9072, + "files": (487, 499), "windows": (350922, 359999)} + + +def scan(csi_path, chunk=4000): + """Chunked scan of the (mmap'd) windows array; returns (nan_mask, big_mask).""" + csi = np.load(csi_path, mmap_mode="r") + n = len(csi) + nan_mask = np.zeros(n, dtype=bool) + big_mask = np.zeros(n, dtype=bool) + for i in range(0, n, chunk): + block = np.asarray(csi[i:i + chunk]) + finite = np.isfinite(block) + nan_mask[i:i + chunk] = (~finite).any(axis=(1, 2)) + big_mask[i:i + chunk] = ( + np.abs(np.where(finite, block, 0)).max(axis=(1, 2)) > 1.5) + if (i // chunk) % 10 == 0: + print(f" scanned {min(i + chunk, n):,}/{n:,} windows " + f"(nan={int(nan_mask.sum()):,} big={int(big_mask.sum()):,})", + flush=True) + return nan_mask, big_mask + + +def describe_files(data_dir, mask): + """Map marked windows to dataset file indices via window_info.npz.""" + info = os.path.join(data_dir, "window_info.npz") + if not os.path.exists(info): + return None + w2f = np.load(info)["window_to_file"] + return np.unique(w2f[mask]) + + +def main(): + parser = argparse.ArgumentParser( + description="Regenerate the corruption masks from a PRISTINE " + "(uncleaned) kagglehub download. See module docstring.") + parser.add_argument("--data-dir", default=os.path.join( + os.path.expanduser("~"), ".cache", "kagglehub", "datasets", "kaka2434", + "wiflow-dataset", "versions", "1", "preprocessed_csi_data"), + help="Directory containing csi_windows.npy (PRISTINE copy)") + parser.add_argument("--out-dir", default=RESULTS, + help="Where to write the two .npy masks") + parser.add_argument("--chunk", type=int, default=4000, + help="Windows per scan chunk (memory/speed tradeoff)") + args = parser.parse_args() + + csi_path = os.path.join(args.data_dir, "csi_windows.npy") + if not os.path.exists(csi_path): + sys.exit(f"csi_windows.npy not found in {args.data_dir}") + + print(f"scanning {csi_path} (chunk={args.chunk}) ...") + nan_mask, big_mask = scan(csi_path, args.chunk) + union = nan_mask | big_mask + print(f"nan: {int(nan_mask.sum()):,} | big: {int(big_mask.sum()):,} | " + f"union: {int(union.sum()):,} of {len(union):,} windows") + + # ---- sanity check: an all-False result means a CLEANED copy ------------ + if not union.any(): + sys.exit( + "ERROR: scan found ZERO corrupted windows.\n" + "\n" + "The pristine Kaggle download (kaka2434/wiflow-dataset v1) is " + "known to contain\n" + "9,072 corrupted windows (NaN/Inf + amplitudes up to 3.4e38) in " + "dataset files\n" + "487-499 (RESULTS.md, reproducibility defect 5). Finding none " + "means this copy\n" + "has almost certainly already been repaired by remote/clean_v2.py " + "(or clean_nan.py),\n" + "which zeroes the corrupted windows IN PLACE -- after that the " + "corruption evidence\n" + "is gone and the masks CANNOT be regenerated from this copy.\n" + "\n" + "Refusing to overwrite the committed ground-truth masks with " + "all-False ones.\n" + "Re-download the dataset (kagglehub.dataset_download(" + "'kaka2434/wiflow-dataset'))\n" + "and point --data-dir at the fresh, uncleaned copy.") + + files = describe_files(args.data_dir, union) + if files is not None: + print(f"marked windows span dataset files {files.min()}-{files.max()}: " + f"{files.tolist()}") + lo, hi = EXPECTED["files"] + if files.min() != lo or files.max() != hi: + print(f"WARNING: expected marked files exactly {lo}-{hi} " + f"(the pristine v1 download); got {files.min()}-{files.max()}. " + f"Different dataset version, or a partially cleaned copy?") + for name, mask, exp in (("nan", nan_mask, EXPECTED["nan"]), + ("big", big_mask, EXPECTED["big"])): + if int(mask.sum()) != exp: + print(f"WARNING: {name} mask has {int(mask.sum()):,} True windows; " + f"the pristine v1 download yields {exp:,}.") + + os.makedirs(args.out_dir, exist_ok=True) + for name, mask in (("nan_windows_mask.npy", nan_mask), + ("big_windows_mask.npy", big_mask)): + out = os.path.join(args.out_dir, name) + np.save(out, mask) + print(f"wrote {out} ({int(mask.sum()):,} True)") + + +if __name__ == "__main__": + main() diff --git a/benchmarks/wiflow-std/onnx_bench.py b/benchmarks/wiflow-std/onnx_bench.py index 5e4c7c14..f9285f66 100644 --- a/benchmarks/wiflow-std/onnx_bench.py +++ b/benchmarks/wiflow-std/onnx_bench.py @@ -21,40 +21,22 @@ import json import os import platform import statistics -import sys import time import traceback import numpy as np import torch -HERE = os.path.dirname(os.path.abspath(__file__)) -UPSTREAM = os.path.join(HERE, "upstream") -RESULTS = os.path.join(HERE, "results") -sys.path.insert(0, UPSTREAM) +from _bench_common import RESULTS, import_upstream, load_wiflow_model -import types # noqa: E402 - -_models_pkg = types.ModuleType("models") -_models_pkg.__path__ = [os.path.join(UPSTREAM, "models")] -sys.modules["models"] = _models_pkg - -from models.pose_model import WiFlowPoseModel # noqa: E402 +import_upstream() # sys.path + models stub + >1GB np.load mmap patch CHECKPOINT = os.path.join(RESULTS, "retrained_best_pose_model.pth") OUT_JSON = os.path.join(RESULTS, "edge_optimization.json") def load_fp32_model(): - state = torch.load(CHECKPOINT, map_location="cpu", weights_only=True) - renames = {"att.": "attention.", "final_conv.": "decoder."} - state = {next((new + k[len(old):] for old, new in renames.items() - if k.startswith(old)), k): v - for k, v in state.items()} - model = WiFlowPoseModel(dropout=0.5) - model.load_state_dict(state, strict=True) - model.eval() - return model + return load_wiflow_model(CHECKPOINT) def try_export(model, path, batch, dynamic, opset=17): @@ -115,6 +97,16 @@ def bench_ort(sess, batch, n_runs): def main(): + import argparse + parser = argparse.ArgumentParser( + description="ONNX export + onnxruntime CPU benchmark for the " + "retrained WiFlow-STD checkpoint (no options; see " + "module docstring). NB: the published " + "retrained_fp32_dynamic.onnx came from the TorchScript " + "exporter; on newer torch the dynamo attempt may succeed " + "first and produce a different (external-data) artifact.") + parser.parse_args() + import onnxruntime model = load_fp32_model() results = { diff --git a/benchmarks/wiflow-std/quantize_bench.py b/benchmarks/wiflow-std/quantize_bench.py index 70939fc6..80d842e1 100644 --- a/benchmarks/wiflow-std/quantize_bench.py +++ b/benchmarks/wiflow-std/quantize_bench.py @@ -28,7 +28,6 @@ import json import os import platform import statistics -import sys import time import numpy as np @@ -36,55 +35,21 @@ import torch import torch.nn as nn from torch.utils.data import DataLoader -HERE = os.path.dirname(os.path.abspath(__file__)) -UPSTREAM = os.path.join(HERE, "upstream") -RESULTS = os.path.join(HERE, "results") -sys.path.insert(0, UPSTREAM) +from _bench_common import HERE, RESULTS, evaluate, import_upstream, load_wiflow_model -# Upstream models/__init__.py is broken as published (imports a name tcn.py -# does not define); register a stub package so it never executes. -import types # noqa: E402 +import_upstream() # sys.path + models stub + >1GB np.load mmap patch -_models_pkg = types.ModuleType("models") -_models_pkg.__path__ = [os.path.join(UPSTREAM, "models")] -sys.modules["models"] = _models_pkg - -import dataset as upstream_dataset # noqa: E402 from dataset import ( # noqa: E402 PreprocessedCSIKeypointsDataset, create_preprocessed_train_val_test_loaders, ) -from models.pose_model import WiFlowPoseModel # noqa: E402 -from utils.metrics import calculate_mpjpe, calculate_pck # noqa: E402 CHECKPOINT = os.path.join(RESULTS, "retrained_best_pose_model.pth") -# csi_windows.npy is ~13 GB; mmap large arrays instead of loading into RAM -# (same trick as eval_repro.py). -_np_load = np.load - - -def _np_load_mmap(path, *a, **kw): - if (isinstance(path, str) and path.endswith(".npy") - and os.path.getsize(path) > 1 << 30 and "mmap_mode" not in kw): - kw["mmap_mode"] = "r" - return _np_load(path, *a, **kw) - - -upstream_dataset.np.load = _np_load_mmap - def load_fp32_model(): - state = torch.load(CHECKPOINT, map_location="cpu", weights_only=True) - # legacy upstream names, harmless no-op on the retrained checkpoint - renames = {"att.": "attention.", "final_conv.": "decoder."} - state = {next((new + k[len(old):] for old, new in renames.items() - if k.startswith(old)), k): v - for k, v in state.items()} - model = WiFlowPoseModel(dropout=0.5) - model.load_state_dict(state, strict=True) - model.eval() - return model + # legacy upstream key remap inside is a harmless no-op on this checkpoint + return load_wiflow_model(CHECKPOINT) def state_dict_size_bytes(model, path): @@ -138,33 +103,6 @@ def build_test_subset(data_dir, subset_size, batch_size=64): return loader, len(clean) -def evaluate(model, loader, dtype=torch.float32, label=""): - totals = {0.2: 0.0, 0.5: 0.0} - total_mpe, n = 0.0, 0 - t0 = time.time() - with torch.no_grad(): - for batch_idx, (bx, by) in enumerate(loader): - out = model(bx.to(dtype)).float() - pck = calculate_pck(out, by, thresholds=[0.2, 0.5]) - mpe = calculate_mpjpe(out, by) - bs = by.size(0) - total_mpe += mpe * bs - for t in totals: - totals[t] += pck[t] * bs - n += bs - if batch_idx % 50 == 0: - print(f" [{label}] batch {batch_idx}: n={n} " - f"pck20={totals[0.2]/n:.4f} mpjpe={total_mpe/n:.4f} " - f"({time.time()-t0:.0f}s)", flush=True) - return { - "samples": n, - "pck@20": totals[0.2] / n, - "pck@50": totals[0.5] / n, - "mpjpe": total_mpe / n, - "wall_seconds": time.time() - t0, - } - - def quantize_int8_dynamic(fp32_model): """torch.ao.quantization.quantize_dynamic on Linear/Conv where supported. Returns (model, report) where report documents what actually quantized.""" @@ -272,7 +210,7 @@ def main(): for name, (model, dtype, _f) in variants.items(): print(f"\n=== accuracy: {name} ===") results["variants"][name]["accuracy"] = evaluate( - model, loader, dtype, label=name) + model, loader, dtype=dtype, label=name) print(json.dumps(results["variants"][name]["accuracy"], indent=2)) # ---- merge into edge_optimization.json --------------------------------- diff --git a/benchmarks/wiflow-std/remote/eval_retrained.py b/benchmarks/wiflow-std/remote/eval_retrained.py index 92b1b17a..7940184d 100644 --- a/benchmarks/wiflow-std/remote/eval_retrained.py +++ b/benchmarks/wiflow-std/remote/eval_retrained.py @@ -3,6 +3,11 @@ Scores the model produced by run.py (train_output/best_pose_model.pth or similar) on the seed-42 test split: full test set AND NaN-free subset (excluding windows that were zero-filled by clean_nan.py — file indices 487-499). + +NOTE: deployed to ruvultra (~/wiflow-std-bench) as a standalone single file, +so it deliberately inlines its helpers. The reference implementations (upstream +import shim, >1GB np.load mmap patch, key-remap loader, canonical evaluate +loop) live in benchmarks/wiflow-std/_bench_common.py — keep copies in sync. """ import json, os, random, sys @@ -10,6 +15,20 @@ import numpy as np import torch from torch.utils.data import DataLoader, Subset +# csi_windows.npy is ~13 GB; mmap large arrays instead of eagerly loading +# ~15 GB into RAM (same patch as _bench_common._np_load_mmap). +_np_load = np.load + + +def _np_load_mmap(path, *a, **kw): + if (isinstance(path, str) and path.endswith('.npy') + and os.path.getsize(path) > 1 << 30 and 'mmap_mode' not in kw): + kw['mmap_mode'] = 'r' + return _np_load(path, *a, **kw) + + +np.load = _np_load_mmap + sys.path.insert(0, os.path.expanduser('~/wiflow-std-bench/upstream')) from dataset import PreprocessedCSIKeypointsDataset, create_preprocessed_train_val_test_loaders from models.pose_model import WiFlowPoseModel diff --git a/benchmarks/wiflow-std/remote/measb/train_measb.py b/benchmarks/wiflow-std/remote/measb/train_measb.py index 5789836f..11555854 100644 --- a/benchmarks/wiflow-std/remote/measb/train_measb.py +++ b/benchmarks/wiflow-std/remote/measb/train_measb.py @@ -54,6 +54,11 @@ Pre-registered protocol (followed exactly): Usage (on ruvultra): nice -n 10 nohup ~/wiflow-std-bench/venv/bin/python train_measb.py > train_measb.log 2>&1 & + +NOTE: deployed to ruvultra as a standalone single file, so it deliberately +inlines its helpers. The reference implementations (upstream import shim, +np.load mmap patch, key-remap loader, canonical evaluate loop) live in +benchmarks/wiflow-std/_bench_common.py — keep copies in sync. """ import json diff --git a/benchmarks/wiflow-std/remote/sweep/run_sweep.py b/benchmarks/wiflow-std/remote/sweep/run_sweep.py index b0312d31..ab406d83 100644 --- a/benchmarks/wiflow-std/remote/sweep/run_sweep.py +++ b/benchmarks/wiflow-std/remote/sweep/run_sweep.py @@ -17,6 +17,11 @@ Usage: nohup venv/bin/python sweep/run_sweep.py > sweep/sweep.log 2>&1 & Idempotent: variants already present in sweep/results.jsonl are skipped. + +NOTE: deployed to ruvultra (~/wiflow-std-bench/sweep) as a standalone file, so +it deliberately inlines its helpers. The reference implementations (upstream +import shim, >1GB np.load mmap patch, key-remap loader, canonical evaluate +loop) live in benchmarks/wiflow-std/_bench_common.py — keep copies in sync. """ import argparse import copy @@ -30,6 +35,20 @@ import numpy as np import torch from torch.utils.data import DataLoader, Subset +# csi_windows.npy is ~13 GB; mmap large arrays instead of eagerly loading +# ~15 GB into RAM (same patch as _bench_common._np_load_mmap). +_np_load = np.load + + +def _np_load_mmap(path, *a, **kw): + if (isinstance(path, str) and path.endswith('.npy') + and os.path.getsize(path) > 1 << 30 and 'mmap_mode' not in kw): + kw['mmap_mode'] = 'r' + return _np_load(path, *a, **kw) + + +np.load = _np_load_mmap + BENCH = os.path.expanduser('~/wiflow-std-bench') SWEEP = os.path.join(BENCH, 'sweep') sys.path.insert(0, os.path.join(BENCH, 'upstream')) diff --git a/benchmarks/wiflow-std/results/big_windows_mask.npy b/benchmarks/wiflow-std/results/big_windows_mask.npy new file mode 100644 index 00000000..56a70171 Binary files /dev/null and b/benchmarks/wiflow-std/results/big_windows_mask.npy differ diff --git a/benchmarks/wiflow-std/results/nan_windows_mask.npy b/benchmarks/wiflow-std/results/nan_windows_mask.npy new file mode 100644 index 00000000..da5e5244 Binary files /dev/null and b/benchmarks/wiflow-std/results/nan_windows_mask.npy differ diff --git a/benchmarks/wiflow-std/ruvector.db b/benchmarks/wiflow-std/ruvector.db new file mode 100644 index 00000000..1a935589 Binary files /dev/null and b/benchmarks/wiflow-std/ruvector.db differ diff --git a/benchmarks/wiflow-std/static_ptq_bench.py b/benchmarks/wiflow-std/static_ptq_bench.py index 1f74eeb6..131c81d3 100644 --- a/benchmarks/wiflow-std/static_ptq_bench.py +++ b/benchmarks/wiflow-std/static_ptq_bench.py @@ -45,10 +45,11 @@ import numpy as np import torch HERE = os.path.dirname(os.path.abspath(__file__)) -RESULTS = os.path.join(HERE, "results") sys.path.insert(0, HERE) +from _bench_common import RESULTS # noqa: E402 # quantize_bench sets up upstream imports + the np.load mmap patch +# (both via _bench_common.import_upstream) from quantize_bench import build_test_subset # noqa: E402 import quantize_bench as qb # noqa: E402 from eval_ort_accuracy import evaluate_ort # noqa: E402 diff --git a/benchmarks/wiflow-std/tiny_edge_bench.py b/benchmarks/wiflow-std/tiny_edge_bench.py index 3d995d9b..67a24e2c 100644 --- a/benchmarks/wiflow-std/tiny_edge_bench.py +++ b/benchmarks/wiflow-std/tiny_edge_bench.py @@ -38,7 +38,8 @@ machinery on the tiny checkpoint: Usage: PYTHONUTF8=1 .venv/Scripts/python.exe tiny_edge_bench.py \ - [--data-dir ] [--subset 10000] [--calib 500] + [--data-dir ] [--subset 10000] [--calib 512] + (--calib must be a multiple of 64; see step 4 above) Writes/merges into results/edge_optimization.json under key "tiny_variant". """ @@ -211,6 +212,13 @@ def main(): parser.add_argument("--out", default=os.path.join(RESULTS, "edge_optimization.json")) args = parser.parse_args() + if args.calib % 64 != 0: + parser.error( + f"--calib must be a multiple of 64 (got {args.calib}): ORT 1.26's " + f"histogram calibration collector np.asarray()'s the per-batch " + f"maxima and crashes on a ragged final batch (calibration batch " + f"size is 64)") + model = load_tiny_model() info = describe(model) print(f"tiny model: {info['params']:,} params, tcn_groups={info['tcn_groups_per_block']}, " diff --git a/v2/crates/wifi-densepose-cli/src/calibrate_api.rs b/v2/crates/wifi-densepose-cli/src/calibrate_api.rs index a8a57180..8be6fb03 100644 --- a/v2/crates/wifi-densepose-cli/src/calibrate_api.rs +++ b/v2/crates/wifi-densepose-cli/src/calibrate_api.rs @@ -39,7 +39,8 @@ use tokio::sync::{mpsc, oneshot, RwLock}; use tower_http::cors::CorsLayer; use wifi_densepose_calibration::extract::{AnchorFeature, Features}; use wifi_densepose_calibration::{ - AnchorLabel, AnchorQualityGate, AnchorRecorder, MixtureOfSpecialists, SpecialistBank, + AnchorLabel, AnchorQualityGate, AnchorRecorder, MixtureOfSpecialists, NodeGeometry, + SpecialistBank, }; use wifi_densepose_core::types::CsiFrame; use wifi_densepose_signal::{BaselineCalibration, CalibrationRecorder}; @@ -207,6 +208,9 @@ struct RoomEnroll { baseline_id: String, fs_hz: f32, anchors: Vec, + /// Transceiver geometry recorded via `POST /enroll/geometry` (ADR-152 + /// §2.1.1); latest recording wins. Snapshotted into the bank at train time. + geometry: Vec, } /// Result of capturing one anchor (`POST /enroll/anchor`). @@ -299,6 +303,7 @@ fn build_router(state: ApiState) -> Router { .route("/api/v1/room/state", get(room_state)) .route("/api/v1/room/train", post(train_room)) .route("/api/v1/enroll/anchor", post(enroll_anchor)) + .route("/api/v1/enroll/geometry", post(enroll_geometry)) .route("/api/v1/enroll/status", get(enroll_status)) .layer(CorsLayer::permissive()) .with_state(state) @@ -670,8 +675,9 @@ async fn descriptor() -> impl IntoResponse { "GET /api/v1/calibration/result": "last finalized baseline summary", "GET /api/v1/calibration/baselines": "list persisted baseline files", "GET /api/v1/room/state?bank=": "live mixture-of-specialists RoomState over the CSI window", - "POST /api/v1/room/train": "{ room_id, baseline_id, anchors[]? } → train + persist a specialist bank (anchors[] optional if enrolled in-server)", + "POST /api/v1/room/train": "{ room_id, baseline_id, anchors[]?, geometry[]? } → train + persist a specialist bank (anchors[]/geometry[] optional if enrolled in-server)", "POST /api/v1/enroll/anchor": "{ room_id, baseline, label, duration_s? } → capture one guided anchor (blocks for the capture)", + "POST /api/v1/enroll/geometry": "{ room_id, geometry: [NodeGeometry…] } → record transceiver geometry for the room (ADR-152 §2.1.1; latest wins)", "GET /api/v1/enroll/status?room=": "enrollment progress (accepted anchors, next, complete)" } })) @@ -740,11 +746,18 @@ struct TrainRequest { baseline_id: String, #[serde(default)] anchors: Vec, + /// Optional transceiver geometry (ADR-152 §2.1.1). Falls back to the + /// geometry recorded in-server via `POST /enroll/geometry`; absent both, + /// the bank trains geometry-free (valid, but no geometry conditioning). + #[serde(default)] + geometry: Vec, } /// Train a per-room specialist bank and persist it as `/.json` /// (the name `room-state` reads back). Uses the posted `anchors` if present, else /// falls back to the in-server enrollment accumulated via `POST /enroll/anchor`. +/// The enrollment's transceiver-geometry snapshot (posted `geometry` or the +/// `POST /enroll/geometry` record) is threaded into the bank (ADR-152 §2.1.1). async fn train_room(State(st): State, Json(req): Json) -> impl IntoResponse { let (anchors, baseline_id) = if !req.anchors.is_empty() { (req.anchors.clone(), req.baseline_id.clone()) @@ -756,11 +769,25 @@ async fn train_room(State(st): State, Json(req): Json) - } } }; + let geometry = if !req.geometry.is_empty() { + req.geometry.clone() + } else { + st.enroll.read().await.get(&req.room_id).map(|re| re.geometry.clone()).unwrap_or_default() + }; let at = (unix_ms() / 1000) as i64; let bank = match SpecialistBank::train(&req.room_id, &baseline_id, &anchors, at) { Ok(b) => b, Err(e) => return (StatusCode::BAD_REQUEST, Json(serde_json::json!({"error": format!("training failed: {e}")}))).into_response(), }; + let bank = if geometry.is_empty() { + eprintln!( + "[calibrate-serve] no transceiver geometry recorded for room '{}' — bank will not support geometry conditioning (ADR-152 §2.1.2)", + req.room_id + ); + bank + } else { + bank.with_geometry(geometry) + }; let name = sanitize_room_id(&req.room_id); let dir = { st.status.read().await.output_dir.clone() }; let path = format!("{dir}/{name}.json"); @@ -777,10 +804,37 @@ async fn train_room(State(st): State, Json(req): Json) - "bank": name, // pass as ?bank= to /room/state "anchor_count": bank.anchor_count, "specialists": kinds, + "geometry_nodes": bank.geometry.len(), "path": path, }))).into_response() } +/// Body for `POST /api/v1/enroll/geometry`. +#[derive(Deserialize)] +struct EnrollGeometryBody { + room_id: String, + /// Per-node transceiver geometry records (ADR-152 §2.1.1). + geometry: Vec, +} + +/// Record the room's transceiver geometry (ADR-152 §2.1.1) into the in-server +/// enrollment; the next `POST /room/train` snapshots it into the bank. A later +/// POST supersedes an earlier one (latest wins), mirroring +/// `EnrollmentSession::record_geometry`. +async fn enroll_geometry(State(st): State, Json(b): Json) -> impl IntoResponse { + if b.geometry.is_empty() { + return (StatusCode::BAD_REQUEST, Json(serde_json::json!({"error":"geometry must be a non-empty array of NodeGeometry records"}))).into_response(); + } + let nodes = b.geometry.len(); + { + let mut map = st.enroll.write().await; + let re = map.entry(b.room_id.clone()).or_insert_with(RoomEnroll::default); + re.geometry = b.geometry; + } + eprintln!("[calibrate-serve] enroll geometry room={} nodes={nodes}", b.room_id); + (StatusCode::OK, Json(serde_json::json!({"room_id": b.room_id, "geometry_nodes": nodes}))).into_response() +} + /// Body for `POST /api/v1/enroll/anchor`. #[derive(Deserialize)] struct EnrollAnchorBody { @@ -1086,6 +1140,59 @@ mod tests { ); } + /// ADR-152 §2.1.1: geometry threads into the trained bank through both API + /// paths — inline in the train request, or recorded via /enroll/geometry — + /// and a geometry-free train still produces a valid (unconditioned) bank. + #[tokio::test] + async fn train_threads_geometry_into_bank() { + let dir = tempfile::tempdir().unwrap(); + let app = build_router(test_state(dir.path().to_str().unwrap())); + let anchors = r#"[ + {"room_id":"g","label":"empty","features":{"mean":1.0,"variance":1.0,"motion":0.1,"breathing_score":0.0,"breathing_hz":0.0,"heart_score":0.0,"heart_hz":0.0}}, + {"room_id":"g","label":"stand_still","features":{"mean":1.0,"variance":10.0,"motion":0.2,"breathing_score":0.0,"breathing_hz":0.0,"heart_score":0.0,"heart_hz":0.0}} + ]"#; + let load_bank = |name: &str| { + let raw = std::fs::read_to_string(dir.path().join(format!("{name}.json"))).unwrap(); + SpecialistBank::from_json(&raw).unwrap() + }; + + // (1) geometry inline in the train request. + let body = format!( + r#"{{"room_id":"g1","baseline_id":"b","anchors":{anchors}, + "geometry":[{{"node_id":1,"position":{{"x_m":0.0,"y_m":0.0,"z_m":1.0}},"method":"tape-measure"}},{{"node_id":2}}]}}"# + ); + assert_eq!(req(app.clone(), "POST", "/api/v1/room/train", Some(&body)).await, StatusCode::OK); + let bank = load_bank("g1"); + assert_eq!(bank.geometry.len(), 2); + assert_eq!(bank.geometry[0].method, "tape-measure"); + assert_eq!(bank.geometry[1].node_id, 2); + + // (2) geometry recorded via /enroll/geometry; train body omits it. + assert_eq!( + req(app.clone(), "POST", "/api/v1/enroll/geometry", + Some(r#"{"room_id":"g2","geometry":[{"node_id":7,"method":"floor-plan"}]}"#)).await, + StatusCode::OK + ); + let body2 = format!(r#"{{"room_id":"g2","baseline_id":"b","anchors":{anchors}}}"#); + assert_eq!(req(app.clone(), "POST", "/api/v1/room/train", Some(&body2)).await, StatusCode::OK); + let bank2 = load_bank("g2"); + assert_eq!(bank2.geometry.len(), 1); + assert_eq!(bank2.geometry[0].node_id, 7); + + // (3) no geometry anywhere → valid geometry-free bank (note logged). + let body3 = format!(r#"{{"room_id":"g3","baseline_id":"b","anchors":{anchors}}}"#); + assert_eq!(req(app.clone(), "POST", "/api/v1/room/train", Some(&body3)).await, StatusCode::OK); + let bank3 = load_bank("g3"); + assert!(bank3.geometry.is_empty()); + assert!(bank3.presence.is_some(), "bank still trains without geometry"); + + // (4) empty geometry array is rejected. + assert_eq!( + req(app, "POST", "/api/v1/enroll/geometry", Some(r#"{"room_id":"g4","geometry":[]}"#)).await, + StatusCode::BAD_REQUEST + ); + } + #[tokio::test] async fn enroll_status_empty_and_bad_label() { let dir = tempfile::tempdir().unwrap(); diff --git a/v2/crates/wifi-densepose-cli/src/room.rs b/v2/crates/wifi-densepose-cli/src/room.rs index 1f010eae..bf60e4b6 100644 --- a/v2/crates/wifi-densepose-cli/src/room.rs +++ b/v2/crates/wifi-densepose-cli/src/room.rs @@ -11,7 +11,7 @@ use std::time::{Duration, Instant, SystemTime, UNIX_EPOCH}; use tokio::net::UdpSocket; use wifi_densepose_calibration::{ Anchor, AnchorLabel, AnchorQualityGate, AnchorRecorder, EnrollmentEvent, EnrollmentSession, - MixtureOfSpecialists, MultiNodeMixture, SpecialistBank, + MixtureOfSpecialists, MultiNodeMixture, NodeGeometry, SpecialistBank, }; use wifi_densepose_calibration::extract::{AnchorFeature, Features}; use wifi_densepose_core::types::CsiFrame; @@ -226,20 +226,50 @@ pub struct TrainRoomArgs { /// Output specialist-bank file. #[arg(long, default_value = "./room-bank.json")] pub output: String, + /// Optional transceiver-geometry file: a JSON array of `NodeGeometry` + /// records (ADR-152 §2.1.1). Recorded into the enrollment session before + /// training so the bank carries the layout it was trained under. + #[arg(long)] + pub geometry: Option, } /// Execute `train-room`. +/// +/// If the enrollment session carries a transceiver-geometry snapshot (recorded +/// at enroll time or supplied here via `--geometry`), it is threaded into the +/// bank (ADR-152 §2.1.1); a geometry-free enrollment still trains a valid bank. pub async fn train_room(args: TrainRoomArgs) -> Result<()> { let raw = std::fs::read_to_string(&args.enrollment) .map_err(|e| anyhow::anyhow!("cannot read {}: {e} — run `enroll` first", args.enrollment))?; - let data: EnrollmentData = + let mut data: EnrollmentData = serde_json::from_str(&raw).map_err(|e| anyhow::anyhow!("invalid enrollment: {e}"))?; if data.anchors.is_empty() { bail!("no accepted anchors in {} — re-run enroll", args.enrollment); } - let bank = SpecialistBank::train(&data.room_id, &data.baseline_id, &data.anchors, now_unix()) + if let Some(path) = &args.geometry { + let graw = std::fs::read_to_string(path) + .map_err(|e| anyhow::anyhow!("cannot read geometry {path}: {e}"))?; + let geometry: Vec = serde_json::from_str(&graw).map_err(|e| { + anyhow::anyhow!("invalid geometry {path}: {e} (expected a JSON array of NodeGeometry records)") + })?; + data.session.record_geometry(geometry, now_unix()); + } + + let mut bank = SpecialistBank::train(&data.room_id, &data.baseline_id, &data.anchors, now_unix()) .map_err(|e| anyhow::anyhow!("training failed: {e}"))?; + match data.session.geometry() { + Some(g) if !g.is_empty() => { + bank = bank.with_geometry(g.to_vec()); + eprintln!( + "[train-room] geometry: {} node(s) snapshotted into the bank (ADR-152 §2.1.1)", + bank.geometry.len() + ); + } + _ => eprintln!( + "[train-room] no transceiver geometry recorded — bank will not support geometry conditioning (ADR-152 §2.1.2)" + ), + } std::fs::write(&args.output, bank.to_json().map_err(|e| anyhow::anyhow!("{e}"))?) .map_err(|e| anyhow::anyhow!("cannot write {}: {e}", args.output))?; @@ -456,3 +486,141 @@ async fn room_watch_multi(args: RoomWatchArgs) -> Result<()> { } Ok(()) } + +#[cfg(test)] +mod tests { + use super::*; + + fn feature(label: AnchorLabel, variance: f32, motion: f32) -> AnchorFeature { + AnchorFeature { + room_id: "t".into(), + label, + features: Features { + mean: 1.0, + variance, + motion, + breathing_score: 0.0, + breathing_hz: 0.0, + heart_score: 0.0, + heart_hz: 0.0, + }, + } + } + + /// Write a minimal valid enrollment file (two anchors, no geometry event). + fn write_enrollment(dir: &std::path::Path) -> String { + let data = EnrollmentData { + room_id: "t".into(), + baseline_id: "base-1".into(), + fs_hz: 15.0, + anchors: vec![ + feature(AnchorLabel::Empty, 1.0, 0.1), + feature(AnchorLabel::StandStill, 10.0, 0.2), + ], + session: EnrollmentSession::new("t", "base-1", 1000), + }; + let path = dir.join("enrollment.json"); + std::fs::write(&path, serde_json::to_string(&data).unwrap()).unwrap(); + path.to_string_lossy().into_owned() + } + + fn trained_bank(out: &std::path::Path) -> SpecialistBank { + SpecialistBank::from_json(&std::fs::read_to_string(out).unwrap()).unwrap() + } + + /// ADR-152 §2.1.1: `--geometry` records into the session and the bank + /// snapshots it — enrollment geometry reaches the trained bank. + #[tokio::test] + async fn train_room_threads_geometry_when_provided() { + let dir = tempfile::tempdir().unwrap(); + let enrollment = write_enrollment(dir.path()); + let geometry = vec![ + NodeGeometry::new(1, "tape-measure").with_position(0.0, 0.0, 1.0), + NodeGeometry::unknown(2), + ]; + let gpath = dir.path().join("geometry.json"); + std::fs::write(&gpath, serde_json::to_string(&geometry).unwrap()).unwrap(); + let out = dir.path().join("bank.json"); + + train_room(TrainRoomArgs { + enrollment, + output: out.to_string_lossy().into_owned(), + geometry: Some(gpath.to_string_lossy().into_owned()), + }) + .await + .unwrap(); + + assert_eq!(trained_bank(&out).geometry, geometry); + } + + /// A geometry-free enrollment still trains a valid bank (optional by + /// design) — it just carries no snapshot. + #[tokio::test] + async fn train_room_without_geometry_yields_geometry_free_bank() { + let dir = tempfile::tempdir().unwrap(); + let enrollment = write_enrollment(dir.path()); + let out = dir.path().join("bank.json"); + + train_room(TrainRoomArgs { + enrollment, + output: out.to_string_lossy().into_owned(), + geometry: None, + }) + .await + .unwrap(); + + let bank = trained_bank(&out); + assert!(bank.geometry.is_empty()); + assert!(bank.presence.is_some(), "bank still trains without geometry"); + } + + /// Geometry recorded at enroll time (in the session event log) is picked up + /// without the `--geometry` flag. + #[tokio::test] + async fn train_room_uses_session_geometry() { + let dir = tempfile::tempdir().unwrap(); + let geometry = vec![NodeGeometry::new(3, "floor-plan").with_position(1.0, 2.0, 1.5)]; + let mut session = EnrollmentSession::new("t", "base-1", 1000); + session.record_geometry(geometry.clone(), 1000); + let data = EnrollmentData { + room_id: "t".into(), + baseline_id: "base-1".into(), + fs_hz: 15.0, + anchors: vec![ + feature(AnchorLabel::Empty, 1.0, 0.1), + feature(AnchorLabel::StandStill, 10.0, 0.2), + ], + session, + }; + let epath = dir.path().join("enrollment.json"); + std::fs::write(&epath, serde_json::to_string(&data).unwrap()).unwrap(); + let out = dir.path().join("bank.json"); + + train_room(TrainRoomArgs { + enrollment: epath.to_string_lossy().into_owned(), + output: out.to_string_lossy().into_owned(), + geometry: None, + }) + .await + .unwrap(); + + assert_eq!(trained_bank(&out).geometry, geometry); + } + + #[tokio::test] + async fn train_room_rejects_invalid_geometry_file() { + let dir = tempfile::tempdir().unwrap(); + let enrollment = write_enrollment(dir.path()); + let gpath = dir.path().join("geometry.json"); + std::fs::write(&gpath, r#"{"not":"an array"}"#).unwrap(); + + let err = train_room(TrainRoomArgs { + enrollment, + output: dir.path().join("bank.json").to_string_lossy().into_owned(), + geometry: Some(gpath.to_string_lossy().into_owned()), + }) + .await + .unwrap_err(); + assert!(err.to_string().contains("invalid geometry"), "{err}"); + } +} diff --git a/v2/crates/wifi-densepose-hardware/src/ieee80211bf/events.rs b/v2/crates/wifi-densepose-hardware/src/ieee80211bf/events.rs new file mode 100644 index 00000000..c366e52d --- /dev/null +++ b/v2/crates/wifi-densepose-hardware/src/ieee80211bf/events.rs @@ -0,0 +1,108 @@ +//! Session FSM I/O types for the 802.11bf sensing model: events in +//! ([`SessionEvent`]), actions out ([`Action`]), close reasons, static +//! configuration, and the state enum. +//! +//! Split from [`super::session`] to keep each file under the ADR-153 +//! 500-line maintainability cap; the canonical public path re-exports +//! these from [`super::session`]. + +use super::messages::{ + CsiReportPayload, SbpRequest, SbpResponse, SbpStatus, SensingMeasurementInstance, + SensingMeasurementReport, SensingMeasurementSetupRequest, SensingMeasurementSetupResponse, + SensingSessionTermination, TerminationReason, +}; +use super::types::{MeasurementInstanceId, SensingCapabilities, SetupStatus, SpecProfile}; + +/// Session FSM states. +#[derive(Debug, Clone, Copy, PartialEq, Eq)] +pub enum SessionState { + Idle, + SetupNegotiating, + Active, + Terminating, +} + +/// Inputs to the session FSM. `Start*` are local commands; `*Received` are +/// frames from the peer; `Timeout`/`InstanceElapsed` are scheduler ticks. +#[derive(Debug, Clone, PartialEq)] +pub enum SessionEvent { + /// Local command (initiator): begin setup negotiation. + StartSetup(SensingMeasurementSetupRequest), + /// Local command (initiator): request sensing-by-proxy from an AP. + StartSbp(SbpRequest), + SetupRequestReceived(SensingMeasurementSetupRequest), + SetupResponseReceived(SensingMeasurementSetupResponse), + SbpRequestReceived(SbpRequest), + SbpResponseReceived(SbpResponse), + /// Scheduler tick: the negotiated periodicity elapsed (the + /// measurement-driving endpoint — initiator or SBP proxy — emits the + /// next measurement-instance trigger). + InstanceElapsed, + /// A sensing receiver captured a measurement for an instance (payload is + /// fed by the transport/bridge — see `OpportunisticCsiBridge`). + MeasurementCaptured { + instance_id: MeasurementInstanceId, + payload: CsiReportPayload, + }, + ReportReceived(SensingMeasurementReport), + /// Generic timeout tick for the current state. + Timeout, + /// Local command: terminate the session. + Terminate(TerminationReason), + TerminationReceived(SensingSessionTermination), +} + +/// Outputs of the session FSM. `Send*`/`TriggerInstance`/`RelaySbpReport` +/// go to the transport; `DeliverReport`/`SessionClosed` go to the local +/// consumer. +#[derive(Debug, Clone, PartialEq)] +pub enum Action { + SendSetupRequest(SensingMeasurementSetupRequest), + SendSetupResponse(SensingMeasurementSetupResponse), + SendSbpRequest(SbpRequest), + SendSbpResponse(SbpResponse), + TriggerInstance(SensingMeasurementInstance), + SendReport(SensingMeasurementReport), + DeliverReport(SensingMeasurementReport), + /// SBP proxy mode: forward a report received from the sensing responder + /// to the SBP client. The transport maps this to a frame toward the + /// client (`SensingFrame::SbpReport`), distinct from `SendReport`, + /// which travels toward the sensing initiator. + RelaySbpReport(SensingMeasurementReport), + SendTermination(SensingSessionTermination), + SessionClosed(CloseReason), +} + +/// Why a session returned to Idle. +#[derive(Debug, Clone, Copy, PartialEq)] +pub enum CloseReason { + SetupRejected(SetupStatus), + SbpRejected(SbpStatus), + Terminated(TerminationReason), + /// Terminating-state quiescence completed (no peer echo required). + Completed, +} + +/// Static configuration for a sensing session. +#[derive(Debug, Clone, PartialEq)] +pub struct SessionConfig { + /// Spec profile this endpoint advertises/accepts. + pub profile: SpecProfile, + /// Capability set used to evaluate inbound setups. + pub capabilities: SensingCapabilities, + /// Consecutive negotiation timeouts before aborting to Idle. + pub max_setup_timeouts: u8, + /// Consecutive missed instances (Active timeouts) before terminating. + pub max_missed_instances: u8, +} + +impl Default for SessionConfig { + fn default() -> Self { + Self { + profile: SpecProfile::Ieee80211Bf2025, + capabilities: SensingCapabilities::sim_full(), + max_setup_timeouts: 3, + max_missed_instances: 5, + } + } +} diff --git a/v2/crates/wifi-densepose-hardware/src/ieee80211bf/messages.rs b/v2/crates/wifi-densepose-hardware/src/ieee80211bf/messages.rs index 4b4cab98..b78231ad 100644 --- a/v2/crates/wifi-densepose-hardware/src/ieee80211bf/messages.rs +++ b/v2/crates/wifi-densepose-hardware/src/ieee80211bf/messages.rs @@ -125,12 +125,38 @@ impl SbpRequest { } /// Status carried by an SBP response. +/// +/// Mirrors [`SetupStatus`] 1:1 (see the `From` impl): an SBP +/// request is validated through the same chain as a direct setup, so every +/// rejection class must survive the proxy translation. +/// `RejectedNotSupported` additionally covers a proxy that lacks the SBP +/// capability itself. #[derive(Debug, Clone, Copy, PartialEq, Eq, Serialize, Deserialize)] pub enum SbpStatus { Accepted, RejectedNotSupported, RejectedUnsupportedParams, + RejectedSetupIdCollision, + RejectedIncompatibleProfile, RejectedByPolicy, + RejectedCapacity, +} + +impl From for SbpStatus { + /// 1:1 mapping from the direct-setup status space, keeping the SBP path + /// on the single `evaluate_setup` validation chain (no SBP-only policy + /// drift or bypass). + fn from(status: SetupStatus) -> Self { + match status { + SetupStatus::Accepted => SbpStatus::Accepted, + SetupStatus::RejectedNotSupported => SbpStatus::RejectedNotSupported, + SetupStatus::RejectedUnsupportedParams => SbpStatus::RejectedUnsupportedParams, + SetupStatus::RejectedSetupIdCollision => SbpStatus::RejectedSetupIdCollision, + SetupStatus::RejectedIncompatibleProfile => SbpStatus::RejectedIncompatibleProfile, + SetupStatus::RejectedByPolicy => SbpStatus::RejectedByPolicy, + SetupStatus::RejectedCapacity => SbpStatus::RejectedCapacity, + } + } } /// Sensing-by-Proxy (SBP) response (proxy AP → requesting STA). diff --git a/v2/crates/wifi-densepose-hardware/src/ieee80211bf/mod.rs b/v2/crates/wifi-densepose-hardware/src/ieee80211bf/mod.rs index 9845dee4..0e23bdc5 100644 --- a/v2/crates/wifi-densepose-hardware/src/ieee80211bf/mod.rs +++ b/v2/crates/wifi-densepose-hardware/src/ieee80211bf/mod.rs @@ -33,12 +33,17 @@ //! measurement instance, CSI-variant report, SBP exchange, termination). //! - [`session`] — deterministic event-driven session FSM: //! `Idle → SetupNegotiating → Active → Terminating → Idle`, with explicit -//! rejection paths and timeout handling. No async, no clocks. +//! rejection paths, timeout handling, single-role enforcement, and the +//! first-class SBP proxy mode. No async, no clocks. +//! - [`events`] — the FSM I/O types ([`events::SessionEvent`], +//! [`events::Action`], close reasons, configuration), re-exported via +//! [`session`]. //! - [`table`] — responder-side setup registry (setup-ID collision and -//! capacity rejection paths). +//! capacity rejection paths, for direct setups and SBP alike). //! - [`transport`] — the [`transport::SensingTransport`] seam, the //! [`transport::SimTransport`] test double, and the ESP32 bridge. +pub mod events; pub mod messages; pub mod session; pub mod table; @@ -68,4 +73,6 @@ mod tests; #[cfg(test)] mod tests_fsm; #[cfg(test)] +mod tests_sbp; +#[cfg(test)] mod testutil; diff --git a/v2/crates/wifi-densepose-hardware/src/ieee80211bf/session.rs b/v2/crates/wifi-densepose-hardware/src/ieee80211bf/session.rs index 4b6f92a7..d6376e9d 100644 --- a/v2/crates/wifi-densepose-hardware/src/ieee80211bf/session.rs +++ b/v2/crates/wifi-densepose-hardware/src/ieee80211bf/session.rs @@ -12,103 +12,38 @@ //! (responder responds with a rejected setup status), setup-ID collision //! ([`super::table::SessionTable`]), and negotiation timeout (typed //! [`BfError::NegotiationTimeout`] + reset to Idle). +//! +//! **Single-role design:** a session is constructed as initiator or responder +//! and keeps that role for its whole lifetime. An initiator-role session +//! receiving a peer's setup or SBP request answers `RejectedNotSupported` +//! instead of accepting — a peer must never be able to hijack a session out +//! of its configured role. Endpoints that play both roles run one session per +//! role (or a [`super::table::SessionTable`] for the responder side). +//! +//! **SBP proxy mode:** a responder session that accepts an SBP request +//! becomes a first-class proxy ([`SensingSession::is_sbp_proxy`]): it drives +//! the standard initiator path toward the actual sensing responder — +//! including re-triggering measurement instances on +//! [`SessionEvent::InstanceElapsed`] — and relays every received report to +//! the SBP client via [`Action::RelaySbpReport`], in addition to local +//! [`Action::DeliverReport`] delivery. +//! +//! Local `Start*` commands issued outside Idle are caller bugs and surface +//! as typed [`BfError::InvalidStateForCommand`]; genuinely ignorable stray +//! frames/ticks remain silent no-ops. The FSM I/O types live in +//! [`super::events`] and are re-exported here. use super::messages::{ - CsiReportPayload, SbpRequest, SbpResponse, SbpStatus, SensingMeasurementInstance, - SensingMeasurementReport, SensingMeasurementSetupRequest, SensingMeasurementSetupResponse, - SensingSessionTermination, TerminationReason, + SbpRequest, SbpResponse, SbpStatus, SensingMeasurementInstance, SensingMeasurementReport, + SensingMeasurementSetupRequest, SensingMeasurementSetupResponse, SensingSessionTermination, + TerminationReason, }; use super::types::{ BfError, MeasurementInstanceId, MeasurementSetupId, MeasurementSetupParams, ReportingConfig, - SensingCapabilities, SensingRole, SetupStatus, SpecProfile, + SensingRole, SetupStatus, }; -/// Session FSM states. -#[derive(Debug, Clone, Copy, PartialEq, Eq)] -pub enum SessionState { - Idle, - SetupNegotiating, - Active, - Terminating, -} - -/// Inputs to the session FSM. `Start*` are local commands; `*Received` are -/// frames from the peer; `Timeout`/`InstanceElapsed` are scheduler ticks. -#[derive(Debug, Clone, PartialEq)] -pub enum SessionEvent { - /// Local command (initiator): begin setup negotiation. - StartSetup(SensingMeasurementSetupRequest), - /// Local command (initiator): request sensing-by-proxy from an AP. - StartSbp(SbpRequest), - SetupRequestReceived(SensingMeasurementSetupRequest), - SetupResponseReceived(SensingMeasurementSetupResponse), - SbpRequestReceived(SbpRequest), - SbpResponseReceived(SbpResponse), - /// Scheduler tick: the negotiated periodicity elapsed (initiator emits - /// the next measurement-instance trigger). - InstanceElapsed, - /// A sensing receiver captured a measurement for an instance (payload is - /// fed by the transport/bridge — see `OpportunisticCsiBridge`). - MeasurementCaptured { - instance_id: MeasurementInstanceId, - payload: CsiReportPayload, - }, - ReportReceived(SensingMeasurementReport), - /// Generic timeout tick for the current state. - Timeout, - /// Local command: terminate the session. - Terminate(TerminationReason), - TerminationReceived(SensingSessionTermination), -} - -/// Outputs of the session FSM. `Send*`/`TriggerInstance` go to the transport; -/// `DeliverReport`/`SessionClosed` go to the local consumer. -#[derive(Debug, Clone, PartialEq)] -pub enum Action { - SendSetupRequest(SensingMeasurementSetupRequest), - SendSetupResponse(SensingMeasurementSetupResponse), - SendSbpRequest(SbpRequest), - SendSbpResponse(SbpResponse), - TriggerInstance(SensingMeasurementInstance), - SendReport(SensingMeasurementReport), - DeliverReport(SensingMeasurementReport), - SendTermination(SensingSessionTermination), - SessionClosed(CloseReason), -} - -/// Why a session returned to Idle. -#[derive(Debug, Clone, Copy, PartialEq)] -pub enum CloseReason { - SetupRejected(SetupStatus), - SbpRejected(SbpStatus), - Terminated(TerminationReason), - /// Terminating-state quiescence completed (no peer echo required). - Completed, -} - -/// Static configuration for a sensing session. -#[derive(Debug, Clone, PartialEq)] -pub struct SessionConfig { - /// Spec profile this endpoint advertises/accepts. - pub profile: SpecProfile, - /// Capability set used to evaluate inbound setups. - pub capabilities: SensingCapabilities, - /// Consecutive negotiation timeouts before aborting to Idle. - pub max_setup_timeouts: u8, - /// Consecutive missed instances (Active timeouts) before terminating. - pub max_missed_instances: u8, -} - -impl Default for SessionConfig { - fn default() -> Self { - Self { - profile: SpecProfile::Ieee80211Bf2025, - capabilities: SensingCapabilities::sim_full(), - max_setup_timeouts: 3, - max_missed_instances: 5, - } - } -} +pub use super::events::{Action, CloseReason, SessionConfig, SessionEvent, SessionState}; /// One sensing session (one measurement setup) on one endpoint. #[derive(Debug, Clone)] @@ -122,6 +57,10 @@ pub struct SensingSession { setup: Option<(MeasurementSetupId, MeasurementSetupParams)>, /// True when this session awaits proxied sensing (SBP client). sbp_client: bool, + /// True when this responder-role session proxies sensing for an SBP + /// client: it drives the initiator path toward the sensing responder + /// and relays received reports back to the client. + sbp_proxy: bool, setup_timeouts: u8, missed_instances: u8, instance_counter: u32, @@ -146,6 +85,7 @@ impl SensingSession { pending_request: None, setup: None, sbp_client: false, + sbp_proxy: false, setup_timeouts: 0, missed_instances: 0, instance_counter: 0, @@ -161,13 +101,20 @@ impl SensingSession { self.role } + /// True when this session is acting as an SBP proxy (accepted via + /// [`SessionEvent::SbpRequestReceived`]); cleared on reset to Idle. + pub fn is_sbp_proxy(&self) -> bool { + self.sbp_proxy + } + pub fn setup_id(&self) -> Option { self.setup.as_ref().map(|(id, _)| *id) } /// Drive the FSM with one event. Protocol-level rejections surface as - /// `Ok` actions (responses to the peer); malformed/adversarial input and - /// negotiation timeout surface as typed `Err` (never a panic). + /// `Ok` actions (responses to the peer); malformed/adversarial input, + /// out-of-state local commands, and negotiation timeout surface as typed + /// `Err` (never a panic). pub fn handle(&mut self, event: SessionEvent) -> Result, BfError> { match self.state { SessionState::Idle => self.handle_idle(event), @@ -212,6 +159,13 @@ impl SensingSession { status, }) }; + // Single-role design (module docs): an initiator-role + // session never accepts a peer's setup request — accepting + // here would let a peer hijack the session into the + // responder path. + if self.role != SensingRole::Responder { + return Ok(vec![response(SetupStatus::RejectedNotSupported)]); + } match self.evaluate_setup(&req) { SetupStatus::Accepted => { self.setup = Some((req.setup_id, req.params.clone())); @@ -223,7 +177,16 @@ impl SensingSession { status => Ok(vec![response(status)]), } } - SessionEvent::SbpRequestReceived(sbp) => Ok(self.handle_sbp_request(sbp)), + SessionEvent::SbpRequestReceived(sbp) => { + // Single-role design: only responder-role sessions proxy. + if self.role != SensingRole::Responder { + return Ok(vec![Action::SendSbpResponse(SbpResponse { + proxy_setup_id: sbp.proxy_setup_id, + status: SbpStatus::RejectedNotSupported, + })]); + } + Ok(self.handle_sbp_request(sbp)) + } // Stray frames/ticks in Idle are ignored, not errors. _ => Ok(vec![]), } @@ -232,6 +195,12 @@ impl SensingSession { /// SBP proxy path: accept the request, then run the *standard initiator /// path* toward the actual sensing responder. No direct sensor coupling — /// the proxied setup is an ordinary `SendSetupRequest` on the transport. + /// + /// Validation is the single [`Self::evaluate_setup`] chain: the proxied + /// setup request is built first and evaluated exactly as a direct setup + /// would be, with the resulting [`SetupStatus`] mapped 1:1 onto + /// [`SbpStatus`] — no SBP-only re-implementation that could drift from + /// (or bypass) the setup policy. fn handle_sbp_request(&mut self, sbp: SbpRequest) -> Vec { let respond = |status| { Action::SendSbpResponse(SbpResponse { @@ -239,29 +208,22 @@ impl SensingSession { status, }) }; + // SBP-specific capability gate; everything else is the setup chain. if !self.config.capabilities.sensing_by_proxy { return vec![respond(SbpStatus::RejectedNotSupported)]; } - if !self.config.profile.accepts(&sbp.profile) { - return vec![respond(SbpStatus::RejectedUnsupportedParams)]; - } - match sbp.validate() { - Err(BfError::SensingDisabledByPolicy) => { - return vec![respond(SbpStatus::RejectedByPolicy)]; - } - Err(_) => return vec![respond(SbpStatus::RejectedUnsupportedParams)], - Ok(()) => {} - } - if self.config.capabilities.evaluate(&sbp.params).is_err() { - return vec![respond(SbpStatus::RejectedUnsupportedParams)]; - } let req = SensingMeasurementSetupRequest { profile: sbp.profile.clone(), setup_id: sbp.proxy_setup_id, params: sbp.params.clone(), }; + match self.evaluate_setup(&req) { + SetupStatus::Accepted => {} + status => return vec![respond(SbpStatus::from(status))], + } self.setup = Some((req.setup_id, req.params.clone())); self.pending_request = Some(req.clone()); + self.sbp_proxy = true; self.setup_timeouts = 0; self.state = SessionState::SetupNegotiating; vec![respond(SbpStatus::Accepted), Action::SendSetupRequest(req)] @@ -362,6 +324,13 @@ impl SensingSession { term.reason, ))]) } + // Local Start* outside Idle is a caller bug — typed error. + SessionEvent::StartSetup(_) | SessionEvent::StartSbp(_) => { + Err(BfError::InvalidStateForCommand { + state: "SetupNegotiating", + }) + } + // Genuinely ignorable stray frames/ticks are no-ops. _ => Ok(vec![]), } } @@ -369,7 +338,13 @@ impl SensingSession { fn handle_active(&mut self, event: SessionEvent) -> Result, BfError> { match event { SessionEvent::InstanceElapsed => { - if self.role == SensingRole::Initiator && !self.sbp_client { + // The measurement-driving endpoint re-triggers here: the + // initiator, or an SBP proxy running the initiator path + // toward the sensing responder. SBP *clients* only consume + // proxied reports and never trigger instances. + let drives_instances = + (self.role == SensingRole::Initiator || self.sbp_proxy) && !self.sbp_client; + if drives_instances { match self.next_instance_record() { Some(instance) => Ok(vec![Action::TriggerInstance(instance)]), None => Ok(vec![]), @@ -387,6 +362,11 @@ impl SensingSession { Some((id, p)) => (*id, p.clone()), None => return Ok(vec![]), }; + // A successful capture means this instance was not missed — + // the missed-instance budget counts *consecutive* misses, + // so it resets here even when threshold-based reporting + // suppresses the report below. + self.missed_instances = 0; let mean = payload.mean_amplitude(); let should_report = match params.reporting { ReportingConfig::EveryInstance => true, @@ -418,7 +398,16 @@ impl SensingSession { }); } self.missed_instances = 0; - Ok(vec![Action::DeliverReport(report)]) + if self.sbp_proxy { + // Proxy mode: deliver to the local consumer *and* relay + // toward the SBP client on the transport. + Ok(vec![ + Action::DeliverReport(report.clone()), + Action::RelaySbpReport(report), + ]) + } else { + Ok(vec![Action::DeliverReport(report)]) + } } SessionEvent::Timeout => { self.missed_instances = self.missed_instances.saturating_add(1); @@ -439,6 +428,12 @@ impl SensingSession { term.reason, ))]) } + // Local Start* outside Idle is a caller bug — typed error. + SessionEvent::StartSetup(_) | SessionEvent::StartSbp(_) => { + Err(BfError::InvalidStateForCommand { state: "Active" }) + } + // Genuinely ignorable stray frames (duplicate setup/SBP traffic) + // are no-ops. _ => Ok(vec![]), } } @@ -456,6 +451,12 @@ impl SensingSession { self.reset(); Ok(vec![Action::SessionClosed(CloseReason::Completed)]) } + // Local Start* outside Idle is a caller bug — typed error. + SessionEvent::StartSetup(_) | SessionEvent::StartSbp(_) => { + Err(BfError::InvalidStateForCommand { + state: "Terminating", + }) + } _ => Ok(vec![]), } } @@ -489,6 +490,7 @@ impl SensingSession { self.pending_request = None; self.setup = None; self.sbp_client = false; + self.sbp_proxy = false; self.setup_timeouts = 0; self.missed_instances = 0; self.instance_counter = 0; diff --git a/v2/crates/wifi-densepose-hardware/src/ieee80211bf/table.rs b/v2/crates/wifi-densepose-hardware/src/ieee80211bf/table.rs index 8b9967c4..064145b5 100644 --- a/v2/crates/wifi-densepose-hardware/src/ieee80211bf/table.rs +++ b/v2/crates/wifi-densepose-hardware/src/ieee80211bf/table.rs @@ -1,10 +1,16 @@ //! Responder-side setup registry for the 802.11bf sensing model — enforces //! the setup-ID-collision and capacity rejection paths a single session //! cannot see on its own (ADR-153 acceptance: duplicate setup ID rejected). +//! Both entry points — direct setups ([`SessionTable::handle_setup_request`]) +//! and sensing-by-proxy ([`SessionTable::handle_sbp_request`]) — share the +//! same guards and the same per-setup session storage. use std::collections::BTreeMap; -use super::messages::{SensingMeasurementSetupRequest, SensingMeasurementSetupResponse}; +use super::messages::{ + SbpRequest, SbpResponse, SbpStatus, SensingMeasurementSetupRequest, + SensingMeasurementSetupResponse, +}; use super::session::{Action, SensingSession, SessionConfig, SessionEvent, SessionState}; use super::types::{BfError, MeasurementSetupId, SetupStatus}; @@ -16,6 +22,9 @@ use super::types::{BfError, MeasurementSetupId, SetupStatus}; pub struct SessionTable { config: SessionConfig, sessions: BTreeMap, + /// Events dropped because no session owned the setup ID (see + /// [`Self::handle_for`]). + unknown_setup_drops: u64, } impl SessionTable { @@ -23,6 +32,7 @@ impl SessionTable { Self { config, sessions: BTreeMap::new(), + unknown_setup_drops: 0, } } @@ -38,6 +48,13 @@ impl SessionTable { self.sessions.get(&setup_id.value()) } + /// Count of events dropped by [`Self::handle_for`] because the setup ID + /// was unknown — lets an AP spot peers addressing setups it never + /// accepted without turning stray frames into errors. + pub fn unknown_setup_drops(&self) -> u64 { + self.unknown_setup_drops + } + /// Route an inbound setup request, rejecting setup-ID collisions and /// capacity overruns before delegating to a responder session. pub fn handle_setup_request( @@ -49,12 +66,10 @@ impl SessionTable { SensingMeasurementSetupResponse { setup_id, status }, )]) }; - if let Some(existing) = self.sessions.get(&req.setup_id.value()) { - if existing.state() != SessionState::Idle { - return reject(req.setup_id, SetupStatus::RejectedSetupIdCollision); - } + if self.is_collision(req.setup_id) { + return reject(req.setup_id, SetupStatus::RejectedSetupIdCollision); } - if self.active_setups() >= self.config.capabilities.max_active_setups as usize { + if self.at_capacity() { return reject(req.setup_id, SetupStatus::RejectedCapacity); } let key = req.setup_id.value(); @@ -64,8 +79,35 @@ impl SessionTable { Ok(actions) } - /// Route any other event to the session owning `setup_id` (no-op if the - /// setup is unknown — stray frames are ignored, not errors). + /// Route an inbound SBP request, rejecting proxy-setup-ID collisions and + /// capacity overruns before delegating to a (new) proxy session — the + /// SBP mirror of [`Self::handle_setup_request`], so a table-driven AP + /// accepts SBP end-to-end instead of silently dropping it. + pub fn handle_sbp_request(&mut self, sbp: SbpRequest) -> Result, BfError> { + let reject = |proxy_setup_id, status| { + Ok(vec![Action::SendSbpResponse(SbpResponse { + proxy_setup_id, + status, + })]) + }; + if self.is_collision(sbp.proxy_setup_id) { + return reject(sbp.proxy_setup_id, SbpStatus::RejectedSetupIdCollision); + } + if self.at_capacity() { + return reject(sbp.proxy_setup_id, SbpStatus::RejectedCapacity); + } + let key = sbp.proxy_setup_id.value(); + let mut session = SensingSession::new_responder(self.config.clone()); + let actions = session.handle(SessionEvent::SbpRequestReceived(sbp))?; + self.sessions.insert(key, session); + Ok(actions) + } + + /// Route any other event to the session owning `setup_id`. + /// + /// Frames addressing an unknown setup are dropped *by design* (stray + /// frames are ignored, not errors), but the drop is observable through + /// [`Self::unknown_setup_drops`]. pub fn handle_for( &mut self, setup_id: MeasurementSetupId, @@ -73,7 +115,22 @@ impl SessionTable { ) -> Result, BfError> { match self.sessions.get_mut(&setup_id.value()) { Some(session) => session.handle(event), - None => Ok(vec![]), + None => { + self.unknown_setup_drops = self.unknown_setup_drops.saturating_add(1); + Ok(vec![]) + } } } + + /// A non-Idle session already owns this setup ID. + fn is_collision(&self, setup_id: MeasurementSetupId) -> bool { + self.sessions + .get(&setup_id.value()) + .is_some_and(|existing| existing.state() != SessionState::Idle) + } + + /// The active-setup budget is exhausted. + fn at_capacity(&self) -> bool { + self.active_setups() >= self.config.capabilities.max_active_setups as usize + } } diff --git a/v2/crates/wifi-densepose-hardware/src/ieee80211bf/tests_fsm.rs b/v2/crates/wifi-densepose-hardware/src/ieee80211bf/tests_fsm.rs index 443d5397..ae698d63 100644 --- a/v2/crates/wifi-densepose-hardware/src/ieee80211bf/tests_fsm.rs +++ b/v2/crates/wifi-densepose-hardware/src/ieee80211bf/tests_fsm.rs @@ -1,7 +1,8 @@ //! ADR-153 acceptance tests — session FSM full cycle, rejection paths, -//! timeout handling, threshold-based reporting, SBP flows, and adversarial -//! no-panic coverage. Type/serde/transport/bridge tests live in -//! [`super::tests`]. All tests are hardware-free (simulation only). +//! timeout handling, threshold-based reporting, single-role enforcement, +//! and adversarial no-panic coverage. SBP flows live in [`super::tests_sbp`]; +//! type/serde/transport/bridge tests in [`super::tests`]. All tests are +//! hardware-free (simulation only). use super::messages::*; use super::session::{ @@ -300,77 +301,72 @@ fn threshold_report_emitted_only_when_threshold_crossed() { assert!(responder.handle(capture(125.0)).unwrap().is_empty()); } -// ---------- SBP ---------- +// ---------- consecutive missed-instance semantics ---------- #[test] -fn sbp_proxy_request_maps_to_standard_responder_path() { - // Proxy AP: accepts the SBP request and initiates an ordinary setup - // toward the sensing responder — no direct sensor coupling. - let mut proxy = SensingSession::new_responder(SessionConfig::default()); - let sbp = SbpRequest { - profile: SpecProfile::Ieee80211Bf2025, - proxy_setup_id: MeasurementSetupId::new(11).unwrap(), - params: params(), +fn missed_instance_budget_is_consecutive_not_cumulative() { + // Review finding 2: a successful measurement must reset the + // missed-instance counter — `max_missed_instances` bounds *consecutive* + // misses (as documented on SessionConfig), not cumulative ones. + let mut responder = SensingSession::new_responder(SessionConfig::default()); // 5 missed max + responder + .handle(SessionEvent::SetupRequestReceived(setup_request(2))) + .unwrap(); + assert_eq!(responder.state(), SessionState::Active); + let capture = || SessionEvent::MeasurementCaptured { + instance_id: MeasurementInstanceId::new(0), + payload: payload(10.0), }; - let actions = proxy.handle(SessionEvent::SbpRequestReceived(sbp)).unwrap(); - let forwarded = match &actions[..] { - [Action::SendSbpResponse(SbpResponse { - status: SbpStatus::Accepted, - .. - }), Action::SendSetupRequest(req)] => req.clone(), - other => panic!("expected SBP accept + setup request, got {other:?}"), - }; - assert_eq!(proxy.state(), SessionState::SetupNegotiating); - assert_eq!(forwarded.setup_id.value(), 11); - // The forwarded request drives a *normal* responder session. - let mut responder = SensingSession::new_responder(SessionConfig::default()); - let actions = responder - .handle(SessionEvent::SetupRequestReceived(forwarded)) - .unwrap(); - let resp = match &actions[..] { - [Action::SendSetupResponse(r)] => *r, - other => panic!("expected accept, got {other:?}"), - }; - assert_eq!(resp.status, SetupStatus::Accepted); - proxy - .handle(SessionEvent::SetupResponseReceived(resp)) - .unwrap(); - assert_eq!(proxy.state(), SessionState::Active); + // Miss 4, then succeed once... + for _ in 0..4 { + assert!(responder.handle(SessionEvent::Timeout).unwrap().is_empty()); + } + let actions = responder.handle(capture()).unwrap(); + assert!(matches!(actions[..], [Action::SendReport(_)])); + + // ...so 4 more misses still leave the session alive. + for _ in 0..4 { + assert!(responder.handle(SessionEvent::Timeout).unwrap().is_empty()); + assert_eq!(responder.state(), SessionState::Active); + } + // The 5th consecutive miss terminates. + let actions = responder.handle(SessionEvent::Timeout).unwrap(); + assert!(matches!( + actions[..], + [Action::SendTermination(SensingSessionTermination { + reason: TerminationReason::Timeout, + .. + })] + )); + assert_eq!(responder.state(), SessionState::Terminating); } +// ---------- single-role enforcement & out-of-state commands ---------- + #[test] -fn sbp_client_flow_and_rejections() { - let mut client = SensingSession::new_initiator(SessionConfig::default()); +fn initiator_role_session_rejects_inbound_setup_and_sbp_requests() { + // Review finding 4a: single-role design — a peer must not be able to + // hijack an initiator-role session into the responder path. + let mut initiator = SensingSession::new_initiator(SessionConfig::default()); + let actions = initiator + .handle(SessionEvent::SetupRequestReceived(setup_request(3))) + .unwrap(); + assert!(matches!( + actions[..], + [Action::SendSetupResponse(SensingMeasurementSetupResponse { + status: SetupStatus::RejectedNotSupported, + .. + })] + )); + assert_eq!(initiator.state(), SessionState::Idle); + let sbp = SbpRequest { profile: SpecProfile::Ieee80211Bf2025, - proxy_setup_id: MeasurementSetupId::new(12).unwrap(), + proxy_setup_id: MeasurementSetupId::new(4).unwrap(), params: params(), }; - let actions = client.handle(SessionEvent::StartSbp(sbp.clone())).unwrap(); - assert!(matches!(actions[..], [Action::SendSbpRequest(_)])); - let accept = SbpResponse { - proxy_setup_id: sbp.proxy_setup_id, - status: SbpStatus::Accepted, - }; - client - .handle(SessionEvent::SbpResponseReceived(accept)) - .unwrap(); - assert_eq!(client.state(), SessionState::Active); - // Proxied report is delivered to the local consumer. - let report = SensingMeasurementReport { - setup_id: sbp.proxy_setup_id, - instance_id: MeasurementInstanceId::new(0), - payload: payload(1.0), - }; - let actions = client.handle(SessionEvent::ReportReceived(report)).unwrap(); - assert!(matches!(actions[..], [Action::DeliverReport(_)])); - - // A proxy without SBP capability rejects. - let mut cfg = SessionConfig::default(); - cfg.capabilities.sensing_by_proxy = false; - let mut no_sbp = SensingSession::new_responder(cfg); - let actions = no_sbp + let actions = initiator .handle(SessionEvent::SbpRequestReceived(sbp)) .unwrap(); assert!(matches!( @@ -380,7 +376,59 @@ fn sbp_client_flow_and_rejections() { .. })] )); - assert_eq!(no_sbp.state(), SessionState::Idle); + assert_eq!(initiator.state(), SessionState::Idle); + assert!(!initiator.is_sbp_proxy()); +} + +#[test] +fn local_start_commands_error_outside_idle() { + // Review finding 4b: StartSetup/StartSbp outside Idle are caller bugs + // and must surface as typed errors, not silent no-ops. + let sbp = SbpRequest { + profile: SpecProfile::Ieee80211Bf2025, + proxy_setup_id: MeasurementSetupId::new(13).unwrap(), + params: params(), + }; + let start_err = |s: &mut SensingSession, expected: SessionState| { + assert!(matches!( + s.handle(SessionEvent::StartSetup(setup_request(8))), + Err(BfError::InvalidStateForCommand { .. }) + )); + assert!(matches!( + s.handle(SessionEvent::StartSbp(sbp.clone())), + Err(BfError::InvalidStateForCommand { .. }) + )); + // The rejected commands must not disturb the session. + assert_eq!(s.state(), expected); + }; + + let mut s = SensingSession::new_initiator(SessionConfig::default()); + s.handle(SessionEvent::StartSetup(setup_request(7))) + .unwrap(); + start_err(&mut s, SessionState::SetupNegotiating); + + s.handle(SessionEvent::SetupResponseReceived( + SensingMeasurementSetupResponse { + setup_id: MeasurementSetupId::new(7).unwrap(), + status: SetupStatus::Accepted, + }, + )) + .unwrap(); + start_err(&mut s, SessionState::Active); + // Genuinely ignorable stray frames remain no-ops in Active. + assert!(s + .handle(SessionEvent::SbpResponseReceived(SbpResponse { + proxy_setup_id: MeasurementSetupId::new(7).unwrap(), + status: SbpStatus::Accepted, + })) + .unwrap() + .is_empty()); + + s.handle(SessionEvent::Terminate( + TerminationReason::InitiatorRequested, + )) + .unwrap(); + start_err(&mut s, SessionState::Terminating); } // ---------- adversarial: no panics anywhere ---------- diff --git a/v2/crates/wifi-densepose-hardware/src/ieee80211bf/tests_sbp.rs b/v2/crates/wifi-densepose-hardware/src/ieee80211bf/tests_sbp.rs new file mode 100644 index 00000000..a80623ca --- /dev/null +++ b/v2/crates/wifi-densepose-hardware/src/ieee80211bf/tests_sbp.rs @@ -0,0 +1,338 @@ +//! ADR-153 sensing-by-proxy (SBP) acceptance tests — proxy lifecycle +//! (re-triggering + report relay), client flow, table-driven AP entry +//! point, and the single-validation-path status mapping. Other FSM tests +//! live in [`super::tests_fsm`]; type/serde/transport/bridge tests in +//! [`super::tests`]. All tests are hardware-free (simulation only). + +use super::messages::*; +use super::session::{ + Action, CloseReason, SensingSession, SessionConfig, SessionEvent, SessionState, +}; +use super::table::SessionTable; +use super::testutil::{params, payload}; +use super::transport::{action_to_frame, frame_to_event, SensingFrame}; +use super::types::*; +use crate::csi_frame::Bandwidth; + +fn sbp_request(id: u8) -> SbpRequest { + SbpRequest { + profile: SpecProfile::Ieee80211Bf2025, + proxy_setup_id: MeasurementSetupId::new(id).unwrap(), + params: params(), + } +} + +#[test] +fn sbp_proxy_request_maps_to_standard_responder_path() { + // Proxy AP: accepts the SBP request and initiates an ordinary setup + // toward the sensing responder — no direct sensor coupling. + let mut proxy = SensingSession::new_responder(SessionConfig::default()); + let actions = proxy + .handle(SessionEvent::SbpRequestReceived(sbp_request(11))) + .unwrap(); + let forwarded = match &actions[..] { + [Action::SendSbpResponse(SbpResponse { + status: SbpStatus::Accepted, + .. + }), Action::SendSetupRequest(req)] => req.clone(), + other => panic!("expected SBP accept + setup request, got {other:?}"), + }; + assert_eq!(proxy.state(), SessionState::SetupNegotiating); + assert_eq!(forwarded.setup_id.value(), 11); + + // The forwarded request drives a *normal* responder session. + let mut responder = SensingSession::new_responder(SessionConfig::default()); + let actions = responder + .handle(SessionEvent::SetupRequestReceived(forwarded)) + .unwrap(); + let resp = match &actions[..] { + [Action::SendSetupResponse(r)] => *r, + other => panic!("expected accept, got {other:?}"), + }; + assert_eq!(resp.status, SetupStatus::Accepted); + proxy + .handle(SessionEvent::SetupResponseReceived(resp)) + .unwrap(); + assert_eq!(proxy.state(), SessionState::Active); +} + +#[test] +fn sbp_client_flow_and_rejections() { + let mut client = SensingSession::new_initiator(SessionConfig::default()); + let sbp = sbp_request(12); + let actions = client.handle(SessionEvent::StartSbp(sbp.clone())).unwrap(); + assert!(matches!(actions[..], [Action::SendSbpRequest(_)])); + let accept = SbpResponse { + proxy_setup_id: sbp.proxy_setup_id, + status: SbpStatus::Accepted, + }; + client + .handle(SessionEvent::SbpResponseReceived(accept)) + .unwrap(); + assert_eq!(client.state(), SessionState::Active); + // Proxied report is delivered to the local consumer. + let report = SensingMeasurementReport { + setup_id: sbp.proxy_setup_id, + instance_id: MeasurementInstanceId::new(0), + payload: payload(1.0), + }; + let actions = client.handle(SessionEvent::ReportReceived(report)).unwrap(); + assert!(matches!(actions[..], [Action::DeliverReport(_)])); + + // A proxy without SBP capability rejects. + let mut cfg = SessionConfig::default(); + cfg.capabilities.sensing_by_proxy = false; + let mut no_sbp = SensingSession::new_responder(cfg); + let actions = no_sbp + .handle(SessionEvent::SbpRequestReceived(sbp)) + .unwrap(); + assert!(matches!( + actions[..], + [Action::SendSbpResponse(SbpResponse { + status: SbpStatus::RejectedNotSupported, + .. + })] + )); + assert_eq!(no_sbp.state(), SessionState::Idle); +} + +#[test] +fn sbp_proxy_full_lifecycle_retriggers_and_relays() { + // Review finding 1: the SBP proxy is a first-class mode — after the + // proxied setup is accepted it keeps driving measurement instances on + // InstanceElapsed (like an initiator) and relays every received report + // to the SBP client in addition to local delivery. + let mut proxy = SensingSession::new_responder(SessionConfig::default()); + + // Accept: SBP response to the client + proxied setup to the responder. + let actions = proxy + .handle(SessionEvent::SbpRequestReceived(sbp_request(21))) + .unwrap(); + let forwarded = match &actions[..] { + [Action::SendSbpResponse(SbpResponse { + status: SbpStatus::Accepted, + .. + }), Action::SendSetupRequest(req)] => req.clone(), + other => panic!("expected SBP accept + setup request, got {other:?}"), + }; + assert!(proxy.is_sbp_proxy()); + + // Responder accepts → proxy Active, instance 0 triggered. + let actions = proxy + .handle(SessionEvent::SetupResponseReceived( + SensingMeasurementSetupResponse { + setup_id: forwarded.setup_id, + status: SetupStatus::Accepted, + }, + )) + .unwrap(); + assert_eq!(proxy.state(), SessionState::Active); + match &actions[..] { + [Action::TriggerInstance(i)] => assert_eq!(i.instance_id.value(), 0), + other => panic!("expected instance 0 trigger, got {other:?}"), + } + + // InstanceElapsed re-triggers instance 1+ (proxy drives the schedule). + let actions = proxy.handle(SessionEvent::InstanceElapsed).unwrap(); + match &actions[..] { + [Action::TriggerInstance(i)] => assert_eq!(i.instance_id.value(), 1), + other => panic!("expected instance 1 trigger, got {other:?}"), + } + + // A report from the sensing responder is delivered locally AND relayed. + let report = SensingMeasurementReport { + setup_id: forwarded.setup_id, + instance_id: MeasurementInstanceId::new(1), + payload: payload(5.0), + }; + let actions = proxy + .handle(SessionEvent::ReportReceived(report.clone())) + .unwrap(); + assert_eq!( + actions, + vec![ + Action::DeliverReport(report.clone()), + Action::RelaySbpReport(report.clone()), + ] + ); + // The relay action maps to a frame toward the SBP client, which + // consumes it through the standard report path. + let frame = action_to_frame(&Action::RelaySbpReport(report.clone())).unwrap(); + assert_eq!(frame, SensingFrame::SbpReport(report.clone())); + assert_eq!( + frame_to_event(frame), + Some(SessionEvent::ReportReceived(report)) + ); + + // Terminate cleanly: notify the responder, quiesce back to Idle. + let actions = proxy + .handle(SessionEvent::Terminate( + TerminationReason::InitiatorRequested, + )) + .unwrap(); + assert!(matches!(actions[..], [Action::SendTermination(_)])); + assert_eq!(proxy.state(), SessionState::Terminating); + let actions = proxy.handle(SessionEvent::Timeout).unwrap(); + assert!(matches!( + actions[..], + [Action::SessionClosed(CloseReason::Completed)] + )); + assert_eq!(proxy.state(), SessionState::Idle); + assert!(!proxy.is_sbp_proxy()); +} + +#[test] +fn session_table_routes_sbp_end_to_end() { + // Review finding 3: the table has a first-class SBP entry point with + // the same collision/capacity guards as direct setups — a table-driven + // AP accepts SBP instead of silently dropping it. + let mut table = SessionTable::new(SessionConfig::default()); + let actions = table.handle_sbp_request(sbp_request(31)).unwrap(); + let forwarded = match &actions[..] { + [Action::SendSbpResponse(SbpResponse { + status: SbpStatus::Accepted, + .. + }), Action::SendSetupRequest(req)] => req.clone(), + other => panic!("expected SBP accept + setup request, got {other:?}"), + }; + let setup_id = forwarded.setup_id; + assert_eq!(table.active_setups(), 1); + assert!(table.session(setup_id).unwrap().is_sbp_proxy()); + + // Proxy-setup-ID collision while the first proxy is live. + let actions = table.handle_sbp_request(sbp_request(31)).unwrap(); + assert!(matches!( + actions[..], + [Action::SendSbpResponse(SbpResponse { + status: SbpStatus::RejectedSetupIdCollision, + .. + })] + )); + + // Drive the proxied negotiation to Active through the table. + let actions = table + .handle_for( + setup_id, + SessionEvent::SetupResponseReceived(SensingMeasurementSetupResponse { + setup_id, + status: SetupStatus::Accepted, + }), + ) + .unwrap(); + assert!(matches!(actions[..], [Action::TriggerInstance(_)])); + assert_eq!( + table.session(setup_id).unwrap().state(), + SessionState::Active + ); + + // Reports relay to the SBP client through the table-owned proxy. + let report = SensingMeasurementReport { + setup_id, + instance_id: MeasurementInstanceId::new(0), + payload: payload(2.0), + }; + let actions = table + .handle_for(setup_id, SessionEvent::ReportReceived(report.clone())) + .unwrap(); + assert!(actions.contains(&Action::RelaySbpReport(report))); + + // Capacity guard mirrors the direct-setup path. + let mut cfg = SessionConfig::default(); + cfg.capabilities.max_active_setups = 1; + let mut small = SessionTable::new(cfg); + small.handle_sbp_request(sbp_request(1)).unwrap(); + let actions = small.handle_sbp_request(sbp_request(2)).unwrap(); + assert!(matches!( + actions[..], + [Action::SendSbpResponse(SbpResponse { + status: SbpStatus::RejectedCapacity, + .. + })] + )); + + // Unknown-setup drops are observable, not silent (finding 3). + assert_eq!(table.unknown_setup_drops(), 0); + let actions = table + .handle_for(MeasurementSetupId::new(99).unwrap(), SessionEvent::Timeout) + .unwrap(); + assert!(actions.is_empty()); + assert_eq!(table.unknown_setup_drops(), 1); +} + +#[test] +fn sbp_validation_shares_setup_chain_with_one_to_one_status_mapping() { + // Review finding 5: SBP requests are validated by building the proxied + // setup request first and running it through the single evaluate_setup + // chain — statuses map 1:1, so no rejection class is folded away and no + // setup policy can be bypassed via SBP. + + // Incompatible profile now surfaces as its own status (the old + // duplicated SBP chain folded it into RejectedUnsupportedParams). + let mut cfg = SessionConfig::default(); + cfg.profile = SpecProfile::VendorExtension("acme".into()); + let mut proxy = SensingSession::new_responder(cfg); + let actions = proxy + .handle(SessionEvent::SbpRequestReceived(sbp_request(41))) + .unwrap(); + assert!(matches!( + actions[..], + [Action::SendSbpResponse(SbpResponse { + status: SbpStatus::RejectedIncompatibleProfile, + .. + })] + )); + + // Consent policy rejection passes through unchanged. + let mut proxy = SensingSession::new_responder(SessionConfig::default()); + let mut sbp = sbp_request(42); + sbp.params.consent = ConsentMode::Disabled; + let actions = proxy.handle(SessionEvent::SbpRequestReceived(sbp)).unwrap(); + assert!(matches!( + actions[..], + [Action::SendSbpResponse(SbpResponse { + status: SbpStatus::RejectedByPolicy, + .. + })] + )); + + // Capability rejection (bandwidth beyond the advertised maximum). + let mut cfg = SessionConfig::default(); + cfg.capabilities.max_bandwidth_mhz = 40; + let mut proxy = SensingSession::new_responder(cfg); + let mut sbp = sbp_request(43); + sbp.params.bandwidth = Bandwidth::Bw80; + let actions = proxy.handle(SessionEvent::SbpRequestReceived(sbp)).unwrap(); + assert!(matches!( + actions[..], + [Action::SendSbpResponse(SbpResponse { + status: SbpStatus::RejectedUnsupportedParams, + .. + })] + )); + + // The status translation itself is exhaustive and 1:1. + let pairs = [ + (SetupStatus::Accepted, SbpStatus::Accepted), + ( + SetupStatus::RejectedNotSupported, + SbpStatus::RejectedNotSupported, + ), + ( + SetupStatus::RejectedUnsupportedParams, + SbpStatus::RejectedUnsupportedParams, + ), + ( + SetupStatus::RejectedSetupIdCollision, + SbpStatus::RejectedSetupIdCollision, + ), + ( + SetupStatus::RejectedIncompatibleProfile, + SbpStatus::RejectedIncompatibleProfile, + ), + (SetupStatus::RejectedByPolicy, SbpStatus::RejectedByPolicy), + (SetupStatus::RejectedCapacity, SbpStatus::RejectedCapacity), + ]; + for (setup, sbp) in pairs { + assert_eq!(SbpStatus::from(setup), sbp); + } +} diff --git a/v2/crates/wifi-densepose-hardware/src/ieee80211bf/transport.rs b/v2/crates/wifi-densepose-hardware/src/ieee80211bf/transport.rs index b57a3652..07c014b4 100644 --- a/v2/crates/wifi-densepose-hardware/src/ieee80211bf/transport.rs +++ b/v2/crates/wifi-densepose-hardware/src/ieee80211bf/transport.rs @@ -39,6 +39,10 @@ pub enum SensingFrame { Report(SensingMeasurementReport), SbpRequest(SbpRequest), SbpResponse(SbpResponse), + /// Proxied measurement report forwarded by an SBP proxy toward its SBP + /// client ([`Action::RelaySbpReport`]) — distinct from [`Self::Report`], + /// which travels toward the sensing initiator. + SbpReport(SensingMeasurementReport), Termination(SensingSessionTermination), } @@ -106,6 +110,7 @@ pub fn action_to_frame(action: &Action) -> Option { Action::SendSbpResponse(resp) => Some(SensingFrame::SbpResponse(*resp)), Action::TriggerInstance(instance) => Some(SensingFrame::InstanceTrigger(*instance)), Action::SendReport(report) => Some(SensingFrame::Report(report.clone())), + Action::RelaySbpReport(report) => Some(SensingFrame::SbpReport(report.clone())), Action::SendTermination(term) => Some(SensingFrame::Termination(*term)), Action::DeliverReport(_) | Action::SessionClosed(_) => None, } @@ -122,6 +127,9 @@ pub fn frame_to_event(frame: SensingFrame) -> Option Some(E::SetupRequestReceived(req)), SensingFrame::SetupResponse(resp) => Some(E::SetupResponseReceived(resp)), SensingFrame::Report(report) => Some(E::ReportReceived(report)), + // The SBP client consumes proxied reports through the standard + // report path (its session is in sbp_client mode). + SensingFrame::SbpReport(report) => Some(E::ReportReceived(report)), SensingFrame::SbpRequest(req) => Some(E::SbpRequestReceived(req)), SensingFrame::SbpResponse(resp) => Some(E::SbpResponseReceived(resp)), SensingFrame::Termination(term) => Some(E::TerminationReceived(term)), diff --git a/v2/crates/wifi-densepose-hardware/src/ieee80211bf/types.rs b/v2/crates/wifi-densepose-hardware/src/ieee80211bf/types.rs index 9801ad76..2ec99051 100644 --- a/v2/crates/wifi-densepose-hardware/src/ieee80211bf/types.rs +++ b/v2/crates/wifi-densepose-hardware/src/ieee80211bf/types.rs @@ -386,6 +386,10 @@ impl SensingCapabilities { #[derive(Debug, Clone, Copy, PartialEq, Eq, Serialize, Deserialize)] pub enum SetupStatus { Accepted, + /// The receiving endpoint does not act as a sensing responder for this + /// request — e.g. an initiator-role session received a setup request + /// (single-role design, see [`crate::ieee80211bf::session`]). + RejectedNotSupported, RejectedUnsupportedParams, RejectedSetupIdCollision, RejectedIncompatibleProfile, diff --git a/v2/crates/wifi-densepose-train/src/error.rs b/v2/crates/wifi-densepose-train/src/error.rs index 2e44805d..3d2c2fcd 100644 --- a/v2/crates/wifi-densepose-train/src/error.rs +++ b/v2/crates/wifi-densepose-train/src/error.rs @@ -436,6 +436,18 @@ pub enum MaeError { crop: usize, }, + /// The mask ratio is not a finite value strictly inside `(0, 1)` — the + /// same rule as [`MaePretrainConfig::validate`]. A NaN ratio must never + /// silently mask zero patches, and ratios ≤ 0 / ≥ 1 degenerate to + /// all-visible / all-masked grids. + /// + /// [`MaePretrainConfig::validate`]: crate::mae::MaePretrainConfig::validate + #[error("Invalid mask ratio {ratio}: must be finite and strictly inside (0, 1)")] + InvalidMaskRatio { + /// The offending ratio. + ratio: f64, + }, + /// A NaN or ±inf CSI value was found; corrupted input must be cleaned /// upstream, never masked over. #[error("Non-finite CSI value {value} at (t={row}, sc={col})")] diff --git a/v2/crates/wifi-densepose-train/src/mae.rs b/v2/crates/wifi-densepose-train/src/mae.rs index dd026d65..366473bb 100644 --- a/v2/crates/wifi-densepose-train/src/mae.rs +++ b/v2/crates/wifi-densepose-train/src/mae.rs @@ -160,6 +160,13 @@ impl MaePretrainConfig { /// Patchify `window` and draw the deterministic random mask in one step, /// using `self.seed`. See [`patchify`] and [`random_mask`]. + /// + /// # Errors + /// + /// Everything [`patchify`] rejects, plus [`MaeError::InvalidMaskRatio`] + /// if `self.mask_ratio` is not finite or outside `(0, 1)` (the + /// [`Self::validate`] rule) — a NaN ratio must never silently mask zero + /// patches. pub fn mask_window( &self, window: &[f32], @@ -167,7 +174,7 @@ impl MaePretrainConfig { subc: usize, ) -> Result<(PatchGrid, MaskIndices), MaeError> { let grid = patchify(window, time, subc, self)?; - let mask = random_mask(grid.n_patches(), self.mask_ratio, self.seed); + let mask = random_mask(grid.n_patches(), self.mask_ratio, self.seed)?; Ok((grid, mask)) } } @@ -337,8 +344,18 @@ fn unpatchify_select(grid: &PatchGrid, keep: Option<&[usize]>, fill: f32) -> Vec /// ([`Xorshift64`]), so the same `(n_patches, mask_ratio, seed)` triple always /// yields the same mask. Both index lists are sorted ascending, disjoint, and /// together cover `0..n_patches`. -#[must_use] -pub fn random_mask(n_patches: usize, mask_ratio: f64, seed: u64) -> MaskIndices { +/// +/// # Errors +/// +/// [`MaeError::InvalidMaskRatio`] if `mask_ratio` is not finite or outside +/// the open interval `(0, 1)` — the same rule as +/// [`MaePretrainConfig::validate`]. Erroring (never clamping) keeps the +/// module's error-not-silent policy: a NaN ratio would otherwise silently +/// mask zero patches and a ratio ≥ 1 would mask everything. +pub fn random_mask(n_patches: usize, mask_ratio: f64, seed: u64) -> Result { + if !mask_ratio.is_finite() || mask_ratio <= 0.0 || mask_ratio >= 1.0 { + return Err(MaeError::InvalidMaskRatio { ratio: mask_ratio }); + } let n_masked = ((mask_ratio * n_patches as f64).round() as usize).min(n_patches); let mut order: Vec = (0..n_patches).collect(); let mut rng = Xorshift64::new(seed); @@ -350,7 +367,7 @@ pub fn random_mask(n_patches: usize, mask_ratio: f64, seed: u64) -> MaskIndices let mut visible: Vec = order[n_masked..].to_vec(); masked.sort_unstable(); visible.sort_unstable(); - MaskIndices { masked, visible } + Ok(MaskIndices { masked, visible }) } // --------------------------------------------------------------------------- diff --git a/v2/crates/wifi-densepose-train/src/model.rs b/v2/crates/wifi-densepose-train/src/model.rs index d06ebd01..ac575e6e 100644 --- a/v2/crates/wifi-densepose-train/src/model.rs +++ b/v2/crates/wifi-densepose-train/src/model.rs @@ -126,7 +126,15 @@ impl WiFiDensePoseModel { tch::no_grad(|| self.forward_impl(amplitude, phase, false)) } - /// Save model weights to a file (tch safetensors / .pt format). + /// Save model weights to a file. The tch `VarStore` dispatches the format + /// on the file extension: `.safetensors` → safetensors, anything else → + /// torch `.pt`. + /// + /// **Platform constraint:** prefer `.safetensors`. The `.pt` path + /// (`_save_parameters`/`_load_parameters`) is broken on Windows with + /// torch 2.11 (GenericDict internal assert on the load roundtrip — see + /// `wiflow_std/model.rs::save_and_load_roundtrip`), which is why + /// [`crate::trainer::Trainer`] writes `.safetensors` checkpoints. /// /// # Errors /// @@ -137,7 +145,8 @@ impl WiFiDensePoseModel { .map_err(|e| TrainError::training_step(format!("save failed: {e}"))) } - /// Load model weights from a file. + /// Load model weights from a file (format dispatched on extension; see + /// the `.pt`-on-Windows caveat on [`Self::save`]). /// /// # Errors /// @@ -983,7 +992,9 @@ mod tests { let mut model = WiFiDensePoseModel::new(&cfg, Device::Cpu); let tmp = tempdir().expect("tempdir"); - let path = tmp.path().join("weights.pt"); + // safetensors, not .pt: this torch build's .pt roundtrip is broken on + // Windows (torch 2.11 GenericDict internal assert). + let path = tmp.path().join("weights.safetensors"); model.save(&path).expect("save should succeed"); model.load(&path).expect("load should succeed"); diff --git a/v2/crates/wifi-densepose-train/src/trainer.rs b/v2/crates/wifi-densepose-train/src/trainer.rs index a022fed0..a9f8d670 100644 --- a/v2/crates/wifi-densepose-train/src/trainer.rs +++ b/v2/crates/wifi-densepose-train/src/trainer.rs @@ -286,7 +286,12 @@ impl Trainer { best_epoch = epoch; patience_counter = 0; - let ckpt_name = format!("best_epoch{epoch:04}_pck{val_pck:.4}.pt"); + // .safetensors, not .pt: VarStore dispatches the format on + // the extension, and this torch build's .pt + // _save_parameters/_load_parameters roundtrip is broken on + // Windows (torch 2.11 GenericDict internal assert — see + // wiflow_std/model.rs save_and_load_roundtrip). + let ckpt_name = format!("best_epoch{epoch:04}_pck{val_pck:.4}.safetensors"); let ckpt_path = self.config.checkpoint_dir.join(&ckpt_name); match self.model.save(&ckpt_path) { @@ -339,8 +344,8 @@ impl Trainer { } } - // Save final model regardless. - let final_ckpt = self.config.checkpoint_dir.join("final.pt"); + // Save final model regardless (.safetensors — see checkpoint note above). + let final_ckpt = self.config.checkpoint_dir.join("final.safetensors"); if let Err(e) = self.model.save(&final_ckpt) { warn!("Failed to save final model: {e}"); } @@ -413,7 +418,8 @@ impl Trainer { .load(path) .map_err(|e| TrainError::checkpoint(e.to_string(), path))?; - // Try to parse the epoch from the filename (e.g. "best_epoch0042_pck0.7842.pt"). + // Try to parse the epoch from the filename, extension-agnostic + // (e.g. "best_epoch0042_pck0.7842.safetensors"). let epoch = path .file_stem() .and_then(|s| s.to_str()) diff --git a/v2/crates/wifi-densepose-train/src/wiflow_std/config.rs b/v2/crates/wifi-densepose-train/src/wiflow_std/config.rs index 2068a47d..7393d846 100644 --- a/v2/crates/wifi-densepose-train/src/wiflow_std/config.rs +++ b/v2/crates/wifi-densepose-train/src/wiflow_std/config.rs @@ -56,6 +56,10 @@ fn default_input_pw_groups() -> usize { 1 } +fn default_min_feature_width() -> usize { + 15 +} + // --------------------------------------------------------------------------- // WiFlowStdConfig // --------------------------------------------------------------------------- @@ -114,9 +118,28 @@ pub struct WiFlowStdConfig { pub attention_groups: usize, /// Number of 2-D keypoints produced. Default: **15** (upstream skeleton); - /// use **17** for RuView's COCO-skeleton ESP32 eval set. + /// use **17** for RuView's COCO-skeleton ESP32 eval set. Only changes the + /// parameter-free final adaptive pool — never the trunk: the stride + /// schedule is governed by [`Self::min_feature_width`], so 15- and + /// 17-keypoint variants share the identical conv graph and weights + /// (matching the validated Python protocol, + /// `benchmarks/wiflow-std/remote/measb/train_measb.py`, which swaps only + /// `avg_pool` and loads the pretrained state_dict `strict=True`). pub keypoints: usize, + /// Floor for the conv encoder's width downsampling: each + /// `AsymmetricConvBlock` halves the width only while the result stays + /// ≥ this value (see [`Self::conv_strides`]). + /// + /// Default: **15** — the upstream constant. Provenance: the reference's + /// four hardcoded stride-2 blocks exist because its 240-channel TCN + /// output halves cleanly four times, 240 / 2⁴ = 15. The compact presets' + /// schedules were derived with this same floor. Override only when + /// designing a new trunk; do **not** couple it to [`Self::keypoints`] — + /// the adaptive pool maps the decoder height to any keypoint count. + #[serde(default = "default_min_feature_width")] + pub min_feature_width: usize, + /// Elementwise dropout probability inside the TCN blocks, in `[0, 1)`. /// Default: **0.5** (the value used by our verified retraining run). pub dropout: f64, @@ -134,6 +157,7 @@ impl Default for WiFlowStdConfig { conv_channels: vec![8, 16, 32, 64], attention_groups: 8, keypoints: 15, + min_feature_width: 15, dropout: 0.5, } } @@ -142,6 +166,12 @@ impl Default for WiFlowStdConfig { impl WiFlowStdConfig { /// Default architecture with a different keypoint count (e.g. 17 for the /// ESP32 COCO-skeleton eval set, ADR-152 §2.2(b)). + /// + /// The trunk is untouched: [`Self::min_feature_width`] stays at the + /// upstream floor of 15, so e.g. `for_keypoints(17)` keeps the trained + /// `[2, 2, 2, 2]` stride schedule (feature width 15) and the adaptive + /// pool maps 15 → 17 — exactly the validated Python protocol + /// (`benchmarks/wiflow-std/remote/measb/train_measb.py`). pub fn for_keypoints(keypoints: usize) -> Self { WiFlowStdConfig { keypoints, @@ -284,6 +314,12 @@ impl WiFlowStdConfig { if self.keypoints == 0 { return Err(ConfigError::invalid_value("keypoints", "must be >= 1")); } + if self.min_feature_width == 0 { + return Err(ConfigError::invalid_value( + "min_feature_width", + "must be >= 1", + )); + } if !self.dropout.is_finite() || !(0.0..1.0).contains(&self.dropout) { return Err(ConfigError::invalid_value( "dropout", @@ -316,16 +352,20 @@ impl WiFlowStdConfig { /// Width stride of each `AsymmetricConvBlock`, derived with the sweep's /// rule (`model_compact.py::compute_strides`): halve the width /// (`w → ceil(w / 2)`, the `(1,3)`-kernel stride-2 output size) only - /// while the result stays ≥ [`Self::keypoints`], so the final adaptive - /// pool never has to duplicate rows. At the upstream default - /// (240 channels, 15 keypoints) this derives `[2, 2, 2, 2]` — the - /// hardcoded upstream schedule, exactly. + /// while the result stays ≥ [`Self::min_feature_width`]. At the upstream + /// default (240 TCN channels, floor 15) this derives `[2, 2, 2, 2]` — + /// the hardcoded upstream schedule, exactly. + /// + /// Deliberately independent of [`Self::keypoints`]: the keypoint count + /// only changes the parameter-free adaptive pool, so retargeting the + /// skeleton (e.g. [`Self::for_keypoints`]`(17)`) keeps the trained graph + /// and the pool maps `feature_width() → keypoints`. pub fn conv_strides(&self) -> Vec { let mut w = self.tcn_output_channels(); let mut strides = Vec::with_capacity(self.conv_channels.len()); for _ in &self.conv_channels { let next = w.div_ceil(2); - if next >= self.keypoints { + if next >= self.min_feature_width { strides.push(2); w = next; } else { @@ -375,7 +415,16 @@ impl WiFlowStdConfig { /// /// Pins the port against the verified reference: the 15-keypoint default /// must equal **2,225,042** (`RESULTS.md` artifact verification). + /// + /// Returns **0** for any config that fails [`Self::validate`]: the + /// formula is only meaningful for buildable architectures (an invalid + /// config would otherwise index an empty `conv_channels` or divide by a + /// zero group count). Call `validate()` first when you need the reason. pub fn param_count(&self) -> usize { + if self.validate().is_err() { + return 0; + } + let mut total = 0; // TCN stack: per-conv groups follow tcn_groups_mode; only the first @@ -593,6 +642,76 @@ mod tests { assert_eq!(WiFlowStdConfig::tiny().feature_width(), 16); } + #[test] + fn for_keypoints_17_keeps_trained_trunk_and_pools_15_to_17() { + // Pin against the validated Python protocol (train_measb.py): K=17 + // swaps only the adaptive pool, never the stride schedule. A derived + // [2, 2, 2, 1]/width-30 graph here would silently diverge from the + // trained [2, 2, 2, 2]/width-15 checkpoint. + let cfg = WiFlowStdConfig::for_keypoints(17); + assert_eq!(cfg.min_feature_width, 15); + assert_eq!(cfg.conv_strides(), [2, 2, 2, 2]); + assert_eq!(cfg.feature_width(), 15); + assert_eq!(cfg.output_shape(1), (1, 17, 2)); + } + + #[test] + fn min_feature_width_override_changes_schedule_as_designed() { + // Raising the floor stops the downsampling earlier (240 → 30). + let cfg = WiFlowStdConfig { + min_feature_width: 30, + ..Default::default() + }; + cfg.validate().expect("floor 30 validates"); + assert_eq!(cfg.conv_strides(), [2, 2, 2, 1]); + assert_eq!(cfg.feature_width(), 30); + + // Lowering it lets a small trunk halve further (tiny: 32 → 8). + let cfg = WiFlowStdConfig { + min_feature_width: 8, + ..WiFlowStdConfig::tiny() + }; + cfg.validate().expect("floor 8 validates"); + assert_eq!(cfg.conv_strides(), [2, 2, 1, 1]); + assert_eq!(cfg.feature_width(), 8); + } + + #[test] + fn rejects_zero_min_feature_width() { + let cfg = WiFlowStdConfig { + min_feature_width: 0, + ..Default::default() + }; + assert!(cfg.validate().is_err()); + } + + #[test] + fn param_count_returns_zero_for_invalid_configs() { + // Documented total behavior: configs that fail validate() yield 0 + // instead of panicking (OOB index / division by zero). + for cfg in [ + WiFlowStdConfig { + conv_channels: vec![], + ..Default::default() + }, + WiFlowStdConfig { + tcn_groups: 0, + ..Default::default() + }, + WiFlowStdConfig { + input_pw_groups: 0, + ..Default::default() + }, + WiFlowStdConfig { + tcn_channels: vec![], + ..Default::default() + }, + ] { + assert!(cfg.validate().is_err(), "precondition: {cfg:?} is invalid"); + assert_eq!(cfg.param_count(), 0, "no panic, returns 0: {cfg:?}"); + } + } + #[test] fn fixed_mode_with_defaults_is_unchanged_by_new_knobs() { // The new fields default to upstream behavior: gcd(c, 20) == 20 for diff --git a/v2/crates/wifi-densepose-train/src/wiflow_std/model.rs b/v2/crates/wifi-densepose-train/src/wiflow_std/model.rs index 6d0bf813..074a0f02 100644 --- a/v2/crates/wifi-densepose-train/src/wiflow_std/model.rs +++ b/v2/crates/wifi-densepose-train/src/wiflow_std/model.rs @@ -142,7 +142,16 @@ impl WiFlowStdModel { tch::no_grad(|| self.forward_impl(csi, false)) } - /// Save model weights (tch `.pt` / safetensors format). + /// Save model weights. The tch `VarStore` dispatches the format on the + /// file extension: `.safetensors` → safetensors, anything else → torch + /// `.pt`. + /// + /// **Platform constraint:** prefer `.safetensors`. The `.pt` path + /// (`_save_parameters`/`_load_parameters`) is broken on Windows with + /// torch 2.11 (GenericDict internal assert on the load roundtrip — see + /// the `save_and_load_roundtrip` test below), and the verified retrained + /// checkpoint is shipped as key-remapped safetensors anyway + /// (`benchmarks/wiflow-std/export_to_safetensors.py`). /// /// # Errors /// @@ -153,7 +162,8 @@ impl WiFlowStdModel { .map_err(|e| TrainError::training_step(format!("save failed: {e}"))) } - /// Load model weights from a file. + /// Load model weights from a file (format dispatched on extension; see + /// the `.pt`-on-Windows caveat on [`Self::save`]). /// /// # Errors /// diff --git a/v2/crates/wifi-densepose-train/tests/test_mae.rs b/v2/crates/wifi-densepose-train/tests/test_mae.rs index 01702144..7d935b95 100644 --- a/v2/crates/wifi-densepose-train/tests/test_mae.rs +++ b/v2/crates/wifi-densepose-train/tests/test_mae.rs @@ -207,21 +207,55 @@ fn mask_count_is_exact_for_default_recipe() { // 54 patches @ 0.80 → round(43.2) = 43 masked, 11 visible. let cfg = MaePretrainConfig::default(); assert_eq!(cfg.num_masked(54), 43); - let mask = random_mask(54, cfg.mask_ratio, cfg.seed); + let mask = random_mask(54, cfg.mask_ratio, cfg.seed).unwrap(); assert_eq!(mask.masked.len(), 43); assert_eq!(mask.visible.len(), 11); } #[test] fn same_seed_same_mask_different_seed_differs() { - let a = random_mask(100, 0.80, 7); - let b = random_mask(100, 0.80, 7); + let a = random_mask(100, 0.80, 7).unwrap(); + let b = random_mask(100, 0.80, 7).unwrap(); assert_eq!(a, b, "same (n, ratio, seed) must reproduce the mask"); - let c = random_mask(100, 0.80, 8); + let c = random_mask(100, 0.80, 8).unwrap(); assert_ne!(a.masked, c.masked, "different seeds must differ"); } +#[test] +fn random_mask_rejects_invalid_ratios() { + // Error-not-silent: NaN must not silently mask 0 patches; ratios outside + // (0, 1) must not degenerate to all-visible / all-masked grids. + for ratio in [ + f64::NAN, + f64::INFINITY, + f64::NEG_INFINITY, + 1.0, + 1.5, + 0.0, + -0.1, + ] { + let err = random_mask(54, ratio, 42).unwrap_err(); + assert!( + matches!(err, MaeError::InvalidMaskRatio { .. }), + "ratio {ratio} must be rejected, got {err:?}" + ); + } +} + +#[test] +fn mask_window_rejects_invalid_ratio_before_masking() { + let cfg = MaePretrainConfig { + mask_ratio: f64::NAN, + ..MaePretrainConfig::default() + }; + let buf = window(90, 54); + assert!(matches!( + cfg.mask_window(&buf, 90, 54), + Err(MaeError::InvalidMaskRatio { .. }) + )); +} + proptest! { /// Exact count, sortedness, range, disjointness, and full coverage hold /// for arbitrary grid sizes, ratios, and seeds. @@ -231,7 +265,7 @@ proptest! { ratio in 0.01f64..0.99, seed in any::(), ) { - let mask = random_mask(n, ratio, seed); + let mask = random_mask(n, ratio, seed).unwrap(); let expected_masked = ((ratio * n as f64).round() as usize).min(n); prop_assert_eq!(mask.masked.len(), expected_masked); prop_assert_eq!(mask.masked.len() + mask.visible.len(), n); @@ -254,7 +288,10 @@ proptest! { /// Determinism by seed for arbitrary inputs. #[test] fn prop_mask_deterministic(n in 1usize..400, seed in any::()) { - prop_assert_eq!(random_mask(n, 0.80, seed), random_mask(n, 0.80, seed)); + prop_assert_eq!( + random_mask(n, 0.80, seed).unwrap(), + random_mask(n, 0.80, seed).unwrap() + ); } /// Round-trip identity for arbitrary divisible window/patch geometries.