workflow.py

import os
import argparse
import sys
import json
import datetime
import textwrap
import xml.dom.minidom
import yaml
import traceback
import chipathlon
import chipathlon.workflow_job
import chipathlon.db
import chipathlon.workflow_module
import random
from pprint import pprint
from Pegasus.DAX3 import *

class Workflow(object):

    def __init__(self, jobhome, run_file, param_file, host, username, password):
        # Initialize db connection
        self.mdb = chipathlon.db.MongoDB(host, username, password)
        # Jobname info & err
        self.jobhome = os.path.abspath(jobhome)
        self.jobname = os.path.basename(os.path.dirname(self.jobhome + "/"))
        self.err = ""
        # Input file info
        self.run_file = run_file
        self.param_file = param_file
        # Dax specific info
        self.dax = ADAG(self.jobname)
        self.executables = {}
        self.files = {}
        self.jobs = {}
        self.deps = {}
        self.workflow_jobs = {}
        self.modules = {}
        return

    def info(self):
        pprint(self.run_data)
        pprint(self.workflow_jobs)
        pprint(self.modules)
        pprint(self.executables)
        return

    def generate(self):
        try:
            self._load_executables()
            self._load_workflow_jobs()
            self._load_modules()
            self._load_runs()
            # All required information is loaded, start queuing jobs
            self._add_download()
            self._add_align()
            # Create pegasus important stuff
            self._create_setup()
            self._add_notify()
            self._create_replica()
            self._create_sites()
            self._create_submit()
            self._write()
        except SystemExit as se:
            print self.err
        if self.err:
            print self.err
            sys.exit(1)
        return

    def _raise(self):
        if self.err:
            raise SystemExit(self.err)
        return

    def _load_executables(self, os_type="linux", arch="x86_64"):
        # Load wrapper scripts for commands that need to be loaded from module
        for root, dirs, files in os.walk("%s/%s" % (os.path.dirname(os.path.realpath(__file__)), chipathlon.conf.job_wrappers)):
            for f in files:
                ex_name = f.split("_")[0]
                self.executables[ex_name] = Executable(name=ex_name, os=os_type, arch=arch)
                self.executables[ex_name].addPFN(PFN("file://%s/%s" % (root, f), "condorpool"))
                self.dax.addExecutable(self.executables[ex_name])
            break
        # Load actual scripts
        for root, dirs, files in os.walk("%s/%s" % (os.path.dirname(os.path.realpath(__file__)), chipathlon.conf.job_scripts)):
            for f in files:
                self.executables[f] = Executable(name=f, os=os_type, arch=arch)
                self.executables[f].addPFN(PFN("file://%s/%s" % (root, f), "condorpool"))
                self.dax.addExecutable(self.executables[f])
            break
        self._raise()
        return

    def _load_workflow_jobs(self):
        for root, dirs, files in os.walk("%s/%s" % (os.path.dirname(os.path.realpath(__file__)), chipathlon.conf.job_params)):
            for f in files:
                ex_name = f.split("_")[0]
                if ex_name in self.executables:
                    yj = chipathlon.workflow_job.WorkflowJob(os.path.join(root, f), self.param_file, self.executables[ex_name])
                    if not yj.err:
                        self.workflow_jobs[yj.jobname] = yj
                    else:
                        self.err += yj.err
                else:
                    print "[INFO] Skipping param file %s, no corresponding executable found." % (f,)
            break
        self._raise()
        return

    def _load_modules(self):
        for root, dirs, files in os.walk(os.path.join(os.path.dirname(os.path.realpath(__file__)), chipathlon.conf.job_modules)):
            for f in files:
                mod = chipathlon.workflow_module.WorkflowModule(os.path.join(root, f), self.workflow_jobs)
                if not mod.err:
                    self.modules[mod.name] = mod
                else:
                    self.err += mod.err
            break
        self._raise()
        return

    def _load_runs(self):
        with open(self.run_file, "r") as rh:
            try:
                self.run_data = yaml.load(rh)
            except yaml.YAMLError as exc:
                self.err += "Error parsing run template file '%s'. %s.\n" % (self.run_file, exc)
        if not self.err:
            checked = []
            for run in self.run_data["runs"]:
                for module_name in self.modules:
                    if isinstance(self.modules[module_name], dict):
                        keys = self.modules[module_name].keys()
                        if len(keys) > 1:
                            if module_name in run:
                                if run[module_name] not in keys:
                                    self.err += "Error parsing run template file '%s'. Module '%s' has invalid entry '%s', should be one of %s.\n" % (self.run_file, module_name, run[module_name], keys)
                            else:
                                self.err += "Error parsing run template file '%s'.  Module '%s' requires an entry, should be one of %s.\n" % (self.run_file, module_name, keys)
                if "experiment" in run:
                    valid, msg, data = self.mdb.get_samples(run["experiment"])
                    if valid:
                        run["data"] = data
                    else:
                        self.err += msg
                    # This will be changed to _check_genomes(run) eventually.
                    # self._check_genomes(run)
                    # UNCOMMENT THIS:
                    if run["align"] not in checked:
                        self._check_grch38(run)
                        checked += run["align"]
                else:
                    self.err += "Error parsing run template file '%s'.  Required key 'experiment' not defined.\n" % (self.run_file,)
        self._raise()
        return

    def _check_grch38(self, run):
        if run["align"] in self.run_data["genomes"]:
            assembly = "grch38.p6"
            gen_prefix = "genome_%s_%s" % (run["align"], assembly)
            if assembly in self.run_data["genomes"][run["align"]]:
                base_file = self.run_data["genomes"][run["align"]][assembly]
                if os.path.isfile(base_file):
                    base_file_name = gen_prefix + base_file.split(".", 1)[1]
                    f = File(base_file_name)
                    f.addPFN(PFN(base_file, "local"))
                    self.files[base_file_name] = f
                    if base_file.split(".", 1)[1] in chipathlon.conf.genomes[run["align"]]["base_file"]:
                        prefix = base_file.split(".", 1)[0]
                        missing = []
                        for ext in chipathlon.conf.genomes[run["align"]]["additional_files"]:
                            if not os.path.isfile(prefix + "." + ext) and not os.path.isfile(base_file + "." + ext):
                                missing.append(ext)
                            else:
                                fnew = File(gen_prefix + "." + ext)
                                if os.path.isfile(prefix + "." + ext):
                                    fnew.addPFN(prefix + "." + ext)
                                else:
                                    fnew.addPFN(base_file + "." + ext)
                                self.files[gen_prefix + "." + ext] = fnew
                                self.dax.addFile(fnew)
                        if len(missing) > 0:
                            self.err += "Genome defined with tool '%s' and assembly '%s' is missing additional_files with extensions %s.\n" % (run["align"], assembly, missing)
                    else:
                        self.err += "Genome defined with tool '%s' and assembly '%s', has invalid extension.  Should be one of %s.\n" % (run["align"], assembly, chipathlon.conf.genomes[run["align"]]["base_file"])
                else:
                    self.err += "Genome defined with tool '%s' and assembly '%s', has non-existant base file '%s'.\n" % (run["align"], assembly, base_file)
            else:
                self.err += "Alignment defined with tool '%s' for assembly '%s', no corresponding genome definition.\n" % (run["align"], assembly)
        else:
            self.err += "Alignment defined with tool '%s', no corresponding genome definition.\n" % (run["align"])
        self._raise()
        return

    def _check_genomes(self, run):
        # NOTE: THIS function should NOT be called.
        # We will only be using GRCH38.p6 genome.
        # Use the _check_grch38 function instead
        valid, msg, assembly = self.mdb.get_assembly(run["experiment"])
        if valid:
            if run["align"] in self.run_data["genomes"]:
                if assembly in self.run_data["genomes"][run["align"]]:
                    base_file = self.run_data["genomes"][run["align"]][assembly]
                    if os.path.isfile(base_file):
                        if base_file.split(".", 1)[1] in chipathlon.conf.genomes[run["align"]]["base_file"]:
                            prefix = base_file.split(".", 1)[0]
                            missing = []
                            for ext in chipathlon.conf.genomes[run["align"]]["additional_files"]:
                                if not os.path.isfile(prefix + "." + ext):
                                    missing.append(ext)
                            if len(missing) > 0:
                                self.err += "Genome defined with tool '%s' and assembly '%s' is missing additional_files with extensions %s.\n" % (run["align"], assembly, missing)
                        else:
                            self.err += "Genome defined with tool '%s' and assembly '%s', has invalid extension.  Should be one of %s.\n" % (run["align"], assembly, chipathlon.conf.genomes[run["align"]]["base_file"])
                    else:
                        self.err += "Genome defined with tool '%s' and assembly '%s', has non-existant base file '%s'.\n" % (run["align"], assembly, base_file)
                else:
                    self.err += "Alignment defined with tool '%s' for assembly '%s', no corresponding genome definition.\n" % (run["align"], assembly)
            else:
                self.err += "Alignment defined with tool '%s', no corresponding genome definition.\n" % (run["align"])
        else:
            self.err += "DB error: %s.\n" % (msg, )
        return

    def _add_download(self):
        # Remember, experiment always paired with randomly selected control,
        # For alignment doesn't really matter, but we set up the groundwork
        # for the future.  We randomly select unless there is an
        # equal number of control and experiment files
        self.input_sets = []
        for run in self.run_data["runs"]:
            if (len(run["data"]["control"]) != len(run["data"]["experiment"])):
                control_data = [random.choice(run["data"]["control"]) for i in range(0, len(run["data"]["experiment"]))]
            else:
                control_data = run["data"]["control"]
            for pair in zip(run["data"]["experiment"], control_data):
                exp_name = "%s_%s.fastq.gz" % (run["experiment"], pair[0]["accession"])
                control_name = "%s_%s.fastq.gz" % (run["experiment"], pair[1]["accession"])
                if exp_name not in self.files and control_name not in self.files:
                    self.input_sets.append(pair)
                    for f in pair:
                        # For each file we create the download job
                        # and output file.  Base output files will be
                        # EXPID_ACCESSION.fastq.gz
                        output_name = "%s_%s.fastq.gz" % (run["experiment"], f["accession"])
                        output_file = File(output_name)
                        self.files[output_name] = output_file
                        job = Job(self.executables["download_fastq.py"])
                        job.uses(output_file, link=Link.OUTPUT, transfer=True)
                        job.addArguments("-u", f["url"], "-p", output_file, "-t http://", "-m", f["content_md5sum"])
                        self.jobs[output_name] = job
                        self.dax.addJob(job)
        return

    def _add_align(self):
        print self.modules["align"].markers
        markers = {}
        markers["read_end"] = "single"
        for run in self.run_data["runs"]:
            input_files = {}
            additional_files = {}
            markers["tool"] = run["align"]
            gen_prefix = "genome_%s_grch38.p6" % (run["align"],)
            input_files["ref_genome"] = "%s.%s" % (gen_prefix, self.run_data["genomes"][run["align"]]["grch38.p6"])
            for ext in chipathlon.conf.genomes[run["align"]]["additional_files"]:
                additional_files["ref_genome." + ext] = "%s.%s" % (gen_prefix, ext)
            for pair in self.input_sets:
                for f in pair:
                    prefix = "%s_%s" % (run["experiment"], f["accession"])
                    input_files["download_1.fastq"] = "%s.fastq.gz" % (prefix,)
                    self.modules["align"].add_jobs(self.dax, self.jobs, self.files, prefix, markers, input_files, additional_files)
        return

    def _create_setup(self):
        """
            Creates the base structure for job submission.  Everything is contained
            within a folder based on the current timestamp.
        """
        self.basepath = self.jobhome + "/" + datetime.datetime.now().strftime("%Y-%m-%d_%H%M%S")
        if not os.path.exists(self.basepath):
            os.makedirs(self.basepath)
        for folder in ["input", "output"]:
            if not os.path.exists(os.path.join(self.basepath, folder)):
                os.makedirs(os.path.join(self.basepath, folder))
        return

    def _add_notify(self):
        # NEED TO DETERMINE HOW TO READ IN THESE VARS
        notify_path = os.path.join(self.basepath, "input/notify.sh")
        with open(notify_path, "w") as wh:
            notify = textwrap.dedent("""\
                    #!/bin/bash
                    %s/notification/email -t %s --report=pegasus-analyzer
            """ % (config["notify"]["pegasus_home"], config["notify"]["email"]))
            wh.write(notify)
            os.chmod(notify_path, 0755)
        self.dax.invoke(When.AT_END, notify_path)
        return

    def _create_replica(self):
        with open(os.path.join(self.basepath, "input/conf.rc"), "w") as wh:
            pegasusrc = textwrap.dedent("""\
                pegasus.catalog.site = XML
                pegasus.catalog.site.file = %s/sites.xml

                pegasus.condor.logs.symlink = false

                pegasus.data.configuration = sharedfs
                """ % (os.path.join(self.basepath, "input"),))
            wh.write(pegasusrc)
        return

    def _create_sites(self):
        with open(os.path.join(self.basepath, "input/sites.xml"), "w") as wh:
            sites = """\
                <sitecatalog xmlns="http://pegasus.isi.edu/schema/sitecatalog" xmlns:xsi="http://www.w3.org/2001/XMLSchema-instance" xsi:schemaLocation="http://pegasus.isi.edu/schema/sitecatalog http://pegasus.isi.edu/schema/sc-4.0.xsd" version="4.0">
                   <site handle="local" arch="x86_64" os="LINUX">
                       <directory type="shared-scratch" path="%s">
                           <file-server operation="all" url="file://%s" />
                       </directory>
                       <directory type="local-storage" path="%s">
                           <file-server operation="all" url="file://%s" />
                       </directory>
                   </site>
                   <site handle="condorpool" arch="x86_64" os="LINUX">
                       <directory type="shared-scratch" path="%s">
                           <file-server operation="all" url="file://%s" />
                       </directory>
                """ % (
                        os.path.join(self.basepath, "work"),
                        os.path.join(self.basepath, "work"),
                        os.path.join(self.basepath, "output"),
                        os.path.join(self.basepath, "output"),
                        os.path.join(self.basepath, "work"),
                        os.path.join(self.basepath, "work")
                    )
            # NEED TO DETERMINE HOW TO READ IN THIS INFO
            for namespace in config["profile"]:
                for key in config["profile"][namespace]:
                    val = ":".join(config["profile"][namespace][key]) if "path" in key.lower() else config["profile"][namespace][key]
                    sites += """\n\t<profile namespace="%s" key="%s">%s</profile> """ % (namespace, key, val)
            sites += "</site></sitecatalog>"
            sites = sites.replace("\n", "")
            wh.write("\n".join([line for line in xml.dom.minidom.parseString(sites).toprettyxml().split('\n') if line.strip()]))
        return

    def _create_submit(self):
        """
            Creates the pegasus submit script.  submit.sh
        """
        with open(os.path.join(self.basepath, "input/submit.sh"), "w") as wh:
            submit = textwrap.dedent("""\
                    #!/bin/bash
                    plan=`pegasus-plan \\
                    --conf "%s" \\
                    --sites "%s" \\
                    --dir "%s" \\
                    --output-site local \\
                    --dax "%s" \\
                    --randomdir \\
                    """ % (self.basepath + "/input/conf.rc", "condorpool", self.basepath + "/work/", self.basepath + "/input/chipathlon.dax"))
            submit += textwrap.dedent("""\
                    --submit`

                    status=`echo "$plan" | grep pegasus-status | tr -s ' '| cut -d ' ' -f 6`
                    echo -e "#!/bin/bash
                    pegasus-status -l $status" > status.sh
                    chmod 744 status.sh

                    remove=`echo "$plan" | grep pegasus-remove | tr -s ' '| cut -d ' ' -f 5`
                    echo -e "#!/bin/bash
                    pegasus-remove $remove" > remove.sh
                    chmod 744 remove.sh

                    echo "$plan"

                    """)
            wh.write(submit)
            os.chmod(self.basepath + "/input/submit.sh", 0755)
        return

    def _write(self):
        with open(self.basepath + "/input/chipathlon.dax", "w") as wh:
            self.dax.writeXML(wh)
        return