subprograms.py 8.71 KB
Newer Older
1 2 3
"""Various clients for system utilities."""
import collections
import logging
4
import xml.etree.ElementTree
5 6
import tornado.process as process
import subprocess
7

8
logger = logging.getLogger('nrm')
9 10 11 12 13
resources = collections.namedtuple("Resources", ["cpus", "mems"])


def logpopen(p, args, stdout, stderr):
    """log popen cmd."""
14 15 16 17
    logger.debug("popen cmd: %r", args)
    logger.debug("popen return code: %s", p.returncode)
    logger.debug("popen stdout: %r", stdout)
    logger.debug("popen, stderr: %r", stderr)
18 19 20 21 22 23 24 25 26 27 28 29 30 31 32 33 34


def bitmask2list(mask):
    """Convert a bitmask to the list of power of 2 set to 1."""
    i = int(mask or '0x0', base=16)
    ret = []
    for j in range(i.bit_length()):
        m = 1 << j
        if (i & m):
            ret.append(j)
    return ret


def list2bitmask(l):
    """Convert a list into a bitmask."""
    m = 0
    for e in l:
35
        m |= 1 << int(e)
36 37 38 39 40 41 42
    return hex(m)


class NodeOSClient(object):

    """Client to argo_nodeos_config."""

43
    def __init__(self, argo_nodeos_config="argo_nodeos_config"):
44
        """Load client configuration."""
45
        self.prefix = argo_nodeos_config
46 47 48 49 50 51 52 53 54 55 56 57 58 59 60 61 62 63 64 65 66 67 68 69 70 71 72 73 74 75 76 77 78 79 80 81 82 83 84 85 86 87 88 89 90 91 92 93 94 95 96 97 98 99 100 101 102 103 104 105 106

    def getavailable(self):
        """Gather available resources."""
        args = [self.prefix, "--show_available_resources=shared:false"]
        p = subprocess.Popen(args, stdout=subprocess.PIPE,
                             stderr=subprocess.PIPE)
        stdout, stderr = p.communicate()
        logpopen(p, args, stdout, stderr)
        # parse the format: first line is threads, then a list as multiline,
        # then nodes, and the same
        cpus = []
        mems = []
        lines = stdout.splitlines()
        splitindex = lines.index('------------Memory nodes------------')
        cpuslines = lines[1:splitindex]
        memlines = lines[splitindex+1:]
        for l in cpuslines:
            cpus.extend(l.split())
        for l in memlines:
            mems.extend(l.split())
        return resources([int(x) for x in cpus], [int(x) for x in mems])

    def create(self, name, params):
        """Create container, according to params."""
        args = [self.prefix]
        cmd = "--create_container="
        cmd += 'name:{0}'.format(name)
        cmd += ' cpus:[{0}]'.format(",".join([str(x) for x in params.cpus]))
        cmd += ' mems:[{0}]'.format(",".join([str(x) for x in params.mems]))
        args.append(cmd)
        p = subprocess.Popen(args, stdout=subprocess.PIPE,
                             stderr=subprocess.PIPE)
        stdout, stderr = p.communicate()
        logpopen(p, args, stdout, stderr)

    def attach(self, name, pid):
        """Attach a pid to a container."""
        args = [self.prefix]
        cmd = '--attach_to_container='
        cmd += 'name:{0}'.format(name)
        cmd += ' pids:[{0}]'.format(pid)
        args.append(cmd)
        p = subprocess.Popen(args, stdout=subprocess.PIPE,
                             stderr=subprocess.PIPE)
        stdout, stderr = p.communicate()
        logpopen(p, args, stdout, stderr)

    def delete(self, name, kill=False):
        """Destroy container."""
        # destroy container
        args = [self.prefix]
        cmd = '--delete_container='
        cmd += 'name:{0}'.format(name)
        if kill:
            cmd += ' kill_content:true'
        args.append(cmd)
        p = subprocess.Popen(args, stdout=subprocess.PIPE,
                             stderr=subprocess.PIPE)
        stdout, stderr = p.communicate()
        logpopen(p, args, stdout, stderr)

107 108 109 110 111 112 113 114 115 116
    def execute(self, name, argv, environ):
        """Execute argv inside container."""
        args = [self.prefix]
        cmd = '--exec='
        cmd += 'name:{0}'.format(name)
        # argo_nodeos_config takes argv as 'arg0 arg1 ...' so we need to merge
        # the arguments into a single list, with single quotes. We also need
        # to escape spaces from arguments before.
        argv = [s.replace(' ', r'\ ') for s in argv]
        cmd += " argv:'"+" ".join(argv)+"'"
117 118 119
        env = ['{0}={1}'.format(envname, envval.replace(' ', r'\ '))
               for envname, envval in environ.items()]
        cmd += " env:'"+" ".join(env)+"'"
120 121 122
        args.append(cmd)
        return process.Subprocess(args, stdin=process.Subprocess.STREAM,
                                  stdout=process.Subprocess.STREAM,
123
                                  stderr=process.Subprocess.STREAM)
