FastAPI-based working frontend (#10)

cacheflow/utils.py

@@ -1,5 +1,6 @@
 import enum
 import random
+import psutil
 
 import numpy as np
 import torch
@@ -26,6 +27,7 @@ class Counter:
     def reset(self) -> None:
         self.counter = 0
 
+
 def set_random_seed(seed: int):
     random.seed(seed)
     np.random.seed(seed)
@@ -35,3 +37,11 @@ def set_random_seed(seed: int):
 
     if model_parallel_is_initialized():
         model_parallel_cuda_manual_seed(seed)
+
+
+def get_gpu_memory(gpu: int = 0) -> int:
+    return torch.cuda.get_device_properties(gpu).total_memory
+
+
+def get_cpu_memory() -> int:
+    return psutil.virtual_memory().total