How to implement bstop command functionality with python script

Steps

#!/usr/bin/env python3
"""
bstop.py - a Python re-implementation of LSF's bstop built on the
IBM Spectrum LSF Python API (https://github.com/IBMSpectrumComputing/lsf-python-api).
bstop suspends unfinished jobs by sending SIGSTOP (sequential jobs) or
SIGTSTP (parallel jobs). This script finds the matching jobs with
lsb_openjobinfo()/lsb_readjobinfo() and suspends them with lsb_signaljob().
Usage (same shape as the real command):
bstop.py [-a] [-J job_name] [-m host] [-q queue] [-u user|all]
[-C reason] [0] [job_ID ...]
bstop.py 314               suspend job 314
bstop.py -u all            suspend the last submitted job in the system
bstop.py -u jsmith 0       suspend all jobs of user jsmith
bstop.py -q night -J "sim*" 0
bstop.py -a                suspend all jobs
Extra switches for this script only:
--dry-run show what would be suspended, send nothing
--mock run against a fake in-memory job list (no LSF needed)
Requires: LSF environment sourced (profile.lsf) and pythonlsf installed.
"""
import argparse
import getpass
import signal
import sys
---- LSF constants (values from lsf/lsbatch.h; the SWIG module normally
---- exports them, so we prefer lsf. and fall back to the numbers).
JOB_STAT_PEND = 0x01
JOB_STAT_PSUSP = 0x02
JOB_STAT_RUN = 0x04
JOB_STAT_SSUSP = 0x08
JOB_STAT_USUSP = 0x10
JOB_STAT_EXIT = 0x20
JOB_STAT_DONE = 0x40
SUSPENDED_MASK = JOB_STAT_PSUSP | JOB_STAT_SSUSP | JOB_STAT_USUSP
FINISHED_MASK = JOB_STAT_EXIT | JOB_STAT_DONE
lsb_openjobinfo option bits
_OPEN_FLAGS = {
"ALL_JOB": 0x0001,
"PEND_JOB": 0x0004,
"SUSP_JOB": 0x0008,
"CUR_JOB": 0x0010, # unfinished jobs only
"LAST_JOB": 0x0020, # only the most recently submitted job
"RUN_JOB": 0x0040,
}
def _flag(lsf, name):
return getattr(lsf, name, _OPEN_FLAGS[name]) if lsf else _OPEN_FLAGS[name]
--------------------------------------------------------------------------
Backends
--------------------------------------------------------------------------
class Job:
"""Minimal view of a jobInfoEnt."""
def __init__(self, job_id, status, user, name="", queue="", hosts=(), nproc=1):
    self.job_id = job_id
    self.status = status
    self.user = user
    self.name = name
    self.queue = queue
    self.hosts = list(hosts)
    self.nproc = nproc

@property
def parallel(self):
    return self.nproc > 1
class LsfBackend:
"""Talks to a real LSF cluster through pythonlsf."""
def __init__(self):
    try:
        from pythonlsf import lsf
    except ImportError:
        sys.exit("bstop.py: cannot import pythonlsf. Source profile.lsf and "
                 "install lsf-python-api, or use --mock.")
    self.lsf = lsf
    if lsf.lsb_init("bstop.py") != 0:
        sys.exit("bstop.py: failed to initialise LSF batch library: "
                 + self._err())

def _err(self):
    sysmsg = getattr(self.lsf, "lsb_sysmsg", None)
    return sysmsg() if sysmsg else "unknown LSF error"

