vllm.distributed.kv_transfer.kv_connector.v1.umbp.runtime.standalone

Standalone UMBP runtime backed by a node-local umbp_standalone_server.

The server owns one DRAM pool per host. Every engine, data-parallel rank, and worker on that host attaches to it over a Unix-domain socket, so KV stored by one engine is reusable by the others and survives engine restarts. Workers register their GPU KV allocations with the server through HIP IPC; the server then copies between those allocations and its pool directly.

Because the server is reachable from the scheduler process, scheduler lookups query it directly instead of going through a worker-side bridge.

Classes:

• StandaloneRuntime –

  Runtime whose objects live in a node-local standalone server.

Functions:

• normalize_endpoint –

  Return the gRPC address of a standalone server socket.

StandaloneRuntime

Bases: IUMBPRuntime

Runtime whose objects live in a node-local standalone server.

Source code in vllm/distributed/kv_transfer/kv_connector/v1/umbp/runtime/standalone.py

class StandaloneRuntime(IUMBPRuntime):
    """Runtime whose objects live in a node-local standalone server."""

    def __init__(self, config: UMBPRuntimeConfig) -> None:
        self._options = _resolve_standalone_options(config)
        self._address = normalize_endpoint(self._options["endpoint"])
        self._max_workers = int(self._options.get("num_workers", 4))
        self._timeout_s = float(self._options.get("timeout_ms", 30000)) / 1000
        self._lookup_timeout_s = (
            float(self._options.get("lookup_timeout_ms", 1000)) / 1000
        )

    @classmethod
    def from_config(cls, config: UMBPRuntimeConfig) -> StandaloneRuntime:
        return cls(config)

    def create_scheduler_handle(
        self,
        namespace: str,
        topology: RankTopology,
        layout: KVLayoutDescriptor,
    ) -> UMBPSchedulerHandle:
        del namespace, topology, layout
        return _StandaloneSchedulerHandle(
            _build_client(self._options),
            self._address,
            call_timeout_s=self._lookup_timeout_s,
            close_timeout_s=self._timeout_s,
        )

    def create_worker_handle(
        self,
        namespace: str,
        topology: RankTopology,
        layout: KVLayoutDescriptor,
    ) -> UMBPWorkerHandle:
        logger.info("UMBP standalone worker attached to %s", self._address)
        return _StandaloneWorkerHandle(
            _build_client(self._options),
            namespace,
            topology,
            lookup_dir="",
            max_workers=self._max_workers,
            timeout_s=self._timeout_s,
            layout=layout,
            serve_lookups=False,
        )

_DaemonTransferPool

Bases: Executor

A thread pool whose threads never hold up interpreter exit.

ThreadPoolExecutor joins its threads at exit, and a call blocked in an unresponsive service cannot be interrupted.

Source code in vllm/distributed/kv_transfer/kv_connector/v1/umbp/runtime/standalone.py

class _DaemonTransferPool(Executor):
    """A thread pool whose threads never hold up interpreter exit.

    ThreadPoolExecutor joins its threads at exit, and a call blocked in an
    unresponsive service cannot be interrupted.
    """

    def __init__(self, max_workers: int, thread_name_prefix: str) -> None:
        if max_workers <= 0:
            raise ValueError("max_workers must be greater than 0")
        self._work: queue.SimpleQueue[tuple[Future, Callable[[], Any]] | None] = (
            queue.SimpleQueue()
        )
        self._lock = threading.Lock()
        self._shutdown = False
        self._threads = [
            threading.Thread(
                target=self._run, name=f"{thread_name_prefix}_{index}", daemon=True
            )
            for index in range(max_workers)
        ]
        for thread in self._threads:
            thread.start()

    def submit(self, fn: Callable, /, *args: Any, **kwargs: Any) -> Future:
        future: Future = Future()
        with self._lock:
            if self._shutdown:
                raise RuntimeError("cannot schedule new futures after shutdown")
            self._work.put((future, lambda: fn(*args, **kwargs)))
        return future

    def _run(self) -> None:
        while (item := self._work.get()) is not None:
            future, call = item
            if not future.set_running_or_notify_cancel():
                continue
            try:
                result = call()
            except BaseException as exc:
                future.set_exception(exc)
            else:
                future.set_result(result)

    def shutdown(self, wait: bool = True, *, cancel_futures: bool = False) -> None:
        with self._lock:
            if self._shutdown:
                return
            self._shutdown = True
            if cancel_futures:
                while True:
                    try:
                        item = self._work.get_nowait()
                    except queue.Empty:
                        break
                    if item is not None:
                        # wait() counts a cancelled future as done only once
                        # it is notified, which _run() never gets to do here.
                        item[0].cancel()
                        item[0].set_running_or_notify_cancel()
            for _ in self._threads:
                self._work.put(None)
        if wait:
            for thread in self._threads:
                thread.join()

