From 24f8273bfc98417078a7788f4e699040f285fff6 Mon Sep 17 00:00:00 2001 From: Aditi Gaur Date: Tue, 9 Jun 2026 14:35:05 -0700 Subject: [PATCH] Save XIDs on healthagent restart. Only discard XID's if they are older than the boot time. --- healthagent/gpu.py | 75 +++++++++++++++++++++++++++++++++++++++++++++- 1 file changed, 74 insertions(+), 1 deletion(-) diff --git a/healthagent/gpu.py b/healthagent/gpu.py index 20137ee..f041fd1 100644 --- a/healthagent/gpu.py +++ b/healthagent/gpu.py @@ -1,8 +1,11 @@ import asyncio +import json import sys import logging import os import shutil +import tempfile +import time from datetime import datetime, timezone from healthagent import epilog,status,healthcheck,prolog from healthagent.scheduler import Scheduler @@ -19,6 +22,10 @@ # At 1s polling, 300 = 5-minute window for delta_gt rate computation. MAX_KEEP_SAMPLES = 300 +# Persistent XID history file — survives healthagent restarts. +_XID_HISTORY_DIR = os.path.join(os.getenv("HEALTHAGENT_DIR", "/opt/healthagent"), "run") +_XID_HISTORY_FILE = os.path.join(_XID_HISTORY_DIR, "xid_history.json") + # Field-specific enrichments: map field ID -> callable(raw_value, field_values) -> list[str]. # Applied *after* generic message formatting to decode field values into # human-readable details. Independent of watch config (defaults or user overrides). @@ -67,7 +74,7 @@ def setup(self): self.dcgmHandle = None self.watch_fields = [] self.gpu_config = [] - self.xid_history = {} + self.xid_history = self._load_xid_history() if self.test_mode: log.info("Running GPU tests in DCGM_TEST_MODE") self.dcgmGroup, self.dcgmHandle = Wrap.connect(grp_name="healthagent_group", test_mode=self.test_mode) @@ -107,6 +114,67 @@ def setup_background_watches(self): # Seed the persistent sample collection (first call with None returns latest values) self._field_collection = self.dcgmGroup.samples.GetAllSinceLastCall_v2(None, self.field_group) + @staticmethod + def _boot_time() -> float: + """Return wall-clock timestamp of last boot (seconds since epoch).""" + with open('/proc/uptime') as f: + uptime_seconds = float(f.read().split()[0]) + return time.time() - uptime_seconds + + @staticmethod + def _load_xid_history() -> dict: + """Load XID history from JSON. Discard the file entirely if it predates last boot.""" + try: + boot_ts = GpuHealthChecks._boot_time() + if os.path.getmtime(_XID_HISTORY_FILE) < boot_ts: + log.debug("XID history file predates last boot, discarding") + os.remove(_XID_HISTORY_FILE) + return {} + with open(_XID_HISTORY_FILE, 'r') as f: + history = json.load(f) + except (FileNotFoundError, json.JSONDecodeError, OSError) as e: + log.debug(f"No previous XID history loaded: {e}") + return {} + + # Convert JSON string keys back to int + try: + restored = { + gpu_id: {int(k): v for k, v in xids.items()} + for gpu_id, xids in history.items() + } + except (ValueError, AttributeError, TypeError) as e: + log.warning(f"Corrupt XID history file, discarding: {e}") + return {} + count = sum(len(v) for v in restored.values()) + if count: + log.info(f"Restored {count} XIDs from previous session") + return restored + + async def _save_xid_history(self): + """Persist XID history to JSON. Scheduled via Scheduler.add_task (fire-and-forget).""" + data = self.xid_history + # Convert int keys to strings for JSON compatibility + serializable = { + gpu_id: {str(k): v for k, v in xids.items()} + for gpu_id, xids in data.items() + } + try: + await asyncio.to_thread(self._write_xid_file, serializable) + except Exception as e: + log.warning(f"Failed to persist XID history: {e}") + + @staticmethod + def _write_xid_file(data: dict): + """Write XID data atomically via unique temp file + rename.""" + fd, tmp = tempfile.mkstemp(dir=_XID_HISTORY_DIR, suffix='.tmp') + try: + with os.fdopen(fd, 'w', encoding='utf-8') as f: + json.dump(data, f) + os.replace(tmp, _XID_HISTORY_FILE) + except BaseException: + os.unlink(tmp) + raise + def __display_gpu_config(self): ## Invoke method to get gpu IDs of the members of the newly-created group @@ -179,6 +247,7 @@ async def handle_policy_violation(self, callbackresp): self.xid_history[gpu_key] = {} if xid_received not in self.xid_history[gpu_key] or timestamp < self.xid_history[gpu_key][xid_received]["timestamp"]: self.xid_history[gpu_key][xid_received] = {"xid": xid_received, "timestamp": timestamp} + Scheduler.add_task(self._save_xid_history) except ValueError as e: log.exception(e) return @@ -195,6 +264,7 @@ def track_fieldsv2(self): XIDs are handled separately — GPU-only, not part of field watches. """ custom_fields = {'error_count': 0, 'warning_count': 0, 'category': set()} + xid_changed = False try: # Accumulate new samples since last call self._field_collection = self.dcgmGroup.samples.GetAllSinceLastCall_v2( @@ -278,6 +348,7 @@ def track_fieldsv2(self): ts_utc = datetime.fromtimestamp(sample.ts / 1_000_000, tz=timezone.utc).strftime("%Y-%m-%dT%H:%M:%S UTC") if xid_num not in self.xid_history[gpu_id] or ts_utc < self.xid_history[gpu_id][xid_num]["timestamp"]: self.xid_history[gpu_id][xid_num] = {"xid": xid_num, "timestamp": ts_utc} + xid_changed = True # Populate report with full XID history for gpu_id, xids in self.xid_history.items(): @@ -302,6 +373,8 @@ def track_fieldsv2(self): if xids: custom_fields['category'].add("XID") + if xid_changed: + Scheduler.add_task(self._save_xid_history) return custom_fields except Exception as e: log.exception(e)