DIRAC 9.0.0a42__py3-none-any.whl → 9.0.7__py3-none-any.whl

DIRAC/Core/Utilities/CGroups2.py

@@ -0,0 +1,328 @@
+#!/usr/bin/env python3
+"""cgroup2 support for DIRAC pilot."""
+
+import os
+import functools
+import subprocess
+from DIRAC import S_OK, S_ERROR, gLogger
+from DIRAC.Core.Utilities.DIRACSingleton import DIRACSingleton
+from DIRAC.Core.Utilities import Subprocess
+
+
+class CG2Manager(metaclass=DIRACSingleton):
+    """A class to manage cgroup2 hierachy for a typical pilot job use-case.
+
+    This creates a group for all of the pilot processes (anything in the
+    group at the start. This is a requirement for controlling the
+    sub-groups (no processes in non-leaf groups).
+
+    A group is then created on request for each "slot" under the pilot,
+    with the requested limits.
+    """
+
+    # Paths used to lookup cgroup info
+    FILE_MOUNTS = "/proc/mounts"
+    FILE_CUR_CGROUP = f"/proc/{os.getpid()}/cgroup"
+    # Control file names within the cgroup2 hierachy
+    CTRL_CONTROLLERS = "cgroup.controllers"
+    CTRL_PROCS = "cgroup.procs"
+    CTRL_SUBTREE = "cgroup.subtree_control"
+    CTRL_MEM_OOM_GROUP = "memory.oom.group"
+    CTRL_MEM_EVENTS = "memory.events"
+    CTRL_MEM_MAX = "memory.max"
+    CTRL_MEM_SWAP_MAX = "memory.swap.max"
+    CTRL_MEM_PEAK = "memory.peak"
+    CTRL_CPU_MAX = "cpu.max"
+    # CPU controller constants
+    # Weight is the max value for 1 CPU core
+    CPU_WEIGHT = 100000
+    # Period is the averaging time in us to apply the limit
+    # The default is 100k and I see no particularly reason this should change
+    CPU_PERIOD = 100000
+    # Name of the group for the existing pilot processes
+    PILOT_GROUP = f"dirac_pilot_{os.getpid()}"
+
+    def __init__(self):
+        """Set-up CGroup2 manager."""
+        # This boolean will be set to True if the cgroups are configured
+        # in the expected way
+        self._ready = False
+        # A counter of number of subgroups created
+        # Used to create unique group names
+        self._subproc_num = 0
+        # Physical path to the starting cgroup for this process
+        # (i.e. the base of our hierachy)
+        self._cgroup_path = None
+        # Logger
+        self.log = gLogger.getSubLogger("CG2Manager")
+
+    @staticmethod
+    def _filter_file(path, filterfcn):
+        """Opens a file and runs filterfcn for each line.
+        If filterfcn returns any value, that value will be returned
+        by this function.
+        Returns None if no line matches.
+        """
+        with open(path, encoding="ascii") as file_in:
+            for line in file_in.readlines():
+                line = line.strip()
+                if res := filterfcn(line):
+                    return res
+        return None
+
+    def _detect_root(self):
+        """Find the cgroup2 filesystem mountpoint on this system.
+        Returns the mountpoint path or None if it isn't found.
+        """
+
+        def filt(line):
+            """Filter function to find the first cgroup2 mount point
+            from a standard /proc/mounts layout file.
+            """
+            parts = line.split(" ")
+            if len(parts) < 3:
+                return None
+            if parts[2] == "cgroup2":
+                return parts[1]
+            return None
+
+        return self._filter_file(self.FILE_MOUNTS, filt)
+
+    def _detect_path(self):
+        """Finds the full physical path to the current cgroup control dir.
+        Sets self._cgroup_path on success.
+        Raises a RuntimeError if the path cannot be determined.
+        """
+
+        def filt(line):
+            """Filter to find the current cgroup2 name for the current
+            process, without the leading /.
+            """
+            if line.startswith("0::/"):
+                return line[4:]
+            return False
+
+        if not (root_path := self._detect_root()):
+            raise RuntimeError("Failed to find cgroup mount point")
+        if not (cur_group := self._filter_file(self.FILE_CUR_CGROUP, filt)):
+            raise RuntimeError("Failed to find current cgroup")
+        self._cgroup_path = os.path.join(root_path, cur_group)
+
+    def _create_group(self, group_name, isolate_oom=True):
+        """Creates a new group.
+        If "isolate_oom" is True, the new group will be decoupled
+        from the parent's OOM group.
+        Raises a RuntimeError if the group cannot be created.
+        """
+        try:
+            os.mkdir(os.path.join(self._cgroup_path, group_name))
+        except PermissionError as err:
+            raise RuntimeError(f"Permission denied creating sub-cgroup '{group_name}'") from err
+        if isolate_oom:
+            self._write_control(group_name, self.CTRL_MEM_OOM_GROUP, "0")
+
+    def _remove_group(self, group_name):
+        """Removes a group."""
+        os.rmdir(os.path.join(self._cgroup_path, group_name))
+
+    def _move_init_procs(self):
+        """Creates the pilot sub-group and moves all of the initial processes
+        from the top group into the new sub-group.
+        Will raise a RuntimeError if any cgroup configuration problem
+        prevents this from completing succesfully.
+        """
+        self._create_group(self.PILOT_GROUP, isolate_oom=False)
+        cur_pids = self._read_control("", self.CTRL_PROCS)
+        self._write_control(self.PILOT_GROUP, self.CTRL_PROCS, cur_pids)
+
+    def _read_control(self, group_name, ctrl_name):
+        """Reads a control value for the given group_name (relative to our base path).
+        The returned value varies depending on the value content:
+          - For a single token value, a string containing that token will be returned.
+          - For a single line value with space-seperated tokens, a list of tokens will be returned.
+          - For a multi-line value (where each line is a token), a list of tokens will be returned.
+        All tokens in the return values are strings.
+        A RuntimeError will be raised if the control cannot be read.
+        """
+        try:
+            with open(
+                os.path.join(self._cgroup_path, group_name, ctrl_name),
+                encoding="ascii",
+            ) as file_in:
+                values = [line.strip() for line in file_in.readlines()]
+                if " " in values and len(values) == 1:
+                    values = values[0].split(" ")
+                if len(values) == 1:
+                    values = values[0]
+                return values
+        except PermissionError as err:
+            raise RuntimeError(f"Access denied reading read control '{group_name}/{ctrl_name}'") from err
+
+    def _write_control(self, group_name, ctrl_name, value):
+        """Writes a control value for a given group_name (relative to our base path).
+        The value can be a string or an iterable of strings. The values should not
+        contain any whitespace characters.
+        A RuntimeError will be raised if the control cannot be set.
+        """
+        try:
+            ctrl_path = os.path.join(self._cgroup_path, group_name, ctrl_name)
+            with open(ctrl_path, "w", encoding="ascii") as file_out:
+                if isinstance(value, str):
+                    value = [value]
+                for arg in value:
+                    file_out.write(f"{arg}\n")
+                    # Flush is critical here as setting multiple values at the same time may fail
+                    file_out.flush()
+        except PermissionError as err:
+            raise RuntimeError(f"Access denied writing control '{group_name}/{ctrl_name}'") from err
+        except OSError as err:
+            # This generally happens if we're trying to set a value that is
+            # considered invalid, for example delegating a controller that isn't enabled
+            # in the first place.
+            raise RuntimeError(f"Error writing control '{group_name}/{ctrl_name}' = {value}") from err
+
+    def _get_oom_count(self, slot_name):
+        """Extracts the OOM counter as an int for the given slot.
+        Returns an int on success, can return a None if the memory.events
+        doesn't contain an oom counter or throws RuntimeError on failure.
+        """
+
+        def filt(line):
+            """Filter to find the oom counter from a memory.events file."""
+            if line.startswith("oom "):
+                return int(line[4:])
+            return False
+
+        mem_events = os.path.join(self._cgroup_path, slot_name, self.CTRL_MEM_EVENTS)
+        return self._filter_file(mem_events, filt)
+
+    def _set_limits(self, group_name, cores=None, memory=None, noswap=False):
+        """Sets the limits for an existing group.
+        See create_slot for a description of the other parameters.
+        This will raise a RuntimeError if appyling any of the limits fail to apply.
+        """
+        if cores:
+            proc_max = int(cores * self.CPU_WEIGHT)
+            self._write_control(group_name, self.CTRL_CPU_MAX, f"{proc_max} {self.CPU_PERIOD}")
+        if memory:
+            self._write_control(group_name, self.CTRL_MEM_MAX, f"{memory}")
+        if noswap:
+            self._write_control(group_name, self.CTRL_MEM_SWAP_MAX, "0")
+
+    def _prepare(self):
+        """Sets up the cgroup tree for the current process.
+        Should be called once, before using any of the other functions in this class.
+
+        Note that this function (specifcally the _move_init_procs call) assumes that
+        the list of processes is static. If the process list changes while this is running,
+        it is likely that this will fail to set things up properly.
+        """
+        self._detect_path()
+        controllers = self._read_control("", self.CTRL_CONTROLLERS)
+        if not controllers:
+            raise RuntimeError("No controllers enabled")
+        for ctrl in ["cpu", "memory"]:
+            if not ctrl in controllers:
+                raise RuntimeError(f"{ctrl} controller not enabled")
+        self._move_init_procs()
+        self._write_control("", self.CTRL_SUBTREE, ["+cpu", "+memory"])
+        self._ready = True
+
+    def _create_slot(self, slot_name, cores=None, memory=None, noswap=False):
+        """Creates a slot for a job with the given slot_name.
+        Cores is a float, number of CPU cores this group may use.
+        Memory is a string or int, either a number of bytes to limit the group RSS,
+        or a string limit with a unit suffix, e.g. "1G" as supported by the cgroup memory
+        controller.
+        If noswap is set to true, the swap memory limit will be set to 0; this is mostly
+        useful for testing (where the system may swap memory instead of triggering an
+        OOM, which may allow a process to use more than the memory limit).
+        This will raise a RuntimeError if setting up the slot fails.
+        """
+        if not self._ready:
+            return
+        self._create_group(slot_name)
+        self._set_limits(slot_name, cores, memory, noswap)
+
+    def _remove_slot(self, slot_name):
+        """Removes a slot with the given name.
+        Can raise usual filesystem OSError if the slot doesn't exist.
+        """
+        if not self._ready:
+            return
+        self._remove_group(slot_name)
+
+    def _setup_subproc(self, slot_name):
+        """A subprocess preexec function for setting up cgroups.
+        This will move te current process into the given cgroup slot.
+        On failure, no error will be reported.
+        """
+        # Threading danger!
+        # There are potential threading issues with preexec functions
+        # They must not hold any locks that the parent process might already
+        # be holding, including ones in standard library functions.
+        # This function should be kept as minimal as possible.
+        try:
+            self._write_control(slot_name, self.CTRL_PROCS, f"{os.getpid()}")
+        except Exception as err:
+            # We can't even really log here as we're in the set-up
+            # context of the new proces
+            pass
+
+    def setUp(self):
+        """Creates the base cgroup tree if possible. Should be called once
+        per process before using systemCall.
+        Returns S_OK/S_ERROR.
+        """
+        try:
+            self._prepare()
+        except Exception as err:
+            # The majority of CGroup failures will be RuntimeError
+            # However we don't want any unexpected failure to crash the upstream module,
+            # We just want to continue without cgroup support instead
+            return S_ERROR(str(err))
+        return S_OK()
+
+    def systemCall(self, *args, **kwargs):
+        """A proxy function for Subprocess.systemCall but will create a cgroup2 slot
+        if the functionality is available. An optional ceParameters dictionary
+        may be included, which will be searched for specific cgroup memory options.
+        Returns the usual S_OK/S_ERROR from Subprocess.systemCall.
+        """
+        preexec_fn = None
+        slot_name = f"subproc_{os.getpid()}_{self._subproc_num}"
+        self._subproc_num += 1
+        if self._ready:
+            self.log.info(f"Creating slot cgroup {slot_name}")
+            cores = None
+            memory = None
+            noswap = False
+            if "ceParameters" in kwargs:
+                if cpuLimit := kwargs["ceParameters"].get("CPULimit", None):
+                    cores = float(cpuLimit)
+                if memoryMB := int(kwargs["ceParameters"].get("MemoryLimitMB", 0)):
+                    memory = memoryMB * 1024 * 1024
+                if kwargs["ceParameters"].get("MemoryNoSwap", "no").lower() in ("yes", "true"):
+                    noswap = True
+            try:
+                self.log.info(f"CGroup Limits, CPU: {cores}, Mem: {memory}, NoSwap: {noswap}")
+                self._create_slot(slot_name, cores=cores, memory=memory, noswap=noswap)
+                preexec_fn = functools.partial(CG2Manager._setup_subproc, self, slot_name)
+            except Exception as err:
+                self.log.warn("Failed to create slot cgroup:", str(err))
+        kwargs["preexec_fn"] = preexec_fn
+        kwargs.pop("ceParameters", None)
+        res = Subprocess.systemCall(*args, **kwargs)
+        if self._ready:
+            self.log.info(f"Removing slot cgroup {slot_name}")
+            try:
+                oom_count = self._get_oom_count(slot_name)
+                if oom_count:
+                    # Child process triggered an OOM
+                    # We can't readily report this upstream (child process will probably
+                    # fail with an error code), so just log it and continue
+                    self.log.info(f"OOM detected from child process (slot {slot_name})")
+                self._remove_slot(slot_name)
+            except Exception as err:
+                self.log.warn(f"Failed to delete slot {slot_name} cgroup:", str(err))
+        return res

