bench: add the GPU benchmark and fold its findings into the docs

docs/source/parallel_execution.rst

@@ -58,19 +58,30 @@ Backends
        ``cores // num_workers`` while active and restored afterwards;
        ``torch_num_threads`` overrides this explicitly.
 
-Measured speedups (12-core CPU, FedAvg on FedSynthetic, 64 clients, 3
-iterations, 2 local epochs, best of 3 runs; see
-``test-files/benchmark_executors.py`` and the committed
-``benchmark_executors_results.data`` to reproduce): about **1.2×** at 2–4
-workers, degrading slightly beyond that as GIL contention dominates. This
-is an observation **on this light configuration** (tiny MLP; single-threaded
-local steps take only tens of milliseconds, so the Python glue code — batch
-loops, optimizer stepping — holds the GIL most of the time), **not** a
-ceiling for the backend in general: with compute-heavier local training
-(larger models, more/bigger batches, GPU offload) or the (planned)
-process-based backend the balance shifts toward parallel execution. The
+Measured speedups (best of 3 runs; see ``test-files/benchmark_executors.py``
+and ``test-files/benchmark_gpu.py`` to reproduce; the committed
+``benchmark_executors_results.data`` holds the CPU run):
+
+- **CPU, tiny model** (MLP, tens of ms per client task; 64 clients, 3
+  iterations): about **1.2×** at 2–4 workers, degrading slightly beyond that
+  as GIL contention dominates. The Python glue code (batch loops, optimizer
+  stepping) holds the GIL most of the time on such workloads.
+- **CPU, compute-heavier model** (1.19M-parameter MLP, 5 local epochs, batch
+  256, 12-core host): **1.40×** at 8 workers — the torch kernels release the
+  GIL, so the benefit grows with the compute share of each client task.
+- **single GPU**: **no benefit from the threads backend** (sequential and
+  threads measure the same within noise): all clients submit kernels to the
+  same device, where they execute serially anyway — the sequential loop
+  already keeps a single GPU fed. (The measured GPU run was itself ~3×
+  faster than the CPU run for this workload.) Scaling beyond one GPU comes
+  from multi-GPU client placement (``ServerConfig.visiable_gpus``) or CUDA
+  streams, not from concurrent submission to one device.
+
+These are observations on the measured configurations, not ceilings: the
 per-backend arithmetic guarantees (fixed aggregation order, identical
-participation sets) hold regardless of the achieved speedup.
+participation sets) hold regardless of the achieved speedup, and the
+(planned) process-based backend targets the CPU workloads where the GIL
+bounds the threads backend.
 
 How it works
 ------------

test-files/benchmark_gpu.py

@@ -0,0 +1,67 @@
+"""GPU/CPU benchmark for the client-executor backends (light footprint).
+
+Runs FedAvg on FedSynthetic(dimension=512) with the mlp_d3 candidate model
+(~1.19M params), comparing sequential vs threads backends on CPU and GPU.
+Best-of-3 wall clock per configuration; results appended to a log file.
+
+Run:
+    python test-files/benchmark_gpu.py
+"""
+
+import sys
+import time
+from pathlib import Path
+
+sys.path.insert(0, str(Path(__file__).parents[1].resolve()))
+
+import torch  # noqa: E402
+
+import fl_sim.algorithms as algorithms  # noqa: E402
+from fl_sim.data_processing.fed_synthetic import FedSynthetic  # noqa: E402
+
+ITERS = 3
+CLIENTS = 64
+EPOCHS = 5
+BATCH = 256
+WORKERS = 8
+REPEATS = 3
+
+
+def run_once(executor: str, num_workers: int, use_gpu: bool) -> float:
+    reg = algorithms.get_algorithm("FedAvg")
+    dataset = FedSynthetic(1, 1, False, max(30, CLIENTS), dimension=512)
+    server = reg["server"](
+        dataset.candidate_models["mlp_d3"],
+        dataset,
+        reg["server_config"](
+            ITERS,
+            CLIENTS,
+            clients_sample_ratio=1.0,
+            verbose=0,
+            visiable_gpus=[0] if use_gpu else [],
+        ),
+        reg["client_config"](BATCH, EPOCHS),
+    )
+    server.config.executor = executor
+    server.config.num_workers = num_workers
+    start = time.perf_counter()
+    server.train_federated()
+    return time.perf_counter() - start
+
+
+def bench(tag: str, use_gpu: bool) -> None:
+    print(f"=== {tag} ({ITERS} iters x {CLIENTS} clients x {EPOCHS} local epochs, best of {REPEATS}) ===", flush=True)
+    base = min(run_once("sequential", 0, use_gpu) for _ in range(REPEATS))
+    print(f"  sequential      {base:8.2f} s   1.00x", flush=True)
+    t = min(run_once("threads", WORKERS, use_gpu) for _ in range(REPEATS))
+    print(f"  threads@{WORKERS}      {t:8.2f} s   {base / t:5.2f}x", flush=True)
+
+
+if __name__ == "__main__":
+    print(f"torch {torch.__version__}, cuda available: {torch.cuda.is_available()}", flush=True)
+    if torch.cuda.is_available():
+        print(f"gpu: {torch.cuda.get_device_name(0)}", flush=True)
+    bench("CPU", use_gpu=False)
+    if torch.cuda.is_available():
+        bench("GPU", use_gpu=True)
+    print("BENCH_GPU_DONE", flush=True)

