daemon/daemon.py

"""Hermes Daemon — 后台训练守护进程."""

import time
import signal
import sys
from pathlib import Path

# Fix Unicode display on Windows (GBK terminal)
if sys.platform == "win32":
    sys.stdout.reconfigure(encoding="utf-8", errors="replace")
    sys.stderr.reconfigure(encoding="utf-8", errors="replace")
from hermes_core.types import HERMES_DATA_DIR
from hermes_core.db import init_db, get_scopes_needing_training, get_active_records
from hermes_core.scheduler import check_idle, get_training_queue, process_queue
from hermes_core.config import get_config


running = True


def _signal_handler(signum, frame):
    global running
    running = False
    print("\nShutting down daemon gracefully...")


def _get_all_users() -> list[str]:
    """扫描所有有数据的用户。"""
    users_dir = Path(get_config().data_dir) / "users"
    if not users_dir.exists():
        return []
    return [d.name for d in users_dir.iterdir() if d.is_dir()]


def main():
    """Daemon 主循环。从配置文件读取扫描间隔等参数。"""
    config = get_config()
    interval = config.scan_interval_seconds

    signal.signal(signal.SIGINT, _signal_handler)
    signal.signal(signal.SIGTERM, _signal_handler)

    print(f"Hermes Daemon v0.1.0 started")
    print(f"  Scan interval: {interval}s")
    print(f"  Data dir: {config.data_dir}")
    print(f"  CPU threshold: {config.cpu_idle_threshold}")
    print(f"  Log: {config.log_file}")

    while running:
        users = _get_all_users()
        for user_id in users:
            if not running:
                break
            try:
                conn = init_db(user_id)
                scopes_needing = get_scopes_needing_training(conn)
                conn.close()

                if not scopes_needing:
                    continue

                queue = get_training_queue(user_id)
                if not queue:
                    continue

                if check_idle(config.cpu_idle_threshold, config.gpu_idle_threshold):
                    print(f"[{time.strftime('%H:%M:%S')}] Idle detected. "
                          f"Training {queue[0]['scope_id']} for user {user_id}")
                    count = process_queue(user_id)
                    if count > 0:
                        print(f"[{time.strftime('%H:%M:%S')}] "
                              f"Training complete for {queue[0]['scope_id']}")
            except Exception as e:
                print(f"[{time.strftime('%H:%M:%S')}] Error processing user {user_id}: {e}")

        for _ in range(interval):
            if not running:
                break
            time.sleep(1)

    print("Hermes Daemon stopped.")


if __name__ == "__main__":
    main()