Skip to content

REF: extract shared Xavier core and Torch backend #5629

New issue

Have a question about this project? Sign up for a free GitHub account to open an issue and contact its maintainers and the community.

Sign up for GitHub

By clicking “Sign up for GitHub”, you agree to our terms of service and privacy statement. We’ll occasionally send you account related emails.

Already on GitHub? Sign in to your account

Merged
qinxuye merged 2 commits into from
Oct 4, 2026
Merged

REF: extract shared Xavier core and Torch backend #5629

Show file tree
Hide file tree
Changes from all commits
Commits
Show all changes
2 commits
Select commit Hold shift + click to select a range
d214612
REF: extract shared Xavier core and Torch backend
qinxuye Oct 4, 2026
7f6b39d
fix: share Xavier transfer actor UID across coordinator and adapter
qinxuye Oct 4, 2026
File filter

Filter by extension

Filter by extension
Conversations
Failed to load comments.
Loading
Jump to
Jump to file
Failed to load files.
Loading
Diff view
Diff view
18 changes: 8 additions & 10 deletions xinference/core/supervisor.py
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -1059,10 +1059,10 @@ async def signal_handler():
signal.SIGTERM, lambda: asyncio.create_task(signal_handler())
)

from ..model.llm.vllm.xavier.block_tracker import VLLMBlockTracker
from ..model.llm.vllm.xavier.collective_manager import CollectiveManager
from ..model.llm.xavier.block_tracker import BlockTracker
from ..model.llm.xavier.collective_manager import CollectiveManager

self._block_tracker_mapping: Dict[str, xo.ActorRefType[VLLMBlockTracker]] = {} # type: ignore
self._block_tracker_mapping: Dict[str, xo.ActorRefType[BlockTracker]] = {} # type: ignore
self._collective_manager_mapping: Dict[ # type: ignore
str, xo.ActorRefType[CollectiveManager]
] = {}
Expand Down Expand Up @@ -3099,7 +3099,7 @@ async def launch_builtin_model(
and model_engine is not None
and model_engine.lower() == "vllm"
)
from ..model.llm.vllm.xavier.transport import validate_gpu_cache_budget
from ..model.llm.xavier.transport import validate_gpu_cache_budget

