sqtop 0.1.0__py3-none-any.whl

sqtop/cli.py

@@ -0,0 +1,74 @@
+"""Entry point for `sqtop`: the dashboard, or a one-shot text report."""
+
+import argparse
+import sys
+
+from . import hosts, probes, profiles, report, slurm, snapshot
+
+
+def find_clusters():
+    """This cluster, if Slurm runs here, then each cluster with an ssh host.
+
+    The local cluster is detected from `scontrol show config`; a configured
+    remote with the same name is skipped, since it is already shown.
+    """
+    local = hosts.Local()
+    config = local.run(slurm.CONFIG_COMMAND)
+    clusters = []
+    local_name = None
+    if config is not None:
+        local_name = slurm.parse_config(config).get("ClusterName") or None
+        profile = profiles.load_profile(local_name or "")
+        clusters.append(snapshot.Cluster(local_name or "local", profile, local))
+    for profile in profiles.remote_profiles():
+        if profile.name != local_name:
+            remote = hosts.Remote(profile.ssh)
+            clusters.append(snapshot.Cluster(profile.name, profile, remote))
+    return clusters
+
+
+def main(argv=None):
+    # No options: clusters are detected and configured in the config, which
+    # Config edits in the dashboard, where the probe request and hourly email
+    # are set too. Parsing still rejects stray arguments and serves --help.
+    config = profiles.short_path(profiles.CONFIG_PATH)
+    argparse.ArgumentParser(
+        prog="sqtop",
+        description=(
+            "Partitions, account/QOS limits and estimated start times for the "
+            "Slurm cluster you are logged in to; clusters with an ssh host in "
+            f"{config} get a tab that connects when opened, Config edits that "
+            "config, and Email sends the dashboard hourly through its [smtp]. "
+            "Prints a text report of this cluster when output is not a terminal."
+        ),
+    ).parse_args(argv)
+
+    try:
+        clusters = find_clusters()
+        mail = profiles.load_mail()
+    except ValueError as error:
+        sys.exit(f"sqtop: {profiles.short_path(profiles.CONFIG_PATH)}: {error}")
+    if not clusters:
+        sys.exit(
+            "sqtop: Slurm is not available here ('scontrol' failed), and no "
+            f"cluster in {config} has an ssh host."
+        )
+    settings = probes.Settings()
+    if mail is not None:
+        settings = probes.Settings(mail_user=mail[0], mail_type=mail[1])
+
+    if not sys.stdout.isatty():
+        # Like the dashboard at startup, the report covers the cluster this
+        # runs on; remote ones only when Slurm is not available here.
+        reported = [c for c in clusters if c.host.ssh is None] or clusters
+        report.print_reports(reported, settings)
+        return
+
+    # Textual is only imported for the dashboard, keeping the report fast.
+    from .tui import SqtopApp
+
+    SqtopApp(clusters, settings).run()
+
+
+if __name__ == "__main__":
+    main()

sqtop/hosts.py

