nodeview-slurm 0.1.0__py3-none-any.whl

nodeview/nodes.py

@@ -0,0 +1,318 @@
+"""Full node readout from `scontrol -d -o show node`.
+
+nodeocc stops at sinfo and misses a lot: the actual GPU model (which here
+lives in the features, not in GRES), the indices of busy GPUs, allocated RAM
+versus truly free RAM, CPU load, partitions, scheduling weight. A single call
+returns all of it, as key=value pairs.
+"""
+from __future__ import annotations
+
+import re
+import subprocess
+from dataclasses import dataclass, field
+from typing import Optional
+
+from .i18n import t
+
+# "NodeName=x CPUAlloc=0 AllocTRES=cpu=12,mem=96G,gres/gpu=3 ..." -> dict.
+# Careful: some values contain '=' (AllocTRES, CfgTRES), so each token is split
+# on its FIRST '=', not on all of them. The only fields with spaces inside are
+# OS= and Reason=: the first isn't needed, the second is extracted separately.
+_REASON = re.compile(r'\bReason=(.*?)(?=\s+\w+=|$)')
+
+
+def _kv(line: str) -> dict:
+    out = {}
+    for token in line.split():
+        if '=' not in token:
+            continue
+        key, value = token.split('=', 1)
+        out.setdefault(key, value)
+    match = _REASON.search(line)
+    if match:
+        out['Reason'] = match.group(1).strip()
+    return out
+
+
+def _int(value, default=0) -> int:
+    try:
+        return int(str(value).strip())
+    except (TypeError, ValueError):
+        return default
+
+
+def _float(value, default=0.0) -> float:
+    try:
+        return float(str(value).strip())
+    except (TypeError, ValueError):
+        return default
+
+
+def _gpu_count(gres: str) -> int:
+    """'gpu:5(S:0),tmpfs:90198M' -> 5 ; 'gpu:(null):3(IDX:1-2,4)' -> 3"""
+    if not gres:
+        return 0
+    for token in gres.split(','):
+        token = token.strip()
+        if not token.startswith('gpu'):
+            continue
+        core = token.split('(')[0]
+        parts = [p for p in core.split(':') if p and p != 'null']
+        for p in reversed(parts):
+            if p.isdigit():
+                return int(p)
+    return 0
+
+
+def _gpu_indices(gres_used: str) -> list:
+    """'gpu:(null):3(IDX:1-2,4)' -> [1, 2, 4]"""
+    if not gres_used:
+        return []
+    # Slurm writes 'gpu:(null):3(IDX:1-2,4)': between 'gpu:' and '(IDX:' sits
+    # the model, which here is (null), so skip it without swallowing the IDX
+    match = re.search(r'gpu:.*?\(IDX:([^)]*)\)', gres_used)
+    if not match:
+        return []
+    out = []
+    for chunk in match.group(1).split(','):
+        chunk = chunk.strip()
+        if chunk in ('', 'N/A'):
+            continue
+        if '-' in chunk:
+            a, b = chunk.split('-', 1)
+            if a.isdigit() and b.isdigit():
+                out.extend(range(int(a), int(b) + 1))
+        elif chunk.isdigit():
+            out.append(int(chunk))
+    return sorted(set(out))
+
+
+def _model(features: str) -> str:
+    """'gpu_RTX_A5000_24G' -> 'RTX A5000 24G'; several features -> first gpu_*"""
+    for feat in (features or '').split(','):
+        feat = feat.strip()
+        if feat.startswith('gpu_'):
+            return feat[4:].replace('_', ' ')
+    first = (features or '').split(',')[0].strip()
+    return t('na') if first in ('', '(null)', 'null') else first
+
+
+def _mem_mb(value: str) -> float:
+    """'96G' / '98304' / '96000M' -> MB"""
+    value = (value or '').strip()
+    if not value:
+        return 0.0
+    mult = {'K': 1 / 1024, 'M': 1, 'G': 1024, 'T': 1024 * 1024}
+    if value[-1].upper() in mult:
+        return _float(value[:-1]) * mult[value[-1].upper()]
+    return _float(value)
+
+
+def _alloc_tres(tres: str) -> dict:
+    """'cpu=12,mem=96G,gres/gpu=3' -> {'cpu': 12.0, 'mem': 98304.0, 'gpu': 3.0}"""
+    out = {}
+    for item in (tres or '').split(','):
+        if '=' not in item:
+            continue
+        key, value = item.split('=', 1)
+        key = key.strip()
+        if key == 'mem':
+            out['mem'] = _mem_mb(value)
+        elif key == 'cpu':
+            out['cpu'] = _float(value)
+        elif key in ('gres/gpu', 'gpu'):
+            out['gpu'] = _float(value)
+    return out
+
+
+@dataclass
+class NodeSpec:
+    """Everything Slurm says about a node."""
+    name: str
+    state: str                  # IDLE, MIXED, ALLOCATED, DOWN, DRAINED...
+    reason: Optional[str]
+    partitions: tuple = ()
+
+    gpu_model: str = field(default_factory=lambda: t('na'))
+    gpus_total: int = 0
+    gpus_used: int = 0
+    gpu_indices_used: tuple = ()
+
+    cpus_total: int = 0
+    cpus_alloc: int = 0
+    cpu_load: float = 0.0
+
+    mem_total_mb: float = 0.0
+    mem_alloc_mb: float = 0.0
+    mem_free_mb: float = 0.0    # truly free (system RAM), not "unallocated"
+
+    sockets: int = 0
+    cores_per_socket: int = 0
+    threads_per_core: int = 0
+    arch: str = ''
+    weight: int = 0
+    tmp_disk_mb: float = 0.0
+    boot_time: str = ''
+    slurm_version: str = ''
+
+    @property
+    def gpus_free(self) -> int:
+        return max(0, self.gpus_total - self.gpus_used)
+
+    @property
+    def cpus_free(self) -> int:
+        return max(0, self.cpus_total - self.cpus_alloc)
+
+    @property
+    def mem_unalloc_mb(self) -> float:
+        return max(0.0, self.mem_total_mb - self.mem_alloc_mb)
+
+    @property
+    def base_state(self) -> str:
+        return self.state.split('+')[0].split('*')[0].upper()
+
+    @property
+    def healthy(self) -> bool:
+        bad = ('DOWN', 'DRAIN', 'DRAINED', 'DRAINING', 'FAIL', 'FAILING',
+               'MAINT', 'INVAL', 'UNKNOWN', 'NOT_RESPONDING')
+        return not any(b in self.state.upper() for b in bad)
+
+    @property
+    def schedulable(self) -> bool:
+        return self.healthy and self.gpus_free > 0
+
+    @property
+    def gpu_load(self) -> float:
+        return self.gpus_used / self.gpus_total if self.gpus_total else 0.0
+
+    @property
+    def cpu_alloc_load(self) -> float:
+        return self.cpus_alloc / self.cpus_total if self.cpus_total else 0.0
+
+    @property
+    def mem_load(self) -> float:
+        return self.mem_alloc_mb / self.mem_total_mb if self.mem_total_mb else 0.0
+
+    @property
+    def cpu_real_load(self) -> float:
+        """Actual load over total CPUs: exposes jobs that request but don't use."""
+        return self.cpu_load / self.cpus_total if self.cpus_total else 0.0
+
+
+def short_name(name: str) -> str:
+    """Same normalization as nodeocc, so joblet node names match."""
+    return name.replace('aimagelab-srv-', '').replace('ailb-login-', '')
+
+
+def _expand_nodelist(spec: str) -> list:
+    """'huber' -> ['huber'] ; 'ailb-login-0[2-3],nico' -> ['ailb-login-02', ...]"""
+    out, buf, depth = [], '', 0
+    for ch in spec:                      # commas inside [] don't split
+        if ch == '[':
+            depth += 1
+        elif ch == ']':
+            depth -= 1
+        if ch == ',' and depth == 0:
+            out.append(buf)
+            buf = ''
+        else:
+            buf += ch
+    if buf:
+        out.append(buf)
+
+    names = []
+    for item in out:
+        item = item.strip()
+        if '[' not in item:
+            if item:
+                names.append(item)
+            continue
+        prefix, rest = item.split('[', 1)
+        body, suffix = rest.split(']', 1) if ']' in rest else (rest, '')
+        for part in body.split(','):
+            if '-' in part:
+                a, b = part.split('-', 1)
+                if a.isdigit() and b.isdigit():
+                    width = len(a)
+                    names.extend(f"{prefix}{i:0{width}d}{suffix}"
+                                 for i in range(int(a), int(b) + 1))
+            elif part:
+                names.append(f"{prefix}{part}{suffix}")
+    return names
+
+
+def read_gpu_owners() -> dict:
+    """{node: {GPU index: user}}.
+
+    `scontrol -d show job` reports, for each job and each node it occupies,
+    the exact indices of the assigned GPUs: it's the only way to know not just
+    which GPUs are in use, but by whom. A single call for all jobs, with -o
+    putting each job on one line; multi-node jobs have several
+    `Nodes=... GRES=...` blocks on the same line, so all of them are scanned.
+    """
+    try:
+        out = subprocess.run(['scontrol', '-d', '-o', 'show', 'job'],
+                             capture_output=True, text=True, timeout=20)
+    except (OSError, subprocess.SubprocessError):
+        return {}
+    if out.returncode != 0:
+        return {}
+
+    owners = {}
+    block = re.compile(r'Nodes=(\S+)\s+CPU_IDs=\S+\s+Mem=\S+\s+GRES=(\S+)')
+    for line in out.stdout.splitlines():
+        if 'JobState=RUNNING' not in line and 'JobState=COMPLETING' not in line:
+            continue
+        who = re.search(r'UserId=([^(\s]+)', line)
+        if not who:
+            continue
+        user = who.group(1)
+        for nodes_spec, gres in block.findall(line):
+            indices = _gpu_indices(gres)
+            if not indices:
+                continue
+            for node in _expand_nodelist(nodes_spec):
+                slot = owners.setdefault(short_name(node), {})
+                for idx in indices:
+                    slot.setdefault(idx, user)
+    return owners
+
+
+def read_nodes() -> list:
+    out = subprocess.run(['scontrol', '-d', '-o', 'show', 'node'],
+                         capture_output=True, text=True, timeout=20)
+    if out.returncode != 0:
+        raise RuntimeError(t('src.scontrol_failed', err=out.stderr.strip()))
+
+    nodes = []
+    for line in out.stdout.splitlines():
+        if not line.strip().startswith('NodeName='):
+            continue
+        d = _kv(line)
+        alloc = _alloc_tres(d.get('AllocTRES', ''))
+        reason = d.get('Reason', '').strip()
+        nodes.append(NodeSpec(
+            name=short_name(d.get('NodeName', '?')),
+            state=d.get('State', '?'),
+            reason=None if reason in ('', 'none', 'N/A') else reason,
+            partitions=tuple(p for p in d.get('Partitions', '').split(',') if p),
+            gpu_model=_model(d.get('AvailableFeatures', '')),
+            gpus_total=_gpu_count(d.get('Gres', '')),
+            gpus_used=int(alloc.get('gpu', _gpu_count(d.get('GresUsed', '')))),
+            gpu_indices_used=tuple(_gpu_indices(d.get('GresUsed', ''))),
+            cpus_total=_int(d.get('CPUTot')),
+            cpus_alloc=_int(d.get('CPUAlloc')),
+            cpu_load=_float(d.get('CPULoad')),
+            mem_total_mb=_float(d.get('RealMemory')),
+            mem_alloc_mb=_float(d.get('AllocMem')),
+            mem_free_mb=_float(d.get('FreeMem')),
+            sockets=_int(d.get('Sockets')),
+            cores_per_socket=_int(d.get('CoresPerSocket')),
+            threads_per_core=_int(d.get('ThreadsPerCore')),
+            arch=d.get('Arch', ''),
+            weight=_int(d.get('Weight')),
+            tmp_disk_mb=_float(d.get('TmpDisk')),
+            boot_time=d.get('BootTime', ''),
+            slurm_version=d.get('Version', ''),
+        ))
+    return nodes

