Commit e1551e4
committed
[BugFix][KVCache] fix cupy device id caching and pickle for _match_result
## Motivation
修复两个 bug:
1. `transfer_manager.py` 中每次调用 `cp.cuda.runtime.getDevice()` 存在隐患,应在初始化时缓存为实例变量,保证后续操作使用一致的设备 ID。
2. `request.py` 的 `__getstate__` 未跳过 `_match_result`,该字段包含 BlockNode 树的父子循环引用,pickle 时会触发 `RecursionError`;同时补充 `__setstate__` 确保 unpickle 后字段恢复为安全默认值。
## Modifications
- `transfer_manager.py`:初始化时调用 `cp.cuda.runtime.getDevice()` 并缓存到 `self._cupy_device_id`,后续 `with cp.cuda.Device(...)` 和日志均使用该缓存值。
- `request.py`:
- `__getstate__` 中将 `_match_result` 加入跳过集合 `_SKIP_KEYS`,避免循环引用导致 pickle 失败。
- 新增 `__setstate__`,unpickle 后将 `_block_hasher` 和 `_match_result` 恢复为 `None`。
## Usage or Command1 parent 4b43eb7 commit e1551e4
2 files changed
Lines changed: 24 additions & 15 deletions
| Original file line number | Diff line number | Diff line change | |
|---|---|---|---|
| |||
92 | 92 | | |
93 | 93 | | |
94 | 94 | | |
95 | | - | |
| 95 | + | |
96 | 96 | | |
97 | 97 | | |
98 | | - | |
| 98 | + | |
99 | 99 | | |
100 | | - | |
| 100 | + | |
101 | 101 | | |
102 | 102 | | |
103 | 103 | | |
| |||
447 | 447 | | |
448 | 448 | | |
449 | 449 | | |
450 | | - | |
451 | 450 | | |
452 | 451 | | |
453 | | - | |
| 452 | + | |
454 | 453 | | |
455 | | - | |
| 454 | + | |
456 | 455 | | |
457 | 456 | | |
458 | 457 | | |
| |||
534 | 533 | | |
535 | 534 | | |
536 | 535 | | |
537 | | - | |
| 536 | + | |
538 | 537 | | |
539 | 538 | | |
540 | 539 | | |
| |||
640 | 639 | | |
641 | 640 | | |
642 | 641 | | |
643 | | - | |
| 642 | + | |
644 | 643 | | |
645 | 644 | | |
646 | 645 | | |
| |||
| Original file line number | Diff line number | Diff line change | |
|---|---|---|---|
| |||
453 | 453 | | |
454 | 454 | | |
455 | 455 | | |
456 | | - | |
| 456 | + | |
| 457 | + | |
| 458 | + | |
| 459 | + | |
| 460 | + | |
| 461 | + | |
457 | 462 | | |
458 | 463 | | |
459 | | - | |
460 | | - | |
461 | | - | |
462 | | - | |
463 | | - | |
| 464 | + | |
464 | 465 | | |
| 466 | + | |
| 467 | + | |
465 | 468 | | |
466 | 469 | | |
467 | | - | |
468 | 470 | | |
469 | 471 | | |
| 472 | + | |
| 473 | + | |
| 474 | + | |
| 475 | + | |
| 476 | + | |
| 477 | + | |
| 478 | + | |
| 479 | + | |
470 | 480 | | |
471 | 481 | | |
472 | 482 | | |
| |||
0 commit comments