config_utils.py

"""Loads configurations from .yaml files and expands environment variables.
"""
import copy
import collections
import glob
import math
import os
import pprint
import sys
import yaml

import toolz as tz

class CmdNotFound(Exception):
    pass

# ## Generalized configuration

def update_w_custom(config, lane_info):
    """Update the configuration for this lane if a custom analysis is specified.
    """
    name_remaps = {"variant": ["SNP calling", "variant", "variant2"],
                   "SNP calling": ["SNP calling", "variant", "variant2"],
                   "variant2": ["SNP calling", "variant", "variant2"]}
    config = copy.deepcopy(config)
    base_name = lane_info.get("analysis")
    if "algorithm" not in config:
        config["algorithm"] = {}
    for analysis_type in name_remaps.get(base_name, [base_name]):
        custom = config.get("custom_algorithms", {}).get(analysis_type)
        if custom:
            for key, val in custom.iteritems():
                config["algorithm"][key] = val
    # apply any algorithm details specified with the lane
    for key, val in lane_info.get("algorithm", {}).iteritems():
        config["algorithm"][key] = val
    # apply any resource details specified with the lane
    for prog, pkvs in lane_info.get("resources", {}).iteritems():
        if prog not in config["resources"]:
            config["resources"][prog] = {}
        for key, val in pkvs.iteritems():
            config["resources"][prog][key] = val
    return config

# ## Retrieval functions

def load_system_config(config_file=None, work_dir=None):
    """Load bcbio_system.yaml configuration file, handling standard defaults.

    Looks for configuration file in default location within
    final base directory from a standard installation. Handles both standard
    installs (galaxy/bcbio_system.yaml) and docker installs (config/bcbio_system.yaml).
    """
    docker_config = _get_docker_config()
    if config_file is None:
        config_file = "bcbio_system.yaml"
    if not os.path.exists(config_file):
        base_dir = get_base_installdir()
        test_config = os.path.join(base_dir, "galaxy", config_file)
        if os.path.exists(test_config):
            config_file = test_config
        else:
            raise ValueError("Could not find input system configuration file %s, "
                             "including inside standard directory %s" %
                             (config_file, os.path.join(base_dir, "galaxy")))
    config = load_config(config_file)
    if docker_config:
        assert work_dir is not None, "Need working directory to merge docker config"
        config_file = os.path.join(work_dir, "%s-merged%s" % os.path.splitext(os.path.basename(config_file)))
        config = _merge_system_configs(config, docker_config, config_file)
    if "algorithm" not in config:
        config["algorithm"] = {}
    config["bcbio_system"] = config_file
    return config, config_file

def get_base_installdir():
    return os.path.normpath(os.path.join(os.path.realpath(sys.executable), os.pardir, os.pardir, os.pardir))

def _merge_system_configs(host_config, container_config, out_file=None):
    """Create a merged system configuration from external and internal specification.
    """
    out = copy.deepcopy(container_config)
    for k, v in host_config.iteritems():
        if k in set(["galaxy_config"]):
            out[k] = v
        elif k == "resources":
            for pname, resources in v.iteritems():
                if not isinstance(resources, dict) and pname not in out[k]:
                    out[k][pname] = resources
                else:
                    for rname, rval in resources.iteritems():
                        if (rname in set(["cores", "jvm_opts", "memory"])
                              or pname in set(["gatk", "mutect"])):
                            if pname not in out[k]:
                                out[k][pname] = {}
                            out[k][pname][rname] = rval
    # Ensure final file is relocatable by mapping back to reference directory
    if "bcbio_system" in out and ("galaxy_config" not in out or not os.path.isabs(out["galaxy_config"])):
        out["galaxy_config"] = os.path.normpath(os.path.join(os.path.dirname(out["bcbio_system"]),
                                                             os.pardir, "galaxy",
                                                             "universe_wsgi.ini"))
    if out_file:
        with open(out_file, "w") as out_handle:
            yaml.safe_dump(out, out_handle, default_flow_style=False, allow_unicode=False)
    return out

