diff --git a/README.md b/README.md index 5678e15..c173fc8 100644 --- a/README.md +++ b/README.md @@ -2712,10 +2712,10 @@ soup quickstart --yes Check your environment for compatibility issues: ```bash -soup doctor +soup doctor [--nccl] ``` -Shows: Python version, GPU availability, system resources (RAM/Disk), all dependency versions, and fix suggestions. +Shows: Python version, GPU availability, system resources (RAM/Disk), all dependency versions, and fix suggestions. Use `--nccl` to measure and check multi-GPU communication bandwidth against expected hardware ceilings. ## Version Info @@ -3408,7 +3408,7 @@ soup tokenizer train --input c.jsonl --vocab-size N Train BPE tokenizer (v0.53. soup bench --p50 --p95 Bench with tail-latency percentiles (v0.53.9) soup bench --backend auto Auto-detect transformers/mlx backend (v0.53.9) soup serve --reasoning-parser deepseek-r1 Strip blocks from responses (v0.53.9) -soup doctor Check environment +soup doctor [--nccl] Check environment (optionally check NCCL bandwidth) soup quickstart [--dry-run] Full demo soup version [--full] [--json] Show version (--full: system info, --json: JSON output) soup --verbose Full traceback on errors diff --git a/soup_cli/commands/doctor.py b/soup_cli/commands/doctor.py index a4da883..9599df6 100644 --- a/soup_cli/commands/doctor.py +++ b/soup_cli/commands/doctor.py @@ -5,6 +5,7 @@ from __future__ import annotations import platform import sys +import typer from rich.console import Console from rich.panel import Panel from rich.table import Table @@ -51,7 +52,11 @@ _MAX_EXCLUSIVE: dict[str, str] = { } -def doctor(): +def doctor( + nccl: bool = typer.Option( + False, "--nccl", help="Measure NCCL bandwidth and check against reference table." + ), +): """Check system dependencies, GPU, and compatibility.""" console.print("[bold]Soup Doctor[/] - checking your environment...\n") @@ -109,8 +114,7 @@ def doctor(): if max_excl and _version_ge(version_str, max_excl): status = f"[red]INCOMPATIBLE (need <{max_excl})[/]" issues.append( - f"Downgrade {pkg_name}: " - f"pip install '{pkg_name}>={min_ver},<{max_excl}'" + f"Downgrade {pkg_name}: pip install '{pkg_name}>={min_ver},<{max_excl}'" ) elif _version_ok(version_str, min_ver): status = "[green]OK[/]" @@ -145,16 +149,21 @@ def doctor(): # Check torchvision + torch compatibility _check_torchvision_compat(issues) + if nccl: + _run_nccl_check() + # Summary if issues: console.print(f"\n[yellow]Found {len(issues)} issue(s):[/]") for issue in issues: console.print(f" [red]>[/] {issue}") - console.print("\n[dim]Fix all: pip install -U " + " ".join( - f"'{pkg_name}>={min_ver}'" - for _, pkg_name, min_ver, required in DEPS - if required - ) + "[/]") + console.print( + "\n[dim]Fix all: pip install -U " + + " ".join( + f"'{pkg_name}>={min_ver}'" for _, pkg_name, min_ver, required in DEPS if required + ) + + "[/]" + ) else: console.print("\n[bold green]All checks passed![/] Your environment is ready.") @@ -185,7 +194,7 @@ def _check_gpu(): name = torch.cuda.get_device_name(idx) mem = torch.cuda.get_device_properties(idx) total_gb = getattr(mem, "total_memory", getattr(mem, "total_mem", 0)) - total_gb = total_gb / (1024 ** 3) + total_gb = total_gb / (1024**3) gpus.append(f" GPU {idx}: [bold]{name}[/] ({total_gb:.1f} GB)") gpu_info = "\n".join(gpus) cuda_ver = torch.version.cuda or "N/A" @@ -293,7 +302,7 @@ def _detect_dual_python_interpreters() -> str: ) -_GB = 1024 ** 3 +_GB = 1024**3 def _get_ram_gb() -> str: @@ -301,6 +310,7 @@ def _get_ram_gb() -> str: # Prefer psutil if installed try: import psutil + return f"{psutil.virtual_memory().total / _GB:.0f} GB" except ImportError: pass @@ -318,9 +328,13 @@ def _get_ram_gb() -> str: elif system == "Darwin": try: import subprocess + res = subprocess.run( ["sysctl", "-n", "hw.memsize"], - capture_output=True, text=True, timeout=5, check=False, + capture_output=True, + text=True, + timeout=5, + check=False, ) if res.returncode == 0: return f"{int(res.stdout.strip()) / _GB:.0f} GB" @@ -387,8 +401,13 @@ def _check_torchvision_compat(issues: list): # Known compatible pairs (torch minor -> torchvision minor) compat = { - "2.6": "0.21", "2.5": "0.20", "2.4": "0.19", - "2.3": "0.18", "2.2": "0.17", "2.1": "0.16", "2.0": "0.15", + "2.6": "0.21", + "2.5": "0.20", + "2.4": "0.19", + "2.3": "0.18", + "2.2": "0.17", + "2.1": "0.16", + "2.0": "0.15", } expected_tv = compat.get(torch_minor) if expected_tv and not tv_minor.startswith(expected_tv): @@ -437,3 +456,108 @@ def _version_ge(installed: str, threshold: str) -> bool: return inst_parts >= thr_parts except (ValueError, AttributeError): return False + + +def _run_nccl_check(): + """Run NCCL bandwidth test if requested.""" + try: + import torch + import torch.distributed as dist + import torch.multiprocessing as mp + + from soup_cli.utils.profiling_v0_43 import nccl_bandwidth_check + from soup_cli.utils.topology import detect_topology + except ImportError: + console.print("\n[yellow]NCCL bandwidth check requires torch and torch.distributed.[/]") + return + + if not torch.cuda.is_available() or not dist.is_available(): + console.print("\n[yellow]NCCL bandwidth check requires CUDA and torch.distributed.[/]") + return + + topo = detect_topology() + gpu_count = topo["gpu_count"] + if gpu_count < 2: + console.print("\n[yellow]NCCL bandwidth requires >=2 GPUs[/]") + return + + name = torch.cuda.get_device_name(0).lower() + if "h100" in name: + gpu = "h100" + elif "a100" in name: + gpu = "a100" + elif "v100" in name: + gpu = "v100" + elif "4090" in name: + gpu = "rtx4090" + elif "3090" in name: + gpu = "rtx3090" + else: + gpu = "unknown" + + link = topo["interconnect"] + + manager = mp.Manager() + return_dict = manager.dict() + + console.print("\n[bold]Measuring NCCL bandwidth (100MB all_reduce)...[/]") + try: + mp.spawn(_nccl_worker, args=(return_dict,), nprocs=2, join=True) + except Exception as e: + console.print(f"[red]Failed to measure NCCL bandwidth:[/] {e}") + return + + if "gb_per_sec" in return_dict: + measured = return_dict["gb_per_sec"] + res = nccl_bandwidth_check(gpu=gpu, link=link, measured_gb_per_sec=measured) + status = res["status"] + if status == "OK": + color = "green" + elif status == "MINOR": + color = "yellow" + else: + color = "red" + + expected = res.get("expected_gb_per_sec") + expected_str = f" vs expected {expected:.1f}" if expected else "" + + console.print( + f" Result ({gpu.upper()} over {link.upper()}): " + f"[{color}]{status}[/] ({measured:.1f} GB/s{expected_str})" + ) + else: + console.print("[red]Failed to capture NCCL bandwidth measurement.[/]") + + +def _nccl_worker(rank: int, return_dict): + import os + import time + + import torch + import torch.distributed as dist + + os.environ["MASTER_ADDR"] = "localhost" + os.environ["MASTER_PORT"] = "29500" + + 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}") + + # Warmup + 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() + + if rank == 0: + elapsed = (end - start) / iters + size_gb = (tensor.element_size() * tensor.numel()) / 1e9 + return_dict["gb_per_sec"] = size_gb / elapsed + + dist.destroy_process_group() diff --git a/tests/test_doctor.py b/tests/test_doctor.py index 6aa3f0a..5081858 100644 --- a/tests/test_doctor.py +++ b/tests/test_doctor.py @@ -112,9 +112,12 @@ def test_doctor_checks_optional_deps(): def test_doctor_missing_dep(): """soup doctor reports missing required dep.""" - with patch("soup_cli.commands.doctor.DEPS", [ - ("nonexistent_fake_pkg_xyz", "nonexistent-pkg", "1.0.0", True), - ]): + with patch( + "soup_cli.commands.doctor.DEPS", + [ + ("nonexistent_fake_pkg_xyz", "nonexistent-pkg", "1.0.0", True), + ], + ): result = runner.invoke(app, ["doctor"]) assert result.exit_code == 0 assert "MISSING" in result.output @@ -122,9 +125,54 @@ def test_doctor_missing_dep(): def test_doctor_outdated_dep(): """soup doctor reports outdated dep.""" - with patch("soup_cli.commands.doctor.DEPS", [ - ("sys", "sys", "999.0.0", True), # sys has no __version__ but import won't fail - ]): + with patch( + "soup_cli.commands.doctor.DEPS", + [ + ("sys", "sys", "999.0.0", True), # sys has no __version__ but import won't fail + ], + ): result = runner.invoke(app, ["doctor"]) assert result.exit_code == 0 # Either outdated or OK (depends on version attr presence) + + +# --- NCCL Check tests --- + + +def test_doctor_nccl_no_gpu(): + """--nccl with <2 GPUs prints a skip message.""" + with ( + patch("torch.cuda.is_available", return_value=True), + patch("torch.distributed.is_available", return_value=True), + patch( + "soup_cli.utils.topology.detect_topology", + return_value={"gpu_count": 1, "nvlink_pairs": 0, "interconnect": "single"}, + ), + ): + result = runner.invoke(app, ["doctor", "--nccl"]) + assert result.exit_code == 0 + assert "NCCL bandwidth requires >=2 GPUs" in result.output + + +def test_doctor_nccl_mocked_success(): + """--nccl with 2 GPUs runs the check and displays result.""" + + # We mock mp.spawn to just set a value in the return_dict instead of actually running processes. + def mock_spawn(func, args, nprocs, join): + return_dict = args[0] + return_dict["gb_per_sec"] = 350.0 # mock value + + with ( + patch("torch.cuda.is_available", return_value=True), + patch("torch.distributed.is_available", return_value=True), + patch( + "soup_cli.utils.topology.detect_topology", + return_value={"gpu_count": 2, "nvlink_pairs": 1, "interconnect": "nvlink"}, + ), + patch("torch.cuda.get_device_name", return_value="NVIDIA H100 80GB HBM3"), + patch("torch.multiprocessing.spawn", side_effect=mock_spawn), + ): + result = runner.invoke(app, ["doctor", "--nccl"]) + assert result.exit_code == 0 + assert "Measuring NCCL bandwidth" in result.output + assert "Result (H100 over NVLINK)" in result.output