monitoring.py

# - HELP ----------------------------------------------------------------------
# https: https://stackoverflow.com/questions/10175812/how-to-generate-a-self-signed-ssl-certificate-using-openssl


# - IMPORTS -------------------------------------------------------------------
import threading
import time
import pandas as pd
from tabulate import tabulate
import datetime
import json
import yaml

from config.config import arg_parser
from server_watcher import ServerWatcher
from ssh_utils.connection_manager import ConnectionManagerViaSSH
from ssh_utils.cmd_manager import ManageCmdsToHosts


# - CONSTANTS -----------------------------------------------------------------
# Global variables
latest_result = "No data yet"
latest_result_detailed = "No data yet"
latest_result_detailed_space = "No data yet"

# Read server aliases
with open('config/server_names.yaml') as f:
    GPU_SERVER_ALIASES = yaml.safe_load(f)['GPU_SERVER_ALIASES']


# - FUNCTIONS -----------------------------------------------------------------
def calculate_average_utilization(server_data, server_alias):
    """Compute some statistics"""
    # Create dicts
    server_stats = {}
    server_stats_detailed = {}
    disk_stats_detailed = {}

    # Iterate over each server in the results
    # to extract GPU, CPU, and Disk Space details
    cpu_info = server_data['cpus']
    disk_space = server_data["disk_space"]

    # GPU and CPU Statistics
    server_stats[server_alias] = {
        "Avg. GPU Memory (%)": server_data["avg_gpu_mem_used"],
        "Avg. GPU Utilization (%)": server_data["avg_gpu_mem_util"],
        "Active Users": list(cpu_info['normalized_cpu_load_wrt_user'].keys()),
        "Total CPU Load (%)": round((cpu_info.get('cpu_load').get('load_5min')
                                     / cpu_info.get('cpu_count')) * 100, 2),
        "Avg Disk Space Used (%)": disk_space["Avg Use %"]
    }
    # Detailled GPU Statistics
    server_stats_detailed[server_alias] = {
        "GPU Memory (MiB)": server_data["gpus_list_gpu_mem_used"],
        "GPU Utilization (%)": server_data["gpus_list_gpu_gpu_util"],
    }
    # Disk Space Statistics
    disk_stats_detailed[server_alias] = {
        "Total Disp Space Used (TB)": disk_space['Tot Used'],
        "Total Disp Space Avail (TB)": disk_space['Tot Avail']
    }

    return server_stats, server_stats_detailed, disk_stats_detailed


def thread_server_moni(START_TIME, HOURS_TO_RUN, cmd_manager, GPU_server_alias,
                       server_watcher, SLEEP, lock, barrier, shared_server_stats,
                       shared_server_stats_detailed, shared_server_disk_detailed):
    """Monitor a GPU server with one Thread"""

    global latest_result, latest_result_detailed, latest_result_detailed_space

    while time.time() - START_TIME < HOURS_TO_RUN:
        # GPU, CPU, and Disk Space Monitoring
        host_output, _ = \
            cmd_manager.execute_cmd_on_host(GPU_server_alias, 'nvidia-smi')
        cpu_stats = \
            [cmd_manager.execute_cmd_on_host(GPU_server_alias, cmd)[0] for cmd in ["free",
                                                                                   "uptime",
                                                                                   "ps -eo user,pcpu",
                                                                                   "top -b -n 1 | awk 'NR>7 {cpu[$2]+=$9} END {for (u in cpu) print u, cpu[u]}'",
                                                                                   "nproc"]
             ]
        disk_space_stats = \
            [cmd_manager.execute_cmd_on_host(GPU_server_alias, cmd)[0] for cmd in ["df | grep -E '^/dev/' | grep -vE 'boot|tmpfs'"]
             ]

        # Thread-safe write to results
        gpac_msg = server_watcher.get_server_statistics(GPU_server_alias,
                                                        gpu_output=host_output,
                                                        cpu_output=cpu_stats,
                                                        disk_space_output=disk_space_stats)

        server_stats, detailled_server_stats, disk_stats_detailed = calculate_average_utilization(gpac_msg, GPU_server_alias)

        with lock:
            shared_server_stats[GPU_server_alias] = \
                server_stats[GPU_server_alias]
            shared_server_stats_detailed[GPU_server_alias] = \
                detailled_server_stats[GPU_server_alias]
            shared_server_disk_detailed[GPU_server_alias] = \
                disk_stats_detailed[GPU_server_alias]

        # Synchronize threads
        try:
            barrier.wait()
        except threading.BrokenBarrierError:
            print(f"Barrier broken in thread {GPU_server_alias}")

        time.sleep(SLEEP)


