Spaces:
Running
Running
File size: 2,630 Bytes
2f85c93 0e323dc 1476939 8cf77a3 bf722a2 8cf77a3 0e323dc bf722a2 0e323dc bf722a2 0e323dc 8cf77a3 bf722a2 0e323dc bf722a2 8cf77a3 bf722a2 8cf77a3 bf722a2 |
1 2 3 4 5 6 7 8 9 10 11 12 13 14 15 16 17 18 19 20 21 22 23 24 25 26 27 28 29 30 31 32 33 34 35 36 37 38 39 40 41 42 43 44 45 46 47 48 49 50 51 52 53 54 55 56 57 58 59 60 61 62 63 64 65 66 67 68 69 70 71 |
from src.manager.utils.singleton import singleton
import torch
import psutil
@singleton
class BudgetManager():
total_resource_budget = 100
current_resource_usage = 0
total_expense_budget = 1000
current_expense = 0
is_budget_initialized = False
def __init__(self):
if not self.is_budget_initialized:
self.total_resource_budget = self.calculate_total_budget()
self.is_budget_initialized = True
def calculate_total_budget(self)-> int:
total_mem = 0
if torch.cuda.is_available():
gpu_index = torch.cuda.current_device()
gpu_name = torch.cuda.get_device_name(gpu_index)
total_vram = torch.cuda.get_device_properties(gpu_index).total_memory
total_mem = total_vram /1024 ** 3
print(f"GPU detected: {gpu_name}")
print(f"Total VRAM: {total_mem:.2f} GB")
else:
mem = psutil.virtual_memory()
total_mem = mem.total/ 1024 ** 3
print("No GPU detected. Using CPU.")
print(f"Total RAM: {total_mem:.2f} GB")
return round((total_mem / 16) * 100)
def get_total_resource_budget(self):
return self.total_resource_budget
def get_current_resource_usage(self):
return self.current_resource_usage
def get_current_remaining_resource_budget(self):
return self.total_resource_budget - self.current_resource_usage
def can_spend_resource(self, cost):
return True if self.current_resource_usage + cost <= self.total_resource_budget else False
def add_to_resource_budget(self, cost):
if not self.can_spend_resource(cost):
raise Exception("No resource budget remaining")
self.current_resource_usage += cost
def remove_from_resource_expense(self, cost):
if self.current_resource_usage - cost < 0:
raise Exception("Not enough resource budget to remove")
self.current_resource_usage -= cost
def get_total_expense_budget(self):
return self.total_expense_budget
def get_current_expense(self):
return self.current_expense
def get_current_remaining_expense_budget(self):
return self.total_expense_budget - self.current_expense
def can_spend_expense(self, cost):
return True if self.current_expense + cost <= self.total_expense_budget else False
def add_to_expense_budget(self, cost):
if not self.can_spend_expense(cost):
raise Exception("No expense budget remaining")
self.current_expense += cost |