data_collector.py - master - nice_jam/ClickHouseVSPostgres

data_collector.py
126 строк3 KB
Filippenko Pavel
requirements.txt
29 июн 2025, 13:20
29 июн 2025, 13:20ccaed22
import time
import datetime
import psutil
import clickhouse_connect
import psycopg2
from pynvml import *

SLEEP_TIMEOUT = 2 # sec

class PCStatistics:
    def __init__(self, timestamp, cpu, ram, gpu) -> None:
        self.timestamp = timestamp
        self.cpu = cpu
        self.ram = ram
        self.gpu = gpu

def get_statistics(gpu_handle):
    timestamp = datetime.datetime.now()
    cpu = psutil.cpu_percent()
    ram = psutil.virtual_memory().percent

    if gpu_handle is not None:
        try:
            gpu = nvmlDeviceGetUtilizationRates(gpu_handle).gpu
        except:
            gpu = 0.0
    else:
        gpu = 0.0

    return PCStatistics(timestamp, cpu, ram, gpu)

def main():
    # --- GPU init ---
    try:
        nvmlInit()
        gpu_handle = nvmlDeviceGetHandleByIndex(0)
    except:
        print("Can't initialize gpu handler for this device")
        gpu_handle = None

    # --- ClickHouse client (with auth) ---
    try:
        ch_client = clickhouse_connect.get_client(
            host='localhost',
            username='ch_user',
            password='ch_pass',
            database='default'
        )
    except Exception as ex:
        print(f"[ERROR] EXCEPTION (connect to ClickHouse): {ex}")
        return 

    # --- PostgreSQL client ---
    try:
        pg_conn = psycopg2.connect(
            host='localhost',
            dbname='postgres',
            user='postgres',
            password='1234'
        )
    except Exception as ex:
        print(f"[ERROR] EXCEPTION (connect to PostgreSQL): {ex}")
        return 
    pg_cur = pg_conn.cursor()

    # Create table in ClickHouse if not exists
    ch_client.command('''
    CREATE TABLE IF NOT EXISTS system_metrics (
        timestamp DateTime,
        cpu_usage Float32,
        gpu_usage Float32,
        ram_usage Float32
    ) ENGINE = MergeTree()
    ORDER BY timestamp
    ''')

    # Create table in PostgreSQL if not exists
    pg_cur.execute('''
    CREATE TABLE IF NOT EXISTS system_metrics (
        timestamp TIMESTAMP,
        cpu_usage REAL,
        gpu_usage REAL,
        ram_usage REAL
    )
    ''')
    pg_conn.commit()

    # --- Collection loop ---
    print("Starting data collection... (press Ctrl+C to stop)")
    try:
        while True:
            pc_statistics = get_statistics(gpu_handle)

            print(f"{pc_statistics.timestamp} | CPU: {pc_statistics.cpu}%, GPU: {pc_statistics.gpu}%, RAM: {pc_statistics.ram}%")

            # ClickHouse insert
            ch_client.insert(
                table='system_metrics',
                data=[[
                    pc_statistics.timestamp,
                    pc_statistics.cpu,
                    pc_statistics.gpu,
                    pc_statistics.ram
                ]],
                column_names=['timestamp', 'cpu_usage', 'gpu_usage', 'ram_usage']
            )

            # PostgreSQL insert
            pg_cur.execute(
                "INSERT INTO system_metrics (timestamp, cpu_usage, gpu_usage, ram_usage) VALUES (%s, %s, %s, %s)",
                (pc_statistics.timestamp, pc_statistics.cpu, pc_statistics.gpu, pc_statistics.ram)
            )
            pg_conn.commit()

            time.sleep(SLEEP_TIMEOUT)
    except KeyboardInterrupt:
        print("\nData collection stopped.")

    finally:
        pg_cur.close()
        pg_conn.close()
        if gpu_handle is not None:
            nvmlShutdown()
            
if __name__ == '__main__':
    main()