# - MAIN ----------------------------------------------------------------------
def main():
    args = arg_parser()
    HOURS_TO_RUN = args.run
    SLEEP = args.sleep
    START_TIME = time.time()

    global latest_result, latest_result_detailed, latest_result_detailed_space

    server_watcher = ServerWatcher()
    host_manager = ConnectionManagerViaSSH(args.password)
    cmd_manager = ManageCmdsToHosts(host_manager)

    shared_server_stats = {}
    shared_server_stats_detailed = {}
    shared_server_disk_detailed = {}
    results_lock = threading.Lock()

    def update_latest_results():

        global latest_result, latest_result_detailed, latest_result_detailed_space

        with results_lock:
            # Update the global variable for Flask in a thread-safe manner
            current_time = datetime.datetime.now().strftime("%Y-%m-%d %H:%M:%S")

            # Convert shared dictionaries to DataFrames
            df_stats = pd.DataFrame.from_dict(shared_server_stats, orient='index')
            df_stats_detailed = pd.DataFrame.from_dict(shared_server_stats_detailed, orient='index')
            df_disk_space_detailed = pd.DataFrame.from_dict(shared_server_disk_detailed, orient='index')

            # Sort the DataFrames by index (server names) alphabetically
            df_stats_sorted = df_stats.sort_index()
            df_stats_detailed_sorted = df_stats_detailed.sort_index()
            df_disk_space_detailed_sorted = df_disk_space_detailed.sort_index()

            # Use tabulate to create the tables
            latest_result = f"Time Stamp: {current_time}\n"  # Add timestamp
            latest_result += tabulate(df_stats_sorted, headers='keys', tablefmt='grid')

            latest_result_detailed = f"Time Stamp: {current_time}\n"  # Add timestamp
            latest_result_detailed += tabulate(df_stats_detailed_sorted, headers='keys', tablefmt='grid')

            latest_result_detailed_space = f"Time Stamp: {current_time}\n"  # Add timestamp
            latest_result_detailed_space += tabulate(df_disk_space_detailed_sorted, headers='keys', tablefmt='grid')

            payload = {
                "stats": latest_result,
                "detailed_stats": latest_result_detailed,
                "disk_stats": latest_result_detailed_space
                }

            with open("data.json", "w") as file:
                json.dump(payload, file, indent=4)

    try:
        # Establish SSH connections
        for GPU_server_alias in GPU_SERVER_ALIASES:
            host_manager.add_connection(GPU_server_alias)
            host_manager.connect_to_server(GPU_server_alias)
        valid_GPU_SERVER = host_manager.check_current_servers_alive()

        # Create a barrier with the action function
        barrier = threading.Barrier(len(valid_GPU_SERVER), action=update_latest_results)
        threads = []

        # Start monitoring threads
        for GPU_server_alias in valid_GPU_SERVER:
            thread = threading.Thread(
                target=thread_server_moni,
                args=(START_TIME, HOURS_TO_RUN, cmd_manager, GPU_server_alias,
                      server_watcher, SLEEP, results_lock, barrier, shared_server_stats,
                      shared_server_stats_detailed, shared_server_disk_detailed)
            )
            threads.append(thread)
            thread.start()

        # Wait for all threads to finish
        for thread in threads:
            thread.join()

    except KeyboardInterrupt:
        print("\nCtrl+C detected. Terminating gracefully...")

    finally:
        # Ensure that all SSH connections are closed on exit
        print("Closing all SSH connections...")
        host_manager.close_all_connections()
        print("SSH connections closed.")
        latest_result = "SSH connections closed"


if __name__ == '__main__':
    main()