patch_megatron / check_node_status.py
xingzhaohu's picture
Add files using upload-large-folder tool
add5224 verified
Raw
History Blame Contribute Delete
5.55 kB
import subprocess
import json
from datetime import datetime
import re
import os
import argparse
node_dict = {
# "c0f6c83c-01.cloud.together.ai": "gpu01",
# "c0f6c83c-02.cloud.together.ai": "gpu02",
# "c0f6c83c-03.cloud.together.ai": "gpu03",
# "c0f6c83c-04.cloud.together.ai": "gpu04",
# "c0f6c83c-05.cloud.together.ai": "gpu05",
# "c0f6c83c-06.cloud.together.ai": "gpu06",
# "c0f6c83c-07.cloud.together.ai": "gpu07",
# "c0f6c83c-08.cloud.together.ai": "gpu08",
# "c0f6c83c-09.cloud.together.ai": "gpu09",
# "c0f6c83c-10.cloud.together.ai": "gpu10",
# "c0f6c83c-11.cloud.together.ai": "gpu11",
# "c0f6c83c-12.cloud.together.ai": "gpu12",
# "c0f6c83c-13.cloud.together.ai": "gpu13",
# "c0f6c83c-14.cloud.together.ai": "gpu14",
# "c0f6c83c-15.cloud.together.ai": "gpu15",
# "c0f6c83c-16.cloud.together.ai": "gpu16",
# "c0f6c83c-17.cloud.together.ai": "gpu17",
# "c0f6c83c-18.cloud.together.ai": "gpu18",
# "c0f6c83c-19.cloud.together.ai": "gpu19",
# "c0f6c83c-20.cloud.together.ai": "gpu20",
"c0f6c83c-21.cloud.together.ai": "gpu21",
"c0f6c83c-22.cloud.together.ai": "gpu22",
"c0f6c83c-23.cloud.together.ai": "gpu23",
"c0f6c83c-24.cloud.together.ai": "gpu24",
"c0f6c83c-25.cloud.together.ai": "gpu25",
"c0f6c83c-26.cloud.together.ai": "gpu26",
"c0f6c83c-27.cloud.together.ai": "gpu27",
"c0f6c83c-28.cloud.together.ai": "gpu28",
"c0f6c83c-33.cloud.together.ai": "gpu33",
"c0f6c83c-34.cloud.together.ai": "gpu34",
"c0f6c83c-35.cloud.together.ai": "gpu35",
"c0f6c83c-36.cloud.together.ai": "gpu36",
"c0f6c83c-37.cloud.together.ai": "gpu37",
"c0f6c83c-38.cloud.together.ai": "gpu38",
}
## 06 07 11 12 13 15 18 20 21 25 26
def run_ssh_command(node_alias, command):
try:
result = subprocess.run(
["ssh", node_alias, command],
capture_output=True,
text=True,
timeout=30
)
if result.returncode != 0:
print(f"[{node_alias}] Command error: {result.stderr.strip()}")
return ""
return result.stdout.strip()
except Exception as e:
print(f"[{node_alias}] SSH command exception: {e}")
return ""
def get_gpu_info(node_alias):
output = run_ssh_command(node_alias, "nvidia-smi")
gpu_info = []
mem_matches = re.findall(r'(\d+)MiB\s*/\s*(\d+)MiB', output)
power_matches = re.findall(r'(\d+)W\s*/\s*(\d+)W', output)
for i, ((used_mem, total_mem), (used_power, max_power)) in enumerate(zip(mem_matches, power_matches)):
gpu_info.append({
'gpu_id': i,
'used_memory_mb': int(used_mem),
'total_memory_mb': int(total_mem),
'used_memory_gb': round(int(used_mem) / 1024, 2),
'total_memory_gb': round(int(total_mem) / 1024, 2),
'power_usage_w': int(used_power),
'power_limit_w': int(max_power),
'power_usage_percent': round((int(used_power) / int(max_power)) * 100, 2)
})
return gpu_info
def get_system_resources(node_alias):
top_output = run_ssh_command(node_alias, "top -bn1")
free_output = run_ssh_command(node_alias, "free -b")
cpu_idle = re.search(r'%Cpu\(s\):.+?(\d+\.\d+)\s+id', top_output)
cpu_idle_percent = float(cpu_idle.group(1)) if cpu_idle else 0.0
memory_info = free_output.split('\n')
mem_line = [line for line in memory_info if 'Mem:' in line][0].split()
total_memory = int(mem_line[1]) / (1024 ** 3)
used_memory = int(mem_line[2]) / (1024 ** 3)
free_memory = int(mem_line[3]) / (1024 ** 3)
available_memory = int(mem_line[6]) / (1024 ** 3)
return {
'cpu_idle_percent': cpu_idle_percent,
'memory_total_gb': round(total_memory, 2),
'memory_used_gb': round(used_memory, 2),
'memory_free_gb': round(free_memory, 2),
'memory_available_gb': round(available_memory, 2),
'memory_free_percent': round((free_memory / total_memory) * 100, 2),
'memory_available_percent': round((available_memory / total_memory) * 100, 2)
}
def main(node_alias, node_id, save_folder_dir):
try:
gpu_info = get_gpu_info(node_alias)
sys_resources = get_system_resources(node_alias)
monitoring_data = {
'timestamp': datetime.now().isoformat(),
'gpu_info': gpu_info,
'system_resources': sys_resources
}
print(f"\n####### Node: {node_alias} ({node_id}) #######")
for gpu in gpu_info:
print(f"GPU {gpu['gpu_id']}: {gpu['used_memory_gb']} GB used / {gpu['total_memory_gb']} GB")
print(f"CPU idle %: {sys_resources['cpu_idle_percent']}")
print(f"Memory free %: {sys_resources['memory_free_percent']}")
os.makedirs(save_folder_dir, exist_ok=True)
filename = f"{save_folder_dir}/{node_id}.json"
with open(filename, 'w') as f:
json.dump(monitoring_data, f, indent=4)
print(f"Saved to {filename}")
except Exception as e:
print(f"[{node_alias}] Monitoring failed: {e}")
if __name__ == "__main__":
parser = argparse.ArgumentParser()
parser.add_argument('--save_dir', type=str, default='tmp_data',
help='Base directory to save monitoring data')
args = parser.parse_args()
timestamp_dir = f"{args.save_dir}/together_monitoring_{datetime.now().strftime('%Y%m%d_%H%M%S')}"
for node_alias, node_id in node_dict.items():
main(node_alias, node_id, timestamp_dir)