| |
| """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 |
| DEFAULT_MAX_RESP = 64 << 10 |
| _OP_SEQ = struct.Struct("<HQ") |
|
|
| _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() |
| _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]] = {} |
| 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: |
| 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() |
|
|
