diff --git a/pyslurm/pyslurm.pyx b/pyslurm/pyslurm.pyx index 103c09b9..98137d5f 100644 --- a/pyslurm/pyslurm.pyx +++ b/pyslurm/pyslurm.pyx @@ -5343,7 +5343,7 @@ cdef class slurmdb_jobs: slurm.xfree(self.job_cond) slurm.slurmdb_connection_close(&self.db_conn) - def get(self, jobids=[], starttime=0, endtime=0): + def get(self, jobids=[], userids=[], starttime=0, endtime=0, flags = None, db_flags = None, clusters = []): u"""Get Slurmdb information about some jobs. Input formats for start and end times: @@ -5371,6 +5371,23 @@ cdef class slurmdb_jobs: slurm.List JOBSList slurm.ListIterator iters = NULL + + if clusters: + self.job_cond.cluster_list = slurm.slurm_list_create(NULL) + for _cluster in clusters: + _cluster = _cluster.encode("UTF-8") + slurm.slurm_addto_char_list_with_case(self.job_cond.cluster_list, _cluster, False) + + if db_flags: + if isinstance(db_flags, int): + self.job_cond.db_flags = db_flags + else: + self.job_cond.db_flags = slurm.SLURMDB_JOB_FLAG_NOTSET + + if flags: + if isinstance(flags, int): + self.job_cond.flags = flags + if jobids: self.job_cond.step_list = slurm.slurm_list_create(NULL) for _jobid in jobids: @@ -5380,6 +5397,15 @@ cdef class slurmdb_jobs: _jobid = _jobid.encode("UTF-8") slurm.slurm_addto_step_list(self.job_cond.step_list, _jobid) + if userids: + self.job_cond.userid_list = slurm.slurm_list_create(NULL) + for _userid in userids: + if isinstance(_userid, int) or isinstance(_userid, long): + _userid = str(_userid).encode("UTF-8") + else: + _userid = _userid.encode("UTF-8") + slurm.slurm_addto_char_list_with_case(self.job_cond.userid_list, _userid, False) + if starttime: self.job_cond.usage_start = slurm.slurm_parse_time(starttime, 1) errno = slurm.slurm_get_errno() @@ -5466,6 +5492,10 @@ cdef class slurmdb_jobs: slurm.slurm_list_iterator_destroy(iters) slurm.slurm_list_destroy(JOBSList) + if clusters: + slurm.slurm_list_destroy(self.job_cond.cluster_list) + if userids: + slurm.slurm_list_destroy(self.job_cond.userid_list) return J_dict # diff --git a/pyslurm/slurm.pxd b/pyslurm/slurm.pxd index 8c222819..758ec518 100644 --- a/pyslurm/slurm.pxd +++ b/pyslurm/slurm.pxd @@ -2608,6 +2608,7 @@ cdef extern from 'slurm/slurmdb.h' nogil: ctypedef struct slurmdb_job_rec_t: char *account + char *admin_comment char *alloc_gres uint32_t alloc_nodes uint32_t array_job_id @@ -2630,8 +2631,11 @@ cdef extern from 'slurm/slurmdb.h' nogil: uint32_t jobid char *jobname uint32_t lft - char *partition + char *mcs_label char *nodes + char *partition + uint32_t pack_job_id + uint32_t pack_job_offset uint32_t priority uint32_t qosid uint32_t req_cpus @@ -2664,6 +2668,7 @@ cdef extern from 'slurm/slurmdb.h' nogil: uint32_t user_cpu_usec char *wckey uint32_t wckeyid + char *work_dir ctypedef struct slurmdb_qos_usage_t: uint32_t accrue_cnt @@ -2904,6 +2909,7 @@ cdef extern char *slurm_get_checkpoint_dir() cdef extern void slurm_sprint_cpu_bind_type(char *string, cpu_bind_type_t cpu_bind_type) cdef extern void slurm_destroy_char(void *object) cdef extern int slurm_addto_step_list(List step_list, char *names) +cdef extern int slurm_addto_char_list_with_case(List char_list, char *names, bool lower_case_noralization) cdef extern time_t slurm_parse_time(char *time_str, int past) cdef extern int slurm_time_str2mins(const_char_ptr string) cdef extern int slurm_time_str2secs(const_char_ptr string) diff --git a/tests/test-slurmdb.py b/tests/test-slurmdb.py new file mode 100644 index 00000000..afca680c --- /dev/null +++ b/tests/test-slurmdb.py @@ -0,0 +1,81 @@ +import datetime +import pwd +import subprocess + +from nose.tools import assert_equals + +import pyslurm + + +def njobs_sacct_jobs(start, end, username=None): + """ + Count the number of jobs reported by sacct + For comparison with the reults of slurmdb_jobs.get + """ + sacctcmd = ['sacct', '-S', start, '-E', end, '-n', '-X'] + if username is not None: + sacctcmd.extend(['-u', username]) + else: + sacctcmd.append('-a') + sacct = subprocess.Popen(sacctcmd, + stdout=subprocess.PIPE, + stderr=None).communicate() + return len(sacct[0].splitlines()) + + +def njobs_slurmdb_jobs_get(start, end, uid=None): + """ + Count the number of jobs reported by slurmdb + """ + if uid is None: + jobs = pyslurm.slurmdb_jobs().get(starttime=start.encode('utf-8'), + endtime=end.encode('utf-8')) + else: + jobs = pyslurm.slurmdb_jobs().get(starttime=start.encode('utf-8'), + endtime=end.encode('utf-8'), + userids=[uid]) + return len(jobs) + + +def get_user(): + """ + Return a list of usernames and their uid numbers + """ + users = subprocess.Popen(['squeue', '-O', 'username', '-h'], + stdout=subprocess.PIPE, + stderr=None).communicate() + for username in users[0].splitlines(): + print(username.decode()) + uid = pwd.getpwnam("{}".format(username.strip().decode())) + yield username.strip().decode(), uid.pw_uid + + +def test_slurmdb_jobs_get(): + """ + Slurmdb: Compare sacct and slurmdb_jobs.get() for all users + """ + starttime = (datetime.datetime.now() - + datetime.timedelta(days=2)).strftime("%Y-%m-%dT00:00:00") + endtime = (datetime.datetime.now() - + datetime.timedelta(days=1)).strftime("%Y-%m-%dT00:00:00") + njobs_pyslurm = njobs_slurmdb_jobs_get(starttime, endtime) + njobs_sacct = njobs_sacct_jobs(starttime, endtime) + assert_equals(njobs_pyslurm, njobs_sacct) + + +def test_slurmdb_jobs_get_byuser(): + """ + Slurmdb: Compare sacct and slurmdb_jobs.get() for individual users + """ + + userlist = list(get_user()) + for user in userlist[:10]: + starttime = (datetime.datetime.now() - + datetime.timedelta(days=2)).strftime("%Y-%m-%dT00:00:00") + endtime = (datetime.datetime.now() - + datetime.timedelta(days=1)).strftime("%Y-%m-%dT00:00:00") + njobs_pyslurm = njobs_slurmdb_jobs_get(starttime, + endtime, + int(user[1])) + njobs_sacct = njobs_sacct_jobs(starttime, endtime, username=user[0]) + assert_equals(njobs_pyslurm, njobs_sacct)