Simple python script to limit I/O bandwidth of a container

docker_io_limt.py

#!/bin/python

""" 
This script uses cgroups to limit a docker container's read/write bandwidth

Usage:
    limit_docker_io.py <container_id> --read=<read_bps> --write=<write_bps>

Options:
    --read=<read_bps>               Read bps 0 for unlimited [default: 0]
    --write=<write_bps>             Write bps 0 for unlimited [default: 0]
"""

import subprocess
import logging
import os.path
from docopt import docopt

logging.basicConfig()
log = logging.getLogger('limit_io')


def get_mounts():

    out = subprocess.check_output("mount", shell=True).split('\n')
    mounts = list()
    for line in out:
        if line:
            elem = line.split()
            mounts.append({
                'device': elem[0],
                'mountpoint': elem[2]
            })

    return mounts



def get_device_major_minor(devname):
    """
    Device name is the last name 
    eg: /dev/xvda will have dename xvda
    """

    dev_info = open("/proc/partitions")
    for line in dev_info:
        if line.strip():
            elem = line.split()
            if devname == elem[3]:
                return elem[0],elem[1]

    return ""

def limit_container_io(container_id, read_bps, write_bps):

    cgroup_path = "/sys/fs/cgroup/blkio/docker/%s" % container_id

    if not os.path.exists(cgroup_path):
        log.error("Cgroup not found")
        return 
    
    mounts = get_mounts()
    device = None
    device = [m['device'] for m in mounts if m['mountpoint'] == '/']
    device = [m['device'] for m in mounts if m['mountpoint'] == '/var/'] or device
    device = [m['device'] for m in mounts if m['mountpoint'] == '/var/lib'] or device
    device = [m['device'] for m in mounts if m['mountpoint'] == '/var/lib/docker']  or device

    if not device:
        log.error("No device found (strange)")
        return
    
    # get the last element from device
    devname = device[0].split('/')[-1]

    # you can limit io on the base device and not on partition ... get that
    # XXX: assumption: works on parition of the form /dev/sda<num> where num is
    # 1-9 will have problems with dev-mapper based devices as they don't follow
    # that convenction

    devname = devname[:-1]
    major,minor = get_device_major_minor(devname)

    if not major or not minor:
        log.error("Major, minor number not found for %s"% devname)
        return

    
    if read_bps > 0:
        read_bps_cg_file = os.path.join(cgroup_path, "blkio.throttle.read_bps_device")
        conf_str = "%s:%s %s" % (major,minor,read_bps)
        subprocess.call("echo %s > %s" % (conf_str, read_bps_cg_file), shell=True)

    if write_bps > 0:
        write_bps_cg_file = os.path.join(cgroup_path, "blkio.throttle.write_bps_device")
        conf_str = "%s:%s %s" % (major,minor,read_bps)
        subprocess.call("echo %s > %s" % (conf_str, read_bps_cg_file), shell=True)


if __name__ == "__main__":

    args = docopt(__doc__)

    limit_container_io(
            args['<container_id>'], 
            args['--write'], 
            args['--read'],
    )