helloparthshah commited on
Commit
608ee17
·
1 Parent(s): c091ece

Update to use sum of both as the total memory

Browse files
Files changed (1) hide show
  1. src/manager/budget_manager.py +10 -8
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 = 1000
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
- total_mem = total_vram /1024 ** 3
25
  print(f"GPU detected: {gpu_name}")
26
- print(f"Total VRAM: {total_mem:.2f} GB")
27
- else:
28
- mem = psutil.virtual_memory()
29
- total_mem = mem.total/ 1024 ** 3
30
- print("No GPU detected. Using CPU.")
31
- print(f"Total RAM: {total_mem:.2f} GB")
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):