diff --git a/xinference/model/llm/vllm/xavier/gpu_transfer.py b/xinference/model/llm/vllm/xavier/gpu_transfer.py index 3f1d15690d..2b7e20b044 100644 --- a/xinference/model/llm/vllm/xavier/gpu_transfer.py +++ b/xinference/model/llm/vllm/xavier/gpu_transfer.py @@ -102,7 +102,12 @@ async def run(self, function, *args): def completed(future): self.tasks.discard(future) if not future.cancelled(): - future.exception() # Observe failures even after repeated cancellation. + error = future.exception() + if error is not None: + logger.warning( + "Xavier GPU transfer task failed", + exc_info=(type(error), error, error.__traceback__), + ) task.add_done_callback(completed) return await finish_before_cancel(task) @@ -156,6 +161,7 @@ async def _close(self): self.caches.clear() self.send_buffer = self.recv_buffer = None self.store._gpu_lru.clear() + self.store._block_sizes.clear() self.store.blocks.clear() self.store.ready.clear() self.store.tiers.clear() diff --git a/xinference/model/llm/vllm/xavier/test/test_gpu_transfer.py b/xinference/model/llm/vllm/xavier/test/test_gpu_transfer.py index f7ca834d2c..539551c5e9 100644 --- a/xinference/model/llm/vllm/xavier/test/test_gpu_transfer.py +++ b/xinference/model/llm/vllm/xavier/test/test_gpu_transfer.py @@ -363,6 +363,7 @@ def copy(layer, keys, tensors): assert list(r.store.blocks) == [1] assert r.store.publish([2], {"K"}) == [] assert r.store.leases == {"2:live": {1}} + assert r.store._block_sizes == {1: 4} # Published BF16 block survives alone. assert r.store.read("K", [1]).tolist() == [[1, -1]] assert_gpu_lru_consistent(r.store) @@ -926,6 +927,7 @@ def fence(*args): assert not r.store.blocks and not r.store.ready and not r.store.tiers assert not r.store.logical_dtypes and not r.store.leases and not r.store.evicted assert not r.store._gpu_lru + assert not r.store._block_sizes assert r.store.counts == {"gpu": 0, "cpu": 0} assert r.send_buffer is None and r.recv_buffer is None assert not r.recv_refs and not r.caches @@ -1284,3 +1286,27 @@ async def test_new_snapshot_resets_reused_layer_lru_shortcut(monkeypatch): assert list(r.store.blocks) == [3, 1, 2] assert list(r.store._gpu_lru) == [3, 2] assert r.store.tiers == {1: "cpu", 2: "gpu", 3: "gpu"} + + +@pytest.mark.asyncio +async def test_cancelled_caller_logs_background_transfer_failure(monkeypatch, caplog): + r = runtime(monkeypatch) + started, finish = asyncio.Event(), asyncio.Event() + + async def fail(): + started.set() + await finish.wait() + raise RuntimeError("background write failed") + + job = asyncio.create_task(r.run(fail)) + await started.wait() + job.cancel() + await asyncio.sleep(0) + assert not job.done() # Cancellation still drains the protected operation. + finish.set() + with pytest.raises(asyncio.CancelledError): + await job + assert "Xavier GPU transfer task failed" in caplog.text + assert "background write failed" in caplog.text + assert not r.tasks + await r.close() diff --git a/xinference/model/llm/vllm/xavier/test/test_tiered_snapshot.py b/xinference/model/llm/vllm/xavier/test/test_tiered_snapshot.py index c8f0967d1b..6573b7a58d 100644 --- a/xinference/model/llm/vllm/xavier/test/test_tiered_snapshot.py +++ b/xinference/model/llm/vllm/xavier/test/test_tiered_snapshot.py @@ -314,3 +314,50 @@ def fail(*args, **kwargs): stage(s, 2) assert not s._gpu_lru assert not s.blocks + + +def test_incremental_size_uses_physical_dtype_and_cleans_failed_slots(monkeypatch): + s = store(gpu=2) + s.stage("K", [1], torch.ones(1, 2, dtype=torch.float32), torch.bfloat16) + assert s._block_sizes == {1: 8} + # Immutable hits must not add the same layer twice. + s.stage("K", [1], torch.ones(1, 2, dtype=torch.float32)) + assert s._block_sizes == {1: 8} + with pytest.raises(ValueError, match="full-block"): + s.stage("V", [1], torch.ones(1, 3)) + assert s._block_sizes == {1: 8} + + def fail(*args, **kwargs): + raise RuntimeError("copy failed") + + monkeypatch.setattr(torch.Tensor, "to", fail) + # A later-layer failure keeps the existing layer and its physical byte count. + with pytest.raises(RuntimeError, match="copy failed"): + s.stage("V", [1], torch.ones(1, 2)) + assert set(s.blocks[1]) == {"K"} + assert s._block_sizes == {1: 8} + assert s._block_sizes[1] == sum( + value.numel() * value.element_size() for value in s.blocks[1].values() + ) + s._drop(1) + assert not s._block_sizes + + with pytest.raises(RuntimeError, match="copy failed"): + stage(s, 2) + assert not s._block_sizes + + +def test_incremental_size_survives_demotion_eviction_and_key_reuse(): + s = store(gpu=1, cpu=1) + for key in [1, 2]: + stage(s, key) + assert s._block_sizes == {1: 16, 2: 16} + stage(s, 3) + assert s._block_sizes == {2: 16, 3: 16} + # Reusing an evicted hash starts accounting from its new content. + s.stage("K", [1], torch.ones(1, 1)) + assert s._block_sizes == {3: 16, 1: 4} + for key, layers in s.blocks.items(): + assert s._block_sizes[key] == sum( + value.numel() * value.element_size() for value in layers.values() + ) diff --git a/xinference/model/llm/vllm/xavier/tiered_snapshot.py b/xinference/model/llm/vllm/xavier/tiered_snapshot.py index ad1738bcd6..3167454824 100644 --- a/xinference/model/llm/vllm/xavier/tiered_snapshot.py +++ b/xinference/model/llm/vllm/xavier/tiered_snapshot.py @@ -41,6 +41,8 @@ def __init__( self.counts = {"gpu": 0, "cpu": 0} self.metrics = dict(demotions=0, gpu_hits=0, cpu_hits=0, skipped=0) self.block_bytes = block_bytes + # Avoid rescanning every previous layer when validating the next one. + self._block_sizes: Dict[int, int] = {} def touch(self, key: int) -> None: self.blocks.move_to_end(key) @@ -52,6 +54,7 @@ def _drop(self, key: int) -> None: self.counts[self.tiers.pop(key)] -= 1 del self.blocks[key] del self.logical_dtypes[key] + del self._block_sizes[key] self.ready.discard(key) self.evicted.add(key) @@ -120,6 +123,7 @@ def _admit(self, key: int, pinned: set) -> bool: return False self.blocks[key] = {} self.logical_dtypes[key] = {} + self._block_sizes[key] = 0 self.tiers[key] = tier if tier == "gpu": self._gpu_lru[key] = None @@ -138,7 +142,7 @@ def stage( existing = self.blocks.get(key, {}) if layer in existing: continue - used = sum(t.numel() * t.element_size() for t in existing.values()) + used = self._block_sizes.get(key, 0) if used + tensor.numel() * tensor.element_size() > self.block_bytes: raise ValueError("Snapshot exceeds configured full-block size") pinned = set().union(*self.leases.values()) if self.leases else set() @@ -155,6 +159,7 @@ def stage( tensor.detach().to(device, copy=True).contiguous() ) self.logical_dtypes[key][layer] = logical_dtype or tensors.dtype + self._block_sizes[key] += tensor.numel() * tensor.element_size() except Exception: # A failed first layer must not consume an empty slot. if not self.blocks[key]: