class _MoriWorkerHandle(UMBPWorkerHandle):
def __init__(
self,
client: _MoriClient,
namespace: str,
topology: RankTopology,
lookup_dir: str,
max_workers: int,
timeout_s: float,
layout: KVLayoutDescriptor | None = None,
lookup_instance: str = "",
) -> None:
self.client = client
self._namespace = namespace
self._topology = topology
self._lookup_dir = lookup_dir
self._lookup_instance = lookup_instance
self._lookup_server: _MoriLookupServer | None = None
self._registered_storages: set[int] = set()
self._gpu_devices: set[int] = set()
self._executor = ThreadPoolExecutor(
max_workers=max_workers,
thread_name_prefix="umbp-embedded-transfer",
)
self._job_lock = threading.Lock()
self._futures: dict[int, Future[TransferJobState]] = {}
self._store_deadlines: dict[int, float] = {}
self._load_deadlines: dict[int, float] = {}
self._timeout_s = timeout_s
self._layout = layout
self._load_layouts: dict[int, _BlockLoadLayout] = {}
def register_buffers(self, kv_caches: dict[str, torch.Tensor]) -> None:
try:
from mori.cpp import MemoryLocationType
except ImportError as exc:
raise RuntimeError("MORI UMBP Python bindings are unavailable") from exc
for cache in kv_caches.values():
storage = cache.untyped_storage()
storage_ptr = storage.data_ptr()
if storage_ptr in self._registered_storages:
continue
location = (
MemoryLocationType.GPU
if cache.device.type == "cuda"
else MemoryLocationType.CPU
)
device = cache.device.index if cache.device.index is not None else -1
if not self.client.register_memory(
storage_ptr, storage.nbytes(), location, device
):
raise RuntimeError(
f"MORI UMBP failed to register KV storage 0x{storage_ptr:x}"
)
self._registered_storages.add(storage_ptr)
if location == MemoryLocationType.GPU and device >= 0:
self._gpu_devices.add(device)
self._register_load_layouts(kv_caches)
if self._lookup_server is None:
self._lookup_server = _MoriLookupServer(
_lookup_socket_path(
self._namespace,
self._topology.local_namespace,
self._lookup_dir,
self._lookup_instance,
),
self,
)
self._lookup_server.start()
def _register_load_layouts(self, kv_caches: dict[str, torch.Tensor]) -> None:
self._load_layouts = {}
if self._layout is None:
return
for group_id in {r.group_id for r in self._layout.regions}:
bases: list[int] = []
strides: list[int] = []
sizes: list[int] = []
offsets: list[int] = []
counts: list[int] = []
for region in self._layout.regions:
if region.group_id != group_id:
continue
cache = kv_caches.get(region.layer_name)
if cache is None or cache.ndim == 0:
break
stride = cache.stride(0) * cache.element_size()
if (
stride <= 0
or stride > region.block_stride
or region.block_stride % stride
or region.block_bytes > stride
):
break
bases.append(cache.data_ptr())
strides.append(stride)
offsets.append(sum(sizes))
sizes.append(region.block_bytes)
counts.append(cache.shape[0])
else:
self._load_layouts[group_id] = _BlockLoadLayout(
np.asarray(bases, dtype=np.uint64),
np.asarray(strides, dtype=np.uint64),
sizes,
offsets,
min(counts),
sum(sizes),
)
def _bulk_load_args(
self, plans: Sequence[BlockTransferPlan]
) -> tuple[_LoadArgs, tuple[int, ...]] | None:
groups: dict[int, list[tuple[int, int]]] = {}
entries: Iterable[tuple[int | None, int]]
if isinstance(plans, (BlockLoadBatch, BlockStoreBatch)):
entries = zip(plans.group_ids, plans.block_ids, strict=True)
keys = plans.keys
else:
if any(p.ranges for p in plans):
return None
entries = ((p.group_id, p.block_id) for p in plans)
keys = [p.key for p in plans]
for index, (group_id, block_id) in enumerate(entries):
if group_id is None:
return None
layout = self._load_layouts.get(group_id)
if layout is None or not 0 <= block_id < layout.num_blocks:
return None
groups.setdefault(group_id, []).append((index, block_id))
pointers: list[list[int]] = [[]] * len(plans)
sizes: list[list[int]] = [[]] * len(plans)
offsets: list[list[int]] = [[]] * len(plans)
plan_bytes = [0] * len(plans)
for group_id, group_entries in groups.items():
layout = self._load_layouts[group_id]
ids = np.asarray([block for _, block in group_entries], dtype=np.uint64)
addresses = (
layout.bases[None, :] + ids[:, None] * layout.strides[None, :]
).tolist()
for (index, _), row in zip(group_entries, addresses, strict=True):
pointers[index] = row
sizes[index] = layout.sizes
offsets[index] = layout.offsets
plan_bytes[index] = layout.num_bytes
return (
(keys, pointers, sizes, offsets),
tuple(plan_bytes),
)
def load_blocks(
self, plans: Sequence[BlockTransferPlan]
) -> TransferJobState | None:
prepared = self._bulk_load_args(plans)
if prepared is None:
return None
args, plan_bytes = prepared
return self._submit_load(
plans if isinstance(plans, BlockLoadBatch) else tuple(plans),
args,
plan_bytes,
)
def batch_exists(self, keys: Sequence[str]) -> Sequence[bool]:
return [bool(value) for value in self.client.batch_exists(keys)]
def clear(self) -> bool:
return bool(self.client.clear())
@staticmethod
def _range_args(plans: Sequence[BlockTransferPlan], *, for_store: bool = True):
keys, object_sizes, pointers, sizes, offsets = [], [], [], [], []
for plan in plans:
plan_pointers, plan_sizes, plan_offsets = [], [], []
object_size = 0
for item in plan.ranges:
plan_pointers.append(item.base_address)
plan_sizes.append(item.length)
plan_offsets.append(item.object_offset)
if for_store:
end = item.object_offset + item.length
if end > object_size:
object_size = end
keys.append(plan.key)
if for_store:
object_sizes.append(object_size)
pointers.append(plan_pointers)
sizes.append(plan_sizes)
offsets.append(plan_offsets)
return keys, object_sizes, pointers, sizes, offsets
def load(self, plans: Sequence[BlockTransferPlan]) -> TransferJobState:
return self._submit_load(tuple(plans))
def _submit_load(
self,
plans: tuple[BlockTransferPlan, ...] | BlockLoadBatch,
range_args: _LoadArgs | None = None,
plan_bytes: tuple[int, ...] | None = None,
) -> TransferJobState:
job = TransferJobState(plans, plan_bytes=plan_bytes)
job.start()
future = self._executor.submit(self._load_sync, job, plans, range_args)
with self._job_lock:
self._futures[id(job)] = future
return job
def _load_sync(
self,
job: TransferJobState,
plans: tuple[BlockTransferPlan, ...] | BlockLoadBatch,
range_args: _LoadArgs | None = None,
) -> TransferJobState:
with self._job_lock:
self._load_deadlines[id(job)] = time.monotonic() + self._timeout_s
if not plans:
job.complete()
return job
if range_args is None:
keys, _, pointers, sizes, offsets = self._range_args(plans, for_store=False)
else:
keys, pointers, sizes, offsets = range_args
request_id = (
plans.request_id
if isinstance(plans, BlockLoadBatch)
else plans[0].request_id
)
started_at = time.monotonic()
logger.debug(
"MORI UMBP range load started request=%s plans=%d ranges=%d bytes=%d",
request_id,
len(plans),
sum(len(plan_sizes) for plan_sizes in sizes),
sum(sum(plan_sizes) for plan_sizes in sizes),
)
results = self.client.batch_get_ranges_into_ptr(keys, pointers, sizes, offsets)
logger.debug(
"MORI UMBP range load returned request=%s plans=%d elapsed=%.3fs",
request_id,
len(plans),
time.monotonic() - started_at,
)
completed = [key for key, ok in zip(keys, results, strict=True) if ok]
failed = [key for key, ok in zip(keys, results, strict=True) if not ok]
if completed:
job.complete(completed)
if failed:
job.fail(failed, "MORI UMBP range load failed")
return job
def store_blocks(
self, plans: Sequence[BlockTransferPlan]
) -> TransferJobState | None:
"""Store whole blocks from registered layouts; None requests ranges."""
prepared = self._bulk_load_args(plans)
if prepared is None:
return None
(keys, pointers, sizes, offsets), plan_bytes = prepared
return self._submit_store(
plans if isinstance(plans, BlockStoreBatch) else tuple(plans),
(keys, list(plan_bytes), pointers, sizes, offsets),
plan_bytes,
)
def store(self, plans: Sequence[BlockTransferPlan]) -> TransferJobState:
return self._submit_store(tuple(plans))
def _submit_store(
self,
plans: tuple[BlockTransferPlan, ...] | BlockStoreBatch,
range_args: _StoreArgs | None = None,
plan_bytes: tuple[int, ...] | None = None,
) -> TransferJobState:
job = TransferJobState(plans, plan_bytes=plan_bytes)
job.start()
ready_events: list[torch.Event] = []
for device in self._gpu_devices:
with torch.device(f"cuda:{device}"):
event = torch.Event()
event.record(torch.accelerator.current_stream())
ready_events.append(event)
future = self._executor.submit(
self._store_sync, job, plans, ready_events, range_args
)
with self._job_lock:
self._futures[id(job)] = future
return job
def _store_sync(
self,
job: TransferJobState,
plans: tuple[BlockTransferPlan, ...] | BlockStoreBatch,
ready_events: list[torch.Event],
range_args: _StoreArgs | None = None,
) -> TransferJobState:
# Timed from here: a store queued behind others has not stalled.
with self._job_lock:
self._store_deadlines[id(job)] = time.monotonic() + self._timeout_s
for event in ready_events:
event.synchronize()
if not plans:
job.complete()
return job
keys, object_sizes, pointers, sizes, offsets = (
self._range_args(plans) if range_args is None else range_args
)
results = self.client.batch_put_ranges_from_ptr(
keys, object_sizes, pointers, sizes, offsets
)
failed_indices = [index for index, ok in enumerate(results) if not ok]
for attempt in range(_STORE_RETRIES):
if not failed_indices:
break
self.client.flush()
time.sleep(_STORE_RETRY_BACKOFF_S * (attempt + 1))
for index in failed_indices.copy():
retry = self.client.batch_put_ranges_from_ptr(
[keys[index]],
[object_sizes[index]],
[pointers[index]],
[sizes[index]],
[offsets[index]],
)
if retry and retry[0]:
results[index] = True
failed_indices.remove(index)
if failed_indices:
logger.warning(
"MORI UMBP store retry %d left %d/%d objects pending",
attempt + 1,
len(failed_indices),
len(plans),
)
if not self.client.flush():
job.fail([plan.key for plan in plans], "MORI UMBP flush failed")
return job
completed = [plan.key for plan, ok in zip(plans, results, strict=True) if ok]
failed = [plan.key for plan, ok in zip(plans, results, strict=True) if not ok]
if completed:
job.complete(completed)
if failed:
job.fail(failed, "MORI UMBP range store failed")
return job
def wait(self, job: TransferJobState) -> TransferJobState:
with self._job_lock:
future = self._futures.get(id(job))
if future is None:
return job
try:
return future.result(timeout=self._timeout_s)
except TimeoutError as exc:
if not future.done():
# The transfer still owns GPU pointers. Abort this engine step
# rather than reporting completion and allowing block reuse.
raise TimeoutError(
"MORI UMBP transfer timed out with buffers still in use"
) from exc
job.fail([plan.key for plan in job.plans], str(exc))
return job
except Exception as exc:
job.fail([plan.key for plan in job.plans], str(exc))
return job
finally:
if future.done():
with self._job_lock:
self._futures.pop(id(job), None)
self._store_deadlines.pop(id(job), None)
self._load_deadlines.pop(id(job), None)
def poll(self, job: TransferJobState) -> TransferJobState | None:
with self._job_lock:
future = self._futures.get(id(job))
if future is None:
if job.status in (
TransferJobStatus.COMPLETED,
TransferJobStatus.FAILED,
TransferJobStatus.CANCELLED,
):
return job
return None
if not future.done():
with self._job_lock:
deadline = self._store_deadlines.get(
id(job), self._load_deadlines.get(id(job))
)
if deadline is not None and time.monotonic() >= deadline:
raise TimeoutError(
"MORI UMBP transfer timed out with buffers still in use"
)
return None
with self._job_lock:
self._futures.pop(id(job), None)
self._store_deadlines.pop(id(job), None)
self._load_deadlines.pop(id(job), None)
try:
return future.result()
except Exception as exc:
job.fail([plan.key for plan in job.plans], str(exc))
return job
def publish(self, job: TransferJobState) -> None:
if job.status.value != "completed":
raise RuntimeError("cannot publish an incomplete MORI UMBP job")
def cancel(self, job: TransferJobState) -> TransferJobState:
with self._job_lock:
future = self._futures.get(id(job))
if future is None:
return job
if future.cancel():
with self._job_lock:
self._futures.pop(id(job), None)
self._store_deadlines.pop(id(job), None)
self._load_deadlines.pop(id(job), None)
job.cancel("preempted")
return job
return self.wait(job)
def take_evicted_keys(self) -> Sequence[str]:
return ()
def close(self) -> None:
self._executor.shutdown(wait=True, cancel_futures=True)
if self._lookup_server is not None:
self._lookup_server.close()
for storage_ptr in self._registered_storages:
self.client.deregister_memory(storage_ptr)
self._registered_storages.clear()
close = getattr(self.client, "close", None)
if callable(close):
close()
else:
self.client.flush()