def _get_docker_config():
    base_dir = get_base_installdir()
    docker_configfile = os.path.join(base_dir, "config", "bcbio_system.yaml")
    if os.path.exists(docker_configfile):
        return load_config(docker_configfile)

def merge_resources(args):
    """Merge docker local resources and global resource specification in a set of arguments.

    Finds the `data` object within passed arguments and updates the resources
    from a local docker configuration if present.
    """
    docker_config = _get_docker_config()
    if not docker_config:
        return args
    else:
        def _update_resources(config):
            config["resources"] = _merge_system_configs(config, docker_config)["resources"]
            return config
        return _update_config(args, _update_resources)

def load_config(config_file):
    """Load YAML config file, replacing environmental variables.
    """
    with open(config_file) as in_handle:
        config = yaml.load(in_handle)
    config = _expand_paths(config)
    # lowercase resource names, the preferred way to specify, for back-compatibility
    newr = {}
    for k, v in config["resources"].iteritems():
        if k.lower() != k:
            newr[k.lower()] = v
    config["resources"].update(newr)
    return config

def _expand_paths(config):
    for field, setting in config.items():
        if isinstance(config[field], dict):
            config[field] = _expand_paths(config[field])
        else:
            config[field] = expand_path(setting)
    return config

def expand_path(path):
    """ Combines os.path.expandvars with replacing ~ with $HOME.
    """
    try:
        return os.path.expandvars(path.replace("~", "$HOME"))
    except AttributeError:
        return path

def get_resources(name, config):
    """Retrieve resources for a program, pulling from multiple config sources.
    """
    return tz.get_in(["resources", name], config,
                     tz.get_in(["resources", "default"], config, {}))

def get_program(name, config, ptype="cmd", default=None):
    """Retrieve program information from the configuration.

    This handles back compatible location specification in input
    YAML. The preferred location for program information is in
    `resources` but the older `program` tag is also supported.
    """
    # support taking in the data dictionary
    config = config.get("config", config)
    try:
        pconfig = config.get("resources", {})[name]
        # If have leftover old
    except KeyError:
        pconfig = {}
    old_config = config.get("program", {}).get(name, None)
    if old_config:
        for key in ["dir", "cmd"]:
            if not key in pconfig:
                pconfig[key] = old_config
    if ptype == "cmd":
        return _get_program_cmd(name, pconfig, config, default)
    elif ptype == "dir":
        return _get_program_dir(name, pconfig)
    else:
        raise ValueError("Don't understand program type: %s" % ptype)

def _get_check_program_cmd(fn):
    def wrap(name, pconfig, config, default):
        program = expand_path(fn(name, pconfig, config, default))
        is_ok = lambda f: os.path.isfile(f) and os.access(f, os.X_OK)
        if is_ok(program): return program

        for adir in os.environ['PATH'].split(":"):
            if is_ok(os.path.join(adir, program)):
                return os.path.join(adir, program)
        # support bioconda installed programs
        if is_ok(os.path.join(os.path.dirname(sys.executable), name)):
            return (os.path.join(os.path.dirname(sys.executable), name))
        # find system bioconda installed programs if using private code install
        bcbio_system = config.get("bcbio_system", None)
        if bcbio_system:
            system_bcbio_path = os.path.join(os.path.dirname(bcbio_system),
                                             os.pardir, "anaconda", "bin", name)
            if is_ok(system_bcbio_path):
                return system_bcbio_path
        raise CmdNotFound(" ".join(map(repr, (fn.func_name, name, pconfig, default))))
    return wrap

@_get_check_program_cmd
def _get_program_cmd(name, pconfig, config, default):
    """Retrieve commandline of a program.
    """
    if pconfig is None:
        return name
    elif isinstance(pconfig, basestring):
        return pconfig
    elif "cmd" in pconfig:
        return pconfig["cmd"]
    elif default is not None:
        return default
    else:
        return name

def _get_program_dir(name, config):
    """Retrieve directory for a program (local installs/java jars).
    """
    if config is None:
        raise ValueError("Could not find directory in config for %s" % name)
    elif isinstance(config, basestring):
        return config
    elif "dir" in config:
        return expand_path(config["dir"])
    else:
        raise ValueError("Could not find directory in config for %s" % name)