124

125 126 127 128 129 130 131 132 133 134 135 136 137 138 139 140 141 142 143 144 145 146 147 148 149 150 151 152 153

class ChrtClient(object):

    """Client to chrt command line wrapper."""

    flags = {'SCHED_OTHER': '--other',
             'SCHED_BATCH': '--batch',
             'SCHED_FIFO': '--fifo',
             'SCHED_IDLE': '--idle',
             'SCHED_RR': '--rr',
             'SCHED_HPC': '--hpc'
             }

    def __init__(self):
        """Load configuration."""
        self.prefix = "chrt"

    def getwrappedcmd(self, params):
        """Return a list of args to prepend to a popen call."""
        args = [self.prefix]
        args.append(self.flags[params.policy])
        args.append(params.priority)
        return args


class HwlocClient(object):

    """Client to hwloc binaries."""

154
    def __init__(self, hwloc="hwloc"):
155
        """Load configuration."""
156
        self.prefix = hwloc
157 158 159

    def info(self):
        """Return list of all cpus and mems."""
160 161
        cmd = self.prefix + "-ls"
        args = [cmd, '--whole-system', '--output-format', 'xml']
162 163 164 165
        p = subprocess.Popen(args, stdout=subprocess.PIPE,
                             stderr=subprocess.PIPE)
        stdout, stderr = p.communicate()
        logpopen(p, args, stdout, stderr)
166
        xmlroot = xml.etree.ElementTree.fromstring(stdout)
167
        ret = resources([], [])
168 169 170 171 172 173 174 175
        for obj in xmlroot.iter('object'):
            if obj.attrib['type'] == "NUMANode":
                ret.mems.append(int(obj.attrib['os_index']))
            if obj.attrib['type'] == "PU":
                ret.cpus.append(int(obj.attrib['os_index']))
        # if there's only one memory node, hwloc doesn't list it
        if not ret.mems:
            ret.mems.append(0)
176 177 178 179 180 181 182 183 184 185 186 187 188 189 190 191 192 193 194 195 196 197 198 199 200 201 202 203 204 205 206 207 208 209 210 211 212 213 214 215 216 217 218 219 220 221 222 223 224 225 226 227 228 229 230 231 232 233 234 235 236 237 238 239 240 241 242 243
        return ret

    def all2fake(self, resources):
        """Convert resource description of the system into fake topology.

        We need that because hwloc barfs on fake numa nodes.
        """
        # easy version: we have as many numa nodes as we have cores
        mems = len(resources.mems)
        cpus = len(resources.mems)
        assert cpus % mems == 0
        pu = cpus // mems
        return "numa: %s pu:%s".format(mems, pu)

    def distrib(self, numprocs, restrict=None, fake=None):
        """Distribute numprocs across the hierarchy."""
        # The original command only reports back cpusets. We do better, by
        # reporting the mems that go with it. This requires some magic, using
        # hwloc-ls to find the numa node associated with a cpuset reported by
        # distrib
        allresources = self.info()
        cmd = [self.prefix + "-distrib"]
        if fake:
            cmd.extend(['-i', self.all2fake(fake)])
        args = cmd + ["--whole-system", "--taskset", str(numprocs)]
        if restrict:
            mask = list2bitmask(restrict.cpus)
            args.extend(['--restrict', mask])
        p = subprocess.Popen(args, stdout=subprocess.PIPE,
                             stderr=subprocess.PIPE)
        stdout, stderr = p.communicate()
        logpopen(p, args, stdout, stderr)
        cpusets = stdout.splitlines()
        dret = {}
        for c in cpusets:
            dret[c] = resources(bitmask2list(c), [])

        # list all resources, and display cpusets too
        # this will give us the memories associated with each cpuset.
        cmd = [self.prefix + "-ls"]
        if fake:
            cmd.extend(['-i', self.all2fake(fake)])
        args = cmd + ["--whole-system", "-c", "--taskset"]
        p = subprocess.Popen(args, stdout=subprocess.PIPE,
                             stderr=subprocess.PIPE)
        stdout, stderr = p.communicate()
        logpopen(p, args, stdout, stderr)
        lines = stdout.splitlines()
        curmem = allresources.mems
        for l in lines:
            pos = l.find('cpuset=')
            if pos != -1:
                c = l[l.find('cpuset='):].lstrip('cpuset=')
                numa = l.find('NUMANode')
                cset = set(bitmask2list(c))
                if numa != -1:
                    uid = int(l.split()[1].lstrip('L#'))
                    curmem = [uid]
                for mask in dret:
                    cs = set(bitmask2list(mask))
                    if cset.issubset(cs):
                        dret[mask].mems.extend(curmem)
        # At this point, we have valid cpusets, but the mems associated are not
        # restricted, and not necessarily the right amount. We need to:
        #    - remove memories for the restricted set
        #    - split each (cpuset, mems) that is too big into a list of memset
        #    choices of the right size
        return dret.values()