_SharedPoolWorkerHandle

Bases: _MoriWorkerHandle

MORI worker handle whose objects are indexed outside this process.

The pool's service can stall or disappear under a running engine, so its failures degrade to misses and recomputation. A transfer that outlives timeout_ms stays pending rather than failing the step, since MORI still uses its blocks: the scheduler keeps a store's source blocks pinned, and keeps a loading request's blocks until the load is reported, even if the request is aborted. New transfers fail while one is blocked. A synchronous wait on a blocked transfer still raises, since its blocks are about to be read or reused.

Source code in vllm/distributed/kv_transfer/kv_connector/v1/umbp/runtime/standalone.py

class _SharedPoolWorkerHandle(_MoriWorkerHandle):
    """MORI worker handle whose objects are indexed outside this process.

    The pool's service can stall or disappear under a running engine, so its
    failures degrade to misses and recomputation. A transfer that outlives
    timeout_ms stays pending rather than failing the step, since MORI still
    uses its blocks: the scheduler keeps a store's source blocks pinned, and
    keeps a loading request's blocks until the load is reported, even if the
    request is aborted. New transfers fail while one is blocked. A synchronous
    wait on a blocked transfer still raises, since its blocks are about to be
    read or reused.
    """

    mode = "shared"

    def __init__(self, *args: Any, **kwargs: Any) -> None:
        super().__init__(*args, **kwargs)
        self._last_refusal_log = float("-inf")
        self._storages: dict[int, torch.UntypedStorage] = {}

    def _create_executor(self, max_workers: int) -> Executor:
        return _DaemonTransferPool(max_workers, f"umbp-{self.mode}-transfer")

    def register_buffers(self, kv_caches: dict[str, torch.Tensor]) -> None:
        super().register_buffers(kv_caches)
        for cache in kv_caches.values():
            storage = cache.untyped_storage()
            self._storages.setdefault(storage.data_ptr(), storage)

    def _blocked_transfers(self) -> int:
        now = time.monotonic()
        with self._job_lock:
            deadlines = (
                *self._store_deadlines.items(),
                *self._load_deadlines.items(),
            )
            futures = dict(self._futures)
        return sum(
            1
            for job_id, deadline in deadlines
            if now >= deadline
            and (future := futures.get(job_id)) is not None
            and not future.done()
        )

    def load(self, plans: Sequence[BlockTransferPlan]) -> TransferJobState:
        refused = self._refuse_while_blocked(plans)
        return refused if refused is not None else super().load(plans)

    def load_blocks(
        self, plans: Sequence[BlockTransferPlan]
    ) -> TransferJobState | None:
        refused = self._refuse_while_blocked(plans)
        return refused if refused is not None else super().load_blocks(plans)

    def store(self, plans: Sequence[BlockTransferPlan]) -> TransferJobState:
        refused = self._refuse_while_blocked(plans)
        return refused if refused is not None else super().store(plans)

    def store_blocks(
        self, plans: Sequence[BlockTransferPlan]
    ) -> TransferJobState | None:
        refused = self._refuse_while_blocked(plans)
        return refused if refused is not None else super().store_blocks(plans)

    def poll(self, job: TransferJobState) -> TransferJobState | None:
        with self._job_lock:
            future = self._futures.get(id(job))
        if future is not None and not future.done():
            if self._blocked_transfers() and self._fail_if_queued(job):
                return job
            return None
        return super().poll(job)

    def wait(self, job: TransferJobState) -> TransferJobState:
        try:
            return super().wait(job)
        except TimeoutError:
            if self._fail_if_queued(job):
                return job
            raise

    def _fail_if_queued(self, job: TransferJobState) -> bool:
        """Fail a job still waiting for a transfer thread; it holds no buffers."""
        with self._job_lock:
            future = self._futures.get(id(job))
        if future is None or not future.cancel():
            return False
        with self._job_lock:
            self._futures.pop(id(job), None)
        job.fail(
            job.keys,
            f"{TRANSFER_NOT_ATTEMPTED}: queued behind a blocked UMBP {self.mode} "
            "transfer",
        )
        return True

    def cancel(self, job: TransferJobState) -> TransferJobState:
        with self._job_lock:
            is_store = id(job) in self._store_deadlines
            future = self._futures.get(id(job))
        if is_store:
            if future is not None and future.cancel():
                with self._job_lock:
                    self._futures.pop(id(job), None)
                    self._store_deadlines.pop(id(job), None)
                job.cancel("preempted")
            # A running store only reads blocks the scheduler keeps pinned.
            return job
        return super().cancel(job)

    def _refuse_while_blocked(
        self, plans: Sequence[BlockTransferPlan]
    ) -> TransferJobState | None:
        """Fail new transfers, before they touch memory, while one is stuck."""
        blocked = self._blocked_transfers()
        if not blocked:
            return None
        job = TransferJobState(
            plans if isinstance(plans, BlockLoadBatch) else tuple(plans)
        )
        job.start()
        job.fail(
            job.keys,
            f"{TRANSFER_NOT_ATTEMPTED}: an earlier UMBP {self.mode} transfer "
            "is still blocked",
        )
        now = time.monotonic()
        if now - self._last_refusal_log >= _FAILURE_LOG_INTERVAL_S:
            self._last_refusal_log = now
            logger.warning(
                "UMBP %s transfers are refused while %d timed-out transfer(s) "
                "are still blocked in MORI",
                self.mode,
                blocked,
            )
        return job

    def publish(self, job: TransferJobState) -> None:
        if job.status != TransferJobStatus.COMPLETED:
            raise RuntimeError("cannot publish an incomplete MORI UMBP job")
        try:
            super().publish(job)
        except Exception as exc:
            logger.warning("UMBP %s publish failed: %s", self.mode, exc)

    def close(self) -> None:
        # The service may already be gone at engine shutdown and its
        # registrations with it, so a failed deregistration is not an error.
        with self._job_lock:
            futures = list(self._futures.values())
            deadlines = [
                *self._store_deadlines.values(),
                *self._load_deadlines.values(),
            ]
        self._executor.shutdown(wait=False, cancel_futures=True)
        # Running transfers get until their own deadlines; stalled ones get none.
        timeout = max([0.0, *(deadline - time.monotonic() for deadline in deadlines)])
        _, in_flight = wait_futures(futures, timeout=timeout)
        if in_flight:
            # MORI's deregister_memory() and close() wait for every in-flight
            # call, and a stalled service never returns one.
            self._abandon(f"{len(in_flight)} transfer(s) are still blocked in MORI")
            return
        # Neither MORI call has a deadline, so an idle client can hang on them too.
        release = threading.Thread(
            target=self._release_client, name=f"umbp-{self.mode}-close", daemon=True
        )
        release.start()
        release.join(self._timeout_s)
        if release.is_alive():
            self._abandon(f"deregistration did not return in {self._timeout_s:g}s")

    def _release_client(self) -> None:
        for storage_ptr in self._registered_storages:
            try:
                self.client.deregister_memory(storage_ptr)
            except Exception as exc:
                logger.warning(
                    "UMBP %s deregister_memory(0x%x) failed: %s",
                    self.mode,
                    storage_ptr,
                    exc,
                )
        self._registered_storages.clear()
        close = getattr(self.client, "close", None)
        try:
            if callable(close):
                close()
        except Exception as exc:
            logger.warning("UMBP %s client close failed: %s", self.mode, exc)
        self._storages.clear()
        # MORI releases the connection only when the client object is destroyed.
        self.client = None  # type: ignore[assignment]

    def _abandon(self, reason: str) -> None:
        _ABANDONED_STORAGES.extend(self._storages.values())
        self._storages.clear()
        logger.warning("UMBP %s client left open at shutdown: %s", self.mode, reason)

