workflow.py

import os
import copy
import argparse
import sys
import json
import datetime
import textwrap
import yaml
import traceback
import chipathlon
import chipathlon.workflow_job
import chipathlon.db
import chipathlon.workflow_module
from chipathlon.run_parser import RunParser
from chipathlon.generators.align_generator import AlignGenerator
from chipathlon.generators.remove_duplicates_generator import RemoveDuplicatesGenerator
from chipathlon.generators.peak_call_generator import PeakCallGenerator
from chipathlon.generators.idr_generator import IdrGenerator
from pprint import pprint
from Pegasus.DAX3 import *

class Workflow(object):
    """
    :param job_home: The base directory to generate the workflow in.
    :type job_home: str
    :param run_file: The run yaml file defining controls, samples, alignment
        tool, peak calling tool, peak calling type, and which samples to
        run idr on.
    :type run_file: str
    :param config_file: The config yaml file defining system variables.
    :type config_file: str
    :param host: The host address of the MongoDB database.
    :type host: str
    :param username: The username to authenticate for MongoDB access.
    :type username: str
    :param password: The password to authenticate for MongoDB access.
    :type password: str
    :param execute_site: The target site to submit jobs to.  These sites should
        be defined in the sites.xml file.
    :type execute_site: str
    :param output_site: The output site to transfer files to.  This site should
        be defined in the sites.xml file.
    :type output_site: str
    :param save_db: Whether or not we want to save results to the database.
        True by default.
    :type save_db: bool
    :param email: An email address to notify when the workflow is finished.
    :type email: str
    :param rewrite: Whether or not to rewrite existing files.  If true, it will
        ignore files in Mongo and recreate them.  If false, it will download
        files based on the latest available completed job.
    :type rewrite: bool
    :param debug: A flag for printing additional information when generating
        workflows.
    :type debug: bool

    Workflow is the highest level class.  It handles interaction between other
    classes as well as handles the initial loading of yaml files and
    executables.
    """

    def __init__(self, job_home, run_file, param_file, properties_file,
                chip_bin, idr_bin, execute_site="local", output_site="local",
                host="localhost", username=None, password=None, save_db=True,
                email=None, rewrite=False, debug=False):
        # debug mode, print out additional information
        self.debug = debug
        # Job site information
        self.execute_site = execute_site
        self.output_site = output_site
        self.save_db = save_db
        self.email = email

        # DB information
        self.username = username
        self.host = host
        self.password = password
        # Initialize db connection
        self.mdb = chipathlon.db.MongoDB(host, username, password)
        # Jobname info & err
        self.job_home = os.path.abspath(job_home)
        self.base_path = os.path.join(self.job_home, datetime.datetime.now().strftime("%Y-%m-%d_%H%M%S"))
        self.errors = []
        # add new results even if they exist
        self.rewrite = rewrite
        # Input file info
        self.run_file = os.path.abspath(run_file)
        self.param_file = os.path.abspath(param_file)
        self.properties_file = os.path.abspath(properties_file)
        # Validate the existance of but not the contents of the
        # pegasus properties file
        if not os.path.isfile(self.properties_file):
            self.errors.append("Provided pegasus properties file '%s' does not exist." % (self.properties_file,))

        self.chip_bin = chip_bin
        self.idr_bin = idr_bin

        # Dax specific info
        self.dax = ADAG(os.path.basename(os.path.dirname(self.job_home + "/")))
        self.executables = {}
        self.files = {}
        self.jobs = {}
        self.workflow_jobs = {}
        self.modules = {}

        # The steps to take when generating jobs
        self.generation_steps = [
            self._load_executables,
            self._load_workflow_jobs,
            self._load_modules,
            self._load_run_data,
            self._load_generators,
            # All required information is loaded, start queuing jobs
            self._create_setup,
            self._generate_jobs,
            # Create pegasus important stuff
            self._add_notify,
            self._create_pegasus_files
        ]
        return

    def info(self):
        """
        A helper function that prints out all currently loaded run information,
        all loaded workflow_jobs, all loaded workflow_modules, and all loaded
        executables.
        """
        pprint(self.run_data)
        pprint(self.workflow_jobs)
        pprint(self.modules)
        pprint(self.executables)
        return

    def generate(self):
        """
        The function that actually generates the workflow!
        """
        for step in self.generation_steps:
            if self.is_valid():
                step()
            else:
                print self.get_error_string()
                break
        return

    def is_valid(self):
        """
        :returns: True if the workflow is valid else False.

        Whether or not the workflow is valid so far!
        """
        return len(self.errors) == 0

    def get_error_string(self):
        """
        :returns: A newline separated string of all the errors.
        """
        return "\n".join(self.errors)

    def _add_file(self, name, path, site):
        f = File(name)
        f.addPFN(PFN(path, site))
        self.files[name] = f
        self.dax.addFile(f)
        return

    def _add_genome_file(self, gen_file_obj):
        self.files[gen_file_obj["name"]] = gen_file_obj["file"]
        self.dax.addFile(gen_file_obj["file"])
        return

    def _add_executable(self, name, path, os_type="linux", arch="x86_64", site="local", installed=True):
        self.executables[name] = Executable(name=name, os=os_type.lower(), arch=arch, installed=installed)
        self.executables[name].addPFN(PFN(os.path.join("file://", path), site))
        self.dax.addExecutable(self.executables[name])
        if self.debug:
            print "[LOAD_EXECUTABLE] %s" % (name,)
        return

    def _load_executables(self, os_type="linux", arch="x86_64"):
        # Load actual executables don't need no wrappers
        for executable in chipathlon.conf.executables:
            self._add_executable(
                executable,
                os.path.join(self.idr_bin if executable == "idr" else self.chip_bin, executable),
                site=self.execute_site,
                arch=arch,
                os_type=os_type
            )
        # 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:
                if not os.path.splitext(f)[1] == '.pyc':
                    self._add_executable(
                        f,
                        os.path.join(self.chip_bin, f),
                        site=self.execute_site,
                        arch=arch,
                        os_type=os_type,
                    )
            break
        # Handle necessary installed scripts
        for cmd in chipathlon.conf.system_commands:
            self._add_executable(
                cmd,
                os.path.join(chipathlon.conf.system_path, cmd),
                site=self.execute_site,
                arch=arch,
                os_type=os_type
	    )
        return

    def _load_workflow_jobs(self):
        # Load in the specified param file for use in the workflow_jobs
        if os.path.isfile(self.param_file):
            with open(self.param_file, "r") as rh:
                try:
                    self.params = yaml.load(rh)
                    if self.params is None:
                        print "[WARNING] Param file is empty.  Using defaults for all jobs."
                        self.params = {}
                except yaml.YAMLError as ye:
                    self.errors.append("Error parsing param yaml file '%s': %s." % (self.param_file, ye))
        else:
            print "[WARNING] Param file does not exist.  Using defaults for all jobs."
            self.params = {}

        # Load each yaml_job as an actual object
        for root, dirs, files in os.walk("%s/%s" % (os.path.dirname(os.path.realpath(__file__)), chipathlon.conf.job_params)):
            for f in files:
                yaml_job = chipathlon.workflow_job.WorkflowJob(os.path.join(root, f), self.params.get(f.split(".")[0]), self.debug)
                if yaml_job.is_valid():
                    if yaml_job.command in self.executables:
                        yaml_job.add_executable(self.executables[yaml_job.command])
                        self.workflow_jobs[yaml_job.job_name] = yaml_job
                        if self.debug:
                            print "[LOAD WORKFLOW JOB] %s" % (yaml_job.job_name,)
                        for raw_file_name in yaml_job.raw_files:
                            if raw_file_name not in self.files:
                                self._add_file(raw_file_name, yaml_job.raw_files[raw_file_name]["path"], "local")
                            # The _add_file actually creates the Pegasus file class,
                            # we want to reference that within the yaml_job for housekeeping
                            yaml_job.raw_files[raw_file_name]["file"] = self.files[raw_file_name]
                    else:
                        print "[WARNING] Skipping param file %s, corresponding executable %s not found." % (f, yaml_job.command)
                else:
                    for error in yaml_job.errors:
                        self.errors.append(error)
        return

    def _load_modules(self):
        # Load each yaml_module as an actual object
        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, self.debug)
                if not mod.err:
                    self.modules[mod.name] = mod
                    if self.debug:
                        print "[LOAD MODULE] %s" % (mod.name,)
                else:
                    self.err += mod.err
            break
        return

    def _load_run_data(self):
        self.run_parser = RunParser(self.run_file, self.mdb)
        if self.run_parser.is_valid():
            for genome in self.run_parser.get_genomes():
                for gen_file_obj in genome.get_all_files():
                    if gen_file_obj["name"] not in self.files:
                        self._add_genome_file(gen_file_obj)
            for i, run in enumerate(self.run_parser.get_runs()):
                if self.debug:
                    print "[LOADING SAMPLES] Run #%s" % (i,)
                run.load_samples(self.mdb)
                if not run.is_valid():
                    self.errors.append("Error with Run #%s: %s" % (i, run.get_error_string()))
            self.runs = self.run_parser.get_runs()
        else:
            self.errors.append(self.run_parser.get_error_string())
        return

    def _load_generators(self):
        self.align_gen = AlignGenerator(
            self.dax, self.jobs, self.files, self.mdb, self.modules["align"],
            self.workflow_jobs, self.base_path, save_db=self.save_db,
            rewrite=self.rewrite, debug=self.debug
        )
        self.remove_dup_gen = RemoveDuplicatesGenerator(
            self.dax, self.jobs, self.files, self.mdb, self.modules["remove_duplicates"],
            self.workflow_jobs, self.base_path, save_db=self.save_db,
            rewrite=self.rewrite, debug=self.debug
        )
        self.peak_call_gen = PeakCallGenerator(
            self.dax, self.jobs, self.files, self.mdb, self.modules["peak_call"],
            self.workflow_jobs, self.base_path, save_db=self.save_db,
            rewrite=self.rewrite, debug=self.debug
        )
        self.idr_gen = IdrGenerator(
            self.dax, self.jobs, self.files, self.mdb, self.modules["idr"],
            self.workflow_jobs, self.base_path, save_db=self.save_db,
            rewrite=self.rewrite, debug=self.debug
        )
        return

    def _generate_jobs(self):
        for run in self.runs:
            if run.file_type == "fastq":
                self.align_gen.generate(run)
            self.remove_dup_gen.generate(run)
            self.peak_call_gen.generate(run)
            if run.idr is not None:
                self.idr_gen.generate(run)
        return

    def _create_setup(self):
        """
        Creates the base structure for job submission.  Everything is contained
        within a folder based on the current timestamp.
        """
        if not os.path.exists(self.base_path):
            os.makedirs(self.base_path)
        for folder in ["input", "input/db_meta"]:
            if not os.path.exists(os.path.join(self.base_path, folder)):
                os.makedirs(os.path.join(self.base_path, folder))
        return

    def _add_notify(self):
        """
        Add the script to email when the workflow is finished.
        """
        if self.email is not None:
            notify_path = os.path.join(self.base_path, "input/notify.sh")
            with open(notify_path, "w") as wh:
                notify = textwrap.dedent("""\
                        #!/bin/bash
                        pegasus-email -t %s --report=pegasus-analyzer
                """ % (self.email,))
                wh.write(notify)
                os.chmod(notify_path, 0755)
            self.dax.invoke(When.AT_END, notify_path)
        return

    def _create_pegasus_files(self):
        self._create_submit()
        self._create_dax()
        return

    def _create_submit(self):
        """
        Creates the pegasus submit script.  Just a wrapper to pegasus-plan
        with the correct arguments passed through.  Additionally creates
        some easy status/remove scripts for the workflow when submitted.
        """
        with open(os.path.join(self.base_path, "input/submit.sh"), "w") as wh:
            submit = textwrap.dedent("""\
                    #!/bin/bash
                    plan=`pegasus-plan \\
                    --conf %s \\
                    --sites %s \\
                    --output-site %s \\
                    --dir %s \\
                    --dax %s \\
                    --randomdir \\
                    """ % (
                        self.properties_file,
                        self.execute_site,
                        self.output_site,
                        os.path.join(self.base_path, "work"),
                        os.path.join(self.base_path, "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(os.path.join(self.base_path, "input/submit.sh"), 0755)
        return

    def _create_dax(self):
        with open(os.path.join(self.base_path, "input/chipathlon.dax"), "w") as wh:
            self.dax.writeXML(wh)
        return