"""Read-only GPU-Hot snapshot collection with an allowlisted public result.""" from __future__ import annotations import json import math import re from datetime import datetime, timezone from typing import Any from urllib.parse import urlsplit import httpx from ..core.config import settings _MAX_RESPONSE_BYTES = 1_048_576 _MAX_FUTURE_SECONDS = 5.0 def _result(status: str, reason: str | None, *, devices=None, observed_at=None) -> dict[str, Any]: return { 'status': status, 'observedAt': observed_at, 'devices': devices or [], 'reason': reason, 'source': 'gpu-hot', } def _number(value: Any, *, minimum: float = 0.0, maximum: float | None = None) -> float | None: if value is None or isinstance(value, bool): return None try: numeric = float(value) except (TypeError, ValueError, OverflowError): return None if not math.isfinite(numeric) or numeric < minimum or (maximum is not None and numeric > maximum): return None return numeric def _source_time(value: Any) -> datetime | None: if not isinstance(value, str): return None try: parsed = datetime.fromisoformat(value.replace('Z', '+00:00')) if parsed.tzinfo is None or parsed.utcoffset() is None: return None return parsed.astimezone(timezone.utc) except (ValueError, OverflowError): return None def _index(value: Any) -> str | None: if isinstance(value, bool) or not isinstance(value, (int, str)): return None text = str(value) if not re.fullmatch(r'[0-9]{1,6}', text): return None return str(int(text)) def normalize_gpu_snapshot(payload, *, now=None, max_age_seconds=30.0): """Keep measured values; never infer zero usage or a source timezone.""" if not isinstance(payload, dict): return _result('error', 'gpu_metrics_invalid_response') raw_devices = payload.get('gpus') if isinstance(raw_devices, dict): entries = list(raw_devices.items()) elif isinstance(raw_devices, list): entries = list(enumerate(raw_devices)) else: return _result('error', 'gpu_metrics_invalid_response') if not entries: return _result('unavailable', 'gpu_metrics_no_devices') current_time = now if isinstance(now, datetime) else datetime.now(timezone.utc) if current_time.tzinfo is None: current_time = current_time.replace(tzinfo=timezone.utc) current_time = current_time.astimezone(timezone.utc) age_limit = _number(max_age_seconds) or 30.0 devices = [] observed_times = [] seen_ids = set() incomplete = False missing_time = False incomplete_values = False for key, raw_device in entries: if not isinstance(raw_device, dict): incomplete = True continue index = _index(raw_device.get('index')) or _index(key) if index is None or index in seen_ids: incomplete = True continue observed_time = _source_time(raw_device.get('timestamp')) if observed_time is not None: age_seconds = (current_time - observed_time).total_seconds() if age_seconds > age_limit or age_seconds < -_MAX_FUTURE_SECONDS: incomplete = True continue observed_times.append(observed_time) else: missing_time = True raw_name = raw_device.get('name') name = re.sub(r'[\x00-\x1f\x7f]', '', raw_name).strip()[:160] if isinstance(raw_name, str) else '' if not name: name = f'GPU {index}' incomplete_values = True utilization = _number(raw_device.get('utilization'), maximum=100.0) memory_used = _number(raw_device.get('memory_used')) memory_total = _number(raw_device.get('memory_total')) if memory_total == 0: memory_total = None if memory_used is not None and memory_total is not None and memory_used > memory_total: memory_used = None temperature = _number(raw_device.get('temperature'), maximum=150.0) values = (utilization, memory_used, memory_total, temperature) if any(value is None for value in values): incomplete_values = True devices.append({ 'id': f'gpu-{index}', 'name': name, 'utilizationPercent': utilization, 'memoryUsedMiB': memory_used, 'memoryTotalMiB': memory_total, 'temperatureC': temperature, }) seen_ids.add(index) if not devices: return _result('unavailable', 'gpu_metrics_no_devices') observed_at = None if observed_times and not missing_time: observed_at = min(observed_times).isoformat().replace('+00:00', 'Z') partial = incomplete or incomplete_values or missing_time reason = None if partial: reason = 'gpu_metrics_timestamp_unavailable' if missing_time and not incomplete and not incomplete_values else 'gpu_metrics_partial' return _result('partial' if partial else 'ok', reason, devices=devices, observed_at=observed_at) class GpuMetricsConnector: def __init__(self, snapshot_url=None, verify_tls=True, timeout_seconds=3.0): self.snapshot_url = settings.gpu_metrics_url if snapshot_url is None else snapshot_url self.verify_tls = verify_tls duration = _number(timeout_seconds) self.timeout_seconds = min(max(duration if duration is not None else 3.0, 0.05), 3.0) def get_metrics(self): if not self.snapshot_url: return _result('missing_config', 'gpu_metrics_missing_config') if not self._valid_url(): return _result('error', 'gpu_metrics_invalid_url') try: with ( httpx.Client(verify=self.verify_tls, timeout=self.timeout_seconds, follow_redirects=False, trust_env=False) as client, client.stream('GET', self.snapshot_url, headers={'Accept': 'application/json'}) as response, ): if response.status_code != 200: return _result('unavailable', 'gpu_metrics_http_error') body = bytearray() for chunk in response.iter_bytes(chunk_size=65_536): if len(body) + len(chunk) > _MAX_RESPONSE_BYTES: return _result('error', 'gpu_metrics_response_too_large') body.extend(chunk) payload = json.loads(body) return normalize_gpu_snapshot(payload) except httpx.HTTPError: return _result('unavailable', 'gpu_metrics_unreachable') except httpx.InvalidURL: return _result('error', 'gpu_metrics_invalid_url') except (ValueError, UnicodeError): return _result('error', 'gpu_metrics_invalid_response') def _valid_url(self) -> bool: if not isinstance(self.snapshot_url, str) or re.search(r'[\s\x00-\x1f\x7f]', self.snapshot_url): return False try: parsed = urlsplit(self.snapshot_url) return ( parsed.scheme in {'http', 'https'} and bool(parsed.hostname) and parsed.username is None and parsed.password is None and not parsed.fragment and (parsed.port is None or 0 < parsed.port <= 65535) ) except ValueError: return False