Spaces:
Running
Running
| from src.manager.utils.singleton import singleton | |
| import torch | |
| import psutil | |
| class BudgetManager(): | |
| total_resource_budget = 100 | |
| current_resource_usage = 0 | |
| total_expense_budget = 10 | |
| current_expense = 0 | |
| is_budget_initialized = False | |
| is_resource_budget_enabled = True | |
| is_expense_budget_enabled = True | |
| def __init__(self): | |
| if not self.is_budget_initialized: | |
| self.total_resource_budget = self.calculate_total_budget() | |
| self.is_budget_initialized = True | |
| def set_resource_budget_status(self, status: bool): | |
| self.is_enabled = status | |
| if status: | |
| print("Budget manager is enabled.") | |
| else: | |
| print("Budget manager is disabled.") | |
| def set_expense_budget_status(self, status: bool): | |
| self.is_expense_budget_enabled = status | |
| if status: | |
| print("Expense budget manager is enabled.") | |
| else: | |
| print("Expense budget manager is disabled.") | |
| def calculate_total_budget(self)-> int: | |
| total_mem = 0 | |
| gpu_mem = 0 | |
| ram_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 | |
| gpu_mem = total_vram /1024 ** 3 | |
| print(f"GPU detected: {gpu_name}") | |
| print(f"Total VRAM: {gpu_mem:.2f} GB") | |
| mem = psutil.virtual_memory() | |
| ram_mem = mem.total/ 1024 ** 3 | |
| print("No GPU detected. Using CPU.") | |
| print(f"Total RAM: {ram_mem:.2f} GB") | |
| total_mem = gpu_mem + ram_mem | |
| 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): | |
| if not self.is_resource_budget_enabled: | |
| return True | |
| return True if self.current_resource_usage + cost <= self.total_resource_budget else False | |
| def add_to_resource_budget(self, cost): | |
| if not self.is_resource_budget_enabled: | |
| return | |
| 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 not self.is_resource_budget_enabled: | |
| return | |
| 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): | |
| if not self.is_expense_budget_enabled: | |
| return True | |
| return True if self.current_expense + cost <= self.total_expense_budget else False | |
| def add_to_expense_budget(self, cost): | |
| if not self.is_expense_budget_enabled: | |
| return | |
| if not self.can_spend_expense(cost): | |
| raise Exception("No expense budget remaining") | |
| self.current_expense += cost |