more _FutureGroup typing

coreweave · Apr 26, 2024 · f200987 · f200987
1 parent a55e119
commit f200987
Show file tree

Hide file tree

Showing 2 changed files with 5 additions and 5 deletions.
diff --git a/tensorizer/_futuregroup.py b/tensorizer/_futuregroup.py
@@ -3,7 +3,7 @@
 
 
 class _FutureGroup:
-    def __init__(self, futures: Sequence[concurrent.futures.Future]):
+    def __init__(self, futures):  # type: (Sequence[_Future]) -> None
         self.futures = futures
 
     def cancel(self) -> bool:

diff --git a/tensorizer/serialization.py b/tensorizer/serialization.py
@@ -4201,7 +4201,7 @@ def _prepare_for_write_opaque(
             w.set_min_file_version_number(OPAQUE_TENSORIZER_VERSION)
 
     @staticmethod
-    def _do_clone(write_spec, dependency: Optional[concurrent.futures.Future]):
+    def _do_clone(write_spec, dependency: Optional[_Future]):
         if dependency is not None:
             dependency.result(_TIMEOUT)
         write_spec.tensor = write_spec.tensor.clone().detach()
@@ -4234,7 +4234,7 @@ def _prepare_for_write_encryption(
                 ]
             )
 
-            clone_tasks = []
+            clone_tasks: List[_Future] = []
             for w in shared_write_specs[1:]:
                 clone_tasks.append(
                     self._computation_pool.submit(
@@ -4401,7 +4401,7 @@ def compute_sha256(
         for w in write_specs:
             old_tensor_data_task = w.tensor_data_task
 
-            hash_tasks = []
+            hash_tasks: List[_Future] = []
             if w.include_crc32:
                 crc32_task = self._computation_pool.submit(
                     compute_crc32, w, old_tensor_data_task
@@ -4418,7 +4418,7 @@ def compute_sha256(
                 self._jobs.extend(hash_tasks)
 
     def _do_encryption(self, write_specs: Sequence[_WriteSpec]) -> None:
-        def encrypt(write_spec, dependency: _Future):
+        def encrypt(write_spec, dependency: Optional[_Future]):
             if dependency is not None:
                 dependency.result(_TIMEOUT)
             try: