feat: detect use 3 card on multi server (rule)

main.py

@@ -6,7 +6,7 @@ from bs4 import BeautifulSoup
 from email.mime.text import MIMEText
 import os
 
-DEBUG_MODE = True
+DEBUG_MODE = False
 
 URL = 'https://www.cmlab.csie.ntu.edu.tw/status-gpu/'
 GPU_LIMIT = 2
@@ -15,12 +15,12 @@ MAIL_CD_HOUR = 6
 MAIL_MESSAGE = '''
 Hi, {}<br>
 <br>
-提醒您，您目前在 {} 伺服器上已經使用了 {} 張 {}，<br>
-依照 CMLab 規定，每人在每台 CML workstation 上至多使用 2 顆 GPU，詳細規定請參閱 https://www.cmlab.csie.ntu.edu.tw/wiki/doku.php?id=workstation:rule <br>
+提醒您，您目前{}，<br>
+依照 CMLab 規定，每人在 CML workstation 上至多使用 2 顆 GPU <br>
 <br>
 為了公平起見，建議您降低 GPU 使用量！<br>
 雖然我們不會直接處理，但若有人檢舉，我們會停止您運行的程式！<br>
-若有其他問題歡迎來信討論！謝謝您的配合！<br>
+若有其他特殊需求請來信說明！謝謝您的配合！<br>
 <br>
 * 信件由網管自動化工具寄出，若為錯誤寄送請忽略此信件。<br>
 <br>
@@ -106,7 +106,7 @@ def get_server_gpu_status() -> list:
         servers.append({
             'name': server_name,
             'users': users,
-            'gpu': gpu_names
+            'gpus': gpu_names
         })
     return servers
 
@@ -114,8 +114,8 @@ 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.
-        Rules now (2024/03/11):
-            - 每人每台僅限使用 2 張卡
+        Rules now (2024/03/22):
+            - 每人僅限使用 2 張卡
 
         I add this function beacause that we may need more filter rules. we can add some rules in this function easily
 
@@ -125,29 +125,28 @@ def filter(server_status: list, limit: int = 2) -> list:
         Output:
             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 = {}
+    usage = {}
 
     for server in server_status:
-        print(server)
-    '''
-    for user in server_status['users']:
-        if user not in counter:
-            counter[user] = 1
+        print(server['name'])
+        for gpu_index, gpu_name in enumerate(server['gpus']):
+            print("    ", gpu_index, gpu_name, server['users'][gpu_index])
+            for user in server['users'][gpu_index]:
+                if user not in usage:
+                    usage[user] = [{'server': server['name'], 'gpu': gpu_name}]
                 else:
-            counter[user] += 1
+                    usage[user].append({'server': server['name'], 'gpu': gpu_name})
+
+        print('-')
+
+    violators = []
+    for user, state in usage.items():
+        if len(state) > limit:
+            violators.append({'username': user, 'usage': state}) 
             
-    for k, v in counter.items():
-        if v > limit:
-            violators.append({
-                'user': k,
-                'gpu': v,
-            })
-    '''
     return violators
 
-def mail_notify(server_name: str, gpu_name: str, violators: list) -> None:
+def mail_notify(violators: list) -> None:
     def check_send(log: dict, username: str) -> bool:
         if DEBUG_MODE:
             return True
@@ -160,21 +159,40 @@ def mail_notify(server_name: str, gpu_name: str, violators: list) -> None:
             else:
                 return False
 
+    def get_usage_msg(usage: list) -> str:
+        server_usage = {}
+        ans = ""
+        for gpu in usage:
+            if gpu['server'] not in server_usage:
+                server_usage[gpu['server']] = { gpu['gpu']: 1 }
+            else:
+                if gpu['gpu'] not in server_usage[gpu['server']]:
+                    server_usage[gpu['server']][gpu['gpu']] = 1
+                else:
+                    server_usage[gpu['server']][gpu['gpu']] += 1
+
+        for server, gpus in server_usage.items():
+            for gpu, count in gpus.items():
+                ans += f"在 {server} 上使用 {gpu} * {count}, "
+        return ans[:-2]
+
+    print(" ===== 寄送 ===== ")
+
     # get last send time
     with open('send_log.json') as fp:
         send_log = json.load(fp)
 
     for violator in violators:
-        username = violator['user']
-        print(f'    {username}')
+        username = violator['username']
 
         if check_send(send_log, username):
             # update log
             send_log[username] = time.time()
             with open('send_log.json', 'w') as fp:
                 json.dump(send_log, fp)
-            usage = violator['gpu']
-            msg = MAIL_MESSAGE.format(username, server_name, usage, gpu_name)
+            usage_msg = get_usage_msg(violator['usage'])
+            print(f'    {username} {usage_msg}')
+            msg = MAIL_MESSAGE.format(username, usage_msg)
             msg = MIMEText(msg, 'html') # 郵件內文
             msg['Subject'] = f'[網管通知] 提醒您已經超過 CMLab GPU 使用限制！({username})' + "(DEBUG MODE)" if DEBUG_MODE else ""
             msg['From'] = 'unix_manager@cmlab.csie.ntu.edu.tw'
@@ -196,10 +214,10 @@ 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)
+    mail_notify(violators)
+
     '''
     for server in server_status:
         print(server['name'], server['gpu'])