@@ -0,0 +1,272 @@
+"""Where Slurm commands run: this login node, or another one over ssh.
+
+Starting an ssh session on a cluster login node can take seconds (shell
+start-up files, MFA-guarded sshd) while the Slurm commands themselves are
+quick, so a `Remote` runs a whole refresh in one session, and streams a whole
+probe grid through another. Sessions reuse an ssh control master.
+
+Every command runs without the terminal and never prompts: sqtop does not
+log in for the user. A cluster that needs a password or MFA is reached through
+a control master the user opened with `ssh` in a terminal of their own.
+"""
+
+import os
+import shlex
+import subprocess
+import tempfile
+from concurrent.futures import ThreadPoolExecutor, as_completed
+from pathlib import Path
+
+# Separates command outputs within one remote session; Slurm never prints it.
+RECORD = "\x1e"
+# Concurrent local probes; each is a will-run test in slurmctld, so stay polite.
+PROBE_WORKERS = 4
+PROBE_TIMEOUT = 30
+# Control sockets when the user's ssh config does not multiplex already.
+CONTROL_DIR = Path(tempfile.gettempdir()) / f"sqtop-{os.getuid()}"
+CONTROL_PERSIST = "10m"
+
+
+class ConnectionLost(Exception):
+    """ssh could not reach the cluster; the message is ssh's own."""
+
+
+def connect_hint(ssh):
+    """How the user gets `ssh` connected, shown with a `ConnectionLost`."""
+    return (
+        "sqtop never asks for a password or MFA: log in with "
+        f"`ssh {ssh}` in another terminal, with ControlMaster set for it in "
+        "~/.ssh/config so sqtop can share that connection."
+    )
+
+
+def _run(args, timeout):
+    """Return (returncode, stdout, stderr); 127 if missing, 124 on timeout."""
+    try:
+        done = subprocess.run(
+            args,
+            capture_output=True,
+            text=True,
+            timeout=timeout,
+            stdin=subprocess.DEVNULL,
+            check=False,
+        )
+    except FileNotFoundError:
+        return 127, "", f"'{args[0]}' not found"
+    except subprocess.TimeoutExpired:
+        return 124, "", f"Timed out after {timeout}s"
+    return done.returncode, done.stdout, done.stderr
+
+
+def _combined(stdout, stderr):
+    return "\n".join(part.strip() for part in (stdout, stderr) if part.strip())
+
+
+class Local:
+    """Run commands on this machine."""
+
+    ssh = None
+
+    def run(self, args, timeout=60):
+        """Return stdout, or None when the command is missing, fails or hangs."""
+        code, stdout, _ = _run(args, timeout)
+        return stdout if code == 0 else None
+
+    def act(self, args, timeout=60):
+        """Run a command that acts, such as `sbatch` or `scancel`:
+        (returncode, output).
+
+        The output has stdout and stderr together, so a rejection says why.
+        """
+        code, stdout, stderr = _run(args, timeout)
+        return code, _combined(stdout, stderr)
+
+    def run_batch(self, commands, timeout=120):
+        """Run commands in parallel; their stdout in order, None for failures."""
+        if not commands:
+            return []
+        with ThreadPoolExecutor(len(commands)) as pool:
+            return list(pool.map(lambda args: self.run(args, timeout), commands))
+
+    def run_stream(self, commands, on_result, should_stop=lambda: False):
+        """Run commands a few at a time, reporting each as it finishes.
+
+        `on_result(index, returncode, output)` gets stdout and stderr together.
+        Once `should_stop()` is true, commands not yet started are dropped.
+        """
+        with ThreadPoolExecutor(PROBE_WORKERS) as pool:
+            futures = {
+                pool.submit(_run, args, PROBE_TIMEOUT): index
+                for index, args in enumerate(commands)
+            }
+            for future in as_completed(futures):
+                if should_stop():
+                    pool.shutdown(cancel_futures=True)
+                    break
+                code, stdout, stderr = future.result()
+                on_result(futures[future], code, _combined(stdout, stderr))
+
+
+class Remote:
+    """Run commands on a cluster's login node through `ssh <destination>`.
+
+    `destination` is anything ssh accepts, such as a `~/.ssh/config` alias.
+    Commands run with the login shell's PATH (read by `detect_commands`),
+    since non-login sessions often lack Slurm's directory.
+    """
+
+    def __init__(self, ssh):
+        self.ssh = ssh
+        self.path = None
+        self._control = None
+
+    def control_options(self):
+        """Reuse a control master: the user's own if configured, else ours."""
+        if self._control is None:
+            _, stdout, _ = _run(["ssh", "-G", self.ssh], 10)
+            settings = dict(
+                line.split(" ", 1) for line in stdout.splitlines() if " " in line
+            )
+            options = ["-o", "ControlMaster=auto"]
+            if settings.get("controlpath", "none") == "none":
+                CONTROL_DIR.mkdir(mode=0o700, exist_ok=True)
+                options += ["-o", f"ControlPath={CONTROL_DIR}/%C"]
+            if settings.get("controlpersist", "no") in ("no", "false", "0"):
+                options += ["-o", f"ControlPersist={CONTROL_PERSIST}"]
+            self._control = options
+        return self._control
+
+    def ssh_command(self, script):
+        """ssh running `script` with sh, never allocating a terminal or prompting."""
+        options = [*self.control_options(), "-T", "-o", "ConnectTimeout=15"]
+        options += ["-o", "BatchMode=yes", "-o", "LogLevel=ERROR"]
+        return ["ssh", *options, self.ssh, shlex.join(["sh", "-c", script])]
+
+    def _prelude(self, login_path=False):
+        if login_path:
+            # A login shell sets up PATH the way an interactive user sees it.
+            return (
+                "p=$(bash -lc 'printf %s \"$PATH\"' </dev/null 2>/dev/null)"
+                ' && [ -n "$p" ] && PATH=$p; export PATH\n'
+            )
+        if self.path:
+            return f"PATH={shlex.quote(self.path)}; export PATH\n"
+        return ""
+
+    def run(self, args, timeout=60):
+        return self.run_batch([args], timeout)[0]
+
+    def act(self, args, timeout=60):
+        """`Local.act` on the cluster; raises `ConnectionLost` if ssh fails."""
+        script = (
+            f"{self._prelude()}{shlex.join(args)} </dev/null 2>&1; "
+            "printf '\\n\\036%s\\n' $?"
+        )
+        code, stdout, stderr = _run(self.ssh_command(script), timeout)
+        output, record, status = stdout.rpartition(RECORD)
+        if not record:
+            raise ConnectionLost(stderr.strip() or f"ssh exited with status {code}")
+        return int(status), output.strip()
+
+    def run_batch(self, commands, timeout=120, login_path=False):
+        """Run commands in parallel in one session; stdout in order, None on failure.
+
+        Raises `ConnectionLost` when ssh itself fails.
+        """
+        if not commands:
+            return []
+        lines = [
+            self._prelude(login_path),
+            "d=$(mktemp -d) || exit 97",
+            "trap 'rm -rf \"$d\"' EXIT",
+        ]
+        for index, args in enumerate(commands):
+            lines.append(
+                f'{{ {shlex.join(args)} >"$d/{index}" 2>/dev/null; '
+                f'echo $? >"$d/{index}.rc"; }} &'
+            )
+        indexes = " ".join(str(index) for index in range(len(commands)))
+        lines += [
+            "wait",
+            f"for i in {indexes}; do "
+            'printf \'\\036%s %s\\n\' "$i" "$(cat "$d/$i.rc")"; cat "$d/$i"; done',
+        ]
+        code, stdout, stderr = _run(self.ssh_command("\n".join(lines)), timeout)
+        if RECORD not in stdout:
+            raise ConnectionLost(stderr.strip() or f"ssh exited with status {code}")
+        outputs = [None] * len(commands)
+        for record in stdout.split(RECORD)[1:]:
+            header, _, output = record.partition("\n")
+            index, _, returncode = header.partition(" ")
+            if returncode.strip() == "0":
+                outputs[int(index)] = output
+        return outputs
+
+    def run_stream(self, commands, on_result, should_stop=lambda: False):
+        """Run commands over a few sessions, reporting each result as it lands.
+
+        Commands are dealt round-robin to up to `PROBE_WORKERS` sessions, each
+        running its share one after another. `on_result(index, returncode,
+        output)` gets stdout and stderr together, as soon as each finishes.
+        Raises `ConnectionLost` when ssh fails before any result.
+        """
+        if not commands:
+            return
+        shares = [
+            list(range(len(commands)))[start::PROBE_WORKERS]
+            for start in range(min(PROBE_WORKERS, len(commands)))
+        ]
+        with ThreadPoolExecutor(len(shares)) as pool:
+            outcomes = list(
+                pool.map(
+                    lambda share: self._stream_session(
+                        share, commands, on_result, should_stop
+                    ),
+                    shares,
+                )
+            )
+        if not any(received for received, _ in outcomes) and not should_stop():
+            raise ConnectionLost(next(error for _, error in outcomes if error))
+
+    def _stream_session(self, share, commands, on_result, should_stop):
+        """Run `commands[i]` for each i in `share` in one session.
+
+        Returns (results received, ssh's error text or None).
+        """
+        lines = [self._prelude()]
+        for index in share:
+            lines.append(
+                f"timeout {PROBE_TIMEOUT} {shlex.join(commands[index])} "
+                f"</dev/null 2>&1; printf '\\n\\036%s %s\\n' {index} $?"
+            )
+        try:
+            process = subprocess.Popen(
+                self.ssh_command("\n".join(lines)),
+                stdin=subprocess.DEVNULL,
+                stdout=subprocess.PIPE,
+                stderr=subprocess.PIPE,
+                text=True,
+            )
+        except FileNotFoundError:
+            return 0, "'ssh' not found"
+        received = 0
+        output = []
+        try:
+            for line in process.stdout:
+                if not line.startswith(RECORD):
+                    output.append(line)
+                    continue
+                index, _, returncode = line[1:].strip().partition(" ")
+                received += 1
+                on_result(int(index), int(returncode), "".join(output).strip())
+                output = []
+                if should_stop():
+                    break
+        finally:
+            if process.poll() is None:
+                process.terminate()
+            process.wait()
+        if received:
+            return received, None
+        error = process.stderr.read().strip()
+        return 0, error or f"ssh exited with status {process.returncode}"