_fail_if_queued(job)

Fail a job still waiting for a transfer thread; it holds no buffers.

Source code in vllm/distributed/kv_transfer/kv_connector/v1/umbp/runtime/standalone.py

def _fail_if_queued(self, job: TransferJobState) -> bool:
    """Fail a job still waiting for a transfer thread; it holds no buffers."""
    with self._job_lock:
        future = self._futures.get(id(job))
    if future is None or not future.cancel():
        return False
    with self._job_lock:
        self._futures.pop(id(job), None)
    job.fail(
        job.keys,
        f"{TRANSFER_NOT_ATTEMPTED}: queued behind a blocked UMBP {self.mode} "
        "transfer",
    )
    return True

_refuse_while_blocked(plans)

Fail new transfers, before they touch memory, while one is stuck.

Source code in vllm/distributed/kv_transfer/kv_connector/v1/umbp/runtime/standalone.py

def _refuse_while_blocked(
    self, plans: Sequence[BlockTransferPlan]
) -> TransferJobState | None:
    """Fail new transfers, before they touch memory, while one is stuck."""
    blocked = self._blocked_transfers()
    if not blocked:
        return None
    job = TransferJobState(
        plans if isinstance(plans, BlockLoadBatch) else tuple(plans)
    )
    job.start()
    job.fail(
        job.keys,
        f"{TRANSFER_NOT_ATTEMPTED}: an earlier UMBP {self.mode} transfer "
        "is still blocked",
    )
    now = time.monotonic()
    if now - self._last_refusal_log >= _FAILURE_LOG_INTERVAL_S:
        self._last_refusal_log = now
        logger.warning(
            "UMBP %s transfers are refused while %d timed-out transfer(s) "
            "are still blocked in MORI",
            self.mode,
            blocked,
        )
    return job

