nodeview-slurm 0.1.0__py3-none-any.whl

nodeview/profile.py

@@ -0,0 +1,284 @@
+"""Everything the cluster knows about you.
+
+This is information you rarely look at (project budgets, disk quotas, account
+expiration) but that matters a lot when it does: an exhausted budget or a full
+quota blocks your jobs without telling you why.
+
+It costs ~1.7 s in total, almost all spent in the three local commands
+`susage`, `squota` and `sexpiration`, which read a separately updated cache.
+Too much for the job refresh loop: it's read at a much slower pace.
+
+Each source is independent and may be missing: on a cluster without the local
+commands the other sections still show up.
+"""
+from __future__ import annotations
+
+import re
+import subprocess
+from dataclasses import dataclass, field
+from datetime import datetime, timedelta
+from typing import Optional
+
+_ANSI = re.compile(r'\x1b\[[0-9;]*m')
+
+HISTORY_DAYS = 30
+
+
+def _run(args, timeout=25) -> Optional[str]:
+    try:
+        out = subprocess.run(args, capture_output=True, text=True, timeout=timeout)
+    except (OSError, subprocess.SubprocessError):
+        return None
+    if out.returncode != 0 and not out.stdout.strip():
+        return None
+    return out.stdout
+
+
+def _parse_tabulate(text: str):
+    """Parses the fixed-width tables printed by the local commands.
+
+    The row of dashes under the header gives the exact column boundaries:
+    more reliable than splitting on spaces, since the headers contain some
+    ("Total usage (%)").
+    """
+    lines = _ANSI.sub('', text or '').splitlines()
+    sep = None
+    for i, line in enumerate(lines):
+        stripped = line.strip()
+        if stripped and set(stripped) <= set('- ') and '-' in stripped:
+            sep = i
+            break
+    if sep is None or sep == 0:
+        return [], []
+
+    spans = [(m.start(), m.end()) for m in re.finditer(r'-+', lines[sep])]
+    if not spans:
+        return [], []
+    spans[-1] = (spans[-1][0], 10_000)          # the last column may overflow
+
+    def cells(line):
+        return [line[a:b].strip() for a, b in spans]
+
+    headers = cells(lines[sep - 1])
+    rows = []
+    for line in lines[sep + 1:]:
+        if not line.strip():
+            break
+        row = cells(line)
+        if any(row):
+            rows.append(row)
+    return headers, rows
+
+
+def _num(value) -> Optional[float]:
+    try:
+        return float(str(value).replace(',', '').strip())
+    except (TypeError, ValueError):
+        return None
+
+
+@dataclass(frozen=True)
+class Budget:
+    """A project's compute budget, in standard hours."""
+    project: str
+    used: Optional[float]
+    total: Optional[float]
+    pct: Optional[float]
+    year_used: Optional[float]
+    year_total: Optional[float]
+    year_pct: Optional[float]
+
+
+@dataclass(frozen=True)
+class DiskQuota:
+    filesystem: str
+    owner: str
+    used_gb: Optional[float]
+    quota_gb: Optional[float]
+    pct: Optional[float]
+
+
+@dataclass(frozen=True)
+class Expiration:
+    date: str
+    days_left: Optional[int]
+
+
+@dataclass(frozen=True)
+class Share:
+    """Fairshare: how much you've used relative to your allotted share."""
+    account: str
+    norm_shares: Optional[float]
+    effective_usage: Optional[float]
+    factor: Optional[float]
+
+
+@dataclass(frozen=True)
+class History:
+    days: int
+    completed: int = 0
+    failed: int = 0
+    cancelled: int = 0
+    timeout: int = 0
+    other: int = 0
+
+    @property
+    def finished(self) -> int:
+        return self.completed + self.failed + self.cancelled + self.timeout + self.other
+
+    @property
+    def failure_rate(self) -> Optional[float]:
+        return (self.failed + self.timeout) / self.finished if self.finished else None
+
+
+@dataclass
+class Profile:
+    budgets: list = field(default_factory=list)
+    quotas: list = field(default_factory=list)
+    expiration: Optional[Expiration] = None
+    shares: list = field(default_factory=list)
+    accounts: dict = field(default_factory=dict)      # account -> list of QOS
+    history: Optional[History] = None
+    taken_at: float = 0.0
+    duration: float = 0.0
+
+    @property
+    def empty(self) -> bool:
+        return not (self.budgets or self.quotas or self.expiration
+                    or self.shares or self.accounts or self.history)
+
+
+def _read_budgets() -> list:
+    headers, rows = _parse_tabulate(_run(['susage']) or '')
+    if not rows:
+        return []
+
+    def find(*needles):
+        for i, h in enumerate(headers):
+            low = h.lower()
+            if all(n in low for n in needles):
+                return i
+        return None
+
+    i_proj = find('project') or 0
+    i_tot = find('total', 'budget')
+    i_use = find('total', 'usage')
+    i_pct = find('total', 'usage', '%')
+    # "Total usage" and "Total usage (%)" look alike: exclude the percentage
+    if i_use is not None and i_use == i_pct:
+        i_use = next((i for i, h in enumerate(headers)
+                      if 'usage' in h.lower() and '%' not in h and 'total' in h.lower()), None)
+    i_ytot = find('budget', '20')
+    i_yuse = next((i for i, h in enumerate(headers)
+                   if 'usage' in h.lower() and '%' not in h and re.search(r'20\d\d', h)), None)
+    i_ypct = next((i for i, h in enumerate(headers)
+                   if 'usage' in h.lower() and '%' in h and re.search(r'20\d\d', h)), None)
+
+    def cell(row, idx):
+        return _num(row[idx]) if idx is not None and idx < len(row) else None
+
+    return [Budget(project=row[i_proj], used=cell(row, i_use), total=cell(row, i_tot),
+                   pct=cell(row, i_pct), year_used=cell(row, i_yuse),
+                   year_total=cell(row, i_ytot), year_pct=cell(row, i_ypct))
+            for row in rows if row[i_proj]]
+
+
+def _read_quotas() -> list:
+    headers, rows = _parse_tabulate(_run(['squota']) or '')
+    if not rows:
+        return []
+    i_gb = [i for i, h in enumerate(headers) if 'gb' in h.lower() and '%' not in h]
+    i_pct = next((i for i, h in enumerate(headers) if '%' in h and 'gb' in h.lower()), None)
+    i_use = i_gb[0] if i_gb else None
+    i_quota = i_gb[1] if len(i_gb) > 1 else None
+
+    out = []
+    for row in rows:
+        if len(row) < 2 or not row[0]:
+            continue
+        out.append(DiskQuota(
+            filesystem=row[0], owner=row[1],
+            used_gb=_num(row[i_use]) if i_use is not None and i_use < len(row) else None,
+            quota_gb=_num(row[i_quota]) if i_quota is not None and i_quota < len(row) else None,
+            pct=_num(row[i_pct]) if i_pct is not None and i_pct < len(row) else None,
+        ))
+    return out
+
+
+def _read_expiration() -> Optional[Expiration]:
+    headers, rows = _parse_tabulate(_run(['sexpiration']) or '')
+    if not rows:
+        return None
+    row = rows[0]
+    date = row[1] if len(row) > 1 else ''
+    days = _num(row[2]) if len(row) > 2 else None
+    return Expiration(date=date, days_left=int(days) if days is not None else None)
+
+
+def _read_shares() -> list:
+    text = _run(['sshare', '-U', '-P', '-n',
+                 '-o', 'Account,NormShares,EffectvUsage,FairShare'])
+    out = []
+    for line in (text or '').splitlines():
+        parts = [p.strip() for p in _ANSI.sub('', line).split('|')]
+        if len(parts) < 4 or not parts[0]:
+            continue
+        out.append(Share(account=parts[0], norm_shares=_num(parts[1]),
+                         effective_usage=_num(parts[2]), factor=_num(parts[3])))
+    return out
+
+
+def _read_accounts() -> dict:
+    import os
+    user = os.environ.get('USER') or ''
+    text = _run(['sacctmgr', '-nP', 'show', 'assoc', f'user={user}',
+                 'format=Account,QOS'])
+    accounts = {}
+    for line in (text or '').splitlines():
+        parts = line.split('|')
+        if len(parts) < 2 or not parts[0].strip():
+            continue
+        accounts[parts[0].strip()] = [q for q in parts[1].split(',') if q.strip()]
+    return accounts
+
+
+def _read_history(days: int = HISTORY_DAYS) -> Optional[History]:
+    since = (datetime.now() - timedelta(days=days)).strftime('%Y-%m-%d')
+    text = _run(['sacct', '-n', '-X', '-P', '-S', since, '-o', 'State'])
+    if text is None:
+        return None
+    counts = {'completed': 0, 'failed': 0, 'cancelled': 0, 'timeout': 0, 'other': 0}
+    for line in text.splitlines():
+        state = line.strip().upper()
+        if not state:
+            continue
+        if state.startswith('COMPLETED'):
+            counts['completed'] += 1
+        elif state.startswith('FAILED') or state.startswith('NODE_FAIL') or state.startswith('OUT_OF'):
+            counts['failed'] += 1
+        elif state.startswith('CANCELLED'):
+            counts['cancelled'] += 1
+        elif state.startswith('TIMEOUT'):
+            counts['timeout'] += 1
+        elif state.startswith('RUNNING') or state.startswith('PENDING'):
+            continue                      # not finished yet, don't count
+        else:
+            counts['other'] += 1
+    return History(days=days, **counts)
+
+
+def read_profile() -> Profile:
+    import time
+    started = time.perf_counter()
+    profile = Profile()
+    # each source on its own: if a command is missing, the others remain
+    for attr, reader in (('budgets', _read_budgets), ('quotas', _read_quotas),
+                         ('expiration', _read_expiration), ('shares', _read_shares),
+                         ('accounts', _read_accounts), ('history', _read_history)):
+        try:
+            setattr(profile, attr, reader())
+        except Exception:
+            pass
+    profile.taken_at = time.time()
+    profile.duration = time.perf_counter() - started
+    return profile