sqtop/mailer.py

@@ -0,0 +1,230 @@
+"""The dashboard's hourly email: the text report of each cluster it shows.
+
+Each email starts with your job counts and what happened to your jobs since
+the previous email (submitted, started, ended), then has the report
+`sqtop` prints when output is not a terminal. Every email has the
+subject "sqtop" and refers to the first one, so mail clients keep them
+in one conversation rather than as separate emails.
+
+The server comes from `[smtp]` in the config (`profiles.load_smtp`). The
+thread and the jobs last reported are kept in `STATE_PATH`, so the
+conversation carries on from one run of the dashboard to the next.
+"""
+
+import datetime
+import email.message
+import email.utils
+import html
+import io
+import json
+import os
+import smtplib
+import ssl
+from pathlib import Path
+
+from . import hosts, report, slurm, views
+
+SUBJECT = "sqtop"
+INTERVAL = 3600
+SSL_PORT = 465
+SMTP_TIMEOUT = 60
+STATE_PATH = (
+    Path(os.environ.get("XDG_STATE_HOME") or Path.home() / ".local" / "state")
+    / "sqtop"
+    / "mail.json"
+)
+
+
+def compose(smtp, text, thread=None):
+    """The email: `text` as plain text and as monospaced HTML.
+
+    `thread` is the first email's Message-ID. Replying to it, under the same
+    subject, keeps every email in one conversation.
+    """
+    message = email.message.EmailMessage()
+    message["Subject"] = SUBJECT
+    message["From"] = email.utils.formataddr((SUBJECT, smtp.sender))
+    message["To"] = smtp.recipient
+    message["Date"] = email.utils.formatdate(localtime=True)
+    domain = smtp.sender.rpartition("@")[2] or None
+    message["Message-ID"] = email.utils.make_msgid(SUBJECT, domain)
+    if thread:
+        message["In-Reply-To"] = thread
+        message["References"] = thread
+    message.set_content(text)
+    # The tables only line up in a fixed-width font, which plain text
+    # does not get everywhere (Gmail's web client, for one).
+    message.add_alternative(
+        f'<pre style="font: 12px/1.4 monospace">{html.escape(text)}</pre>',
+        subtype="html",
+    )
+    return message
+
+
+def send(smtp, message):
+    """Deliver `message` over SSL (port 465) or STARTTLS.
+
+    A password is never sent over a connection that is not encrypted.
+    """
+    context = ssl.create_default_context()
+    if smtp.port == SSL_PORT:
+        connection = smtplib.SMTP_SSL(
+            smtp.server, smtp.port, timeout=SMTP_TIMEOUT, context=context
+        )
+    else:
+        connection = smtplib.SMTP(smtp.server, smtp.port, timeout=SMTP_TIMEOUT)
+    with connection:
+        if smtp.port != SSL_PORT:
+            connection.ehlo()
+            if connection.has_extn("starttls"):
+                connection.starttls(context=context)
+                connection.ehlo()
+            elif smtp.password:
+                raise smtplib.SMTPException(
+                    f"{smtp.server}:{smtp.port} does not offer STARTTLS; "
+                    "not sending the password unencrypted"
+                )
+        if smtp.password:
+            connection.login(smtp.sender, smtp.password)
+        connection.send_message(message)
+
+
+def job_states(jobs):
+    """What the next email compares with: each job's state and name."""
+    return {job["id"]: {"state": job["state"], "name": job["name"]} for job in jobs}
+
+
+def job_changes(previous, jobs, ended=None):
+    """One line per job whose state changed since `previous` (`job_states`).
+
+    `ended` maps jobs that left the queue to their (state, exit code) from
+    sacct; without an entry such a job only "left the queue".
+    """
+    ended = ended or {}
+    changes = []
+    for job in jobs:
+        before = previous.get(job["id"], {}).get("state")
+        if before == job["state"]:
+            continue
+        label = f"{job['id']} ({job['name']})"
+        if job["state"] == "RUNNING":
+            nodes, _ = views.job_nodes(job)
+            changes.append(f"{label} started on {nodes}")
+        elif before is None and job["state"] == "PENDING":
+            text = f"{label} submitted"
+            if (reason := views.job_reason(job)) != "-":
+                text += f" ({reason})"
+            if (start := views.job_start_time(job, compact=True)) != "-":
+                text += f", expected to start {start}"
+            changes.append(text)
+        else:
+            changes.append(f"{label} is now {job['state']}")
+    current = {job["id"] for job in jobs}
+    for job_id, job in previous.items():
+        if job_id in current:
+            continue
+        label = f"{job_id} ({job['name']})"
+        if job_id not in ended:
+            changes.append(f"{label} left the queue")
+            continue
+        state, code = ended[job_id]
+        exit_code = "" if code in ("", "0:0") else f", exit code {code}"
+        changes.append(f"{label} ended: {state}{exit_code}")
+    return changes
+
+
+def final_states(host, ids):
+    """How jobs that left the queue ended, from sacct; {} when it fails."""
+    # A pending array (`12_[1-9]`) is not an id sacct takes in a list.
+    ids = [job_id for job_id in ids if job_id.replace("_", "").isdigit()]
+    if not ids:
+        return {}
+    try:
+        output = host.run(slurm.sacct_command(ids))
+    except hosts.ConnectionLost:
+        return {}
+    return {} if output is None else slurm.parse_sacct_states(output)
+
+
+def digest(targets, state, now):
+    """The email's text, and the jobs the next email compares with.
+
+    `targets` pairs each cluster with the probe request to report. The text
+    has a line of job counts per cluster, each followed by the changes since
+    the email `state` records, then each cluster's report.
+    """
+    seen = state.get("jobs", {})
+    since = views.parse_time(state.get("sent_at"))
+    if since is not None:
+        since = f"{since:%H:%M}" if since.date() == now.date() else f"{since:%F %H:%M}"
+    latest = dict(seen)
+    heads, reports = [], []
+    for cluster, settings in targets:
+        output = io.StringIO()
+        snap = report.print_cluster(cluster, settings, out=output)
+        reports.append(output.getvalue())
+        if snap is None or snap.jobs is None:
+            heads.append(f"{cluster.name}: cannot read your jobs")
+            continue
+        counts = slurm.job_counts(snap.jobs)
+        head = (
+            f"{cluster.name}: {counts['running']} running, {counts['pending']} pending"
+        )
+        previous = seen.get(cluster.name)
+        latest[cluster.name] = job_states(snap.jobs)
+        if previous is None or since is None:
+            heads.append(head)
+            continue
+        current = {job["id"] for job in snap.jobs}
+        gone = [job_id for job_id in previous if job_id not in current]
+        changes = job_changes(previous, snap.jobs, final_states(cluster.host, gone))
+        if not changes:
+            heads.append(f"{head}; no changes since {since}")
+            continue
+        heads.append(f"{head}. Since {since}:")
+        heads.extend(f"  {change}" for change in changes)
+    separator = f"{report.SEPARATOR}\n\n"
+    return "\n".join(heads) + "\n\n" + separator.join(reports), latest
+
+
+def load_state(path):
+    """What the last email recorded, or {} before the first one."""
+    try:
+        state = json.loads(Path(path).read_text())
+    except (FileNotFoundError, ValueError):
+        return {}
+    return state if isinstance(state, dict) else {}
+
+
+def save_state(path, state):
+    path = Path(path)
+    path.parent.mkdir(parents=True, exist_ok=True)
+    temporary = path.with_name(path.name + ".tmp")
+    temporary.write_text(json.dumps(state, indent=1))
+    temporary.replace(path)
+
+
+def send_report(targets, smtp, path=None, now=None):
+    """Email the report of `targets` once and return the message.
+
+    Raises OSError (including `smtplib.SMTPException`) when it cannot be
+    sent. The state is saved only after the email is sent, so the next email
+    still lists every change since the last one that went out.
+    """
+    path = STATE_PATH if path is None else path
+    now = now or datetime.datetime.now()
+    state = load_state(path)
+    text, jobs = digest(targets, state, now)
+    threads = state.get("threads", {})
+    message = compose(smtp, text, threads.get(smtp.recipient))
+    send(smtp, message)
+    threads.setdefault(smtp.recipient, message["Message-ID"])
+    save_state(
+        path,
+        {
+            "threads": threads,
+            "sent_at": now.isoformat(timespec="seconds"),
+            "jobs": jobs,
+        },
+    )
+    return message