Coverage for cli / slurm.py: 100.00%
254 statements
« prev ^ index » next coverage.py v7.13.5, created at 2026-10-28 16:19 +0000
1"""Business logic behind ``gco slurm`` (cli/commands/slurm_cmd.py).
3The command module owns the Click surface; this module owns everything that
4can be tested without a terminal: reading a batch script and its ``#SBATCH``
5directives into the API's submission body, Slurm's relative time windows,
6the human renderings of the API's answers, and the ``kubectl exec`` argv for
7the login pod (``shell`` / ``output``) that the cluster tunnel reaches.
9The API half lives in gco/services/api_routes/slurm.py; every request here
10is pinned to one region because each region runs its own Slurm cluster.
11"""
13from __future__ import annotations
15import re
16import shlex
17import subprocess # nosec B404 # kubectl runs as a validated argv list, never through a shell
18from collections.abc import Mapping, Sequence
19from datetime import UTC, datetime
20from pathlib import Path
21from typing import Any
23from gco.slurm_config import SLURM_LOGIN_NAME, SLURM_NAMESPACE
25from . import cluster_ui
27#: The Slinky LoginSet GCO deploys and its namespace; ``shell`` and ``output``
28#: exec into its Deployment (the operator names it after the LoginSet) because
29#: it mounts the same /home (EFS) and /scratch (FSx) the workers do and
30#: carries every Slurm client. ``login`` is the operator's container name.
31LOGIN_NAMESPACE = SLURM_NAMESPACE
32LOGIN_DEPLOYMENT = SLURM_LOGIN_NAME
33LOGIN_CONTAINER = "login"
35#: ``--since`` suffixes to Slurm's relative time units (``now-7days``).
36_DURATION_UNITS = {"s": "seconds", "m": "minutes", "h": "hours", "d": "days", "w": "weeks"}
37_DURATION_RE = re.compile(r"^(?P<count>\d{1,6})(?P<unit>[smhdw])$")
38#: Absolute Slurm times (``parse_time``): YYYY-MM-DD[THH:MM[:SS]], HH:MM[:SS], MM/DD[/YY], epoch.
39_SLURM_TIME_RE = re.compile(
40 r"^(?:\d{4}-\d{2}-\d{2}(?:[T ]\d{2}:\d{2}(?::\d{2})?)?|\d{2}:\d{2}(?::\d{2})?"
41 r"|\d{2}/\d{2}(?:/\d{2,4})?|\d{9,10}|now(?:[+-]\d+(?:seconds|minutes|hours|days|weeks))?)$"
42)
43_SBATCH_RE = re.compile(r"^#SBATCH\s+(?P<body>.+?)\s*$")
44_SAFE_PATH_RE = re.compile(r"^/[A-Za-z0-9._/@+%:=,-]{0,1023}$")
45MAX_SCRIPT_BYTES = 64 * 1024
48class SlurmCliError(ValueError):
49 """A request that cannot be shaped: bad script, bad window, bad path."""
52# ---------------------------------------------------------------------------
53# Time windows
54# ---------------------------------------------------------------------------
57def slurm_time_window(since: str | None, start: str | None, end: str | None) -> dict[str, str]:
58 """Translate ``--since 7d`` / ``--start`` / ``--end`` into slurmdb query parameters.
60 ``since`` is a relative window (``30m``, ``24h``, ``7d``, ``2w``) rendered
61 as Slurm's own ``now-7days`` form so slurmdbd does the arithmetic in the
62 cluster's clock; ``start``/``end`` are passed through when they look like
63 a Slurm time. ``since`` and ``start`` are alternatives.
64 """
65 if since and start:
66 raise SlurmCliError("pass either --since or --start, not both")
67 params: dict[str, str] = {}
68 if since:
69 match = _DURATION_RE.match(since.strip().lower())
70 if match is None:
71 raise SlurmCliError(
72 f"invalid --since {since!r}: use a count and unit such as 30m, 24h, 7d or 2w"
73 )
74 count, unit = int(match["count"]), _DURATION_UNITS[match["unit"]]
75 if count == 0:
76 raise SlurmCliError("--since must be a positive window")
77 params["start"] = f"now-{count}{unit}"
78 elif start:
79 params["start"] = _slurm_time(start, "--start")
80 if end:
81 params["end"] = _slurm_time(end, "--end")
82 return params
85def _slurm_time(value: str, flag: str) -> str:
86 text = value.strip()
87 if not _SLURM_TIME_RE.match(text):
88 raise SlurmCliError(
89 f"invalid {flag} {value!r}: use YYYY-MM-DD[THH:MM[:SS]], HH:MM, an epoch, "
90 "or Slurm's relative form such as now-2days"
91 )
92 return text
95# ---------------------------------------------------------------------------
96# Batch scripts
97# ---------------------------------------------------------------------------
100def read_batch_script(path: Path) -> str:
101 """Read a batch script, refusing empty, oversized or shebang-less files."""
102 try:
103 raw = path.read_bytes()
104 except OSError as exc:
105 raise SlurmCliError(f"cannot read {path}: {exc.strerror or exc}") from exc
106 if not raw.strip():
107 raise SlurmCliError(f"{path} is empty")
108 if len(raw) > MAX_SCRIPT_BYTES:
109 raise SlurmCliError(
110 f"{path} is {len(raw)} bytes; batch scripts are limited to {MAX_SCRIPT_BYTES} "
111 "(put the payload in a container image or on /home or /scratch)"
112 )
113 try:
114 text = raw.decode("utf-8")
115 except UnicodeDecodeError as exc:
116 raise SlurmCliError(f"{path} is not UTF-8 text") from exc
117 if not text.lstrip().startswith("#!"):
118 raise SlurmCliError(f"{path} must start with an interpreter line such as #!/bin/bash")
119 return text
122#: ``#SBATCH`` options the CLI understands, mapped to the API body's fields.
123#: slurmrestd does not read ``#SBATCH`` lines (sbatch the client does), so the
124#: CLI applies them as defaults the explicit flags override; anything else in
125#: the script is left for Slurm to ignore, as sbatch would with unknown options.
126_SBATCH_OPTIONS: dict[str, str] = {
127 "-J": "name",
128 "--job-name": "name",
129 "-p": "partition",
130 "--partition": "partition",
131 "-N": "nodes",
132 "--nodes": "nodes",
133 "-n": "tasks",
134 "--ntasks": "tasks",
135 "-c": "cpus_per_task",
136 "--cpus-per-task": "cpus_per_task",
137 "--gpus-per-node": "gpus_per_node",
138 "--mem": "memory_mb",
139 "-t": "time_limit",
140 "--time": "time_limit",
141 "-D": "working_directory",
142 "--chdir": "working_directory",
143 "--comment": "comment",
144 "--exclusive": "exclusive",
145}
146_INT_FIELDS = {"nodes", "tasks", "cpus_per_task", "gpus_per_node"}
149def sbatch_directives(script: str) -> dict[str, Any]:
150 """Collect the ``#SBATCH`` directives the API body can carry.
152 Reads the leading comment block the way sbatch does (directives stop at
153 the first non-comment, non-blank line). Unknown options are ignored.
154 ``--mem`` accepts Slurm's K/M/G/T suffixes and lands in MiB; ``--gpus-per-node``
155 accepts ``N`` or ``type:N``; ``--exclusive`` is a flag.
156 """
157 found: dict[str, Any] = {}
158 for line in script.splitlines()[1:]:
159 stripped = line.strip()
160 if not stripped:
161 continue
162 if not stripped.startswith("#"):
163 break
164 match = _SBATCH_RE.match(stripped)
165 if match is None:
166 continue
167 try:
168 tokens = shlex.split(match["body"], comments=True)
169 except ValueError:
170 continue
171 index = 0
172 while index < len(tokens):
173 token = tokens[index]
174 option, value = token, None
175 if token.startswith("--") and "=" in token:
176 option, value = token.split("=", 1)
177 elif len(token) > 2 and token[0] == "-" and token[1] != "-":
178 option, value = token[:2], token[2:]
179 field = _SBATCH_OPTIONS.get(option)
180 if field is None:
181 index += 1
182 continue
183 if field == "exclusive":
184 found[field] = True
185 index += 1
186 continue
187 if value is None:
188 index += 1
189 if index >= len(tokens):
190 break
191 value = tokens[index]
192 index += 1
193 found[field] = _directive_value(field, value)
194 return found
197def _directive_value(field: str, value: str) -> Any:
198 if field in _INT_FIELDS:
199 text = value.split(":")[-1] if field == "gpus_per_node" else value
200 try:
201 return int(text)
202 except ValueError as exc:
203 raise SlurmCliError(
204 f"#SBATCH {field.replace('_', '-')} {value!r} is not a number"
205 ) from exc
206 if field == "memory_mb":
207 return memory_mb(value)
208 return value
211def memory_mb(value: str) -> int:
212 """Slurm's ``--mem`` (``4000``, ``4G``, ``512M``, ``1T``) in MiB."""
213 match = re.fullmatch(r"(\d+)([KkMmGgTt]?)", value.strip())
214 if match is None:
215 raise SlurmCliError(f"invalid memory {value!r}: use MiB, or a K/M/G/T suffix such as 64G")
216 amount, unit = int(match[1]), match[2].upper() or "M"
217 factor = {"K": 1 / 1024, "M": 1, "G": 1024, "T": 1024 * 1024}[unit]
218 return max(int(amount * factor), 1)
221def submission_body(
222 script: str,
223 *,
224 flags: Mapping[str, Any],
225 environment: Sequence[str] = (),
226) -> dict[str, Any]:
227 """Compose the ``POST /api/v1/slurm/jobs`` body: directives, then flags on top."""
228 body: dict[str, Any] = {"script": script}
229 body.update(sbatch_directives(script))
230 for key, value in flags.items():
231 if value is None or (key == "exclusive" and value is False):
232 continue
233 body[key] = value
234 env: dict[str, str] = {}
235 for item in environment:
236 name, separator, value = item.partition("=")
237 if not separator or not re.fullmatch(r"[A-Za-z_][A-Za-z0-9_]*", name):
238 raise SlurmCliError(f"invalid --env {item!r}: use NAME=value")
239 env[name] = value
240 if env:
241 body["environment"] = env
242 return body
245# ---------------------------------------------------------------------------
246# Human renderings
247# ---------------------------------------------------------------------------
250def _when(epoch: Any) -> str:
251 if not isinstance(epoch, int) or epoch <= 0:
252 return "-"
253 return datetime.fromtimestamp(epoch, tz=UTC).strftime("%Y-%m-%d %H:%M:%S")
256def _join(states: Any) -> str:
257 if isinstance(states, list):
258 return ",".join(str(s) for s in states) or "-"
259 return str(states or "-")
262def render_status(status: Mapping[str, Any], region: str) -> str:
263 lines = [f"\n Slurm — {region} (slurmrestd {status.get('api_version', '?')})", " " + "-" * 70]
264 for controller in status.get("controller") or []:
265 state = "UP" if controller.get("responding") else "DOWN"
266 latency = controller.get("latency_us")
267 suffix = f" ({latency} µs)" if isinstance(latency, int) else ""
268 lines.append(
269 f" Controller {controller.get('hostname')}: {state}{suffix} [{controller.get('mode')}]"
270 )
271 nodes = status.get("nodes") or {}
272 jobs = status.get("jobs") or {}
273 scaled_in = nodes.get("scaled_in") or 0
274 node_line = f" Nodes: {nodes.get('total', 0)} {_counts(nodes.get('by_state'))}"
275 if scaled_in:
276 node_line += f" + {scaled_in} scaled in (records kept until the next scale-up)"
277 lines.append(node_line)
278 lines.append(f" Jobs: {jobs.get('total', 0)} {_counts(jobs.get('by_state'))}")
279 lines.append("")
280 lines.extend(_partition_table(status.get("partitions") or []))
281 return "\n".join(lines)
284def _counts(by_state: Any) -> str:
285 if not isinstance(by_state, dict) or not by_state:
286 return ""
287 return "(" + ", ".join(f"{state} {count}" for state, count in sorted(by_state.items())) + ")"
290def _partition_table(partitions: Sequence[Mapping[str, Any]]) -> list[str]:
291 lines = [
292 f" {'PARTITION':<14} {'STATE':<12} {'NODES':>5} {'CPUS':>5} {'MAX TIME':>10} NODELIST"
293 ]
294 lines.append(" " + "-" * 70)
295 for partition in partitions:
296 name = str(partition.get("name") or "-")
297 if partition.get("default"):
298 name += "*"
299 max_time = partition.get("max_time_minutes")
300 lines.append(
301 f" {name:<14} {_join(partition.get('state')):<12} "
302 f"{partition.get('total_nodes') if partition.get('total_nodes') is not None else '-':>5} "
303 f"{partition.get('total_cpus') if partition.get('total_cpus') is not None else '-':>5} "
304 f"{(str(max_time) + 'm') if isinstance(max_time, int) else 'infinite':>10} "
305 f"{partition.get('node_names') or '-'}"
306 )
307 if not partitions:
308 lines.append(" (no partitions)")
309 return lines
312def render_partitions(partitions: Sequence[Mapping[str, Any]], region: str) -> str:
313 return "\n".join([f"\n Slurm partitions — {region}", *_partition_table(partitions), ""])
316def render_nodes(nodes: Sequence[Mapping[str, Any]], region: str) -> str:
317 lines = [
318 f"\n Slurm nodes — {region} ({len(nodes)})",
319 f" {'NODE':<28} {'STATE':<18} {'CPUS':>9} {'MEMORY MB':>17} {'GRES':<20} PARTITIONS",
320 " " + "-" * 110,
321 ]
322 for node in nodes:
323 cpus = f"{node.get('cpus_allocated') or 0}/{node.get('cpus') or 0}"
324 memory = f"{node.get('memory_allocated_mb') or 0}/{node.get('real_memory_mb') or 0}"
325 gres = str(node.get("gres_used") or node.get("gres") or "-")
326 state = "SCALED-IN" if node.get("scaled_in") else _join(node.get("state"))
327 lines.append(
328 f" {node.get('name') or '-'!s:<28} {state:<18} {cpus:>9} "
329 f"{memory:>17} {gres[:20]:<20} {','.join(node.get('partitions') or []) or '-'}"
330 )
331 if node.get("reason") and not node.get("scaled_in"):
332 lines.append(f" reason: {node['reason']}")
333 if not nodes:
334 lines.append(" (no nodes — autoscaled partitions start at zero workers)")
335 lines.append("")
336 return "\n".join(lines)
339def render_queue(jobs: Sequence[Mapping[str, Any]], total: Any, region: str) -> str:
340 lines = [
341 f"\n Slurm queue — {region} ({len(jobs)} of {total} shown)",
342 f" {'JOBID':>8} {'NAME':<24} {'PARTITION':<10} {'STATE':<12} {'NODES':>5} {'SUBMITTED':<19} REASON/NODELIST",
343 " " + "-" * 110,
344 ]
345 for job in jobs:
346 tail = job.get("nodes") or job.get("state_reason") or "-"
347 lines.append(
348 f" {job.get('job_id') or '-'!s:>8} {str(job.get('name') or '-')[:24]:<24} "
349 f"{str(job.get('partition') or '-')[:10]:<10} {_join(job.get('state'))[:12]:<12} "
350 f"{job.get('node_count') if job.get('node_count') is not None else '-':>5} "
351 f"{_when(job.get('submit_time')):<19} {tail}"
352 )
353 if not jobs:
354 lines.append(" (queue is empty)")
355 lines.append("")
356 return "\n".join(lines)
359def _state_line(job: Mapping[str, Any]) -> str:
360 """``PENDING — Resources`` while the job waits; the bare state otherwise.
362 Slurm's state reason says why a job is still pending. Both daemons keep
363 reporting the last one afterwards (``ReqNodeNotAvail`` on a COMPLETED job,
364 ``None`` on a running one), where it is noise.
365 """
366 states = job.get("state")
367 state = _join(states)
368 reason = job.get("state_reason")
369 pending = "PENDING" in (states if isinstance(states, list) else [states])
370 if pending and reason and reason != "None":
371 return f"{state} — {reason}"
372 return state
375def render_job(job: Mapping[str, Any], region: str) -> str:
376 exit_code = job.get("exit_code") or {}
377 rows = [
378 ("Job", f"{job.get('job_id')} ({job.get('name')})"),
379 ("Region", region),
380 (
381 "Source",
382 "slurmdbd (finished; slurmctld has forgotten it)"
383 if job.get("source") == "slurmdbd"
384 else "slurmctld",
385 ),
386 ("State", _state_line(job)),
387 ("User / account", f"{job.get('user') or '-'} / {job.get('account') or '-'}"),
388 ("Partition", str(job.get("partition") or "-")),
389 (
390 "Nodes",
391 f"{job.get('node_count') if job.get('node_count') is not None else '-'} {job.get('nodes') or ''}".strip(),
392 ),
393 ("Requested", str(job.get("tres_requested") or "-")),
394 ("Allocated", str(job.get("tres_allocated") or "-")),
395 ("Submitted", _when(job.get("submit_time"))),
396 ("Started", _when(job.get("start_time"))),
397 ("Ended", _when(job.get("end_time"))),
398 (
399 "Time limit",
400 f"{job.get('time_limit_minutes')} min" if job.get("time_limit_minutes") else "infinite",
401 ),
402 (
403 "Exit",
404 f"{_join(exit_code.get('status'))} rc={exit_code.get('return_code')}"
405 if exit_code
406 else "-",
407 ),
408 ("Comment", str(job.get("comment") or "-")),
409 ("Working dir", str(job.get("working_directory") or "-")),
410 ("Stdout", str(job.get("standard_output") or "-")),
411 ("Stderr", str(job.get("standard_error") or "-")),
412 ]
413 width = max(len(label) for label, _value in rows)
414 return "\n".join(["", *(f" {label:<{width}} {value}" for label, value in rows), ""])
417def render_accounting(jobs: Sequence[Mapping[str, Any]], total: Any, region: str) -> str:
418 lines = [
419 f"\n Slurm accounting — {region} ({len(jobs)} of {total} shown)",
420 f" {'JOBID':>8} {'NAME':<22} {'USER':<10} {'PARTITION':<10} {'STATE':<12} {'ELAPSED':>9} {'EXIT':>6} {'STARTED':<19} COMMENT",
421 " " + "-" * 120,
422 ]
423 for job in jobs:
424 exit_code = job.get("exit_code") or {}
425 elapsed = job.get("elapsed_seconds")
426 lines.append(
427 f" {job.get('job_id') or '-'!s:>8} {str(job.get('name') or '-')[:22]:<22} "
428 f"{str(job.get('user') or '-')[:10]:<10} {str(job.get('partition') or '-')[:10]:<10} "
429 f"{_join(job.get('state'))[:12]:<12} {_elapsed(elapsed):>9} "
430 f"{exit_code.get('return_code') if exit_code.get('return_code') is not None else '-'!s:>6} "
431 f"{_when(job.get('start_time')):<19} {str(job.get('comment') or '')[:40]}"
432 )
433 if not jobs:
434 lines.append(" (no accounting records in this window)")
435 lines.append("")
436 return "\n".join(lines)
439def _elapsed(seconds: Any) -> str:
440 if not isinstance(seconds, int) or seconds < 0:
441 return "-"
442 hours, remainder = divmod(seconds, 3600)
443 minutes, secs = divmod(remainder, 60)
444 return f"{hours:02d}:{minutes:02d}:{secs:02d}"
447# ---------------------------------------------------------------------------
448# Login pod access (kubectl exec over the cluster tunnel)
449# ---------------------------------------------------------------------------
452def login_exec_command(
453 args: Sequence[str],
454 *,
455 interactive: bool = False,
456 server: str | None = None,
457 tls_server_name: str | None = None,
458) -> list[str]:
459 """``kubectl exec`` into the login Deployment with the tunnel's flags."""
460 cmd = ["kubectl", "exec"]
461 if interactive:
462 cmd.append("-it")
463 cmd += [
464 f"deploy/{LOGIN_DEPLOYMENT}",
465 "-n",
466 LOGIN_NAMESPACE,
467 "-c",
468 LOGIN_CONTAINER,
469 *cluster_ui.kubectl_server_flags(server, tls_server_name),
470 "--",
471 *args,
472 ]
473 return cmd
476def validate_remote_path(path: str) -> str:
477 """An absolute path the login pod may read (``cat`` / ``tail``), with no shell metacharacters."""
478 if not _SAFE_PATH_RE.match(path) or ".." in path.split("/"):
479 raise SlurmCliError(
480 f"refusing to read {path!r}: expected an absolute path on the login pod"
481 )
482 return path
485def run_foreground(cmd: Sequence[str]) -> int:
486 """Run kubectl with inherited stdio (an interactive shell, a tail -f); return its exit code."""
487 try:
488 completed = subprocess.run( # nosemgrep: dangerous-subprocess-use-audit - argv built by login_exec_command; list form, no shell=True
489 list(cmd), check=False
490 )
491 except FileNotFoundError as exc:
492 raise RuntimeError(
493 "kubectl not found. Install kubectl and ensure it's on your PATH."
494 ) from exc
495 return int(completed.returncode)
498def run_capture(cmd: Sequence[str], *, timeout: float = 120.0) -> str:
499 """Run kubectl and return its stdout; a failure carries kubectl's stderr."""
500 try:
501 completed = subprocess.run( # nosemgrep: dangerous-subprocess-use-audit - argv built by login_exec_command; list form, no shell=True
502 list(cmd), capture_output=True, text=True, check=False, timeout=timeout
503 )
504 except FileNotFoundError as exc:
505 raise RuntimeError(
506 "kubectl not found. Install kubectl and ensure it's on your PATH."
507 ) from exc
508 except subprocess.TimeoutExpired as exc:
509 raise RuntimeError(f"Timed out after {timeout:.0f}s running {cmd[0]} {cmd[1]}") from exc
510 if completed.returncode != 0:
511 raise RuntimeError(
512 f"kubectl exec failed ({completed.returncode}): {(completed.stderr or '').strip()}"
513 )
514 return completed.stdout