dev_1.0.0_yuyr_2：重新提交 PR，增加 master/agent 以及系统集成测试 (#17)

Reviewed-on: #17 Reviewed-by: sundapeng <sundp@mail.zgclab.edu.cn> Reviewed-by: xuxt <xuxt@zgclab.edu.cn>
2025-10-11 15:04:46 +08:00
parent 8a38d3d0b2
commit 1e5e91b193
105 changed files with 6096 additions and 107 deletions
--- a/src/agent/app/init.py
+++ b/src/agent/app/init.py
--- a/src/agent/app/client.py
+++ b/src/agent/app/client.py
@@ -0,0 +1,60 @@
+from __future__ import annotations
+
+import json
+from typing import Any, Dict, Optional
+
+import requests
+
+from .log import get_logger
+
+LOGGER = get_logger("argus.agent.client")
+
+
+class MasterAPIError(Exception):
+    def __init__(self, message: str, status_code: int, payload: Optional[Dict[str, Any]] = None) -> None:
+        super().__init__(message)
+        self.status_code = status_code
+        self.payload = payload or {}
+
+
+class AgentClient:
+    def __init__(self, base_url: str, *, timeout: int = 10) -> None:
+        self._base_url = base_url.rstrip("/")
+        self._timeout = timeout
+        self._session = requests.Session()
+
+    def register_node(self, body: Dict[str, Any]) -> Dict[str, Any]:
+        """调用 master 注册接口，返回节点对象。"""
+        url = f"{self._base_url}/api/v1/master/nodes"
+        response = self._session.post(url, json=body, timeout=self._timeout)
+        return self._parse_response(response, "Failed to register node")
+
+    def update_status(self, node_id: str, body: Dict[str, Any]) -> Dict[str, Any]:
+        """上报健康信息，由 master 更新 last_report。"""
+        url = f"{self._base_url}/api/v1/master/nodes/{node_id}/status"
+        response = self._session.put(url, json=body, timeout=self._timeout)
+        return self._parse_response(response, "Failed to update node status")
+
+    def _parse_response(self, response: requests.Response, error_prefix: str) -> Dict[str, Any]:
+        content_type = response.headers.get("Content-Type", "")
+        payload: Dict[str, Any] | None = None
+        if "application/json" in content_type:
+            try:
+                payload = response.json()
+            except json.JSONDecodeError:
+                LOGGER.warning("Response contained invalid JSON", extra={"status": response.status_code})
+
+        if response.status_code >= 400:
+            message = payload.get("error") if isinstance(payload, dict) else response.text
+            raise MasterAPIError(
+                f"{error_prefix}: {message}",
+                status_code=response.status_code,
+                payload=payload if isinstance(payload, dict) else None,
+            )
+
+        if payload is None:
+            try:
+                payload = response.json()
+            except json.JSONDecodeError as exc:
+                raise MasterAPIError("Master returned non-JSON payload", response.status_code) from exc
+        return payload
--- a/src/agent/app/collector.py
+++ b/src/agent/app/collector.py
@@ -0,0 +1,111 @@
+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(),
+    }
+    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() -> int:
+    """采集 GPU 数量，如无法探测则默认为 0。"""
+    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"
--- a/src/agent/app/config.py
+++ b/src/agent/app/config.py
@@ -0,0 +1,74 @@
+from __future__ import annotations
+
+import os
+import socket
+from dataclasses import dataclass
+from pathlib import Path
+from typing import Final
+
+from .version import VERSION
+
+DEFAULT_REPORT_INTERVAL_SECONDS: Final[int] = 60
+
+
+@dataclass(frozen=True)
+class AgentConfig:
+    hostname: str
+    node_file: str
+    version: str
+    master_endpoint: str
+    report_interval_seconds: int
+    health_dir: str
+    request_timeout_seconds: int = 10
+
+
+def _normalise_master_endpoint(value: str) -> str:
+    value = value.strip()
+    if not value:
+        raise ValueError("MASTER_ENDPOINT environment variable is required")
+    if not value.startswith("http://") and not value.startswith("https://"):
+        value = f"http://{value}"
+    return value.rstrip("/")
+
+
+def _read_report_interval(raw_value: str | None) -> int:
+    if raw_value is None or raw_value.strip() == "":
+        return DEFAULT_REPORT_INTERVAL_SECONDS
+    try:
+        interval = int(raw_value)
+    except ValueError as exc:
+        raise ValueError("REPORT_INTERVAL_SECONDS must be an integer") from exc
+    if interval <= 0:
+        raise ValueError("REPORT_INTERVAL_SECONDS must be positive")
+    return interval
+
+
+def _resolve_hostname() -> str:
+    return os.environ.get("AGENT_HOSTNAME") or socket.gethostname()
+
+
+def load_config() -> AgentConfig:
+    """从环境变量推导配置，移除了外部配置文件依赖。"""
+
+    hostname = _resolve_hostname()
+    node_file = f"/private/argus/agent/{hostname}/node.json"
+    health_dir = f"/private/argus/agent/{hostname}/health/"
+
+    master_endpoint_env = os.environ.get("MASTER_ENDPOINT")
+    if master_endpoint_env is None:
+        raise ValueError("MASTER_ENDPOINT environment variable is not set")
+    master_endpoint = _normalise_master_endpoint(master_endpoint_env)
+
+    report_interval_seconds = _read_report_interval(os.environ.get("REPORT_INTERVAL_SECONDS"))
+
+    Path(node_file).parent.mkdir(parents=True, exist_ok=True)
+    Path(health_dir).mkdir(parents=True, exist_ok=True)
+
+    return AgentConfig(
+        hostname=hostname,
+        node_file=node_file,
+        version=VERSION,
+        master_endpoint=master_endpoint,
+        report_interval_seconds=report_interval_seconds,
+        health_dir=health_dir,
+    )
--- a/src/agent/app/health_reader.py
+++ b/src/agent/app/health_reader.py
@@ -0,0 +1,32 @@
+from __future__ import annotations
+
+import json
+from pathlib import Path
+from typing import Any, Dict
+
+from .log import get_logger
+
+LOGGER = get_logger("argus.agent.health")
+
+
+def read_health_directory(path: str) -> Dict[str, Any]:
+    """读取目录中所有 <prefix>-*.json 文件并返回 JSON 映射。"""
+    result: Dict[str, Any] = {}
+    directory = Path(path)
+    if not directory.exists():
+        LOGGER.debug("Health directory does not exist", extra={"path": str(directory)})
+        return result
+
+    for health_file in sorted(directory.glob("*.json")):
+        if "-" not in health_file.stem:
+            LOGGER.debug("Skipping non-prefixed health file", extra={"file": health_file.name})
+            continue
+        try:
+            with health_file.open("r", encoding="utf-8") as handle:
+                content = json.load(handle)
+            result[health_file.stem] = content
+        except json.JSONDecodeError as exc:
+            LOGGER.warning("Failed to parse health file", extra={"file": health_file.name, "error": str(exc)})
+        except OSError as exc:
+            LOGGER.warning("Failed to read health file", extra={"file": health_file.name, "error": str(exc)})
+    return result
--- a/src/agent/app/log.py
+++ b/src/agent/app/log.py
@@ -0,0 +1,18 @@
+from __future__ import annotations
+
+import logging
+import os
+
+
+_LOG_FORMAT = "%(asctime)s %(levelname)s %(name)s - %(message)s"
+
+
+def setup_logging() -> None:
+    level_name = os.environ.get("AGENT_LOG_LEVEL", "INFO").upper()
+    level = getattr(logging, level_name, logging.INFO)
+    logging.basicConfig(level=level, format=_LOG_FORMAT)
+
+
+def get_logger(name: str) -> logging.Logger:
+    setup_logging()
+    return logging.getLogger(name)
--- a/src/agent/app/main.py
+++ b/src/agent/app/main.py
@@ -0,0 +1,163 @@
+from __future__ import annotations
+
+import signal
+import time
+from datetime import datetime, timezone
+from typing import Optional
+
+from .client import AgentClient, MasterAPIError
+from .collector import collect_metadata
+from .config import AgentConfig, load_config
+from .health_reader import read_health_directory
+from .log import get_logger, setup_logging
+from .state import clear_node_state, load_node_state, save_node_state
+
+LOGGER = get_logger("argus.agent")
+
+
+def _current_timestamp() -> str:
+    return datetime.now(timezone.utc).replace(microsecond=0).isoformat().replace("+00:00", "Z")
+
+
+class StopSignal:
+    def __init__(self) -> None:
+        self._stop = False
+
+    def set(self, *_args) -> None:  # type: ignore[override]
+        self._stop = True
+
+    def is_set(self) -> bool:
+        return self._stop
+
+
+def main(argv: Optional[list[str]] = None) -> int:  # noqa: ARG001 - 保留签名以兼容入口调用
+    setup_logging()
+
+    stop_signal = StopSignal()
+    signal.signal(signal.SIGTERM, stop_signal.set)
+    signal.signal(signal.SIGINT, stop_signal.set)
+
+    try:
+        config = load_config()
+    except Exception as exc:
+        LOGGER.error("Failed to load configuration", extra={"error": str(exc)})
+        return 1
+
+    LOGGER.info(
+        "Agent starting",
+        extra={
+            "hostname": config.hostname,
+            "master_endpoint": config.master_endpoint,
+            "node_file": config.node_file,
+        },
+    )
+
+    client = AgentClient(config.master_endpoint, timeout=config.request_timeout_seconds)
+
+    node_state = load_node_state(config.node_file) or {}
+    node_id = node_state.get("id")
+
+    # 与 master 建立注册关系（支持重注册），失败则重试
+    register_response = _register_with_retry(client, config, node_id, stop_signal)
+    if register_response is None:
+        LOGGER.info("Registration aborted due to shutdown signal")
+        return 0
+
+    node_id = register_response.get("id")
+    if not node_id:
+        LOGGER.error("Master did not return node id; aborting")
+        return 1
+    save_node_state(config.node_file, register_response)
+
+    LOGGER.info("Entering status report loop", extra={"node_id": node_id})
+    _status_loop(client, config, node_id, stop_signal)
+    return 0
+
+
+def _register_with_retry(
+    client: AgentClient,
+    config: AgentConfig,
+    node_id: Optional[str],
+    stop_signal: StopSignal,
+):
+    backoff = 5
+    while not stop_signal.is_set():
+        payload = {
+            "name": config.hostname,
+            "type": "agent",
+            "meta_data": collect_metadata(config),
+            "version": config.version,
+        }
+        if node_id:
+            payload["id"] = node_id
+
+        try:
+            response = client.register_node(payload)
+            LOGGER.info("Registration successful", extra={"node_id": response.get("id")})
+            save_node_state(config.node_file, response)
+            return response
+        except MasterAPIError as exc:
+            if exc.status_code == 404 and node_id:
+                LOGGER.warning(
+                    "Master does not recognise node id; clearing local node state",
+                    extra={"node_id": node_id},
+                )
+                clear_node_state(config.node_file)
+                node_id = None
+            elif exc.status_code == 500 and node_id:
+                # id 与 name 不匹配通常意味着配置异常，记录但继续重试
+                LOGGER.error(
+                    "Master rejected node due to id/name mismatch; will retry",
+                    extra={"node_id": node_id},
+                )
+            else:
+                LOGGER.error("Registration failed", extra={"status_code": exc.status_code, "error": str(exc)})
+            time.sleep(min(backoff, 60))
+            backoff = min(backoff * 2, 60)
+        except Exception as exc:  # pragma: no cover - defensive
+            LOGGER.exception("Unexpected error during registration", extra={"error": str(exc)})
+            time.sleep(min(backoff, 60))
+            backoff = min(backoff * 2, 60)
+    return None
+
+
+def _status_loop(
+    client: AgentClient,
+    config: AgentConfig,
+    node_id: str,
+    stop_signal: StopSignal,
+) -> None:
+    interval = config.report_interval_seconds
+    while not stop_signal.is_set():
+        timestamp = _current_timestamp()
+        health_payload = read_health_directory(config.health_dir)
+        body = {
+            "timestamp": timestamp,
+            "health": health_payload,
+        }
+        try:
+            response = client.update_status(node_id, body)
+            LOGGER.info(
+                "Status report succeeded",
+                extra={"node_id": node_id, "health_keys": list(health_payload.keys())},
+            )
+            save_node_state(config.node_file, response)
+        except MasterAPIError as exc:
+            # 保持循环继续执行，等待下一次重试
+            LOGGER.error(
+                "Failed to report status",
+                extra={"status_code": exc.status_code, "error": str(exc)},
+            )
+        except Exception as exc:  # pragma: no cover - defensive
+            LOGGER.exception("Unexpected error during status report", extra={"error": str(exc)})
+
+        for _ in range(interval):
+            if stop_signal.is_set():
+                break
+            time.sleep(1)
+
+    LOGGER.info("Stop signal received; exiting status loop")
+
+
+if __name__ == "__main__":
+    sys.exit(main())
--- a/src/agent/app/state.py
+++ b/src/agent/app/state.py
@@ -0,0 +1,44 @@
+from __future__ import annotations
+
+import json
+import os
+import tempfile
+from pathlib import Path
+from typing import Any, Dict, Optional
+
+from .log import get_logger
+
+LOGGER = get_logger("argus.agent.state")
+
+
+def load_node_state(path: str) -> Optional[Dict[str, Any]]:
+    """读取本地 node.json，容器重启后沿用之前的 ID。"""
+    try:
+        with open(path, "r", encoding="utf-8") as handle:
+            return json.load(handle)
+    except FileNotFoundError:
+        return None
+    except json.JSONDecodeError as exc:
+        LOGGER.warning("node.json is invalid JSON; ignoring", extra={"error": str(exc)})
+        return None
+
+
+def save_node_state(path: str, data: Dict[str, Any]) -> None:
+    """原子化写入 node.json，避免并发读取坏数据。"""
+    directory = Path(path).parent
+    directory.mkdir(parents=True, exist_ok=True)
+    with tempfile.NamedTemporaryFile("w", dir=directory, delete=False, encoding="utf-8") as tmp:
+        json.dump(data, tmp, separators=(",", ":"))
+        tmp.flush()
+        os.fsync(tmp.fileno())
+        temp_path = tmp.name
+    os.replace(temp_path, path)
+
+
+def clear_node_state(path: str) -> None:
+    try:
+        os.remove(path)
+    except FileNotFoundError:
+        return
+    except OSError as exc:
+        LOGGER.warning("Failed to remove node state file", extra={"error": str(exc), "path": path})
--- a/src/agent/app/version.py
+++ b/src/agent/app/version.py
@@ -0,0 +1,69 @@
+from __future__ import annotations
+
+import os
+import sys
+from pathlib import Path
+from typing import Optional
+
+import importlib.metadata
+
+try:
+    import tomllib
+except ModuleNotFoundError:  # pragma: no cover
+    import tomli as tomllib  # type: ignore[no-redef]
+
+
+def _candidate_paths() -> list[Path]:
+    paths = []
+    bundle_dir: Optional[str] = getattr(sys, "_MEIPASS", None)
+    if bundle_dir:
+        paths.append(Path(bundle_dir) / "pyproject.toml")
+    paths.append(Path(__file__).resolve().parent.parent / "pyproject.toml")
+    paths.append(Path(__file__).resolve().parent / "pyproject.toml")
+    paths.append(Path.cwd() / "pyproject.toml")
+    return paths
+
+
+def _read_from_pyproject() -> Optional[str]:
+    for path in _candidate_paths():
+        if not path.exists():
+            continue
+        try:
+            with path.open("rb") as handle:
+                data = tomllib.load(handle)
+        except (OSError, tomllib.TOMLDecodeError):
+            continue
+        project = data.get("project")
+        if isinstance(project, dict):
+            version = project.get("version")
+            if isinstance(version, str):
+                return version
+        tool = data.get("tool")
+        if isinstance(tool, dict):
+            argus_cfg = tool.get("argus")
+            if isinstance(argus_cfg, dict):
+                version = argus_cfg.get("version")
+                if isinstance(version, str):
+                    return version
+    return None
+
+
+def _detect_version() -> str:
+    try:
+        return importlib.metadata.version("argus-agent")
+    except importlib.metadata.PackageNotFoundError:
+        pass
+    override = os.environ.get("AGENT_VERSION_OVERRIDE")
+    if override:
+        return override
+    fallback = _read_from_pyproject()
+    if fallback:
+        return fallback
+    return "0.0.0"
+
+
+VERSION: str = _detect_version()
+
+
+def get_version() -> str:
+    return VERSION