def get_jar(base_name, dname):
    """Retrieve a jar in the provided directory
    """
    jars = glob.glob(os.path.join(expand_path(dname), "%s*.jar" % base_name))

    if len(jars) == 1:
        return jars[0]
    elif len(jars) > 1:
        raise ValueError("Found multiple jars for %s in %s. Need single jar: %s" %
                         (base_name, dname, jars))
    else:
        raise ValueError("Could not find java jar %s in %s" %
                         (base_name, dname))

# ## Retrieval and update to configuration from arguments

def is_std_config_arg(x):
    return isinstance(x, dict) and "algorithm" in x and "resources" in x and "files" not in x

def is_nested_config_arg(x):
    return isinstance(x, dict) and "config" in x and is_std_config_arg(x["config"])

def get_algorithm_config(xs):
    """Flexibly extract algorithm configuration for a sample from any function arguments.
    """
    if isinstance(xs, dict):
        xs = [xs]
    for x in xs:
        if is_std_config_arg(x):
            return x["algorithm"]
        elif is_nested_config_arg(x):
            return x["config"]["algorithm"]
        elif isinstance(x, (list, tuple)) and is_nested_config_arg(x[0]):
            return x[0]["config"]["algorithm"]
    raise ValueError("Did not find algorithm configuration in items: {0}"
                     .format(pprint.pformat(xs)))

def get_dataarg(args):
    """Retrieve the world 'data' argument from a set of input parameters.
    """
    for i, arg in enumerate(args):
        if is_nested_config_arg(arg):
            return i, arg
        elif is_std_config_arg(arg):
            return i, {"config": arg}
        elif isinstance(arg, (list, tuple)) and is_nested_config_arg(arg[0]):
            return i, arg[0]
    raise ValueError("Did not find configuration or data object in arguments: %s" % args)

def add_cores_to_config(args, cores_per_job, parallel=None):
    """Add information about available cores for a job to configuration.
    Ugly hack to update core information in a configuration dictionary.
    """
    def _update_cores(config):
        config["algorithm"]["num_cores"] = int(cores_per_job)
        if parallel:
            parallel.pop("view", None)
            config["parallel"] = parallel
        return config
    return _update_config(args, _update_cores)

def _update_config(args, update_fn):
    """Update configuration, nested in argument list, with the provided update function.
    """
    new_i = None
    for i, arg in enumerate(args):
        if (is_std_config_arg(arg) or is_nested_config_arg(arg) or
              (isinstance(arg, (list, tuple)) and is_nested_config_arg(arg[0]))):
            new_i = i
            break
    if new_i is None:
        raise ValueError("Could not find configuration in args: %s" % str(args))

    new_arg = args[new_i]
    if is_nested_config_arg(new_arg):
        new_arg["config"] = update_fn(copy.deepcopy(new_arg["config"]))
    elif is_std_config_arg(new_arg):
        new_arg = update_fn(copy.deepcopy(new_arg))
    elif isinstance(arg, (list, tuple)) and is_nested_config_arg(new_arg[0]):
        new_arg_first = new_arg[0]
        new_arg_first["config"] = update_fn(copy.deepcopy(new_arg_first["config"]))
        new_arg = [new_arg_first] + new_arg[1:]
    else:
        raise ValueError("Unexpected configuration dictionary: %s" % new_arg)
    args = list(args)[:]
    args[new_i] = new_arg
    return args

def convert_to_bytes(mem_str):
    """Convert a memory specification, potentially with M or G, into bytes.
    """
    if str(mem_str)[-1].upper().endswith("G"):
        return int(mem_str[:-1]) * 1024 * 1024
    elif str(mem_str)[-1].upper().endswith("M"):
        return int(mem_str[:-1]) * 1024
    else:
        return int(mem_str)

