| 1 | # SPDX-License-Identifier: Apache-2.0 |
| 2 | # SPDX-FileCopyrightText: Copyright contributors to the vLLM project |
| 3 | """GPU-to-NIC net-device mapping for RDMA transports (UCX, NVSHMEM, ...). |
| 4 | |
| 5 | Shared by both UniProcExecutor (TP=1) and MultiprocExecutor (TP>1). |
| 6 | All transport-specific env vars and sysfs lookups live here so executor |
| 7 | files only need to call ``set_worker_net_device(local_rank, vllm_config)``. |
| 8 | |
| 9 | Requires two env vars set together: |
| 10 | - ``VLLM_GPU_NIC_PCIE_MAPPING`` -- comma-separated GPU_BDF=NIC_BDF pairs. |
| 11 | - ``VLLM_NIC_SELECTION_VARS`` -- comma-separated list of env vars to set, |
| 12 | each optionally suffixed (e.g. ``UCX_NET_DEVICES:1,NCCL_IB_HCA:1``). |
| 13 | """ |
| 14 | |
| 15 | import os |
| 16 | from pathlib import Path |
| 17 | |
| 18 | import vllm.envs as envs |
| 19 | from vllm.config import VllmConfig |
| 20 | from vllm.logger import init_logger |
| 21 | from vllm.platforms import current_platform |
| 22 | |
| 23 | logger = init_logger(__name__) |
| 24 | |
| 25 | |
| 26 | def normalize_pci(addr: str) -> tuple[int, int, int, int]: |
| 27 | """Parse PCI BDF/domain-bus-device-function into comparable ints (all hex). |
| 28 | |
| 29 | Supported shapes: |
| 30 | - ``domain:bus:dev.fn`` -- domain width varies (e.g. ``00000001:00:00.0``, |
| 31 | ``0001:00:00.0``, ``0000:3f:00.0``). |
| 32 | - ``bus:dev.fn`` -- domain **0** (e.g. ``01:00.0``, ``40:00.0``). |
| 33 | |
| 34 | Function suffix is hex (typically ``0``--``7``). Raises ``ValueError`` if malformed. |
| 35 | """ |
| 36 | s = addr.strip().lower().replace(" ", "") |
| 37 | if s.startswith("0x"): |
| 38 | s = s[2:] |
| 39 | if "." not in s: |
| 40 | raise ValueError(f"invalid PCI BDF (missing function suffix): {addr!r}") |
| 41 | body, fn_s = s.rsplit(".", 1) |
| 42 | if not fn_s or any(c not in "0123456789abcdef" for c in fn_s): |
| 43 | raise ValueError(f"invalid PCI function in BDF: {addr!r}") |
| 44 | fn = int(fn_s, 16) |
| 45 | if fn > 0xFF: |
| 46 | raise ValueError(f"PCI function out of range: {addr!r}") |
| 47 | |
| 48 | parts = body.split(":") |
| 49 | if len(parts) == 2: |
| 50 | domain = 0 |
| 51 | bus = int(parts[0], 16) |
| 52 | device = int(parts[1], 16) |
| 53 | elif len(parts) == 3: |
| 54 | domain = int(parts[0], 16) |
| 55 | bus = int(parts[1], 16) |
| 56 | device = int(parts[2], 16) |
| 57 | else: |
| 58 | raise ValueError( |
| 59 | f"invalid PCI BDF (want domain:bus:dev.fn or bus:dev.fn): {addr!r}" |
| 60 | ) |
| 61 | |
| 62 | if bus > 0xFF or device > 0x1F: |
| 63 | raise ValueError(f"PCI bus or device out of range: {addr!r}") |
| 64 | return (domain, bus, device, fn) |
| 65 | |
| 66 | |
| 67 | def parse_gpu_nic_mapping( |
| 68 | raw: str, |
| 69 | ) -> dict[tuple[int, int, int, int], tuple[int, int, int, int]]: |
| 70 | out: dict[tuple[int, int, int, int], tuple[int, int, int, int]] = {} |
| 71 | for segment in raw.split(","): |
| 72 | segment = segment.strip() |
| 73 | if not segment: |
| 74 | continue |
| 75 | if "=" not in segment: |
| 76 | raise ValueError( |
| 77 | "VLLM_GPU_NIC_PCIE_MAPPING: expected comma-separated" |
| 78 | f" gpu_bdf=nic_bdf pairs; ambiguous segment: {segment!r}" |
| 79 | ) |
| 80 | gpu_s, nic_s = segment.split("=", 1) |
| 81 | gpu_key = normalize_pci(gpu_s.strip()) |
| 82 | nic_val = normalize_pci(nic_s.strip()) |
| 83 | out[gpu_key] = nic_val |
| 84 | return out |
| 85 | |
| 86 | |
| 87 | def rdma_name_for_nic_pci(nic_pci: tuple[int, int, int, int]) -> str: |
| 88 | """Map NIC PCI BDF to sysfs RDMA name (mlx5_*, ibp*, ...). |
| 89 | |
| 90 | Under ``/sys/class/infiniband/<name>/``, ``device`` is a **symlink** to the PCI |
| 91 | device directory (e.g. ``.../0101:00:00.0``). We take ``Path(...).resolve().name`` |
| 92 | as the BDF string. |
| 93 | |
| 94 | ``VLLM_GPU_NIC_PCIE_MAPPING`` NIC keys must **normalize** (via ``normalize_pci``) |
| 95 | to the same tuple as this basename. |
| 96 | """ |
| 97 | ib = Path("/sys/class/infiniband") |
| 98 | if not ib.is_dir(): |
| 99 | raise RuntimeError("/sys/class/infiniband not found or not a directory") |
| 100 | names = sorted(p.name for p in ib.iterdir() if p.is_dir()) |
| 101 | for name in names: |
| 102 | dev_link = ib / name / "device" |
| 103 | if not dev_link.exists(): |
| 104 | continue |
| 105 | try: |
| 106 | resolved = dev_link.resolve() |
| 107 | except OSError: |
| 108 | continue |
| 109 | pci_name = resolved.name |
| 110 | try: |
| 111 | if normalize_pci(pci_name) == nic_pci: |
| 112 | return name |
| 113 | except ValueError: |
| 114 | continue |
| 115 | raise RuntimeError( |
| 116 | f"No /sys/class/infiniband device for NIC PCI {nic_pci}; have entries: {names}" |
| 117 | ) |
| 118 | |
| 119 | |
| 120 | def parse_nic_selection_vars(raw: str) -> list[tuple[str, str]]: |
| 121 | """Parse ``VLLM_NIC_SELECTION_VARS`` into ``(env_var_name, suffix)`` pairs. |
| 122 | |
| 123 | Each entry is ``VAR_NAME`` or ``VAR_NAME:<suffix>``. The colon and |
| 124 | everything after it is appended verbatim to the RDMA device name. |
| 125 | """ |
| 126 | result: list[tuple[str, str]] = [] |
| 127 | for entry in raw.split(","): |
| 128 | entry = entry.strip() |
| 129 | if not entry: |
| 130 | continue |
| 131 | if ":" in entry: |
| 132 | var_name, suffix = entry.split(":", 1) |
| 133 | result.append((var_name, ":" + suffix)) |
| 134 | else: |
| 135 | result.append((entry, "")) |
| 136 | return result |
| 137 | |
| 138 | |
| 139 | def set_worker_gpu_nic_mapping(local_rank: int) -> None: |
| 140 | """Set NIC selection env vars from VLLM_GPU_NIC_PCIE_MAPPING for a worker. |
| 141 | |
| 142 | Which env vars are set is controlled by ``VLLM_NIC_SELECTION_VARS``. |
| 143 | """ |
| 144 | raw = envs.VLLM_GPU_NIC_PCIE_MAPPING.strip() |
| 145 | if not raw: |
| 146 | return |
| 147 | selection_raw = envs.VLLM_NIC_SELECTION_VARS.strip() |
| 148 | selection_vars = parse_nic_selection_vars(selection_raw) |
| 149 | mapping = parse_gpu_nic_mapping(raw) |
| 150 | pci_by_index = current_platform.get_all_gpu_pci_bus_ids() |
| 151 | # Translate CUDA-relative local_rank to the physical device index, |
| 152 | # which accounts for CUDA_VISIBLE_DEVICES narrowing (e.g. DP sharding). |
| 153 | physical_id = current_platform.device_id_to_physical_device_id(local_rank) |
| 154 | if physical_id not in pci_by_index: |
| 155 | raise RuntimeError( |
| 156 | f"No GPU PCI for physical device index {physical_id} " |
| 157 | f"(local_rank={local_rank}) in map " |
| 158 | f"(have indices {sorted(pci_by_index.keys())})" |
| 159 | ) |
| 160 | gpu_bdf = pci_by_index[physical_id] |
| 161 | gpu_key = normalize_pci(gpu_bdf) |
| 162 | if gpu_key not in mapping: |
| 163 | keys_fmt = ", ".join( |
| 164 | f"{d:04x}:{b:02x}:{dev:02x}.{fn}" |
| 165 | for d, b, dev, fn in sorted(mapping.keys()) |
| 166 | ) |
| 167 | raise RuntimeError( |
| 168 | f"No VLLM_GPU_NIC_PCIE_MAPPING entry for GPU PCI {gpu_bdf} " |
| 169 | f"(worker local_rank={local_rank}); mapped GPUs: {keys_fmt}" |
| 170 | ) |
| 171 | nic_pci = mapping[gpu_key] |
| 172 | rdma_dev = rdma_name_for_nic_pci(nic_pci) |
| 173 | |
| 174 | set_vars: list[str] = [] |
| 175 | for var_name, suffix in selection_vars: |
| 176 | value = f"{rdma_dev}{suffix}" |
| 177 | existing = os.environ.get(var_name, "").strip() |
| 178 | if existing: |
| 179 | value = f"{value},{existing}" |
| 180 | os.environ[var_name] = value |
| 181 | set_vars.append(f"{var_name}={value}") |
| 182 | |
| 183 | nic_fmt = f"{nic_pci[0]:04x}:{nic_pci[1]:02x}:{nic_pci[2]:02x}.{nic_pci[3]}" |
| 184 | logger.info( |
| 185 | "GPU rank %s (PCIe addr %s) mapped to NIC %s (PCIe addr %s) via env vars: %s", |
| 186 | local_rank, |
| 187 | gpu_bdf, |
| 188 | rdma_dev, |
| 189 | nic_fmt, |
| 190 | ", ".join(set_vars), |
| 191 | ) |
| 192 | |
| 193 | |
| 194 | def _dp_adjusted_local_rank(tp_local_rank: int, vllm_config: VllmConfig) -> int: |
| 195 | """Compute the node-wide GPU index accounting for data parallelism. |
| 196 | |
| 197 | On CUDA-alike platforms without env-var device isolation (the common |
| 198 | MP-backend path), the worker sees *all* GPUs on the node and selects |
| 199 | its device via ``torch.accelerator.set_device_index()`` using:: |
| 200 | |
| 201 | dp_local_rank * tp_pp_world_size + tp_local_rank |
| 202 | |
| 203 | This mirrors the adjustment in ``Worker.init_device()`` so we resolve |
| 204 | the correct GPU PCI address *before* the CUDA device is initialised. |
| 205 | """ |
| 206 | pc = vllm_config.parallel_config |
| 207 | if ( |
| 208 | pc.distributed_executor_backend not in ("ray", "external_launcher") |
| 209 | and pc.data_parallel_backend != "ray" |
| 210 | and pc.nnodes_within_dp == 1 |
| 211 | ): |
| 212 | dp_local_rank = pc.data_parallel_rank_local |
| 213 | if dp_local_rank is None: |
| 214 | dp_local_rank = pc.data_parallel_index |
| 215 | tp_pp_world_size = pc.pipeline_parallel_size * pc.tensor_parallel_size |
| 216 | return dp_local_rank * tp_pp_world_size + tp_local_rank |
| 217 | return tp_local_rank |
| 218 | |
| 219 | |
| 220 | def set_worker_net_device(local_rank: int, vllm_config: VllmConfig) -> None: |
| 221 | """Top-level entry point for both UniProcExecutor and MultiprocExecutor. |
| 222 | |
| 223 | Sets NIC selection env vars from ``VLLM_GPU_NIC_PCIE_MAPPING`` and |
| 224 | ``VLLM_NIC_SELECTION_VARS`` if present; no-op otherwise. |
| 225 | """ |
| 226 | has_pcie_mapping = bool(envs.VLLM_GPU_NIC_PCIE_MAPPING.strip()) |
| 227 | has_selection_vars = bool(envs.VLLM_NIC_SELECTION_VARS.strip()) |
| 228 | if has_pcie_mapping and not has_selection_vars: |
| 229 | raise RuntimeError( |
| 230 | "VLLM_GPU_NIC_PCIE_MAPPING is set but VLLM_NIC_SELECTION_VARS " |
| 231 | "is not; both must be set together." |
| 232 | ) |
| 233 | if has_selection_vars and not has_pcie_mapping: |
| 234 | raise RuntimeError( |
| 235 | "VLLM_NIC_SELECTION_VARS is set but VLLM_GPU_NIC_PCIE_MAPPING " |
| 236 | "is not; both must be set together." |
| 237 | ) |
| 238 | # No-op when neither env var is present. |
| 239 | if not has_pcie_mapping and not has_selection_vars: |
| 240 | return |
| 241 | # Both env vars are present, so set the NIC selection env vars. |
| 242 | adjusted_rank = _dp_adjusted_local_rank(local_rank, vllm_config) |
| 243 | set_worker_gpu_nic_mapping(adjusted_rank) |