No more future annotations

src/exo/__main__.py

@@ -1,5 +1,3 @@
-from __future__ import annotations
-
 import sys
 from collections.abc import Sequence
 from multiprocessing import freeze_support

src/exo/master/adapters/ollama.py

@@ -1,5 +1,3 @@
-from __future__ import annotations
-
 import json
 from collections.abc import AsyncGenerator
 from typing import Any

src/exo/shared/tracing.py

@@ -1,5 +1,3 @@
-from __future__ import annotations
-
 import json
 import time
 from collections import defaultdict

src/exo/shared/types/ollama_api.py

@@ -1,5 +1,3 @@
-from __future__ import annotations
-
 import time
 from typing import Any, Literal
 

src/exo/utils/info_gatherer/nvml.py

@@ -1,6 +1,4 @@
 # pyright: reportMissingImports=false
-from __future__ import annotations
-
 from exo.shared.types.profiling import SystemPerformanceProfile
 from exo.utils.pydantic_ext import TaggedModel
 
@@ -72,7 +70,9 @@ def gather_nvidia_metrics() -> NvmlMetrics | None:
         nvmlShutdown()
 
         gpu_load_fraction = min(total_gpu_power / _GPU_POWER_MAX, 1.0)
-        estimated_cpu_power = _CPU_POWER_IDLE + (_CPU_POWER_MAX - _CPU_POWER_IDLE) * gpu_load_fraction
+        estimated_cpu_power = (
+            _CPU_POWER_IDLE + (_CPU_POWER_MAX - _CPU_POWER_IDLE) * gpu_load_fraction
+        )
 
         return NvmlMetrics(
             system_profile=SystemPerformanceProfile(

src/exo/worker/engines/kv_cache.py

@@ -1,5 +1,3 @@
-from __future__ import annotations
-
 from collections.abc import Iterator, Sequence
 from copy import deepcopy
 from dataclasses import dataclass
@@ -100,7 +98,7 @@ class TorchKVCache:
             if isinstance(layer, (KVLayerState, RotatingKVLayerState))
         ]
 
-    def detach_cpu(self) -> TorchKVCache:
+    def detach_cpu(self) -> "TorchKVCache":
         layers: list[LayerState] = []
         for layer in self.layers:
             if isinstance(layer, KVLayerState):
@@ -134,7 +132,7 @@ class TorchKVCache:
             torch.cuda.synchronize()
         return TorchKVCache(layers, list(self.token_offset_per_group))
 
-    def trim_to(self, num_tokens: int) -> TorchKVCache:
+    def trim_to(self, num_tokens: int) -> "TorchKVCache":
         trimmed = TorchKVCache(list(self.layers), list(self.token_offset_per_group))
         trimmed._num_tokens = num_tokens
         return trimmed
@@ -149,7 +147,7 @@ class TorchKVCache:
         cache: Sequence[
             KVCache | RotatingKVCache | QuantizedKVCache | ArraysCache | CacheList
         ],
-    ) -> TorchKVCache:
+    ) -> "TorchKVCache":
         layers: list[LayerState] = []
         for c in cache:
             if isinstance(c, RotatingKVCache):
@@ -230,7 +228,7 @@ class TorchKVCache:
         layer_to_group: list[int],
         num_tokens: int,
         token_offset_per_group: list[int] | None = None,
-    ) -> TorchKVCache:
+    ) -> "TorchKVCache":
         block_tables = [
             torch.tensor(ids, dtype=torch.long) for ids in block_ids_per_group
         ]

src/exo/worker/engines/mlx/cache.py

@@ -1,5 +1,3 @@
-from __future__ import annotations
-
 import os
 from copy import deepcopy
 from typing import TYPE_CHECKING

src/exo/worker/engines/vllm/growable_cache.py

@@ -1,5 +1,3 @@
-from __future__ import annotations
-
 from typing import TYPE_CHECKING
 
 import torch

src/exo/worker/runner/llm_inference/batch_generator.py

@@ -1,5 +1,3 @@
-from __future__ import annotations
-
 import itertools
 import time
 from abc import ABC, abstractmethod