|
18 | 18 | import struct |
19 | 19 | from abc import ABC, abstractmethod |
20 | 20 | from concurrent.futures import ThreadPoolExecutor |
21 | | -from typing import Any, Callable, Optional, TypeAlias, Union |
| 21 | +from typing import Any, Callable, Optional, TypeAlias |
22 | 22 |
|
23 | 23 | import torch |
24 | 24 | from torch import Tensor |
@@ -94,7 +94,7 @@ def __init__(self, config: dict): |
94 | 94 | logger.info("YuanrongStorageClient: Create DsTensorClient to connect with yuanrong-datasystem backend!") |
95 | 95 |
|
96 | 96 | @staticmethod |
97 | | - def init(config: dict) -> Union["StorageStrategy", None]: |
| 97 | + def init(config: dict) -> Optional["StorageStrategy"]: |
98 | 98 | """Initialize only if NPU and torch_npu are available.""" |
99 | 99 | torch_npu_imported: bool = True |
100 | 100 | try: |
@@ -468,10 +468,12 @@ def _route_to_strategies( |
468 | 468 | ) -> dict[StorageStrategy, list[int]]: |
469 | 469 | """Groups item indices by the first strategy that supports them. |
470 | 470 |
|
471 | | - Used to route keys/values/custom_meta to storage backends by grouped indexes. |
| 471 | + Used to route data to storage strategies by grouped indexes. |
472 | 472 |
|
473 | 473 | Args: |
474 | | - items: A list of items (e.g., values for put, or custom_meta strings for get/clear). |
| 474 | + items: A list used to distinguish which storage strategy the data is routed to. |
| 475 | + e.g., route <keys, values> for put based on types of values, |
| 476 | + or route <keys, Optional[shapes], Optional[dtypes]> for get/clear based on strategy_tag. |
475 | 477 | The order must correspond to the original keys. |
476 | 478 | selector: A function that determines whether a strategy supports an item. |
477 | 479 | Signature: `(strategy: StorageStrategy, item: Any) -> bool`. |
|
0 commit comments