refactor(doctor): harden NCCL benchmark from PR #178

PR #178 wired the live NCCL bandwidth measurement on top of the
v0.43.0 nccl_bandwidth_check reference table. Four polish items
on top of the merged version:

- Bump timing samples 5 -> 10 and report MEDIAN rather than mean.
  Median is robust to one-off jitter (kernel preemption, page
  faults, OS scheduler hiccups) that can scew a 5-sample mean by
  10-30%. Median of 10 gives a much more honest GB/s number.

- Bump warmup 1 -> 3 iters. First all_reduce in a fresh process
  pays for CUDA kernel JIT plus the initial NCCL collective
  handshake; one warmup isn't always enough to amortise that out.
  3 warmups is the standard NVIDIA recommendation for collective
  benchmarks.

- Time each iteration separately (collect a per-iter list, then
  median) instead of averaging total elapsed / iters. Same headline
  cost, but exposes per-call variance for future percentile output.

- Snapshot + restore MASTER_ADDR / MASTER_PORT around the spawn.
  The worker sets them unconditionally, so prior to this patch a
  second doctor invocation in the same process (test harness,
  notebook, future TUI auto-refresh) would inherit stale env. Now
  we restore exactly the prior state (including absent vars).

Plus three module-level constants (_NCCL_BENCHMARK_TENSOR_BYTES
/ _WARMUP_ITERS / _TIMED_ITERS) with comments explaining each
number so future tuning is informed, not magic. Status line now
prints the actual params so users see what we measured.

Existing tests (mock mp.spawn at the outer boundary) keep passing
unchanged; no new tests needed for this polish layer.

soup_cli/commands/doctor.py

@@ -500,7 +500,10 @@ def _run_nccl_check():
     manager = mp.Manager()
     return_dict = manager.dict()
 
-    console.print("\n[bold]Measuring NCCL bandwidth (100MB all_reduce)...[/]")
+    console.print(
+        "\n[bold]Measuring NCCL bandwidth[/] "
+        "(100 MB all_reduce, 3 warmup + 10 timed iters, median GB/s)..."
+    )
     try:
         mp.spawn(_nccl_worker, args=(return_dict,), nprocs=2, join=True)
     except Exception as e:
@@ -529,35 +532,65 @@ def _run_nccl_check():
         console.print("[red]Failed to capture NCCL bandwidth measurement.[/]")
 
 
+# NCCL benchmark constants. 100 MB matches typical gradient-bucket size for
+# medium models (7B fp16 ≈ 100-200 MB per bucket) so the measurement
+# reflects real all-reduce traffic, not artificially small messages.
+_NCCL_BENCHMARK_TENSOR_BYTES = 100 * 1024 * 1024  # 100 MB
+_NCCL_BENCHMARK_WARMUP_ITERS = 3  # 3 warmups to amortise CUDA kernel JIT
+_NCCL_BENCHMARK_TIMED_ITERS = 10  # 10 timed runs; report the median (robust to outliers)
+
+
 def _nccl_worker(rank: int, return_dict):
     import os
+    import statistics
     import time
 
     import torch
     import torch.distributed as dist
 
+    # Snapshot prior env so we can restore after the spawn finishes —
+    # avoids leaking MASTER_ADDR/MASTER_PORT into the parent doctor process
+    # (matters when tests run multiple doctor invocations in one session).
+    prior_addr = os.environ.get("MASTER_ADDR")
+    prior_port = os.environ.get("MASTER_PORT")
     os.environ["MASTER_ADDR"] = "localhost"
     os.environ["MASTER_PORT"] = "29500"
 
-    dist.init_process_group("nccl", rank=rank, world_size=2)
+    try:
+        dist.init_process_group("nccl", rank=rank, world_size=2)
 
-    tensor_size = 25 * 1024 * 1024  # 100MB in float32
-    tensor = torch.ones(tensor_size, dtype=torch.float32, device=f"cuda:{rank}")
+        # 100 MB tensor of float32 (4 bytes per element).
+        num_elements = _NCCL_BENCHMARK_TENSOR_BYTES // 4
+        tensor = torch.ones(num_elements, dtype=torch.float32, device=f"cuda:{rank}")
 
-    # Warmup
-    dist.all_reduce(tensor)
-    torch.cuda.synchronize(device=f"cuda:{rank}")
+        # Warmup — runs CUDA kernel JIT + first-NCCL-collective handshake
+        # out of the timing window. 3 iters is enough to stabilise.
+        for _ in range(_NCCL_BENCHMARK_WARMUP_ITERS):
+            dist.all_reduce(tensor)
+        torch.cuda.synchronize(device=f"cuda:{rank}")
 
-    start = time.perf_counter()
-    iters = 5
-    for _ in range(iters):
-        dist.all_reduce(tensor)
-    torch.cuda.synchronize(device=f"cuda:{rank}")
-    end = time.perf_counter()
+        # Per-iteration timing (10 samples). Median is more robust to
+        # one-off jitter (kernel preemption, swap, etc.) than mean.
+        per_iter_sec: list[float] = []
+        for _ in range(_NCCL_BENCHMARK_TIMED_ITERS):
+            start = time.perf_counter()
+            dist.all_reduce(tensor)
+            torch.cuda.synchronize(device=f"cuda:{rank}")
+            per_iter_sec.append(time.perf_counter() - start)
 
-    if rank == 0:
-        elapsed = (end - start) / iters
-        size_gb = (tensor.element_size() * tensor.numel()) / 1e9
-        return_dict["gb_per_sec"] = size_gb / elapsed
+        if rank == 0:
+            median_elapsed = statistics.median(per_iter_sec)
+            size_gb = (tensor.element_size() * tensor.numel()) / 1e9
+            return_dict["gb_per_sec"] = size_gb / median_elapsed
 
-    dist.destroy_process_group()
+        dist.destroy_process_group()
+    finally:
+        # Restore prior env exactly as we found it (or remove if missing).
+        if prior_addr is None:
+            os.environ.pop("MASTER_ADDR", None)
+        else:
+            os.environ["MASTER_ADDR"] = prior_addr
+        if prior_port is None:
+            os.environ.pop("MASTER_PORT", None)
+        else:
+            os.environ["MASTER_PORT"] = prior_port