gpu_cache_bytes = validate_gpu_cache_budget(
kwargs.pop("xavier_gpu_cache_bytes", None), enable_xavier, replica
Expand Down Expand Up @@ -3266,15 +3266,13 @@ async def _launch_model():
nonlocal download_hub
try:
if enable_xavier:
from ..model.llm.vllm.xavier.block_tracker import VLLMBlockTracker
from ..model.llm.vllm.xavier.collective_manager import (
CollectiveManager,
)
from ..model.llm.xavier.block_tracker import BlockTracker
from ..model.llm.xavier.collective_manager import CollectiveManager

self._block_tracker_mapping[model_uid] = await xo.create_actor(
VLLMBlockTracker,
BlockTracker,
address=self.address,
uid=f"{VLLMBlockTracker.default_uid()}-{model_uid}",
uid=f"{BlockTracker.default_uid()}-{model_uid}",
)
self._collective_manager_mapping[model_uid] = await xo.create_actor(
CollectiveManager,
Expand Down
62 changes: 8 additions & 54 deletions xinference/model/llm/vllm/xavier/actor_loop.py
View file
Open in desktop
Original file line number Diff line number Diff line change
@@ -1,60 +1,14 @@
# Copyright 2022-2026 Xinference Holdings Pte. Ltd
# Licensed under the Apache License, Version 2.0.
"""Share the synchronous actor RPC loop among connectors in an EngineCore thread."""
"""Compatibility alias for the shared Xavier actor_loop module."""

import asyncio
import logging
import os
import threading
import sys
from typing import TYPE_CHECKING

logger = logging.getLogger(__name__)
_SHUTDOWN_TIMEOUT = 5.0
_local = threading.local()
from ...xavier import actor_loop as _implementation

if TYPE_CHECKING:
from ...xavier.actor_loop import acquire_actor_loop as acquire_actor_loop
from ...xavier.actor_loop import release_actor_loop as release_actor_loop

def acquire_actor_loop() -> asyncio.AbstractEventLoop:
# Never reuse a parent's event loop after fork, or a different thread's loop.
if getattr(_local, "pid", None) != os.getpid() or not getattr(_local, "users", 0):
_local.loop = asyncio.new_event_loop()
_local.pid = os.getpid()
_local.users = 0
_local.users += 1
return _local.loop


def release_actor_loop(loop: asyncio.AbstractEventLoop) -> None:
if (
getattr(_local, "pid", None) != os.getpid()
or getattr(_local, "loop", None) is not loop
):
return
_local.users -= 1
if _local.users:
return
try:
# xoscar listener tasks close their clients in cancellation cleanup.
pending = asyncio.all_tasks(loop)
for task in pending:
task.cancel()

async def drain():
if pending:
await asyncio.gather(*pending, return_exceptions=True)
await loop.shutdown_asyncgens()

cleanup = loop.create_task(drain())
_, unfinished = loop.run_until_complete(
asyncio.wait({cleanup}, timeout=_SHUTDOWN_TIMEOUT)
)
if unfinished:
logger.warning("Timed out draining Xavier actor loop during shutdown")
cleanup.cancel()
# Give cancellation one turn, without waiting on unresponsive peers.
loop.run_until_complete(asyncio.sleep(0))
elif not cleanup.cancelled():
cleanup.result()
finally:
try:
loop.close()
finally:
del _local.loop
sys.modules[__name__] = _implementation
153 changes: 9 additions & 144 deletions xinference/model/llm/vllm/xavier/block_tracker.py
View file
Open in desktop
Original file line number Diff line number Diff line change
@@ -1,149 +1,14 @@
# Copyright 2022-2026 Xinference Holdings Pte. Ltd
#
# Licensed under the Apache License, Version 2.0 (the "License");
# you may not use this file except in compliance with the License.
# You may obtain a copy of the License at
#
# http://www.apache.org/licenses/LICENSE-2.0
#
# Unless required by applicable law or agreed to in writing, software
# distributed under the License is distributed on an "AS IS" BASIS,
# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
# See the License for the specific language governing permissions and
# limitations under the License.
import random
from typing import Dict, List, Optional, Set, Tuple
# Licensed under the Apache License, Version 2.0.
"""Compatibility alias for the shared Xavier block_tracker module."""

import xoscar as xo
import sys
from typing import TYPE_CHECKING

from ...xavier import block_tracker as _implementation

class VLLMBlockTracker(xo.StatelessActor):
@classmethod
def default_uid(cls):
return f"vllm-block-tracker-actor"
if TYPE_CHECKING:
from ...xavier.block_tracker import BlockTracker as BlockTracker
from ...xavier.block_tracker import VLLMBlockTracker as VLLMBlockTracker

def __init__(self):
super().__init__()
# engine -> hash -> (rank, block_id)
self._hash_to_rank_and_block_id: Dict[int, Dict[int, Set[Tuple[int, int]]]] = {} # type: ignore
# engine -> rank -> (hash, block_id)
self._rank_to_hash_and_block_id: Dict[int, Dict[int, Set[Tuple[int, int]]]] = {} # type: ignore
self._unavailable_ranks: Set[int] = set() # type: ignore

def register_blocks(
self, virtual_engine: int, block_infos: List[Tuple[int, int]], rank: int
):
if virtual_engine not in self._hash_to_rank_and_block_id:
self._hash_to_rank_and_block_id[virtual_engine] = {}
hash_to_rank_and_block_id = self._hash_to_rank_and_block_id[virtual_engine]

if virtual_engine not in self._rank_to_hash_and_block_id:
self._rank_to_hash_and_block_id[virtual_engine] = {}
rank_to_hash_and_block_id = self._rank_to_hash_and_block_id[virtual_engine]
if rank not in rank_to_hash_and_block_id:
rank_to_hash_and_block_id[rank] = set()
hash_and_block_id = rank_to_hash_and_block_id[rank]

# Replace-on-register: a vLLM block id can be reused for new content
# (its KV cache slot is reallocated), so before registering the new
# (hash, block_id) mappings, drop any stale (old_hash, block_id) this
# rank still holds for the same block ids. Without this a stale hash
# keeps pointing at a block that now holds different KV, and a later
# query on that hash would load the wrong KV. This also bounds the
# tracker to one entry per (rank, block_id).
new_block_ids = {block_id for _, block_id in block_infos}
new_pairs = set(block_infos)
stale = [
(stale_hash, stale_block_id)
for stale_hash, stale_block_id in hash_and_block_id
if stale_block_id in new_block_ids
and (stale_hash, stale_block_id) not in new_pairs
]
for stale_hash, stale_block_id in stale:
hash_and_block_id.discard((stale_hash, stale_block_id))
stale_entries = hash_to_rank_and_block_id.get(stale_hash)
if stale_entries is not None:
stale_entries.discard((rank, stale_block_id))
if not stale_entries:
hash_to_rank_and_block_id.pop(stale_hash, None)

for hash_content, block_id in block_infos:
hash_to_rank_and_block_id.setdefault(hash_content, set()).add(
(rank, block_id)
)
hash_and_block_id.add((hash_content, block_id))

def query_blocks(
self,
virtual_engine: int,
hash_contents: List[Tuple[int, int]],
exclude_rank: Optional[int] = None,
) -> Dict[int, Set[Tuple[int, int, int]]]:
if virtual_engine not in self._hash_to_rank_and_block_id:
return {}
hash_to_rank_and_block_id = self._hash_to_rank_and_block_id[virtual_engine]
remote: Dict[int, Set[Tuple[int, int, int]]] = {}
for hash_content, _id in hash_contents:
if (
hash_content in hash_to_rank_and_block_id
) and hash_to_rank_and_block_id[hash_content]:
# exclude ranks that are in the recovery process or the caller's own rank
rank_and_block_id = [
(r, b)
for r, b in hash_to_rank_and_block_id[hash_content]
if r not in self._unavailable_ranks and r != exclude_rank
]
if rank_and_block_id:
# TODO: Randomly select here, and try to distribute requests as evenly as possible.
# There may be better methods in the future.
rank, block_id = random.choice(rank_and_block_id)
if rank not in remote:
remote[rank] = {
(hash_content, block_id, _id),
}
else:
remote[rank].add((hash_content, block_id, _id))
return remote

def unregister_block(self, virtual_engine: int, rank: int, block_id: int):
self.unregister_blocks(virtual_engine, rank, [block_id])

def unregister_blocks(self, virtual_engine: int, rank: int, block_ids: List[int]):
entries = self._rank_to_hash_and_block_id.get(virtual_engine, {}).get(rank)
hashes = self._hash_to_rank_and_block_id.get(virtual_engine, {})
if not entries:
return
removed_ids = set(block_ids)
removed = {
(content_hash, block_id)
for content_hash, block_id in entries
if block_id in removed_ids
}
entries.difference_update(removed)
for content_hash, block_id in removed:
locations = hashes.get(content_hash)
if locations is not None:
locations.discard((rank, block_id))
if not locations:
hashes.pop(content_hash, None)

def unregister_rank(self, rank: int):
"""
This rank is in the recovery process, and its query results will be excluded.
"""
self._unavailable_ranks.add(rank)

def register_rank(self, rank: int):
"""
After recovery is successful, clear all stale data of the rank and mark the rank as available.
"""
for _, rank_to_hash_and_block_id in self._rank_to_hash_and_block_id.items():
rank_to_hash_and_block_id.pop(rank, None)

for _, hash_to_rank_and_block_id in self._hash_to_rank_and_block_id.items():
for _, rank_and_block_id in hash_to_rank_and_block_id.items():
to_delete = [(r, b) for r, b in rank_and_block_id if r == rank]
if to_delete:
rank_and_block_id.difference_update(to_delete)

self._unavailable_ranks.discard(rank)
sys.modules[__name__] = _implementation
96 changes: 8 additions & 88 deletions xinference/model/llm/vllm/xavier/collective.py
View file
Open in desktop
Original file line number Diff line number Diff line change
@@ -1,93 +1,13 @@
# Copyright 2022-2026 Xinference Holdings Pte. Ltd
#
# Licensed under the Apache License, Version 2.0 (the "License");
# you may not use this file except in compliance with the License.
# You may obtain a copy of the License at
#
# http://www.apache.org/licenses/LICENSE-2.0
#
# Unless required by applicable law or agreed to in writing, software
# distributed under the License is distributed on an "AS IS" BASIS,
# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
# See the License for the specific language governing permissions and
# limitations under the License.
import logging
from typing import List, Optional
# Licensed under the Apache License, Version 2.0.
"""Compatibility alias for the shared Xavier collective module."""

from .transport import get_transport_host
import sys
from typing import TYPE_CHECKING

logger = logging.getLogger(__name__)
from ...xavier import collective as _implementation

if TYPE_CHECKING:
from ...xavier.collective import CollectiveRank as CollectiveRank

class CollectiveRank:
def __init__(
self,
rank: int,
world_size: int,
rank_address: str,
store_address: str,
store_port: int,
world_addresses: List[str],
):
self._rank = rank
self._world_size = world_size
self._rank_address = rank_address
self._world_addresses = world_addresses
self._store_address = store_address
self._store_port = store_port
self._device = None
self._tcp_store = None
self._context = None

def init_rank(self):
from xoscar.collective import xoscar_pygloo as xp

logger.debug(
f"init_rank called: rank={self._rank}, world_size={self._world_size}, "
f"rank_address={self._rank_address}, store_address={self._store_address}, "
f"store_port={self._store_port}, world_addresses={self._world_addresses}"
)

self._context = xp.rendezvous.Context(rank=self._rank, size=self._world_size)

attr = xp.transport.tcp.attr(get_transport_host(self._rank_address))
self._device = xp.transport.tcp.CreateDevice(attr)

opt = xp.rendezvous.TCPStoreOptions()
opt.port = self._store_port
opt.numWorkers = self._world_size
opt.isServer = self._rank == 0
opt.waitWorkers = False

if self._store_address is None:
raise ValueError(
f"store_address is None for rank {self._rank}. "
f"This usually means the rank 0 model was not properly initialized before other ranks."
)
if self._store_port is None:
raise ValueError(
f"store_port is None for rank {self._rank}. "
f"This usually means the rank 0 model was not properly initialized before other ranks."
)

self._tcp_store = xp.rendezvous.TCPStore(self._store_address, opt)
if self._world_addresses:
self.connect_full_mesh()

def connect_full_mesh(
self, prefix: Optional[str] = None, world_addresses: Optional[List[str]] = None
):
from xoscar.collective import xoscar_pygloo as xp

assert self._device is not None
assert self._tcp_store is not None
assert self._context is not None
if world_addresses is not None:
self._world_addresses = world_addresses
prefix_store = xp.rendezvous.PrefixStore(
prefix or str(self._world_size), self._tcp_store
)
self._context.connectFullMesh(prefix_store, self._device)
logger.debug(
f"Rank {self._rank} arrives successfully, world addresses: {self._world_addresses}"
)
sys.modules[__name__] = _implementation
Loading
Loading