194 lines
7.3 KiB
Python
194 lines
7.3 KiB
Python
"""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
|