issues 153 and 150

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
 
 #

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)

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)