nodeview/palette.py

@@ -0,0 +1,325 @@
+"""Color palette.
+
+Color here carries information, not decoration. Four distinct jobs, never
+mixed:
+
+- **categorical** -> identity: who occupies a GPU. Fixed order, never cycled:
+  the eighth user is the last with a color of their own; from the ninth on
+  they fall into "others". Color never travels alone, the name is always next
+  to it.
+- **sequential** -> quantity: how full a node is. A single hue, light to dark.
+  No rainbows.
+- **status** -> condition: free / nearly full / full / broken. Reserved, never
+  reused as "series 5", always paired with an icon and a word.
+- **ink** -> text and borders. Numbers and labels never take the series
+  color.
+
+The values come from the visualization guide's reference palette and were run
+through the validator for both backgrounds (lightness band, minimum chroma,
+color-blindness separation, contrast): all checks pass. On a light background
+three hues stay below 3:1 contrast, which is why user names are always
+written next to the color.
+
+Terminals without truecolor support get the 256- or 16-color approximation:
+rich handles that on its own.
+"""
+from __future__ import annotations
+
+import colorsys
+from dataclasses import dataclass
+
+
+@dataclass(frozen=True)
+class Palette:
+    name: str
+
+    # identity: fixed order, 8 slots
+    categorical: tuple
+    other: str
+
+    # quantity: a single hue, from light to dark
+    sequential: tuple
+
+    # status: reserved
+    good: str
+    warning: str
+    serious: str
+    critical: str
+
+    # ink
+    ink: str
+    ink_soft: str
+    ink_muted: str
+    rule: str
+
+    def slot(self, index: int) -> str:
+        """Color of slot `index`; past the eighth, the neutral 'others' color."""
+        if index < 0 or index >= len(self.categorical):
+            return self.other
+        return self.categorical[index]
+
+    def fill(self, ratio: float) -> str:
+        """Step of the sequential ramp for a quantity between 0 and 1."""
+        if not self.sequential:
+            return self.ink
+        steps = len(self.sequential)
+        idx = min(steps - 1, max(0, int(round(ratio * (steps - 1)))))
+        return self.sequential[idx]
+
+
+DARK = Palette(
+    name='dark',
+    categorical=('#3987e5', '#d95926', '#199e70', '#c98500',
+                 '#d55181', '#008300', '#9085e9', '#e66767'),
+    other='#898781',
+    sequential=('#cde2fb', '#9ec5f4', '#6da7ec', '#3987e5', '#256abf', '#184f95'),
+    good='#0ca30c',
+    warning='#fab219',
+    serious='#ec835a',
+    critical='#d03b3b',
+    ink='#ffffff',
+    ink_soft='#c3c2b7',
+    ink_muted='#898781',
+    rule='#383835',
+)
+
+LIGHT = Palette(
+    name='light',
+    categorical=('#2a78d6', '#eb6834', '#1baf7a', '#eda100',
+                 '#e87ba4', '#008300', '#4a3aa7', '#e34948'),
+    other='#898781',
+    sequential=('#cde2fb', '#9ec5f4', '#6da7ec', '#3987e5', '#256abf', '#184f95'),
+    good='#006300',
+    warning='#b37400',
+    serious='#c4562a',
+    critical='#c22a2a',
+    ink='#0b0b0b',
+    ink_soft='#52514e',
+    ink_muted='#898781',
+    rule='#c3c2b7',
+)
+
+_RAMP = ('#cde2fb', '#9ec5f4', '#6da7ec', '#3987e5', '#256abf', '#184f95')
+
+# The active palette. The app swaps it when you change theme; --once uses the
+# dark one, which is right for the vast majority of terminals.
+_active = DARK
+
+
+def active() -> Palette:
+    return _active
+
+
+def use(palette: Palette) -> None:
+    global _active
+    _active = palette
+
+
+def for_theme(dark: bool) -> Palette:
+    """Fallback when there is no theme to measure."""
+    return DARK if dark else LIGHT
+
+
+def default() -> Palette:
+    """Palette for static mode: the terminal background can't be queried, so
+    textual-dark's is assumed."""
+    try:
+        from textual.theme import BUILTIN_THEMES
+        return from_theme(BUILTIN_THEMES['textual-dark'])
+    except Exception:
+        return DARK
+
+
+# Shape on top of color. Eight hues can't all be told apart: measured in
+# pairs, the worst drops to ΔE 1.6 under deuteranopia, and on the GPU dots any
+# two users can end up side by side. Shape is the secondary encoding the guide
+# calls for in these cases.
+#
+# The pairing isn't arbitrary: of all the ways to pair eight slots into four
+# groups, it's the one that maximizes the minimum separation between colors
+# sharing a shape. That way the worst pair with the same shape rises to
+# ΔE 11.5, and the weak pairs always get different shapes.
+SHAPES = ("●", "●", "◆", "■", "■", "◆", "▲", "▲")
+SHAPE_OTHER = "●"
+SHAPE_FREE = "○"
+
+
+class Occupants:
+    """Assigns identity colors to users.
+
+    Fixed, stable order: whoever holds the most GPUs gets the lowest slot, and
+    you always get the first. Beyond the first eight you fall into "others":
+    never generate a ninth hue.
+    """
+
+    LIMIT = 8
+
+    def __init__(self, counts: dict, me: str | None = None) -> None:
+        ranked = [u for u, _ in sorted(counts.items(), key=lambda kv: (-kv[1], kv[0]))]
+        if me in ranked:
+            ranked.remove(me)
+            ranked.insert(0, me)
+        self._slots = {user: i for i, user in enumerate(ranked[:self.LIMIT])}
+        self.me = me
+        self.overflow = ranked[self.LIMIT:]
+
+    def style(self, user: str) -> str:
+        pal = active()
+        return pal.slot(self._slots[user]) if user in self._slots else pal.other
+
+    def shape(self, user: str) -> str:
+        """The user's glyph: the encoding that holds up even without color."""
+        idx = self._slots.get(user)
+        return SHAPES[idx] if idx is not None else SHAPE_OTHER
+
+    def named(self) -> list:
+        """(user, color, shape) in slot order, for the legend."""
+        return [(u, self.style(u), self.shape(u))
+                for u, _ in sorted(self._slots.items(), key=lambda kv: kv[1])]
+
+
+# --------------------------------------------------------------------------
+# adapting to the actual background
+# --------------------------------------------------------------------------
+# A fixed palette isn't enough: Textual themes have very different
+# backgrounds, and "dark" doesn't say how dark. Nord's panel is #434C5E,
+# textual-dark's is #242F38: on one the dark ramp vanishes, on the other it
+# doesn't. So the data hues stay the validated ones, but the surroundings (ink,
+# rules) and the ramp direction are derived from the actual background.
+
+def _rgb(value: str) -> tuple:
+    value = str(value).lstrip('#')[:6]
+    return tuple(int(value[i:i + 2], 16) for i in (0, 2, 4))
+
+
+def _hex(rgb: tuple) -> str:
+    return '#%02x%02x%02x' % tuple(max(0, min(255, int(round(c)))) for c in rgb)
+
+
+def _luminance(value: str) -> float:
+    """WCAG relative luminance."""
+    def channel(c):
+        c = c / 255
+        return c / 12.92 if c <= 0.04045 else ((c + 0.055) / 1.055) ** 2.4
+    r, g, b = (channel(c) for c in _rgb(value))
+    return 0.2126 * r + 0.7152 * g + 0.0722 * b
+
+
+def contrast(a: str, b: str) -> float:
+    la, lb = _luminance(a), _luminance(b)
+    hi, lo = max(la, lb), min(la, lb)
+    return (hi + 0.05) / (lo + 0.05)
+
+
+def blend(a: str, b: str, t: float) -> str:
+    """t=1 all a, t=0 all b."""
+    ra, rb = _rgb(a), _rgb(b)
+    return _hex(tuple(x * t + y * (1 - t) for x, y in zip(ra, rb)))
+
+
+def toward(ink: str, surface: str, target: float) -> str:
+    """The color closest to the background that still reaches `target` contrast.
+
+    Blending by a fixed percentage doesn't work: some themes (solarized)
+    already have low-contrast text, and dimming it further makes it unreadable.
+    Instead this searches for the right point, and if the theme's text doesn't
+    reach the target it is used as is.
+    """
+    if contrast(ink, surface) <= target:
+        return ink
+    lo, hi = 0.0, 1.0
+    for _ in range(16):
+        mid = (lo + hi) / 2
+        if contrast(blend(ink, surface, mid), surface) >= target:
+            hi = mid
+        else:
+            lo = mid
+    return blend(ink, surface, hi)
+
+
+def lift(color: str, surface: str, target: float = 3.0) -> str:
+    """Brings `color` to `target` contrast against the background, keeping its hue.
+
+    Status colors are reserved and recognized by hue: green free, red broken.
+    On a mid-lightness background like nord's (#434C5E) both drown, but
+    changing their hue would turn them into something else. So only lightness
+    moves, lighter on a dark background and darker on a light one, until the
+    threshold is crossed.
+    """
+    if contrast(color, surface) >= target:
+        return color
+    r, g, b = (c / 255 for c in _rgb(color))
+    h, l, sat = colorsys.rgb_to_hls(r, g, b)
+    up = _luminance(surface) < 0.4          # dark background -> lighten
+    lo, hi = (l, 1.0) if up else (0.0, l)
+
+    best = color
+    for _ in range(20):
+        mid = (lo + hi) / 2
+        cand = _hex(tuple(c * 255 for c in colorsys.hls_to_rgb(h, mid, sat)))
+        if contrast(cand, surface) >= target:
+            best = cand
+            if up:
+                hi = mid
+            else:
+                lo = mid
+        else:
+            if up:
+                lo = mid
+            else:
+                hi = mid
+    return best
+
+
+def _ramp_for(surface: str) -> tuple:
+    """Ramp ordered so that 'fuller' is always 'more visible'.
+
+    The guide's rule is a single hue, monotonic in lightness; the direction,
+    though, depends on the background. On a dark background a ramp that ends
+    dark makes a full meter look duller than a half-empty one: that's a
+    misreading of the data, not an aesthetic quirk. Steps too close to the
+    background are dropped and the rest sorted by increasing contrast.
+    """
+    steps = [c for c in _RAMP if contrast(c, surface) >= 1.35]
+    if len(steps) < 4:
+        steps = sorted(_RAMP, key=lambda c: contrast(c, surface), reverse=True)[:4]
+    return tuple(sorted(steps, key=lambda c: contrast(c, surface)))
+
+
+def from_theme(theme) -> Palette:
+    """Builds the palette for a Textual theme by measuring its background."""
+    try:
+        resolved = theme.to_color_system().generate()
+    except Exception:
+        resolved = {}
+    surface = str(resolved.get('panel') or resolved.get('surface')
+                  or getattr(theme, 'panel', None) or '#242F38')
+    ink = str(resolved.get('foreground') or getattr(theme, 'foreground', None) or '#e0e0e0')
+    # ansi-* themes leave colors to the terminal and expose no hex values:
+    # there's no background to measure there, so fall back to the theme's label
+    measurable = surface.startswith('#') and ink.startswith('#')
+    if not measurable:
+        dark_theme = bool(getattr(theme, 'dark', True))
+        surface = '#242F38' if dark_theme else '#D0D0D0'
+        ink = '#e0e0e0' if dark_theme else '#0b0b0b'
+
+    # where there is a real background, it decides, not the "dark" label
+    dark = _luminance(surface) < 0.25
+    base = DARK if dark else LIGHT
+
+    return Palette(
+        name=f"{getattr(theme, 'name', '?')} ({'dark' if dark else 'light'})",
+        categorical=base.categorical,
+        other=toward(ink, surface, 2.5),
+        sequential=_ramp_for(surface),
+        good=lift(base.good, surface),
+        warning=lift(base.warning, surface),
+        serious=lift(base.serious, surface),
+        critical=lift(base.critical, surface),
+        ink=ink,
+        # targets: readable secondary text, distinguishable labels,
+        # rules that are visible but don't grab the eye
+        ink_soft=toward(ink, surface, 4.5),
+        ink_muted=toward(ink, surface, 3.0),
+        rule=toward(ink, surface, 1.6),
+    )