_StandaloneSchedulerHandle

Bases: UMBPSchedulerHandle

Answers lookups from the shared server.

A lookup that fails, or gets no reply within lookup_timeout_ms, is reported as a miss, so a stalled or unreachable server degrades to recomputation instead of stalling the scheduler. MORI puts no deadline on these calls, so each runs on a daemon thread; until one that overran its deadline returns, later calls fail at once instead of queueing behind it.

Source code in vllm/distributed/kv_transfer/kv_connector/v1/umbp/runtime/standalone.py

class _StandaloneSchedulerHandle(UMBPSchedulerHandle):
    """Answers lookups from the shared server.

    A lookup that fails, or gets no reply within lookup_timeout_ms, is
    reported as a miss, so a stalled or unreachable server degrades to
    recomputation instead of stalling the scheduler. MORI puts no deadline on
    these calls, so each runs on a daemon thread; until one that overran its
    deadline returns, later calls fail at once instead of queueing behind it.
    """

    def __init__(
        self,
        client: Any,
        address: str,
        call_timeout_s: float,
        close_timeout_s: float,
    ) -> None:
        self._client = client
        self._address = address
        self._call_timeout_s = call_timeout_s
        self._close_timeout_s = close_timeout_s
        self._calls = _DaemonTransferPool(1, "umbp-standalone-lookup")
        self._call_lock = threading.Lock()
        self._unanswered: Future | None = None
        self._last_failure_log = float("-inf")

    def _call(self, method: str, *args: Any) -> Any:
        with self._call_lock:
            if self._unanswered is not None and not self._unanswered.done():
                raise TimeoutError("an earlier call to the server has not returned")
            future = self._calls.submit(getattr(self._client, method), *args)
        try:
            return future.result(timeout=self._call_timeout_s)
        except TimeoutError:
            if future.done():
                raise
            with self._call_lock:
                self._unanswered = future
            raise TimeoutError(
                f"{method} got no reply within {self._call_timeout_s:g}s"
            ) from None

    def lookup(self, keys: Sequence[str]) -> Sequence[bool]:
        if not keys:
            return []
        try:
            result = [bool(value) for value in self._call("batch_exists", list(keys))]
        except Exception as exc:
            self._log_failure("lookup", exc)
            return [False] * len(keys)
        if len(result) != len(keys):
            self._log_failure(
                "lookup", ValueError(f"{len(result)} results for {len(keys)} keys")
            )
            return [False] * len(keys)
        return result

    def clear(self) -> bool:
        try:
            cleared = bool(self._call("clear"))
        except Exception as exc:
            self._log_failure("clear", exc)
            return False
        if cleared:
            logger.warning(
                "Cleared the UMBP standalone pool at %s; this removes KV for "
                "every engine attached to it",
                self._address,
            )
        else:
            logger.warning("UMBP standalone clear at %s failed", self._address)
        return cleared

    def close(self) -> None:
        self._calls.shutdown(wait=False, cancel_futures=True)
        with self._call_lock:
            unanswered = self._unanswered is not None and not self._unanswered.done()
        if unanswered:
            # MORI's close() waits for every in-flight call.
            self._leave_open("a call to the server never returned")
            return
        release = threading.Thread(
            target=self._release_client, name="umbp-standalone-close", daemon=True
        )
        release.start()
        release.join(self._close_timeout_s)
        if release.is_alive():
            self._leave_open(f"close did not return in {self._close_timeout_s:g}s")

    def _release_client(self) -> None:
        close = getattr(self._client, "close", None)
        try:
            if callable(close):
                close()
        except Exception as exc:
            logger.warning("UMBP standalone client close failed: %s", exc)
        self._client = None

    def _leave_open(self, reason: str) -> None:
        self._client = None
        logger.warning(
            "UMBP standalone scheduler client for %s left open at shutdown: %s",
            self._address,
            reason,
        )

    def _log_failure(self, operation: str, exc: Exception) -> None:
        now = time.monotonic()
        if now - self._last_failure_log >= _FAILURE_LOG_INTERVAL_S:
            self._last_failure_log = now
            logger.warning(
                "UMBP standalone %s against %s failed; treating as a miss: %s",
                operation,
                self._address,
                exc,
            )

normalize_endpoint(endpoint)

Return the gRPC address of a standalone server socket.

The data plane depends on file-descriptor and IPC-handle passing over a Unix-domain socket, so only socket paths are accepted.

Source code in vllm/distributed/kv_transfer/kv_connector/v1/umbp/runtime/standalone.py

def normalize_endpoint(endpoint: str) -> str:
    """Return the gRPC address of a standalone server socket.

    The data plane depends on file-descriptor and IPC-handle passing over a
    Unix-domain socket, so only socket paths are accepted.
    """
    if endpoint.startswith(_UNIX_SCHEME):
        path = endpoint[len(_UNIX_SCHEME) :]
    elif "://" in endpoint:
        raise ValueError(
            "standalone UMBP endpoint must be a Unix socket path or unix:// "
            f"address, got {endpoint!r}"
        )
    else:
        path = endpoint
    if not path.startswith("/"):
        raise ValueError(
            f"standalone UMBP endpoint must be an absolute path, got {endpoint!r}"
        )
    return _UNIX_SCHEME + path