def find_jobs(self, job_id, job_name, user, queue, host, last_only):
    lsf = self.lsf
    options = _flag(lsf, "CUR_JOB")  # unfinished jobs only
    if last_only:
        options |= _flag(lsf, "LAST_JOB")

    count = lsf.lsb_openjobinfo(job_id, job_name or "", user or "",
                                queue or "", host or "", options)
    if count < 0:
        sys.exit("bstop.py: " + self._err())

    jobs = []
    more = lsf.new_intp()
    lsf.intp_assign(more, count)
    try:
        for _ in range(count):
            info = lsf.lsb_readjobinfo(more)
            if info is None:
                break
            submit = getattr(info, "submit", None)
            nproc = getattr(submit, "numProcessors", 1) or 1
            # exHosts is a raw char** in the SWIG wrapper and is not
            # indexable from Python; host filtering is done by LSF itself
            # (hostName arg of lsb_openjobinfo), so this list is optional.
            hosts = []
            try:
                n_ex = getattr(info, "numExHosts", 0)
                ex_hosts = getattr(info, "exHosts", None)
                if ex_hosts is not None and n_ex:
                    hosts = [ex_hosts[i] for i in range(n_ex)]
            except (TypeError, IndexError):
                hosts = []
            jobs.append(Job(
                job_id=info.jobId,
                status=info.status,
                user=info.user,
                name=getattr(submit, "jobName", "") or getattr(info, "jName", ""),
                queue=getattr(submit, "queue", ""),
                hosts=hosts,
                nproc=nproc,
            ))
    finally:
        lsf.lsb_closejobinfo()  # always close to avoid leaking the connection
    return jobs

def signal_job(self, job_id, sig):
    """Returns (ok, error_message)."""
    if self.lsf.lsb_signaljob(job_id, int(sig)) == 0:
        return True, ""
    return False, self._err()
class MockBackend:
"""In-memory fake cluster so the script can be tried without LSF."""
def __init__(self):
    me = getpass.getuser()
    self.jobs = [
        Job(101, JOB_STAT_RUN, me, "sim_a", "normal", ["hostA"], 1),
        Job(102, JOB_STAT_RUN, me, "sim_b", "night", ["hostB"], 8),
        Job(103, JOB_STAT_PEND, me, "post", "normal", [], 1),
        Job(104, JOB_STAT_USUSP, me, "held", "normal", ["hostA"], 1),
        Job(105, JOB_STAT_RUN, "jsmith", "other", "night", ["hostC"], 1),
        Job(106, JOB_STAT_DONE, me, "old", "normal", ["hostA"], 1),
    ]

def find_jobs(self, job_id, job_name, user, queue, host, last_only):
    import fnmatch
    me = getpass.getuser()
    out = []
    for j in self.jobs:
        if j.status & FINISHED_MASK:
            continue
        if job_id and j.job_id != job_id:
            continue
        if job_name and not fnmatch.fnmatch(j.name, job_name):
            continue
        if user != "all" and j.user != (user or me):
            if not job_id:  # explicit job IDs may belong to anyone
                continue
        if queue and j.queue != queue:
            continue
        if host and host not in j.hosts:
            continue
        out.append(j)
    if last_only and out:
        out = [max(out, key=lambda j: j.job_id)]
    return out

def signal_job(self, job_id, sig):
    for j in self.jobs:
        if j.job_id == job_id:
            if j.status & (JOB_STAT_RUN | JOB_STAT_PEND):
                j.status = JOB_STAT_USUSP if j.status & JOB_STAT_RUN else JOB_STAT_PSUSP
            return True, ""
    return False, "No matching job found"
