lmcache--lmcache
wehub-resource-sync
76 行
1.9 KiB
Python
76 行
1.9 KiB
Python
# SPDX-License-Identifier: Apache-2.0
|
|
# Standard
|
|
import abc
|
|
import time
|
|
|
|
# Third Party
|
|
import torch
|
|
|
|
# First Party
|
|
from lmcache.logging import init_logger
|
|
from lmcache.utils import _lmcache_nvtx_annotate
|
|
|
|
logger = init_logger(__name__)
|
|
|
|
|
|
class Serializer(metaclass=abc.ABCMeta):
|
|
@abc.abstractmethod
|
|
def to_bytes(self, t: torch.Tensor) -> bytes:
|
|
"""
|
|
Serialize a pytorch tensor to bytes. The serialized bytes should contain
|
|
both the data and the metadata (shape, dtype, etc.) of the tensor.
|
|
|
|
Input:
|
|
t: the input pytorch tensor, can be on any device, in any shape,
|
|
with any dtype
|
|
|
|
Returns:
|
|
bytes: the serialized bytes
|
|
"""
|
|
raise NotImplementedError
|
|
|
|
|
|
class SerializerDebugWrapper(Serializer):
|
|
def __init__(self, s: Serializer):
|
|
self.s = s
|
|
|
|
def to_bytes(self, t: torch.Tensor) -> bytes:
|
|
start = time.perf_counter()
|
|
bs = self.s.to_bytes(t)
|
|
end = time.perf_counter()
|
|
|
|
logger.debug(f"Serialization took {end - start:.2f} seconds")
|
|
return bs
|
|
|
|
|
|
class Deserializer(metaclass=abc.ABCMeta):
|
|
def __init__(self, dtype):
|
|
self.dtype = dtype
|
|
|
|
@abc.abstractmethod
|
|
def from_bytes(self, bs: bytes) -> torch.Tensor:
|
|
"""
|
|
Deserialize a pytorch tensor from bytes.
|
|
|
|
Input:
|
|
bytes: a stream of bytes
|
|
|
|
Output:
|
|
torch.Tensor: the deserialized pytorch tensor
|
|
"""
|
|
raise NotImplementedError
|
|
|
|
|
|
class DeserializerDebugWrapper(Deserializer):
|
|
def __init__(self, d: Deserializer):
|
|
self.d = d
|
|
|
|
@_lmcache_nvtx_annotate
|
|
def from_bytes(self, t: bytes) -> torch.Tensor:
|
|
start = time.perf_counter()
|
|
ret = self.d.from_bytes(t)
|
|
end = time.perf_counter()
|
|
|
|
logger.debug(f"Deserialization took {(end - start) * 1000:.2f} ms")
|
|
return ret
|