diff --git a/main.py b/main.py index ff8fc40..aecb514 100644 --- a/main.py +++ b/main.py @@ -77,30 +77,40 @@ def get_server_gpu_status() -> list: soup = BeautifulSoup(box, 'html.parser') # get server name server_name = soup.find('span', class_='f7').text.replace(' ', '') - - # get GPU name - gpu_name = soup.find('span', class_='f4').text.replace(' ', '') + gpu_names = [] users = [] gpus = box.split('\n') - for i in gpus[1:]: - soup = BeautifulSoup(i, 'html.parser') + for i in gpus[:]: + if 'f6' in i: # if this line is for a single GPU informations + soup = BeautifulSoup(i, 'html.parser') - # get all users who is using this server's GPU resources. - user_obj = soup.find('span', class_='f0') - if user_obj != None: - username = user_obj.text - if username != 'gdm': - users.append(username) + # get all users who is using this server's GPU resources. + user_objs = soup.findAll('span', class_='f0') + user_in_this_gpu = set() + + # get users on this GPU + for user_obj in user_objs: + if user_obj != None: + username = user_obj.text + if username != 'gdm': + user_in_this_gpu.add(username) + + # get gpu names + gpu_name = soup.find('span', class_='f4').text.replace(' ', '') + + # log gpu name & users on this GPU + gpu_names.append(gpu_name) + users.append(list(user_in_this_gpu)) servers.append({ 'name': server_name, 'users': users, - 'gpu': gpu_name + 'gpu': gpu_names }) return servers -def filter(server_status: dict, limit: int = 2) -> list: +def filter(server_status: list, limit: int = 2) -> list: ''' You can set some rule here. it will return the a list which contains the username that exceed the limit. @@ -110,14 +120,18 @@ def filter(server_status: dict, limit: int = 2) -> list: I add this function beacause that we may need more filter rules. we can add some rules in this function easily Input: - server_status (dict): e.g. {'name': 'cml5', ...} + server_status (list): from get_server_gpu_status() limit (int): GPU limit (default: 2) Output: - violators (list): a violator list(set) e.g. [ {'user': 'snsd0805', 'gpu': 3}] means that the violator 'snsd0805' is using 3 GPU. + violators (list): a violator list(set) e.g. [ {'user': 'snsd0805', 'usage': [{'server': 'cml5', 'gpu': 'NVDIA3090', 'number': 3}...] }] means that the violator 'snsd0805' is using 3 GPU. ''' violators = [] counter = {} + + for server in server_status: + print(server) + ''' for user in server_status['users']: if user not in counter: counter[user] = 1 @@ -130,6 +144,7 @@ def filter(server_status: dict, limit: int = 2) -> list: 'user': k, 'gpu': v, }) + ''' return violators def mail_notify(server_name: str, gpu_name: str, violators: list) -> None: @@ -181,9 +196,15 @@ def mail_notify(server_name: str, gpu_name: str, violators: list) -> None: if __name__ == '__main__': server_status = get_server_gpu_status() + for i in server_status: + print(i) + + violators = filter(server_status, 2) + ''' for server in server_status: print(server['name'], server['gpu']) violators = filter(server, 2) print(violators) mail_notify(server['name'], server['gpu'], violators) print("=" * 20) + '''