Spaces:
Running
Running
Commit
·
608ee17
1
Parent(s):
c091ece
Update to use sum of both as the total memory
Browse files
src/manager/budget_manager.py
CHANGED
@@ -6,7 +6,7 @@ import psutil
|
|
6 |
class BudgetManager():
|
7 |
total_resource_budget = 100
|
8 |
current_resource_usage = 0
|
9 |
-
total_expense_budget =
|
10 |
current_expense = 0
|
11 |
is_budget_initialized = False
|
12 |
|
@@ -17,18 +17,20 @@ class BudgetManager():
|
|
17 |
|
18 |
def calculate_total_budget(self)-> int:
|
19 |
total_mem = 0
|
|
|
|
|
20 |
if torch.cuda.is_available():
|
21 |
gpu_index = torch.cuda.current_device()
|
22 |
gpu_name = torch.cuda.get_device_name(gpu_index)
|
23 |
total_vram = torch.cuda.get_device_properties(gpu_index).total_memory
|
24 |
-
|
25 |
print(f"GPU detected: {gpu_name}")
|
26 |
-
print(f"Total VRAM: {
|
27 |
-
|
28 |
-
|
29 |
-
|
30 |
-
|
31 |
-
|
32 |
return round((total_mem / 16) * 100)
|
33 |
|
34 |
def get_total_resource_budget(self):
|
|
|
6 |
class BudgetManager():
|
7 |
total_resource_budget = 100
|
8 |
current_resource_usage = 0
|
9 |
+
total_expense_budget = 10
|
10 |
current_expense = 0
|
11 |
is_budget_initialized = False
|
12 |
|
|
|
17 |
|
18 |
def calculate_total_budget(self)-> int:
|
19 |
total_mem = 0
|
20 |
+
gpu_mem = 0
|
21 |
+
ram_mem = 0
|
22 |
if torch.cuda.is_available():
|
23 |
gpu_index = torch.cuda.current_device()
|
24 |
gpu_name = torch.cuda.get_device_name(gpu_index)
|
25 |
total_vram = torch.cuda.get_device_properties(gpu_index).total_memory
|
26 |
+
gpu_mem = total_vram /1024 ** 3
|
27 |
print(f"GPU detected: {gpu_name}")
|
28 |
+
print(f"Total VRAM: {gpu_mem:.2f} GB")
|
29 |
+
mem = psutil.virtual_memory()
|
30 |
+
ram_mem = mem.total/ 1024 ** 3
|
31 |
+
print("No GPU detected. Using CPU.")
|
32 |
+
print(f"Total RAM: {ram_mem:.2f} GB")
|
33 |
+
total_mem = gpu_mem + ram_mem
|
34 |
return round((total_mem / 16) * 100)
|
35 |
|
36 |
def get_total_resource_budget(self):
|