Coverage for cli / commands / slurm_cmd.py: 100.00%
253 statements
« prev ^ index » next coverage.py v7.13.5, created at 2026-10-28 16:19 +0000
1"""GCO Slurm command group — the Slinky Slurm cluster through the GCO API.
3Provides the ``gco slurm`` sub-commands:
5* ``status`` / ``partitions`` / ``nodes`` — the cluster as slurmctld sees it
6 (``sinfo``), per region.
7* ``queue`` / ``job`` / ``submit`` / ``cancel`` — the live queue, one job, an
8 ``sbatch`` submission from a local script (its ``#SBATCH`` directives are
9 honoured; flags override them) and ``scancel``.
10* ``accounting`` — slurmdbd job history (``sacct``) over a time window.
11* ``shell`` / ``output`` — an interactive shell on the login pod, or a job's
12 output file read from it, over the PRIVATE EKS API endpoint
13 (``--via-ssm`` tunnels through an SSM-managed instance, like
14 ``gco monitoring open``).
16Every API command is pinned to one region (``--region``, else the root
17``--region`` / configured default) because each region runs its own Slurm
18cluster. See docs/SLURM_OPERATOR.md and docs/API.md#slurm.
19"""
21from __future__ import annotations
23import sys
24from pathlib import Path
25from typing import Any
27import click
29from .. import cluster_ui
30from ..config import GCOConfig
31from ..output import confirm, get_output_formatter
33pass_config = click.make_pass_decorator(GCOConfig, ensure=True)
35_REGION_HELP = "Region whose Slurm cluster to use (default: the root --region / configured default)"
38def _slurm_region(config: Any, region: str | None) -> str:
39 """The region a Slurm request is pinned to; Slurm clusters are per region."""
40 chosen = region or config.default_region
41 if not isinstance(chosen, str) or not chosen.strip():
42 raise click.UsageError(
43 "Slurm commands need a region: pass --region or set one with `gco -r REGION`"
44 )
45 return chosen.strip()
48def _call(config: Any, method: str, path: str, region: str, **kwargs: Any) -> dict[str, Any]:
49 from ..aws_client import get_aws_client
51 result: dict[str, Any] = get_aws_client(config).call_api(
52 method=method, path=path, region=region, **kwargs
53 )
54 return result
57def _fail(formatter: Any, verb: str, exc: Exception) -> None:
58 formatter.print_error(f"Failed to {verb}: {exc}")
59 sys.exit(1)
62@click.group()
63@pass_config
64def slurm(config: Any) -> None:
65 """Slurm cluster (cdk.json helm.slurm): queue, submit, accounting, login shell."""
68# ---------------------------------------------------------------------------
69# Cluster views
70# ---------------------------------------------------------------------------
73@slurm.command("status")
74@click.option("--region", "-r", help=_REGION_HELP)
75@pass_config
76def slurm_status(config: Any, region: str | None) -> None:
77 """Controller reachability, partitions, node and job counts.
79 Examples:
80 gco slurm status -r us-east-1
81 gco -r us-east-1 slurm status --output json
82 """
83 from ..slurm import render_status
85 formatter = get_output_formatter(config)
86 target = _slurm_region(config, region)
87 try:
88 result = _call(config, "GET", "/api/v1/slurm/status", target)
89 except Exception as exc:
90 _fail(formatter, "read Slurm status", exc)
91 if config.output_format != "table":
92 formatter.print(result)
93 return
94 print(render_status(result, target))
97@slurm.command("partitions")
98@click.option("--region", "-r", help=_REGION_HELP)
99@pass_config
100def slurm_partitions(config: Any, region: str | None) -> None:
101 """Partitions with their state, time limit and node counts (sinfo).
103 Examples:
104 gco slurm partitions -r us-east-1
105 """
106 from ..slurm import render_partitions
108 formatter = get_output_formatter(config)
109 target = _slurm_region(config, region)
110 try:
111 result = _call(config, "GET", "/api/v1/slurm/partitions", target)
112 except Exception as exc:
113 _fail(formatter, "list Slurm partitions", exc)
114 if config.output_format != "table":
115 formatter.print(result)
116 return
117 print(render_partitions(result.get("partitions") or [], target))
120@slurm.command("nodes")
121@click.option("--region", "-r", help=_REGION_HELP)
122@click.option("--partition", "-p", help="Only nodes in this partition")
123@pass_config
124def slurm_nodes(config: Any, region: str | None, partition: str | None) -> None:
125 """Slurm nodes (the slurmd worker pods) with state, resources and gres.
127 Examples:
128 gco slurm nodes -r us-east-1
129 gco slurm nodes -r us-east-1 --partition gpu
130 """
131 from ..slurm import render_nodes
133 formatter = get_output_formatter(config)
134 target = _slurm_region(config, region)
135 try:
136 result = _call(
137 config, "GET", "/api/v1/slurm/nodes", target, params={"partition": partition}
138 )
139 except Exception as exc:
140 _fail(formatter, "list Slurm nodes", exc)
141 if config.output_format != "table":
142 formatter.print(result)
143 return
144 print(render_nodes(result.get("nodes") or [], target))
147# ---------------------------------------------------------------------------
148# Jobs
149# ---------------------------------------------------------------------------
152@slurm.command("queue")
153@click.option("--region", "-r", help=_REGION_HELP)
154@click.option("--state", "-s", help="Only jobs in this state (PENDING, RUNNING, ...)")
155@click.option("--partition", "-p", help="Only jobs in this partition")
156@click.option("--name", "-n", help="Only jobs with exactly this name")
157@click.option("--limit", "-l", default=200, type=click.IntRange(1, 5000), help="Maximum jobs")
158@pass_config
159def slurm_queue(
160 config: Any,
161 region: str | None,
162 state: str | None,
163 partition: str | None,
164 name: str | None,
165 limit: int,
166) -> None:
167 """The live queue as slurmctld sees it (squeue).
169 Examples:
170 gco slurm queue -r us-east-1
171 gco slurm queue -r us-east-1 --state PENDING --partition gpu
172 """
173 from ..slurm import render_queue
175 formatter = get_output_formatter(config)
176 target = _slurm_region(config, region)
177 try:
178 result = _call(
179 config,
180 "GET",
181 "/api/v1/slurm/jobs",
182 target,
183 params={
184 "state": state.upper() if state else None,
185 "partition": partition,
186 "name": name,
187 "limit": str(limit),
188 },
189 )
190 except Exception as exc:
191 _fail(formatter, "list Slurm jobs", exc)
192 if config.output_format != "table":
193 formatter.print(result)
194 return
195 print(render_queue(result.get("jobs") or [], result.get("total", 0), target))
198@slurm.command("job")
199@click.argument("job_id", type=int)
200@click.option("--region", "-r", help=_REGION_HELP)
201@pass_config
202def slurm_job(config: Any, job_id: int, region: str | None) -> None:
203 """One job as slurmctld sees it (scontrol show job).
205 Examples:
206 gco slurm job 42 -r us-east-1
207 """
208 from ..slurm import render_job
210 formatter = get_output_formatter(config)
211 target = _slurm_region(config, region)
212 try:
213 result = _call(config, "GET", f"/api/v1/slurm/jobs/{job_id}", target)
214 except Exception as exc:
215 _fail(formatter, f"read Slurm job {job_id}", exc)
216 if config.output_format != "table":
217 formatter.print(result)
218 return
219 print(render_job(result.get("job") or {}, target))
222@slurm.command("submit")
223@click.argument("script", type=click.Path(exists=True, dir_okay=False, path_type=Path))
224@click.option("--region", "-r", help=_REGION_HELP)
225@click.option(
226 "--name", "-J", help="Job name (default: the script's #SBATCH --job-name, else gco-job)"
227)
228@click.option("--partition", "-p", help="Partition (default: the cluster default partition)")
229@click.option("--nodes", "-N", type=click.IntRange(1, 1024), help="Node count")
230@click.option("--tasks", "-n", type=click.IntRange(1, 65536), help="Task count (--ntasks)")
231@click.option("--cpus-per-task", "-c", type=click.IntRange(1, 1024), help="CPUs per task")
232@click.option("--gpus-per-node", type=click.IntRange(0, 64), help="GPUs per node")
233@click.option("--mem", "memory", help="Memory per node: MiB or K/M/G/T suffix (e.g. 64G)")
234@click.option("--time", "-t", "time_limit", help="Wall time: minutes, HH:MM:SS or D-HH:MM:SS")
235@click.option("--chdir", "-D", "working_directory", help="Working directory on the worker")
236@click.option(
237 "--env",
238 "-e",
239 "environment",
240 multiple=True,
241 metavar="NAME=VALUE",
242 help="Extra environment (repeatable)",
243)
244@click.option("--comment", help="Free-text comment stored with the job")
245@click.option("--exclusive", is_flag=True, default=False, help="Whole-node allocation")
246@pass_config
247def slurm_submit(
248 config: Any,
249 script: Path,
250 region: str | None,
251 name: str | None,
252 partition: str | None,
253 nodes: int | None,
254 tasks: int | None,
255 cpus_per_task: int | None,
256 gpus_per_node: int | None,
257 memory: str | None,
258 time_limit: str | None,
259 working_directory: str | None,
260 environment: tuple[str, ...],
261 comment: str | None,
262 exclusive: bool,
263) -> None:
264 """Submit a batch script (sbatch) through the GCO API.
266 The script's leading #SBATCH directives (--job-name, --partition, --nodes,
267 --ntasks, --cpus-per-task, --gpus-per-node, --mem, --time, --chdir,
268 --comment, --exclusive) are applied first; flags override them. The job
269 runs as Slurm's own user with your caller identity in its comment.
271 Examples:
272 gco slurm submit hello.sbatch -r us-east-1
273 gco slurm submit train.sbatch -r us-east-1 -p gpu -N 2 --gpus-per-node 8 -t 04:00:00
274 """
275 from ..slurm import SlurmCliError, memory_mb, read_batch_script, submission_body
277 formatter = get_output_formatter(config)
278 target = _slurm_region(config, region)
279 try:
280 body = submission_body(
281 read_batch_script(script),
282 flags={
283 "name": name,
284 "partition": partition,
285 "nodes": nodes,
286 "tasks": tasks,
287 "cpus_per_task": cpus_per_task,
288 "gpus_per_node": gpus_per_node,
289 "memory_mb": memory_mb(memory) if memory else None,
290 "time_limit": time_limit,
291 "working_directory": working_directory,
292 "comment": comment,
293 "exclusive": exclusive,
294 },
295 environment=environment,
296 )
297 except SlurmCliError as exc:
298 _fail(formatter, f"prepare {script.name}", exc)
299 try:
300 result = _call(config, "POST", "/api/v1/slurm/jobs", target, body=body)
301 except Exception as exc:
302 _fail(formatter, f"submit {script.name}", exc)
303 formatter.print_success(
304 f"Submitted job {result.get('job_id')} ({result.get('name')}) to "
305 f"{result.get('partition') or 'the default partition'} in {target}"
306 )
307 if result.get("message"):
308 formatter.print_info(str(result["message"]))
309 formatter.print(result)
312@slurm.command("cancel")
313@click.argument("job_id", type=int)
314@click.option("--region", "-r", help=_REGION_HELP)
315@click.option("--signal", "-s", help="Signal to send instead of the default (e.g. SIGUSR1)")
316@click.option("--yes", "-y", "assume_yes", is_flag=True, help="Skip the confirmation prompt")
317@pass_config
318def slurm_cancel(
319 config: Any, job_id: int, region: str | None, signal: str | None, assume_yes: bool
320) -> None:
321 """Cancel a job (scancel).
323 Examples:
324 gco slurm cancel 42 -r us-east-1
325 gco slurm cancel 42 -r us-east-1 --signal SIGUSR1 -y
326 """
327 formatter = get_output_formatter(config)
328 target = _slurm_region(config, region)
329 if not assume_yes:
330 confirm(f"Cancel Slurm job {job_id} in {target}?", abort=True)
331 try:
332 result = _call(
333 config,
334 "DELETE",
335 f"/api/v1/slurm/jobs/{job_id}",
336 target,
337 params={"signal": signal.upper() if signal else None},
338 )
339 except Exception as exc:
340 _fail(formatter, f"cancel Slurm job {job_id}", exc)
341 formatter.print_success(f"Cancelled job {job_id} in {target}")
342 formatter.print(result)
345@slurm.command("accounting")
346@click.option("--region", "-r", help=_REGION_HELP)
347@click.option(
348 "--since", help="Relative window (30m, 24h, 7d, 2w); default 1d unless --start is given"
349)
350@click.option("--start", help="Window start as a Slurm time (2026-10-01T00:00:00, now-2days)")
351@click.option("--end", help="Window end as a Slurm time (default: now)")
352@click.option("--user", "-u", help="Only this Slurm user")
353@click.option("--partition", "-p", help="Only this partition")
354@click.option("--state", "-s", help="Only these states (COMPLETED, FAILED,...; comma-separated)")
355@click.option("--name", "-n", "job_name", help="Only jobs with exactly this name")
356@click.option("--limit", "-l", default=200, type=click.IntRange(1, 5000), help="Maximum rows")
357@pass_config
358def slurm_accounting(
359 config: Any,
360 region: str | None,
361 since: str | None,
362 start: str | None,
363 end: str | None,
364 user: str | None,
365 partition: str | None,
366 state: str | None,
367 job_name: str | None,
368 limit: int,
369) -> None:
370 """Job history from slurmdbd (sacct).
372 Examples:
373 gco slurm accounting -r us-east-1
374 gco slurm accounting -r us-east-1 --since 7d --state COMPLETED,FAILED
375 gco slurm accounting -r us-east-1 --start 2026-10-01 --end 2026-10-08 --user slurm
376 """
377 from ..slurm import SlurmCliError, render_accounting, slurm_time_window
379 formatter = get_output_formatter(config)
380 target = _slurm_region(config, region)
381 try:
382 window = slurm_time_window(since or (None if start else "1d"), start, end)
383 except SlurmCliError as exc:
384 _fail(formatter, "read Slurm accounting", exc)
385 try:
386 result = _call(
387 config,
388 "GET",
389 "/api/v1/slurm/accounting/jobs",
390 target,
391 params={
392 **window,
393 "user": user,
394 "partition": partition,
395 "state": state.upper() if state else None,
396 "job_name": job_name,
397 "limit": str(limit),
398 },
399 )
400 except Exception as exc:
401 _fail(formatter, "read Slurm accounting", exc)
402 if config.output_format != "table":
403 formatter.print(result)
404 return
405 print(render_accounting(result.get("jobs") or [], result.get("total", 0), target))
408# ---------------------------------------------------------------------------
409# Login pod
410# ---------------------------------------------------------------------------
413@slurm.command("shell")
414@cluster_ui.tunnel_options
415@pass_config
416def slurm_shell(
417 config: Any,
418 region: str | None,
419 via_ssm: str | None,
420 bastion_ttl_minutes: int,
421 assume_yes: bool,
422) -> None:
423 """Open an interactive shell on the Slurm login pod (sbatch, squeue, sacct, srun).
425 Runs `kubectl exec -it` against the private EKS API endpoint; the pod
426 mounts /home (EFS) and, when enabled, /scratch (FSx for Lustre). Table
427 output only: an interactive terminal cannot be captured as JSON.
429 Examples:
430 gco slurm shell --via-ssm auto -y
431 gco slurm shell -r us-west-2 --via-ssm i-0123456789abcdef0
432 """
433 from ..cluster_tunnel import open_api_server_tunnel, resolve_region
434 from ..kubectl_helpers import update_kubeconfig
435 from ..slurm import LOGIN_DEPLOYMENT, LOGIN_NAMESPACE, login_exec_command, run_foreground
437 formatter = get_output_formatter(config)
438 if config.output_format != "table":
439 formatter.print_error("gco slurm shell is interactive; run it without --output json/yaml")
440 sys.exit(1)
441 target_region = resolve_region(config, region)
442 cluster = f"{config.project_name}-{target_region}"
443 try:
444 update_kubeconfig(cluster, target_region)
445 with open_api_server_tunnel(
446 formatter,
447 cluster=cluster,
448 region=target_region,
449 via_ssm=via_ssm,
450 bastion_ttl_minutes=bastion_ttl_minutes,
451 assume_yes=assume_yes,
452 ) as session:
453 cmd = login_exec_command(
454 ["bash", "-l"],
455 interactive=True,
456 server=session.server,
457 tls_server_name=session.tls_server_name,
458 )
459 formatter.print_info(
460 f"Opening a shell on {LOGIN_NAMESPACE}/{LOGIN_DEPLOYMENT} in {target_region} "
461 "(exit the shell to close the tunnel)..."
462 )
463 code = run_foreground(cmd)
464 except (RuntimeError, ValueError) as exc:
465 formatter.print_error(str(exc))
466 sys.exit(1)
467 if code != 0:
468 sys.exit(code)
471@slurm.command("output")
472@click.argument("job_id", type=int)
473@cluster_ui.tunnel_options
474@click.option("--stderr", "use_stderr", is_flag=True, help="Read the job's stderr file instead")
475@click.option("--follow", "-f", is_flag=True, help="Follow the file (tail -f) until Ctrl-C")
476@click.option("--tail", "tail_lines", type=click.IntRange(1, 100000), help="Only the last N lines")
477@pass_config
478def slurm_output(
479 config: Any,
480 job_id: int,
481 region: str | None,
482 via_ssm: str | None,
483 bastion_ttl_minutes: int,
484 assume_yes: bool,
485 use_stderr: bool,
486 follow: bool,
487 tail_lines: int | None,
488) -> None:
489 """Print a job's output file from the login pod.
491 Looks the job up through the API for its stdout/stderr path (the
492 default `slurm-<jobid>.out` under the job's working directory), then
493 reads it over the private EKS endpoint with `kubectl exec`.
495 Examples:
496 gco slurm output 42 -r us-east-1 --via-ssm auto -y
497 gco slurm output 42 -r us-east-1 --via-ssm i-0123456789abcdef0 --follow
498 """
499 from ..cluster_tunnel import open_api_server_tunnel, resolve_region
500 from ..kubectl_helpers import update_kubeconfig
501 from ..slurm import (
502 SlurmCliError,
503 login_exec_command,
504 run_capture,
505 run_foreground,
506 validate_remote_path,
507 )
509 formatter = get_output_formatter(config)
510 if follow and config.output_format != "table":
511 formatter.print_error("--follow streams a terminal; run it without --output json/yaml")
512 sys.exit(1)
513 target_region = resolve_region(config, region)
514 try:
515 job = _call(config, "GET", f"/api/v1/slurm/jobs/{job_id}", target_region).get("job") or {}
516 except Exception as exc:
517 _fail(formatter, f"read Slurm job {job_id}", exc)
518 key = "standard_error" if use_stderr else "standard_output"
519 try:
520 path = validate_remote_path(str(job.get(key) or ""))
521 except SlurmCliError as exc:
522 _fail(formatter, f"locate the {key.replace('_', ' ')} of job {job_id}", exc)
523 reader = (
524 ["tail", "-f", path]
525 if follow
526 else (["tail", "-n", str(tail_lines), path] if tail_lines else ["cat", path])
527 )
528 cluster = f"{config.project_name}-{target_region}"
529 try:
530 update_kubeconfig(cluster, target_region)
531 with open_api_server_tunnel(
532 formatter,
533 cluster=cluster,
534 region=target_region,
535 via_ssm=via_ssm,
536 bastion_ttl_minutes=bastion_ttl_minutes,
537 assume_yes=assume_yes,
538 ) as session:
539 cmd = login_exec_command(
540 reader, server=session.server, tls_server_name=session.tls_server_name
541 )
542 if follow:
543 # Streams until Ctrl-C; the tunnel closes on the way out and
544 # tail's exit status is the command's (SystemExit passes the
545 # except clause below untouched).
546 formatter.print_info(f"Following {path} (Ctrl-C to stop)...")
547 code = run_foreground(cmd)
548 if code != 0:
549 sys.exit(code)
550 return
551 text = run_capture(cmd)
552 except (RuntimeError, ValueError) as exc:
553 formatter.print_error(str(exc))
554 sys.exit(1)
555 if config.output_format != "table":
556 formatter.print({"job_id": job_id, "region": target_region, "path": path, "content": text})
557 return
558 sys.stdout.write(text)
559 if text and not text.endswith("\n"):
560 sys.stdout.write("\n")