DIRAC/Core/Utilities/ClassAd/ClassAdLight.py

@@ -2,294 +2,8 @@
     Condor ClassAd library.
 """
 
+# Import from DIRACCommon for backward compatibility
+from DIRACCommon.Core.Utilities.ClassAd.ClassAdLight import ClassAd
 
-class ClassAd:
-    def __init__(self, jdl):
-        """ClassAd constructor from a JDL string"""
-        self.contents = {}
-        result = self.__analyse_jdl(jdl)
-        if result:
-            self.contents = result
-
-    def __analyse_jdl(self, jdl, index=0):
-        """Analyse one [] jdl enclosure"""
-
-        jdl = jdl.strip()
-
-        # Strip all the blanks first
-        # temp = jdl.replace(' ','').replace('\n','')
-        temp = jdl
-
-        result = {}
-
-        if temp[0] != "[" or temp[-1] != "]":
-            print("Invalid JDL: it should start with [ and end with ]")
-            return result
-
-        # Parse the jdl string now
-        body = temp[1:-1]
-        index = 0
-        namemode = 1
-        valuemode = 0
-        while index < len(body):
-            if namemode:
-                ind = body.find("=", index)
-                if ind != -1:
-                    name = body[index:ind]
-                    index = ind + 1
-                    valuemode = 1
-                    namemode = 0
-                else:
-                    break
-            elif valuemode:
-                ind1 = body.find("[", index)
-                ind2 = body.find(";", index)
-                if ind1 != -1 and ind1 < ind2:
-                    value, newind = self.__find_subjdl(body, ind1)
-                elif ind1 == -1 and ind2 == -1:
-                    value = body[index:]
-                    newind = len(body)
-                else:
-                    if index == ind2:
-                        return {}
-                    else:
-                        value = body[index:ind2]
-                        newind = ind2 + 1
-
-                result[name.strip()] = value.strip().replace("\n", "")
-                index = newind
-                valuemode = 0
-                namemode = 1
-
-        return result
-
-    def __find_subjdl(self, body, index):
-        """Find a full [] enclosure starting from index"""
-        result = ""
-        if body[index] != "[":
-            return (result, 0)
-
-        depth = 0
-        ind = index
-        while depth < 10:
-            ind1 = body.find("]", ind + 1)
-            ind2 = body.find("[", ind + 1)
-            if ind2 != -1 and ind2 < ind1:
-                depth += 1
-                ind = ind2
-            else:
-                if depth > 0:
-                    depth -= 1
-                    ind = ind1
-                else:
-                    result = body[index : ind1 + 1]
-                    if body[ind1 + 1] == ";":
-                        return (result, ind1 + 2)
-                    return result, 0
-
-        return result, 0
-
-    def insertAttributeInt(self, name, attribute):
-        """Insert a named integer attribute"""
-
-        self.contents[name] = str(attribute)
-
-    def insertAttributeBool(self, name, attribute):
-        """Insert a named boolean attribute"""
-
-        if attribute:
-            self.contents[name] = "true"
-        else:
-            self.contents[name] = "false"
-
-    def insertAttributeString(self, name, attribute):
-        """Insert a named string attribute"""
-
-        self.contents[name] = '"' + str(attribute) + '"'
-
-    def insertAttributeVectorString(self, name, attributelist):
-        """Insert a named string list attribute"""
-
-        tmp = ['"' + x + '"' for x in attributelist]
-        tmpstr = ",".join(tmp)
-        self.contents[name] = "{" + tmpstr + "}"
-
-    def insertAttributeVectorInt(self, name, attributelist):
-        """Insert a named string list attribute"""
-
-        tmp = [str(x) for x in attributelist]
-        tmpstr = ",".join(tmp)
-        self.contents[name] = "{" + tmpstr + "}"
-
-    def insertAttributeVectorStringList(self, name, attributelist):
-        """Insert a named list of string lists"""
-
-        listOfLists = []
-        for stringList in attributelist:
-            # tmp = map ( lambda x : '"' + x + '"', stringList )
-            tmpstr = ",".join(stringList)
-            listOfLists.append("{" + tmpstr + "}")
-        self.contents[name] = "{" + ",".join(listOfLists) + "}"
-
-    def lookupAttribute(self, name):
-        """Check the presence of the given attribute"""
-
-        return name in self.contents
-
-    def set_expression(self, name, attribute):
-        """Insert a named expression attribute"""
-
-        self.contents[name] = str(attribute)
-
-    def get_expression(self, name):
-        """Get expression corresponding to a named attribute"""
-
-        if name in self.contents:
-            if isinstance(self.contents[name], int):
-                return str(self.contents[name])
-            return self.contents[name]
-        return ""
-
-    def isAttributeList(self, name):
-        """Check if the given attribute is of the List type"""
-        attribute = self.get_expression(name).strip()
-        return attribute.startswith("{")
-
-    def getListFromExpression(self, name):
-        """Get a list of strings from a given expression"""
-
-        tempString = self.get_expression(name).strip()
-        listMode = False
-        if tempString.startswith("{"):
-            tempString = tempString[1:-1]
-            listMode = True
-
-        tempString = tempString.replace(" ", "").replace("\n", "")
-        if tempString.find("{") < 0:
-            if not listMode:
-                tempString = tempString.replace('"', "")
-                if not tempString:
-                    return []
-                return tempString.split(",")
-
-        resultList = []
-        while tempString:
-            if tempString.find("{") == 0:
-                end = tempString.find("}")
-                resultList.append(tempString[: end + 1])
-                tempString = tempString[end + 1 :]
-                if tempString.startswith(","):
-                    tempString = tempString[1:]
-            elif tempString.find('"') == 0:
-                end = tempString[1:].find('"')
-                resultList.append(tempString[1 : end + 1])
-                tempString = tempString[end + 2 :]
-                if tempString.startswith(","):
-                    tempString = tempString[1:]
-            else:
-                end = tempString.find(",")
-                if end < 0:
-                    resultList.append(tempString.replace('"', "").replace(" ", ""))
-                    break
-                else:
-                    resultList.append(tempString[:end].replace('"', "").replace(" ", ""))
-                    tempString = tempString[end + 1 :]
-
-        return resultList
-
-    def getDictionaryFromSubJDL(self, name):
-        """Get a dictionary of the JDL attributes from a subsection"""
-
-        tempList = self.get_expression(name)[1:-1]
-        resDict = {}
-        for item in tempList.split(";"):
-            if len(item.split("=")) == 2:
-                resDict[item.split("=")[0].strip()] = item.split("=")[1].strip().replace('"', "")
-            else:
-                return {}
-
-        return resDict
-
-    def deleteAttribute(self, name):
-        """Delete a named attribute"""
-
-        if name in self.contents:
-            del self.contents[name]
-            return 1
-        return 0
-
-    def isOK(self):
-        """Check the JDL validity - to be defined"""
-
-        if self.contents:
-            return 1
-        return 0
-
-    def asJDL(self):
-        """Convert the JDL description into a string"""
-
-        result = []
-        for name, value in sorted(self.contents.items()):
-            if value[0:1] == "{":
-                result += [4 * " " + name + " = \n"]
-                result += [8 * " " + "{\n"]
-                strings = value[1:-1].split(",")
-                for st in strings:
-                    result += [12 * " " + st.strip() + ",\n"]
-                result[-1] = result[-1][:-2]
-                result += ["\n" + 8 * " " + "};\n"]
-            elif value[0:1] == "[":
-                tempad = ClassAd(value)
-                tempjdl = tempad.asJDL() + ";"
-                lines = tempjdl.split("\n")
-                result += [4 * " " + name + " = \n"]
-                for line in lines:
-                    result += [8 * " " + line + "\n"]
-
-            else:
-                result += [4 * " " + name + " = " + str(value) + ";\n"]
-        if result:
-            result[-1] = result[-1][:-1]
-        return "[ \n" + "".join(result) + "\n]"
-
-    def getAttributeString(self, name):
-        """Get String type attribute value"""
-        value = ""
-        if self.lookupAttribute(name):
-            value = self.get_expression(name).replace('"', "")
-        return value
-
-    def getAttributeInt(self, name):
-        """Get Integer type attribute value"""
-        value = None
-        if self.lookupAttribute(name):
-            try:
-                value = int(self.get_expression(name).replace('"', ""))
-            except Exception:
-                value = None
-        return value
-
-    def getAttributeBool(self, name):
-        """Get Boolean type attribute value"""
-        if not self.lookupAttribute(name):
-            return False
-
-        value = self.get_expression(name).replace('"', "")
-        return value.lower() == "true"
-
-    def getAttributeFloat(self, name):
-        """Get Float type attribute value"""
-        value = None
-        if self.lookupAttribute(name):
-            try:
-                value = float(self.get_expression(name).replace('"', ""))
-            except Exception:
-                value = None
-        return value
-
-    def getAttributes(self) -> list[str]:
-        """Get the list of all the attribute names
-
-        :return: list of names as strings
-        """
-        return list(self.contents)
+# Re-export for backward compatibility
+__all__ = ["ClassAd"]