wip

pyslurm/db/assoc.pxd

@@ -45,6 +45,7 @@ from pyslurm.db.tres cimport (
     merge_tres_str,
     tres_ids_to_names,
     TrackableResources,
+    TrackableResourceLimits,
 )
 from pyslurm.db.connection cimport Connection
 from pyslurm.utils cimport cstr
@@ -70,6 +71,15 @@ cdef class Association:
         QualitiesOfService qos_data
         TrackableResources tres_data
 
+    cdef public:
+        group_tres
+        group_tres_mins
+        group_tres_run_mins
+        max_tres_mins_per_job
+        max_tres_run_mins_per_user
+        max_tres_per_job
+        max_tres_per_node
+
     @staticmethod
     cdef Association from_ptr(slurmdb_assoc_rec_t *in_ptr)
 

pyslurm/db/assoc.pyx

@@ -73,6 +73,7 @@ cdef class Associations(dict):
             assoc = Association.from_ptr(<slurmdb_assoc_rec_t*>assoc_ptr.data)
             assoc.qos_data = qos_data
             assoc.tres_data = tres_data
+            assoc._parse_tres()
             assoc_dict[assoc.id] = assoc
 
         return assoc_dict
@@ -194,6 +195,23 @@ cdef class Association:
         wrap.ptr = in_ptr
         return wrap
 
