snsd0805/gpu-provider
1
0
Fork 0
Code Issues Pull Requests Packages Projects Releases Wiki Activity

feat: get GPU name in the node

Browse Source
This commit is contained in:
Ting-Jun Wang 2024-05-29 18:34:44 +08:00
parent 68123809bb
commit 847b6f8814
Signed by: snsd0805
GPG Key ID: 48D331A3D6160354
1 changed files with 11 additions and 3 deletions
Show Stats Download Patch File Download Diff File Expand all files Collapse all files

14
src/node_manager.py
View File

@ -1,8 +1,8 @@
import threading import threading
from src.communication import ServiceExplorationModule, ClusterCommunicationModule from src.communication import ServiceExplorationModule, ClusterCommunicationModule
import torch
import time import time
class NodeManager(): class NodeManager():
def __init__(self, host, port): def __init__(self, host, port):
self.status = 'none' self.status = 'none'
@ -10,8 +10,8 @@ class NodeManager():
{'explanation': 'Add another node into our cluster', 'function': 'add_node'}, {'explanation': 'Add another node into our cluster', 'function': 'add_node'},
{'explanation': 'Exit', 'function': 'exit'}, {'explanation': 'Exit', 'function': 'exit'},
] ]
self.GPU = 'RTX 4090' self.get_GPU_info()
self.GPU_num = 1 print(f"You have {self.GPU} * {self.GPU_num}")
# start Cluster Communication Module # start Cluster Communication Module
# let the nodes in the cluster can communicate # let the nodes in the cluster can communicate
@ -23,6 +23,14 @@ class NodeManager():
time.sleep(2) time.sleep(2)
def get_GPU_info(self):
self.GPU_num = torch.cuda.device_count()
assert self.GPU_num > 0, "Your computer doesn't have GPU resource"
self.GPU = torch.cuda.get_device_name(0)
for i in range(self.GPU_num):
assert torch.cuda.get_device_name(i) == self.GPU, "Please provide same type of GPUs."
def start_service(self): def start_service(self):
communication_thread = threading.Thread(target=self.cluster_communication_module.listen) communication_thread = threading.Thread(target=self.cluster_communication_module.listen)
communication_thread.daemon = True communication_thread.daemon = True