-
Notifications
You must be signed in to change notification settings - Fork 0
Expand file tree
/
Copy paththrottle.py
More file actions
121 lines (102 loc) · 3.98 KB
/
Copy paththrottle.py
File metadata and controls
121 lines (102 loc) · 3.98 KB
1
2
3
4
5
6
7
8
9
10
11
12
13
14
15
16
17
18
19
20
21
22
23
24
25
26
27
28
29
30
31
32
33
34
35
36
37
38
39
40
41
42
43
44
45
46
47
48
49
50
51
52
53
54
55
56
57
58
59
60
61
62
63
64
65
66
67
68
69
70
71
72
73
74
75
76
77
78
79
80
81
82
83
84
85
86
87
88
89
90
91
92
93
94
95
96
97
98
99
100
101
102
103
104
105
106
107
108
109
110
111
112
113
114
115
116
117
118
119
120
121
import os
import subprocess
import sys
from dataclasses import dataclass
import simple_parsing
from sarray.utils import console, err_console
@dataclass
class ThrottleConfig:
jobid: int = simple_parsing.field(positional=True, metavar="JOBID")
max_tasks: int = simple_parsing.field(
alias=["-n", "--max", "--max-tasks"], metavar="N"
)
kill: bool = simple_parsing.field(default=False, alias=["-k", "--kill"])
requeue: bool = simple_parsing.field(default=False, alias=["-r", "--requeue"])
def scontrol(*args: str) -> subprocess.CompletedProcess:
return subprocess.run(["scontrol", *args], capture_output=True, text=True)
def parse_scontrol_fields(line: str) -> dict[str, str]:
"""Parse a scontrol -o line into a key→value dict."""
fields = {}
for token in line.strip().split():
if "=" in token:
k, _, v = token.partition("=")
fields[k] = v
return fields
def cmd_throttle(config: ThrottleConfig):
# 1. Fetch job info (one record per line with -o)
result = scontrol("show", "-o", "job", str(config.jobid))
if result.returncode != 0:
err_console.print(f"[bold red]Error:[/] job {config.jobid} not found.")
sys.exit(1)
records = [
parse_scontrol_fields(line)
for line in result.stdout.splitlines()
if line.strip()
]
if not records:
err_console.print(f"[bold red]Error:[/] job {config.jobid} not found.")
sys.exit(1)
first = records[0]
# Check ownership
uid_field = first.get("UserId", "") # format: "username(uid)"
if f"({os.getuid()})" not in uid_field:
err_console.print(
f"[bold red]Error:[/] job {config.jobid} does not belong to you."
)
sys.exit(1)
# Check it's a job array
if "ArrayJobId" not in first:
err_console.print(f"[bold red]Error:[/] job {config.jobid} is not a job array.")
sys.exit(1)
# 2. Update throttle
result = scontrol(
"update", f"JobId={config.jobid}", f"ArrayTaskThrottle={config.max_tasks}"
)
if result.returncode != 0:
err_console.print(f"[bold red]scontrol error:[/]\n{result.stderr.strip()}")
sys.exit(result.returncode)
console.print(
f"[green]Throttle updated to [bold]{config.max_tasks}[/] "
f"for job [bold]{config.jobid}[/].[/]"
)
# 3. Act on excess running tasks if requested
if config.kill and config.requeue:
err_console.print(
"[bold red]Error:[/] --kill and --requeue are mutually exclusive."
)
sys.exit(1)
if not config.kill and not config.requeue:
return
running = sorted(
[r for r in records if r.get("JobState") == "RUNNING"],
key=lambda r: r.get("StartTime", ""),
reverse=False,
)
excess = running[config.max_tasks :]
if not excess:
console.print("[dim]No excess running tasks to act on.[/]")
return
if config.requeue:
console.print(f"[yellow]Requeueing {len(excess)} excess running task(s)...[/]")
for r in excess:
jid = r["JobId"]
display_id = f"{r['ArrayJobId']}_{r['ArrayTaskId']}"
res = scontrol("requeue", jid)
if res.returncode != 0:
err_console.print(
f"[bold red]Failed to requeue {display_id}:[/] {res.stderr.strip()}"
)
else:
console.print(f" [dim]requeued {display_id}[/]")
if config.kill:
console.print(f"[yellow]Cancelling {len(excess)} excess running task(s)...[/]")
for r in excess:
jid = r["JobId"]
display_id = f"{r['ArrayJobId']}_{r['ArrayTaskId']}"
res = subprocess.run(["scancel", jid], capture_output=True, text=True)
if res.returncode != 0:
err_console.print(
f"[bold red]Failed to cancel {display_id}:[/] {res.stderr.strip()}"
)
else:
console.print(f" [dim]cancelled {display_id}[/]")