test/test_nodes.py

@@ -94,7 +94,9 @@ def train(self) -> None:
 
 
 def test_aggregate_results_from_json_log():
-    json_log_file = list((Path(__file__).parents[1].resolve() / "test-files").glob("*.json"))[0]
+    json_log_file = (
+        Path(__file__).parents[1].resolve() / "test-files" / "FedAvg-FedProxFEMNIST-CNNFEMnist_Tiny_06-04_23-25.json"
+    )
     for part, metric in product(["train", "val"], ["acc", "loss"]):
         curve = Server.aggregate_results_from_json_log(json_log_file, part=part, metric=metric)
         assert len(curve) > 0

test/test_utils.py

@@ -367,7 +367,9 @@ def test_compute_clustering_metrics():
 def test_aggregate_results_from_json_log():
     from fl_sim.utils.misc import aggregate_results_from_json_log
 
-    json_log_file = list((Path(__file__).parents[1].resolve() / "test-files").glob("*.json"))[0]
+    json_log_file = (
+        Path(__file__).parents[1].resolve() / "test-files" / "FedAvg-FedProxFEMNIST-CNNFEMnist_Tiny_06-04_23-25.json"
+    )
 
     # dict input
     d = json.loads(json_log_file.read_text())

test/test_viz.py

@@ -33,7 +33,7 @@ def test_find_log_files():
 def test_get_config_from_log():
     from fl_sim.utils.viz import get_config_from_log
 
-    json_log = list(logdir.glob("*.json"))[0]
+    json_log = logdir / "FedAvg-FedProxFEMNIST-CNNFEMnist_Tiny_06-04_23-25.json"
     config = get_config_from_log(json_log)
     assert isinstance(config, dict) and len(config) > 0
 
@@ -47,7 +47,7 @@ def test_get_curves_and_labels_from_log():
 
     from fl_sim.utils.viz import get_curves_and_labels_from_log
 
-    json_log = list(logdir.glob("*.json"))[0]
+    json_log = logdir / "FedAvg-FedProxFEMNIST-CNNFEMnist_Tiny_06-04_23-25.json"
     curves, stems = get_curves_and_labels_from_log(json_log, part="train", x_metric="step", y_metric="acc")
     assert len(curves) == len(stems) == 1
     assert set(curves[0].keys()) == {"x", "y"}
@@ -67,7 +67,7 @@ def test_plot_curves_and_mean_curve():
 
     from fl_sim.utils.viz import get_curves_and_labels_from_log, plot_curves, plot_mean_curve_with_error_bounds
 
-    json_log = list(logdir.glob("*.json"))[0]
+    json_log = logdir / "FedAvg-FedProxFEMNIST-CNNFEMnist_Tiny_06-04_23-25.json"
 
     fig, ax = plot_curves(json_log, part="train", x_metric="step", y_metric="acc")
     assert ax.get_ylabel() == "train acc"