+    def _parse_tres(self):
+        cdef TrackableResources tres = self.tres_data
+        self.group_tres = TrackableResourceLimits.from_ids(
+                self.ptr.grp_tres, tres)
+        self.group_tres_mins = TrackableResourceLimits.from_ids(
+                self.ptr.grp_tres_mins, tres)
+        self.group_tres_run_mins = TrackableResourceLimits.from_ids(
+                self.ptr.grp_tres_mins, tres)
+        self.max_tres_mins_per_job = TrackableResourceLimits.from_ids(
+                self.ptr.max_tres_mins_pj, tres)
+        self.max_tres_run_mins_per_user = TrackableResourceLimits.from_ids(
+                self.ptr.max_tres_run_mins, tres)
+        self.max_tres_per_job = TrackableResourceLimits.from_ids(
+                self.ptr.max_tres_pj, tres)
+        self.max_tres_per_node = TrackableResourceLimits.from_ids(
+                self.ptr.max_tres_pn, tres)
+
     def as_dict(self):
         """Database Association information formatted as a dictionary.
 
@@ -204,19 +222,20 @@ cdef class Association:
 
     def _validate_tres(self):
         self.tres_data = TrackableResources.load(name_is_key=False)
-        self.group_tres = tres_names_to_ids(self.group_tres, self.tres_data) 
-        self.group_tres_mins = tres_names_to_ids(
-                self.group_tres_mins, self.tres_data) 
-        self.group_tres_run_mins = tres_names_to_ids(
-                self.group_tres_run_mins, self.tres_data) 
-        self.max_tres_mins_per_job = tres_names_to_ids(
-                self.max_tres_mins_per_job, self.tres_data) 
-        self.max_tres_run_mins_per_user = tres_names_to_ids(
-                self.max_tres_run_mins_per_user, self.tres_data) 
-        self.max_tres_per_job = tres_names_to_ids(
-                self.max_tres_per_job, self.tres_data) 
-        self.max_tres_per_node = tres_names_to_ids(
-                self.max_tres_per_node, self.tres_data) 
+        cstr.from_dict(&self.ptr.grp_tres,
+                       self.group_tres._validate(self.tres_data))
+        cstr.from_dict(&self.ptr.grp_tres_mins,
+                       self.group_tres_mins._validate(self.tres_data))
+        cstr.from_dict(&self.ptr.grp_tres_run_mins,
+                       self.group_tres_run_mins._validate(self.tres_data))
+        cstr.from_dict(&self.ptr.max_tres_mins_pj,
+                       self.max_tres_mins_per_job._validate(self.tres_data))
+        cstr.from_dict(&self.ptr.max_tres_run_mins,
+                       self.max_tres_run_mins_per_user._validate(self.tres_data))
+        cstr.from_dict(&self.ptr.max_tres_pj,
+                       self.max_tres_per_job._validate(self.tres_data))
+        cstr.from_dict(&self.ptr.max_tres_pn,
+                       self.max_tres_per_node._validate(self.tres_data))
 
     @staticmethod
     def load(name):
@@ -274,30 +293,6 @@ cdef class Association:
     def group_submit_jobs(self, val):
         self.ptr.grp_submit_jobs = u32(val, zero_is_noval=False)
 
-    @property
-    def group_tres(self):
-        return tres_ids_to_names(self.ptr.grp_tres, self.tres_data)
-
-    @group_tres.setter
-    def group_tres(self, val):
-        cstr.from_dict(&self.ptr.grp_tres, val)
-
-    @property
-    def group_tres_mins(self):
-        return tres_ids_to_names(self.ptr.grp_tres_mins, self.tres_data)
-
-    @group_tres_mins.setter
-    def group_tres_mins(self, val):
-        cstr.from_dict(&self.ptr.grp_tres_mins, val)
-
-    @property
-    def group_tres_run_mins(self):
-        return tres_ids_to_names(self.ptr.grp_tres_run_mins, self.tres_data)
-
-    @group_tres_run_mins.setter
-    def group_tres_run_mins(self, val):
-        cstr.from_dict(&self.ptr.grp_tres_run_mins, val)
-
     @property
     def group_wall_time(self):
         return u32_parse(self.ptr.grp_wall, zero_is_noval=False)
@@ -346,38 +341,6 @@ cdef class Association:
     def max_submit_jobs(self, val):
         self.ptr.max_submit_jobs = u32(val, zero_is_noval=False)
 
-    @property
-    def max_tres_mins_per_job(self):
-        return tres_ids_to_names(self.ptr.max_tres_mins_pj, self.tres_data)
-
-    @max_tres_mins_per_job.setter
-    def max_tres_mins_per_job(self, val):
-        cstr.from_dict(&self.ptr.max_tres_mins_pj, val)
-
-    @property
-    def max_tres_run_mins_per_user(self):
-        return tres_ids_to_names(self.ptr.max_tres_run_mins, self.tres_data)
-
-    @max_tres_run_mins_per_user.setter
-    def max_tres_run_mins_per_user(self, val):
-        cstr.from_dict(&self.ptr.max_tres_run_mins, val)
-
-    @property
-    def max_tres_per_job(self):
-        return tres_ids_to_names(self.ptr.max_tres_pj, self.tres_data)
-
-    @max_tres_per_job.setter
-    def max_tres_per_job(self, val):
-        cstr.from_dict(&self.ptr.max_tres_pj, val)
-
-    @property
-    def max_tres_per_node(self):
-        return tres_ids_to_names(self.ptr.max_tres_pn, self.tres_data)
-
-    @max_tres_per_node.setter
-    def max_tres_per_node(self, val):
-        cstr.from_dict(&self.ptr.max_tres_pn, val)
-
     @property
     def max_wall_time_per_job(self):
         return u32_parse(self.ptr.max_wall_pj, zero_is_noval=False)
@@ -424,6 +387,7 @@ cdef class Association:
 
     @qos.setter
     def qos(self, val):
+        # TODO: must be ids
         make_char_list(&self.ptr.qos_list, val)
 
     @property

pyslurm/db/tres.pxd

@@ -45,6 +45,24 @@ cdef merge_tres_str(char **tres_str, typ, val)
 cdef tres_ids_to_names(char *tres_str, TrackableResources tres_data)
 
 
+cdef class TrackableResourceLimits:
+
+    cdef public:
+        cpu
+        mem
+        energy
+        node
+        billing
+        fs
+        vmem
+        pages
+        gres
+        license
+
+    @staticmethod
+    cdef from_ids(char *tres_id_str, TrackableResources tres_data)
+
+
 cdef class TrackableResourceFilter:
     cdef slurmdb_tres_cond_t *ptr
 

pyslurm/db/tres.pyx

@@ -25,6 +25,85 @@
 from pyslurm.utils.uint import *
 from pyslurm.constants import UNLIMITED
 from pyslurm.core.error import RPCError
+from pyslurm.utils.helpers import instance_to_dict
+from pyslurm.utils import cstr
+import json
+
+
+TRES_TYPE_DELIM = "/"
+
+
+cdef class TrackableResourceLimits:
+
+    def __init__(self, **kwargs):
+        self.fs = {}
+        self.gres = {}
+        self.license = {}
+
+        for k, v in kwargs.items():
+            if TRES_TYPE_DELIM in k:
+                typ, name = self._unflatten_tres(k)
+                cur_val = getattr(self, typ)
+
+                if not isinstance(cur_val, dict):
+                    raise ValueError(f"TRES Type {typ} cannot have a name "
+                                     f"({name}). Invalid Value: {typ}/{name}")
+
+                cur_val.update({name : int(v)})
+                setattr(self, typ, cur_val)
+            else:
+                setattr(self, k, v)
+
+    @staticmethod
+    cdef from_ids(char *tres_id_str, TrackableResources tres_data):
+        tres_list = tres_ids_to_names(tres_id_str, tres_data)
+        if not tres_list:
+            return None
+
+        cdef TrackableResourceLimits out = TrackableResourceLimits()
+
+        for tres in tres_list:
+            typ, name, cnt = tres
+            cur_val = getattr(out, typ, slurm.NO_VAL64)
+            if cur_val != slurm.NO_VAL64:
+                if isinstance(cur_val, dict):
+                    cur_val.update({name : cnt})
+                    setattr(out, typ, cur_val)
+                else:
+                    setattr(out, typ, cnt)
+
+        return out
+
+    def _validate(self, TrackableResources tres_data):
+        id_dict = tres_names_to_ids(self.as_dict(flatten_limits=True),
+                                    tres_data)
+        return id_dict
+
+    def _unflatten_tres(self, type_and_name):
+        typ, name = type_and_name.split(TRES_TYPE_DELIM, 1)
+        return typ, name
+
+    def _flatten_tres(self, typ, vals):
+        cdef dict out = {}
+        for name, cnt in vals.items():
+            out[f"{typ}{TRES_TYPE_DELIM}{name}"] = cnt
+
+        return out
+
+    def as_dict(self, flatten_limits=False):
+        cdef dict inst_dict = instance_to_dict(self)
+
+        if flatten_limits:
+            vals = inst_dict.pop("fs")
+            inst_dict.update(self._flatten_tres("fs", vals))
+
+            vals = inst_dict.pop("license")
+            inst_dict.update(self._flatten_tres("license", vals))
+
+            vals = inst_dict.pop("gres")
+            inst_dict.update(self._flatten_tres("gres", vals))
+
+        return inst_dict
 
 
 cdef class TrackableResourceFilter:
@@ -162,7 +241,7 @@ cdef class TrackableResource:
     def type_and_name(self):
         type_and_name = self.type
         if self.name:
-            type_and_name = f"{type_and_name}/{self.name}"
+            type_and_name = f"{type_and_name}{TRES_TYPE_DELIM}{self.name}"
 
         return type_and_name
 
@@ -206,24 +285,22 @@ cdef merge_tres_str(char **tres_str, typ, val):
 
 cdef tres_ids_to_names(char *tres_str, TrackableResources tres_data):
     if not tres_str:
-        return {}
+        return None
 
     cdef:
         dict tdict = cstr.to_dict(tres_str)
-        dict out = {}
+        list out = []
 
     if not tres_data:
-        return tdict
+        return None
 
     for tid, cnt in tdict.items():
         if isinstance(tid, str) and tid.isdigit():
             _tid = int(tid)
             if _tid in tres_data:
-                out[tres_data[_tid].type_and_name] = cnt
-                continue
-
-        # If we can't find the TRES ID in our data, return it raw.
-        out[tid] = cnt
+                out.append(
+                    (tres_data[_tid].type, tres_data[_tid].name, int(cnt))
+                )
 
     return out