add pyi file (#3)

hpcaitech · Jul 5, 2022 · 738a0c0 · 738a0c0
1 parent b14800a
commit 738a0c0
Showing 1 changed file with 11 additions and 0 deletions.
diff --git a/colo_nvme/_C/__init__.pyi b/colo_nvme/_C/__init__.pyi
@@ -0,0 +1,11 @@
+from torch import Tensor
+from typing import Callable, Optional
+
+
+class Offloader:
+    def __init__(self, filename: str, n_entries: int, backend: str = 'uring') -> None: ...
+    def write(self, tensor: Tensor, key: str, callback: Optional[Callable[[], None]] = None) -> None: ...
+    def read(self, tensor: Tensor, key: str, callback: Optional[Callable[[], None]] = None) -> None: ...
+    def sync_write_events(self) -> None: ...
+    def sync_read_events(self) -> None: ...
+    def synchronize(self) -> None: ...