--------------------------------------------------------------------------
bstop logic
--------------------------------------------------------------------------
def parse_args(argv):
p = argparse.ArgumentParser(
prog="bstop.py",
description="Suspend unfinished LSF jobs (SIGSTOP for sequential, "
"SIGTSTP for parallel jobs).",
)
p.add_argument("-a", action="store_true", help="suspend all jobs")
p.add_argument("-J", dest="job_name", metavar="job_name",
help="only jobs with this name (wildcard * allowed)")
p.add_argument("-m", dest="host", metavar="host_name",
help="only jobs dispatched to this host")
p.add_argument("-q", dest="queue", metavar="queue_name",
help="only jobs in this queue")
p.add_argument("-u", dest="user", metavar="user_name|all",
help="only jobs of this user, or 'all'")
p.add_argument("-C", dest="reason", metavar="suspend_reason",
help="suspend reason (max 4095 chars)")
p.add_argument("-V", action="store_true", help="print version and exit")
p.add_argument("--dry-run", action="store_true",
help="list the jobs that would be suspended, send nothing")
p.add_argument("--mock", action="store_true",
help="use a fake in-memory cluster instead of real LSF")
p.add_argument("job_ids", nargs="*", metavar="job_ID",
help="job IDs to suspend; 0 means every job matching the filters")
return p.parse_args(argv)
def suspend_signal(job):
return signal.SIGTSTP if job.parallel else signal.SIGSTOP
def main(argv=None):
args = parse_args(argv)
if args.V:
    print("bstop.py - Python simulation of LSF bstop (lsf-python-api)")
    return 0

if args.reason and len(args.reason) > 4095:
    print("bstop.py: suspend reason is longer than 4095 characters", file=sys.stderr)
    return 1

ids = []
for token in args.job_ids:
    try:
        ids.append(int(token))
    except ValueError:
        print(f"bstop.py: Illegal job ID: {token}", file=sys.stderr)
        return 1

has_filter = any([args.job_name, args.host, args.queue, args.user])
if not (ids or args.a or has_filter):
    print("bstop.py: specify a job ID or one of -a, -J, -m, -q, -u",
          file=sys.stderr)
    return 1

backend = MockBackend() if args.mock else LsfBackend()

if args.reason and not args.mock:
    print("bstop.py: warning: -C reason is not passed through "
          "lsb_signaljob and will be ignored", file=sys.stderr)

# Work out the list of (jobId) queries.
#   explicit IDs (not 0)        -> one query per ID
#   0 or -a                     -> every job that matches the filters
#   filters only, no 0/-a       -> only the last submitted matching job
all_matching = args.a or 0 in ids
explicit = [i for i in ids if i != 0]

user = "all" if args.a else args.user

jobs = []
if explicit:
    for jid in explicit:
        found = backend.find_jobs(jid, args.job_name, "all", args.queue,
                                  args.host, last_only=False)
        if not found:
            print(f"Job <{jid}>: No matching job found", file=sys.stderr)
        jobs.extend(found)
if all_matching or not explicit:
    jobs.extend(backend.find_jobs(0, args.job_name, user, args.queue,
                                  args.host, last_only=not all_matching))

if not jobs and not explicit:
    print("No matching job found", file=sys.stderr)
    return 1

seen, rc = set(), 0
for job in jobs:
    if job.job_id in seen:
        continue
    seen.add(job.job_id)

    if job.status & FINISHED_MASK:
        print(f"Job <{job.job_id}>: Job has already finished", file=sys.stderr)
        rc = 1
        continue
    if job.status & JOB_STAT_USUSP:
        print(f"Job <{job.job_id}>: Job has already been stopped", file=sys.stderr)
        rc = 1
        continue

    sig = suspend_signal(job)
    if args.dry_run:
        print(f"Job <{job.job_id}> would be stopped ({sig.name})")
        continue

    ok, err = backend.signal_job(job.job_id, sig)
    if ok:
        print(f"Job <{job.job_id}> is being stopped")
    else:
        print(f"Job <{job.job_id}>: {err}", file=sys.stderr)
        rc = 1
return rc
if name == "main":
sys.exit(main())

Test result

[egoadmin@gen2 ~]$ bsub sleep 30
Job <418> is submitted to default queue <normal>.
[egoadmin@gen2 examples]$ python3 bstop.py 418
Job <418> is being stopped
[egoadmin@gen2 ~]$ bjobs
JOBID   USER    STAT  QUEUE      FROM_HOST   EXEC_HOST   JOB_NAME   SUBMIT_TIME
418     egoadmi USUSP normal     gen2.fyre.i gen3.fyre.i sleep 30   Nov 11 15:41