PERF: avoid repeated Xavier snapshot size scans

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()

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()

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()
+        )

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]: