File size: 8,842 Bytes
e1acc15 | 1 2 3 4 5 6 7 8 9 10 11 12 13 14 15 16 17 18 19 20 21 22 23 24 25 26 27 28 29 30 31 32 33 34 35 36 37 38 39 40 41 42 43 44 45 46 47 48 49 50 51 52 53 54 55 56 57 58 59 60 61 62 63 64 65 66 67 68 69 70 71 72 73 74 75 76 77 78 79 80 81 82 83 84 85 86 87 88 89 90 91 92 93 94 95 96 97 98 99 100 101 102 103 104 105 106 107 108 109 110 111 112 113 114 115 116 117 118 119 120 121 122 123 124 125 126 127 128 129 130 131 132 133 134 135 136 137 138 139 140 141 142 143 144 145 146 147 148 149 150 151 152 153 154 155 156 157 158 159 160 161 162 163 164 165 166 167 168 169 170 171 172 173 174 175 176 177 178 179 180 181 182 183 184 185 186 187 188 189 190 191 192 193 194 195 196 197 198 199 200 201 202 203 204 205 | # SPDX-License-Identifier: Apache-2.0
"""Adapter to lane O1's transport ``ofl`` (offload/transport/API.md, commit 41cf2f7+).
Endpoint forms (``bind`` = this side's direct-link IPv4; it also selects the RDMA device
and GID, API.md "Selecting device / port / GID"):
o1rdma://10.202.28.1:18620?bind=10.202.28.2 RDMA RC (the fast path)
o1tcp://10.202.28.1:18620?bind=10.202.28.2 ofl's TCP backend
Mapping onto the contract (CONTRACT 5): one contract message = one ofl request.
ofl op = AppHeader.op (1..5; ofl's op 0 is its built-in echo)
ofl generation = AppHeader.seq (echoed by ofl; the contract's own echo is checked too)
ofl payload = the whole contract message (AppHeader + body)
ofl response = the whole contract reply frame, RDMA-written into the registered
reply buffer (one range)
The ofl library and its ctypes binding are loaded from VLLM_DSPARK_REMOTE_OFL_DIR (a
directory holding libofl.so and ofl.py, bind-mounted into the container), else from
OFL_LIB's directory. Timeouts break the ofl connection by design (the QP goes to error
before the call returns, so late data never lands); the proxy then reconnects.
"""
from __future__ import annotations
import ctypes
import importlib.util
import os
import struct
import sys
from .transport import Endpoint, RpcClient, TransportError, TransportTimeout
SERVICE = "dspark-v1"
DEFAULT_MAX_REQ = 16 << 20 # >= 32 + 192 + 512 rows x 30,720 B = 15.73 MB (CONTRACT 8.1)
DEFAULT_MAX_RESP = 64 << 10 # STEP replies are < 1 KiB; STATS JSON fits
_OP_SEQ = struct.Struct("<HQ") # AppHeader.op (offset 6) and .seq (offset 8)
_ofl = None
def load_ofl():
"""Import O1's ofl.py (and through it libofl.so) once."""
global _ofl
if _ofl is not None:
return _ofl
d = os.environ.get("VLLM_DSPARK_REMOTE_OFL_DIR") or ""
if not d and os.environ.get("OFL_LIB"):
d = os.path.dirname(os.environ["OFL_LIB"])
if not d:
raise TransportError(
"o1 endpoint: set VLLM_DSPARK_REMOTE_OFL_DIR to the directory with libofl.so "
"and ofl.py (lane O1, offload/transport)")
py = os.path.join(d, "ofl.py")
so = os.path.join(d, "libofl.so")
if not os.path.exists(py) or not os.path.exists(so):
raise TransportError(f"o1 endpoint: {py} or {so} missing")
os.environ.setdefault("OFL_LIB", so)
spec = importlib.util.spec_from_file_location("dsv41_ofl", py)
if spec is None or spec.loader is None:
raise TransportError(f"cannot import {py}")
mod = importlib.util.module_from_spec(spec)
sys.modules["dsv41_ofl"] = mod
spec.loader.exec_module(mod)
mod.lib() # load + declare now: fail at startup, not on the first round
_ofl = mod
return mod
def _addr(buf) -> tuple[int, int]:
"""Address and length of a contiguous writable buffer (bytearray, numpy, memoryview)."""
if hasattr(buf, "ctypes") and hasattr(buf, "nbytes"):
if not buf.flags.c_contiguous:
raise TransportError("buffer must be contiguous")
return int(buf.ctypes.data), int(buf.nbytes)
mv = memoryview(buf)
if mv.readonly:
raise TransportError("buffer must be writable (ctypes needs its address)")
if not mv.contiguous:
raise TransportError("buffer must be contiguous")
c = (ctypes.c_char * mv.nbytes).from_buffer(mv.cast("B"))
return ctypes.addressof(c), mv.nbytes
class O1Client(RpcClient):
def __init__(self, ep: Endpoint):
if ep.scheme not in ("o1tcp", "o1rdma"):
raise TransportError(f"O1Client cannot serve {ep.scheme}://")
self.ofl = load_ofl()
self.ep = ep
self.backend = self.ofl.BACKEND_RDMA if ep.scheme == "o1rdma" else self.ofl.BACKEND_TCP
self.local_ip = ep.options.get("bind")
if not self.local_ip:
raise TransportError(
f"endpoint {ep}: ?bind=<this side's direct-link IPv4> is required "
"(it selects the RDMA device and GID)")
self.max_req = int(ep.options.get("max_req", DEFAULT_MAX_REQ))
self.max_resp = int(ep.options.get("max_resp", DEFAULT_MAX_RESP))
self.gid_index = int(ep.options.get("gid", -1))
self.rdma_device = ep.options.get("dev")
self._c = None
self._regions: dict[int, tuple[int, int, int]] = {} # id(buf) -> (addr, len, region)
self._keep: list[object] = []
self._res = None
self._seg = None
@property
def connected(self) -> bool:
return self._c is not None
def connect(self, timeout_s: float) -> None:
self.close()
try:
self._c = self.ofl.Client(
self.backend, self.ep.host, self.ep.port, self.local_ip, service=SERVICE,
max_req=self.max_req, max_resp=self.max_resp, inflight=1,
rdma_device=self.rdma_device, gid_index=self.gid_index,
connect_timeout_ms=max(1, int(timeout_s * 1000)))
except self.ofl.OflError as e:
self._c = None
if e.code == -3:
raise TransportTimeout(f"ofl connect {self.ep}: {e}") from e
raise TransportError(f"ofl connect {self.ep}: {e}") from e
self._regions.clear()
self._keep.clear()
self._res = self.ofl.Result()
self._seg = (self.ofl._Seg * 1)()
def _region_for(self, buf) -> tuple[int, int]:
addr, n = _addr(buf)
hit = self._regions.get(id(buf))
if hit is not None and hit[0] == addr and hit[1] == n:
return hit[2], n
if n < self.max_resp:
raise TransportError(f"reply buffer of {n} B < max_resp {self.max_resp}")
try:
region = self._c.reg(buf)
except self.ofl.OflError as e:
raise TransportError(f"ofl reg: {e}") from e
self._regions[id(buf)] = (addr, n, region)
self._keep.append(buf)
return region, n
def request(self, msg, resp_buf, timeout_s: float) -> int:
if self._c is None:
raise TransportError("not connected")
addr, n = _addr(msg)
if n > self.max_req:
raise TransportError(f"message of {n} B > ofl max_req {self.max_req}")
op, seq = _OP_SEQ.unpack_from(memoryview(msg).cast("B"), 6)
return self._call(addr, n, op, seq, resp_buf, timeout_s)
def request_parts(self, parts, resp_buf, timeout_s: float) -> int:
"""The parts are copied straight into ofl's send slot (ofl_client_payload_buf), and the call then skips its own
payload copy: the same single copy request() costs, with no joined buffer in between."""
c = self._c
if c is None:
raise TransportError("not connected")
spans = [_addr(p) for p in parts]
n = sum(k for _, k in spans)
if n > self.max_req:
raise TransportError(f"message of {n} B > ofl max_req {self.max_req}")
if not spans or spans[0][1] < _OP_SEQ.size + 6:
raise TransportError("the first part must hold the AppHeader")
op, seq = _OP_SEQ.unpack_from(memoryview(parts[0]).cast("B"), 6)
pl, cap = c.payload_buf()
if not pl:
self.close()
raise TransportError("ofl connection is broken (no payload slot)")
if n > cap:
raise TransportError(f"message of {n} B > ofl payload slot {cap} B")
off = 0
for a, k in spans:
ctypes.memmove(pl + off, a, k)
off += k
return self._call(pl, n, op, seq, resp_buf, timeout_s)
def _call(self, addr: int, n: int, op: int, seq: int, resp_buf, timeout_s: float) -> int:
c = self._c
region, cap = self._region_for(resp_buf)
seg = self._seg
seg[0].region, seg[0].offset, seg[0].len = region, 0, min(cap, self.max_resp)
L = self.ofl.lib()
rc = L.ofl_client_call(c.h, op, seq, addr, n, seg, 1, ctypes.byref(self._res),
int(timeout_s * 1e9))
if rc != 0:
msg_txt = (L.ofl_client_error(c.h) or b"").decode(errors="replace")
name = self.ofl.ERRORS.get(rc, str(rc))
if rc == -3: # ETIMEDOUT: ofl already broke the connection
self.close()
raise TransportTimeout(f"ofl call: {name}: {msg_txt}")
self.close()
raise TransportError(f"ofl call: {name}: {msg_txt}")
if self._res.generation != seq:
self.close()
raise TransportError(f"ofl generation echo {self._res.generation} != {seq}")
return int(self._res.resp_len)
def close(self) -> None:
if self._c is not None:
try:
self._c.close()
finally:
self._c = None
self._regions.clear()
self._keep.clear()
|