crunch_scripts/arvados_samtools.py

   1 import arvados
   2 import re
   3 import os
   4 import sys
   5 import fcntl
   6 import subprocess
   7
   8 samtools_path = None
   9
  10 def samtools_install_path():
  11     global samtools_path
  12     if samtools_path:
  13         return samtools_path
  14     samtools_path = arvados.util.tarball_extract(
  15         tarball = arvados.current_job()['script_parameters']['samtools_tgz'],
  16         path = 'samtools')
  17
  18     # build "samtools" binary
  19     lockfile = open(os.path.split(samtools_path)[0] + '.samtools-make.lock',
  20                     'w')
  21     fcntl.flock(lockfile, fcntl.LOCK_EX)
  22     arvados.util.run_command(['make', '-j16'], cwd=samtools_path)
  23     lockfile.close()
  24
  25     return samtools_path
  26
  27 def samtools_binary():
  28     return os.path.join(samtools_install_path(), 'samtools')
  29
  30 def run(command, command_args, **kwargs):
  31     execargs = [samtools_binary(),
  32                 command]
  33     execargs += command_args
  34     sys.stderr.write("%s.run: exec %s\n" % (__name__, str(execargs)))
  35     arvados.util.run_command(
  36         execargs,
  37         cwd=arvados.current_task().tmpdir,
  38         stdin=kwargs.get('stdin', subprocess.PIPE),
  39         stderr=kwargs.get('stderr', sys.stderr),
  40         stdout=kwargs.get('stdout', sys.stderr))
  41
  42 def one_task_per_bam_file(if_sequence=0, and_end_task=True):
  43     if if_sequence != arvados.current_task()['sequence']:
  44         return
  45     job_input = arvados.current_job()['script_parameters']['input']
  46     cr = arvados.CollectionReader(job_input)
  47     bam = {}
  48     bai = {}
  49     for s in cr.all_streams():
  50         for f in s.all_files():
  51             if re.search(r'\.bam$', f.name()):
  52                 bam[s.name(), f.name()] = f
  53             elif re.search(r'\.bai$', f.name()):
  54                 bai[s.name(), f.name()] = f
  55     for ((s_name, f_name), bam_f) in bam.items():
  56         bai_f = bai.get((s_name, re.sub(r'bam$', 'bai', f_name)), None)
  57         task_input = bam_f.as_manifest()
  58         if bai_f:
  59             task_input += bai_f.as_manifest()
  60         new_task_attrs = {
  61             'job_uuid': arvados.current_job()['uuid'],
  62             'created_by_job_task_uuid': arvados.current_task()['uuid'],
  63             'sequence': if_sequence + 1,
  64             'parameters': {
  65                 'input': task_input
  66                 }
  67             }
  68         arvados.api().job_tasks().create(body=new_task_attrs).execute()
  69     if and_end_task:
  70         arvados.api().job_tasks().update(uuid=arvados.current_task()['uuid'],
  71                                          body={'success':True}
  72                                          ).execute()
  73         exit(0)