TellReadJob.py

from os.path import join
from .Job import Job, KISSLoader
from .PipelineError import JobFailedError
import logging
from jinja2 import Environment
from .Pipeline import Pipeline
from .PipelineError import PipelineError
from metapool import load_sample_sheet
from os import walk
import re
from collections import defaultdict


logging.basicConfig(level=logging.DEBUG)


class TellReadJob(Job):
    def __init__(self, run_dir, output_path, sample_sheet_path, queue_name,
                 node_count, wall_time_limit, jmem, modules_to_load,
                 qiita_job_id, reference_base,
                 reference_map, sing_script_path, cores_per_task):
        """
        ConvertJob provides a convenient way to run bcl-convert or bcl2fastq
        on a directory BCL files to generate Fastq files.
        :param run_dir: The 'run' directory that contains BCL files.
        :param output_path: Path where all pipeline-generated files live.
        :param sample_sheet_path: The path to a sample-sheet.
        :param queue_name: The name of the Torque queue to use for processing.
        :param node_count: The number of nodes to request.
        :param wall_time_limit: A hard time limit (in min) to bound processing.
        :param jmem: String representing total memory limit for entire job.
        :param modules_to_load: A list of Linux module names to load
        :param qiita_job_id: identify Torque jobs using qiita_job_id
        :param reference_base: None
        :param reference_map: None
        :param cores_per_task: (Optional) # of CPU cores per node to request.
        """
        super().__init__(run_dir,
                         output_path,
                         'TellReadJob',
                         [],
                         1,
                         modules_to_load=modules_to_load)

        self.sample_sheet_path = sample_sheet_path
        self._file_check(self.sample_sheet_path)
        metadata = self._process_sample_sheet()
        self.sample_ids = metadata['sample_ids']
        self.queue_name = queue_name
        self.node_count = node_count
        self.wall_time_limit = wall_time_limit
        self.cores_per_task = cores_per_task

        self.reference_base = reference_base
        self.reference_map = reference_map

        # raise an Error if jmem is not a valid floating point value.
        self.jmem = str(int(jmem))
        self.qiita_job_id = qiita_job_id
        self.jinja_env = Environment(loader=KISSLoader('templates'))
        self.sing_script_path = sing_script_path

        sheet = load_sample_sheet(self.sample_sheet_path)
        lane = sheet.samples[0].Lane

        # force self.lane_number to be int. raise an Error if it's not.
        tmp = int(lane)
        if tmp < 1 or tmp > 8:
            raise ValueError(f"'{tmp}' is not a valid lane number")
        self.lane_number = tmp

        self.job_name = (f"{self.qiita_job_id}-tellread")

    def run(self, callback=None):
        job_script_path = self._generate_job_script()

        # everything is in the job script so there are no additional params.
        params = []

        try:
            self.job_info = self.submit_job(job_script_path,
                                            job_parameters=' '.join(params),
                                            exec_from=None,
                                            callback=callback)

            logging.debug(f'TellReadJob Job Info: {self.job_info}')
        except JobFailedError as e:
            # When a job has failed, parse the logs generated by this specific
            # job to return a more descriptive message to the user.
            # TODO: We need more examples of failed jobs before we can create
            #  a parser for the logs.
            info = self.parse_logs()
            # prepend just the message component of the Error.
            info.insert(0, str(e))
            raise JobFailedError('\n'.join(info))

        self.mark_job_completed()

        logging.debug(f'TellReadJob {self.job_info["job_id"]} completed')

    def _process_sample_sheet(self):
        sheet = load_sample_sheet(self.sample_sheet_path)

        if not sheet.validate_and_scrub_sample_sheet():
            s = "Sample sheet %s is not valid." % self.sample_sheet_path
            raise PipelineError(s)

        header = sheet.Header
        chemistry = header['chemistry']

        if header['Assay'] not in Pipeline.assay_types:
            s = "Assay value '%s' is not recognized." % header['Assay']
            raise PipelineError(s)

        sample_ids = []
        for sample in sheet.samples:
            sample_ids.append((sample['Sample_ID'],
                               sample['Sample_Project'],
                               sample['barcode_id']))

        bioinformatics = sheet.Bioinformatics

        # reorganize the data into a list of dictionaries, one for each row.
        # the ordering of the rows will be preserved in the order of the list.
        lst = bioinformatics.to_dict('records')

        # human-filtering jobs are scoped by project. Each job requires
        # particular knowledge of the project.
        return {'chemistry': chemistry,
                'projects': lst,
                'sample_ids': sample_ids}

    def _generate_job_script(self):
        job_script_path = join(self.output_path, 'tellread_test.sbatch')
        template = self.jinja_env.get_template("tellread.sbatch")

        # generate a comma separated list of sample-ids from the tuples stored
        # in self.sample_ids.

        # NB: Proposed sample-sheets will have traditional Sample_ID and
        # Sample_Name columns as well as a new value named barcode_id. It's
        # this column that will contain the 'C50n' values needed to be
        # supplied to tellread. Later we will use this mapping to rename the
        # files from C50n...fastq.gz to sample-name...fastq.gz.
        samples = ','.join([id[2] for id in self.sample_ids])

        # since we haven't included support for reference_map yet, whenever a
        # reference is not included, the mapping against the list of sample_ids
        # is ['NONE', 'NONE', ..., 'NONE'].
        refs = ','.join(['NONE' for _ in self.sample_ids])

        extra = ""

        with open(job_script_path, mode="w", encoding="utf-8") as f:
            f.write(template.render({
                "job_name": "tellread",
                "wall_time_limit": self.wall_time_limit,
                "mem_in_gb": self.jmem,
                "node_count": self.node_count,
                "cores_per_task": self.cores_per_task,
                "queue_name": self.queue_name,
                "sing_script_path": self.sing_script_path,
                "modules_to_load": ' '.join(self.modules_to_load),
                "lane": f"s_{self.lane_number}",
                # NB: Note that we no longer create a sub-directory under the
                # working directory for TellRead to create all its output
                # folders and files. This means it is creating folders and
                # files in the same directory that has our sbatch script and
                # logs directory. Currently there are no name collisions,
                # however.
                "output": self.output_path,
                "rundir_path": self.root_dir,
                "samples": samples,
                "refs": refs,
                "extra": extra
            }))

        return job_script_path

    def audit(self):
        # this overriden audit method does not need sample-ids passed as a
        # parameter because this job is already aware of what samples should
        # be present and more importantly, how they map to barcode_ids.
        def map_barcode_id_to_sample_id(barcode_id):
            for s_id, _, b_id in self.sample_ids:
                if barcode_id == b_id:
                    return s_id

        corrected = defaultdict(list)
        for root, dirs, files in walk(join(self.output_path, 'Full')):
            for _file in files:
                m = re.match(r"TellReadJob_(.\d)_(C\d\d\d).fastq.gz.corrected."
                             r"err_barcode_removed.fastq", _file)
                if m:
                    read, barcode_id = m.groups(1)
                    corrected[barcode_id].append(read)

        # a sample was processed successfully if all three expected reads are
        # present.
        failed = []
        for barcode_id in corrected:
            # we expect only the following read/orientations.
            if not set(corrected[barcode_id]) == {'I1', 'R1', 'R2'}:
                failed.append(map_barcode_id_to_sample_id(barcode_id))

        return sorted(failed)