[#1] agent模块开发

src/agent/app/collector.py

@@ -0,0 +1,114 @@
+from __future__ import annotations
+
+import os
+import re
+import socket
+import subprocess
+from pathlib import Path
+from typing import Any, Dict
+
+from .config import AgentConfig
+from .log import get_logger
+
+LOGGER = get_logger("argus.agent.collector")
+
+_HOSTNAME_PATTERN = re.compile(r"^([^-]+)-([^-]+)-([^-]+)-.*$")
+
+
+def collect_metadata(config: AgentConfig) -> Dict[str, Any]:
+    """汇总节点注册需要的静态信息。"""
+    hostname = config.hostname
+    env, user, instance = _parse_hostname(hostname)
+    meta = {
+        "hostname": hostname,
+        "ip": _detect_ip_address(),
+        "env": env,
+        "user": user,
+        "instance": instance,
+        "cpu_number": _detect_cpu_count(),
+        "memory_in_bytes": _detect_memory_bytes(),
+        "gpu_number": _detect_gpu_count(config),
+    }
+    return meta
+
+
+def _parse_hostname(hostname: str) -> tuple[str, str, str]:
+    """按照约定的 env-user-instance 前缀拆解主机名。"""
+    match = _HOSTNAME_PATTERN.match(hostname)
+    if not match:
+        LOGGER.warning("Hostname does not match expected pattern", extra={"hostname": hostname})
+        return "", "", ""
+    return match.group(1), match.group(2), match.group(3)
+
+
+def _detect_cpu_count() -> int:
+    count = os.cpu_count()
+    return count if count is not None else 0
+
+
+def _detect_memory_bytes() -> int:
+    """优先读取 cgroup 限额，失败时退回 /proc/meminfo。"""
+    cgroup_path = Path("/sys/fs/cgroup/memory.max")
+    try:
+        raw = cgroup_path.read_text(encoding="utf-8").strip()
+        if raw and raw != "max":
+            return int(raw)
+    except FileNotFoundError:
+        LOGGER.debug("cgroup memory.max not found, falling back to /proc/meminfo")
+    except ValueError:
+        LOGGER.warning("Failed to parse memory.max, falling back", extra={"value": raw})
+
+    try:
+        with open("/proc/meminfo", "r", encoding="utf-8") as handle:
+            for line in handle:
+                if line.startswith("MemTotal:"):
+                    parts = line.split()
+                    if len(parts) >= 2:
+                        return int(parts[1]) * 1024
+    except FileNotFoundError:
+        LOGGER.error("/proc/meminfo not found; defaulting memory to 0")
+    return 0
+
+
+def _detect_gpu_count(config: AgentConfig) -> int:
+    """采集 GPU 数量，可被配置覆盖。"""
+    if config.gpu_number_override is not None:
+        return config.gpu_number_override
+
+    try:
+        proc = subprocess.run(
+            ["nvidia-smi", "-L"],
+            check=False,
+            stdout=subprocess.PIPE,
+            stderr=subprocess.PIPE,
+            text=True,
+            timeout=5,
+        )
+    except FileNotFoundError:
+        LOGGER.debug("nvidia-smi not available; assuming 0 GPUs")
+        return 0
+    except subprocess.SubprocessError as exc:
+        LOGGER.warning("nvidia-smi invocation failed", extra={"error": str(exc)})
+        return 0
+
+    if proc.returncode != 0:
+        LOGGER.debug("nvidia-smi returned non-zero", extra={"stderr": proc.stderr.strip()})
+        return 0
+
+    count = sum(1 for line in proc.stdout.splitlines() if line.strip())
+    return count
+
+
+def _detect_ip_address() -> str:
+    """尝试通过 UDP socket 获得容器出口 IP，失败则回退解析主机名。"""
+    try:
+        with socket.socket(socket.AF_INET, socket.SOCK_DGRAM) as sock:
+            sock.connect(("8.8.8.8", 80))
+            return sock.getsockname()[0]
+    except OSError:
+        LOGGER.debug("UDP socket trick failed; falling back to hostname lookup")
+    try:
+        return socket.gethostbyname(socket.gethostname())
+    except OSError:
+        LOGGER.warning("Unable to resolve hostname to IP; defaulting to 127.0.0.1")
+        return "127.0.0.1"