def adjust_memory(val, magnitude, direction="increase", out_modifier=""):
    """Adjust memory based on number of cores utilized.
    """
    modifier = val[-1:]
    amount = int(val[:-1])
    if direction == "decrease":
        new_amount = amount / float(magnitude)
        # dealing with a specifier like 1G, need to scale to Mb
        if new_amount < 1:
            if modifier.upper().startswith("G"):
                new_amount = (amount * 1024) / magnitude
                modifier = "M" + modifier[1:]
            else:
                raise ValueError("Unexpected decrease in memory: %s by %s" % (val, magnitude))
        amount = int(new_amount)
    elif direction == "increase":
        # for increases with multiple cores, leave small percentage of
        # memory for system to maintain process running resource and
        # avoid OOM killers
        adjuster = 0.91
        amount = int(math.ceil(amount * (adjuster * magnitude)))
    if out_modifier.upper().startswith("G") and modifier.upper().startswith("M"):
        modifier = out_modifier
        amount = int(math.floor(amount / 1024.0))
    if out_modifier.upper().startswith("M") and modifier.upper().startswith("G"):
        modifier = out_modifier
        modifier = int(amount * 1024)
    return "{amount}{modifier}".format(amount=amount, modifier=modifier)

def adjust_opts(in_opts, config):
    """Establish JVM opts, adjusting memory for the context if needed.

    This allows using less or more memory for highly parallel or multicore
    supporting processes, respectively.
    """
    memory_adjust = config["algorithm"].get("memory_adjust", {})
    out_opts = []
    for opt in in_opts:
        if opt.startswith("-Xmx") or (opt.startswith("-Xms") and memory_adjust.get("direction") == "decrease"):
            arg = opt[:4]
            opt = "{arg}{val}".format(arg=arg,
                                      val=adjust_memory(opt[4:],
                                                        memory_adjust.get("magnitude", 1),
                                                        memory_adjust.get("direction")))
        out_opts.append(opt)
    return out_opts

# specific program usage

def use_vqsr(algs):
    """Processing uses GATK's Variant Quality Score Recalibration.
    """
    vqsr_callers = set(["gatk", "gatk-haplotype"])
    vqsr_sample_thresh = 50
    vqsr_supported = collections.defaultdict(int)
    coverage_intervals = set([])
    for alg in algs:
        callers = alg.get("variantcaller", "gatk")
        if isinstance(callers, basestring):
            callers = [callers]
        if not callers:  # no variant calling, no VQSR
            continue
        if "vqsr" in alg.get("tools_off", []):  # VQSR turned off
            continue
        for c in callers:
            if c in vqsr_callers:
                vqsr_supported[c] += 1
                coverage_intervals.add(alg.get("coverage_interval", "exome").lower())
    if len(vqsr_supported) > 0:
        num_samples = max(vqsr_supported.values())
        if "genome" in coverage_intervals or num_samples >= vqsr_sample_thresh:
            return True
    return False

def use_snpeff(algs):
    """Processing uses snpEff. Avoids memory requirements if not used.
    """
    return any(alg.get("effects", "snpeff") == "snpeff" and alg.get("variantcaller") for alg in algs)

def use_bcbio_variation_recall(algs):
    """Processing uses bcbio-variation-recall. Avoids core requirement if not used.
    """
    for alg in algs:
        jointcaller = alg.get("jointcaller", [])
        if not isinstance(jointcaller, (tuple, list)):
            jointcaller = [jointcaller]
        for caller in jointcaller:
            if caller not in set(["gatk-haplotype-joint", None, False]):
                return True
    return False

## functions for navigating through the standard galaxy directory of files

def get_rRNA_interval(genome_dir):
    return os.path.join(genome_dir, "rnaseq", "rRNA.interval_list")

def get_transcript_refflat(genome_dir):
    return os.path.join(genome_dir, "rnaseq", "ref-transcripts.refFlat")

def get_rRNA_sequence(genome_dir):
    return os.path.join(genome_dir, "rnaseq", "rRNA.fa")

def program_installed(program, data):
    """
    returns True if the path to a program can be found
    """
    try:
        path = get_program(program, data)
    except CmdNotFound:
        return False
    return True