def make_decomp_stage(self, ready_to_decomp): """ :param args: :param ready_to_decomp: generator of (id_and_count, column0, column1, ..., [:rest of input]) :return: a generator of [ agd_read_handle, num_records, first_ordinal, record_id, id_and_count, {rest of input} ] """ ready_to_decomp = sanitize_generator(ready_to_decomp) num_columns = len(self.columns) # to_agd_reader = just the columns # pass_around_agd_reader = (id_and_count, rest, of, input, ...) to_agd_reader, pass_around_agd_reader = zip(*( (rtd[1:1+num_columns], (rtd[0],)+tuple(rtd[1+num_columns:])) for rtd in ready_to_decomp )) def gen_timestamps(): for group in pass_around_agd_reader: with tf.control_dependencies((group[0],)): yield gate.unix_timestamp(name="align_head_timestamp") reader_kwargs = {} timestamps = [] if self.log_goodput: timestamps.extend(gen_timestamps()) assert len(timestamps) == len(ready_to_decomp) # control dependencies have to be an iterable reader_kwargs["control_ops"] = tuple((a,) for a in timestamps) # [output_buffer_handles], num_records, first_ordinal, record_id; in order, for each column group in upstream_tensorz multi_column_gen = tuple(pipeline.agd_reader_multi_column_pipeline(upstream_tensorz=to_agd_reader, verify=self.deep_verify, name="align_reader", **reader_kwargs)) # around = num_records, first_ordinal, record_id for each group to_assembler, around_assembler = zip(*( (a[:2], a[1:]) for a in multi_column_gen )) assembler_kwargs = {} if self.log_goodput: log_event_ops = [ (gate.log_events( # single element tuple because that's how tf.control_dependencies works item_names=("id","time","ordinal","record_id"), components=(in_id, timestamp, ordinal, record_id), event_name="align_head", directory=self.log_directory, name="align_head_event_logger" ),) for in_id, timestamp, ordinal, record_id in zip( (slice_id(a[0]) for a in pass_around_agd_reader), timestamps, (b[2] for b in multi_column_gen), (b[3] for b in multi_column_gen) ) ] assembler_kwargs["control_deps"] = log_event_ops # each element is an agd_reads handle agd_assembled_reads = pipeline.agd_read_assembler(upstream_tensors=to_assembler, include_meta=False, **assembler_kwargs) for agd_read, around_assembler_group, around_reader_group in zip(agd_assembled_reads, around_assembler, pass_around_agd_reader): yield (agd_read,) + tuple(around_assembler_group) + tuple(around_reader_group)
def make_sort_pipeline(self, args, input_gen, buf_pool, bufpair_pool): ready_to_process = pipeline.join( upstream_tensors=input_gen, parallel=args.sort_process_parallel, capacity=4, # multiplied by some factor? multi=True, name="ready_to_process") # need to unpack better here multi_column_gen = list( pipeline.agd_reader_multi_column_pipeline( upstream_tensorz=ready_to_process, buffer_pool=buf_pool)) # [ [base qual meta result], num_recs, first_ord, record_id ] chunks_and_recs = [] for chunks, num_recs, first_ord, record_id in multi_column_gen: entry = [] for chunk in chunks: entry.append(chunk) entry.append(num_recs) chunks_and_recs.append(entry) ready = tf.train.batch_join(chunks_and_recs, batch_size=args.column_grouping, allow_smaller_final_batch=True, name="chunk_batcher") name_queue = pipeline.join([name_generator("intermediate_file")], parallel=args.sort_parallel, capacity=4, multi=False, name="inter_file_gen_q") #bpp = persona_ops.buffer_pair_pool(size=0, bound=False, name="local_read_buffer_pair_pool") if args.order_by == location_value: sorter = persona_ops.agd_sort else: sorter = persona_ops.agd_sort_metadata sorters = [] for i in range(args.sort_parallel): #b, q, m, r, num = ready num = ready[-1] r = ready[0] # the sort predicate column must be first cols = tf.stack(ready[1:-1]) superchunk_matrix, num_recs = sorter(buffer_pair_pool=bufpair_pool, results_handles=r, column_handles=cols, num_records=num, name="local_read_agd_sort") # super chunk is r, b, q, m sorters.append([superchunk_matrix, num_recs, name_queue[i]]) return sorters
def import_sga_local(in_queue, argsj, outdir=None, parallel_parse=1, feature="NFAT", path="."): manifest = argsj.dataset if 'reference' not in manifest: raise Exception( "No reference data in manifest {}. Unaligned BAM not yet supported. Please align dataset first." .format(args.dataset)) """ key: tensor with chunk key string local_directory: the "base path" from which these should be read column_grouping_factor: the number of keys to put together parallel_parse: the parallelism for processing records (decomp) """ ref_lens = [] ref_seqs = [] for contig in manifest['reference_contigs']: ref_lens.append(contig['length']) ref_seqs.append(contig['name']) parallel_key_dequeue = tuple(in_queue.dequeue() for _ in range(parallel_parse)) result_chunks = pipeline.local_read_pipeline( upstream_tensors=parallel_key_dequeue, columns=['results']) result_chunk_list = [list(c) for c in result_chunks] parsed_results = pipeline.agd_reader_multi_column_pipeline( upstream_tensorz=result_chunk_list) parsed_results_list = list(parsed_results) parsed_result = pipeline.join(parsed_results_list, parallel=1, capacity=8, multi=True)[0] result_buf, num_results, first_ord, record_id = parsed_result result_buf = tf.unstack(result_buf)[0] result = persona_ops.import_sga(results_handle=result_buf, num_records=num_results, ref_sequences=ref_seqs, ref_seq_sizes=ref_lens, feature=feature, path=path, name="importsgaop") return result
def make_decomp_stage(self, ready_to_decomp): """ :param args: :param ready_to_decomp: generator of (id_and_count, [:rest of input], {column_stack} ) :return: a generator of [ id_and_count, [output_buffer_handles], num_records, first_ordinal, record_id, { rest of input } ] """ ready_to_decomp = sanitize_generator(ready_to_decomp) num_columns = len(self.extended_columns) # to_agd_reader = just the columns # pass_around_agd_reader = (id_and_count, rest, of, input, ...) to_agd_reader, pass_around_agd_reader = zip( *((rtd[-1], (rtd[0], ) + tuple(rtd[1:-1])) for rtd in ready_to_decomp)) def gen_timestamps(): for group in pass_around_agd_reader: idc = group[0] with tf.control_dependencies((idc, )): ts = gate.unix_timestamp(name="sort_head_timestamp") event_log_op = gate.log_events(item_names=("id", "time"), components=(slice_id(idc), ts), event_name="sort_head", directory=self.log_directory, name="sort_head_event_logger") yield event_log_op reader_kwargs = {} timestamps = [] if self.log_goodput: timestamps.extend(gen_timestamps()) assert len(timestamps) == len(ready_to_decomp) # control dependencies have to be an iterable reader_kwargs["control_ops"] = tuple((a, ) for a in timestamps) # [output_buffer_handles], num_records, first_ordinal, record_id; in order, for each column group in upstream_tensorz repack = [c == base_extension for c in self.extended_columns] multi_column_gen = tuple( pipeline.agd_reader_multi_column_pipeline( upstream_tensorz=to_agd_reader, verify=self.deep_verify, repack=repack, name="align_reader", **reader_kwargs)) for pass_around, generated in zip(pass_around_agd_reader, multi_column_gen): yield (pass_around[0], ) + tuple(generated) + tuple( pass_around[1:])
def agd_mark_duplicates_local(in_queue, outdir=None, parallel_parse=1, parallel_write=1, parallel_compress=1): """ key: tensor with chunk key string local_directory: the "base path" from which these should be read column_grouping_factor: the number of keys to put together parallel_parse: the parallelism for processing records (decomp) """ parallel_key_dequeue = tuple(in_queue.dequeue() for _ in range(parallel_parse)) result_chunks = pipeline.local_read_pipeline(upstream_tensors=parallel_key_dequeue, columns=['results']) result_chunk_list = [ list(c) for c in result_chunks ] parsed_results = pipeline.agd_reader_multi_column_pipeline(upstream_tensorz=result_chunk_list) parsed_results_list = list(parsed_results) parsed_result = pipeline.join(parsed_results_list, parallel=1, capacity=8, multi=True)[0] # result_buf, num_recs, first_ord, record_id #parsed_results = tf.contrib.persona.persona_in_pipe(key=key, dataset_dir=local_directory, columns=["results"], parse_parallel=parallel_parse, #process_parallel=1) print(parsed_result) result_buf, num_results, first_ord, record_id = parsed_result result_buf = tf.unstack(result_buf)[0] print(result_buf) bpp = persona_ops.buffer_pair_pool(size=0, bound=False, name="output_buffer_pair_pool") result_out = persona_ops.agd_mark_duplicates(results_handle=result_buf, num_records=num_results, buffer_pair_pool=bpp, name="markdupsop") result_to_write = pipeline.join([result_out, num_results, first_ord, record_id], parallel=parallel_write, capacity=8, multi=False) compressed = compress_pipeline(result_to_write, parallel_compress) written = _make_writers(compressed_batch=list(compressed), output_dir=outdir, write_parallelism=parallel_write) recs = list(written) all_written_keys = pipeline.join(recs, parallel=1, capacity=8, multi=False) return all_written_keys
def agd_flagstat_local(in_queue, outdir=None, parallel_parse=1, parallel_write=1, parallel_compress=1): """ key: tensor with chunk key string local_directory: the "base path" from which these should be read column_grouping_factor: the number of keys to put together parallel_parse: the parallelism for processing records (decomp) """ parallel_key_dequeue = tuple(in_queue.dequeue() for _ in range(parallel_parse)) result_chunks = pipeline.local_read_pipeline( upstream_tensors=parallel_key_dequeue, columns=['results']) result_chunk_list = [list(c) for c in result_chunks] parsed_results = pipeline.agd_reader_multi_column_pipeline( upstream_tensorz=result_chunk_list) parsed_results_list = list(parsed_results) parsed_result = pipeline.join(parsed_results_list, parallel=1, capacity=8, multi=True)[0] # print(parsed_result) result_buf, num_results, first_ord, record_id = parsed_result result_buf = tf.unstack(result_buf)[0] # print(result_buf) result_out = persona_ops.agd_flagstat(results_handle=result_buf, num_records=num_results, name="flagstat") return result_out
def export_bam(in_queue, args): manifest = args.dataset if 'reference' not in manifest: raise Exception( "No reference data in manifest {}. Unaligned BAM not yet supported. Please align dataset first." .format(args.dataset)) #bp_handle = persona_ops.buffer_pool(size=10, bound=False, name="buf_pool") #mmap_pool = persona_ops.m_map_pool(size=10, bound=False, name="file_mmap_buffer_pool") columns = ["base", "qual", "metadata", "results"] num_secondary = 0 for column in manifest['columns']: if 'secondary' in column: columns.append(column) secondary += 1 print("BAM output using columns: {}".format(columns)) # TODO provide option for reading from Ceph result_chunks = pipeline.local_read_pipeline( upstream_tensors=[in_queue.dequeue()], columns=columns) result_chunk_list = [list(c) for c in result_chunks] to_parse = pipeline.join(upstream_tensors=result_chunk_list, parallel=args.parallel_parse, multi=True, capacity=8) parsed_results = pipeline.agd_reader_multi_column_pipeline( upstream_tensorz=to_parse) parsed_results_list = list(parsed_results) parsed_result = pipeline.join(parsed_results_list, parallel=1, capacity=8, multi=True)[0] # base, qual, meta, result, [secondary], num_recs, first_ord, record_id handles = parsed_result[0] bases = handles[0] quals = handles[1] meta = handles[2] # give a matrix of all the result columns results = tf.stack(handles[3:]) num_recs = parsed_result[1] first_ord = parsed_result[2] if args.output_path == "": output_path = manifest['name'] + ".bam" else: output_path = args.output_path ref_lens = [] ref_seqs = [] for contig in manifest['reference_contigs']: ref_lens.append(contig['length']) ref_seqs.append(contig['name']) sort = manifest['sort'] if 'sort' in manifest else 'unsorted' pg_id = "personaAGD" # TODO get from manifest read_group = manifest['name'] agd_to_bam = persona_ops.agd_output_bam(results_handle=results, bases_handle=bases, qualities_handle=quals, metadata_handle=meta, num_records=num_recs, path=output_path, ref_sequences=ref_seqs, ref_seq_sizes=ref_lens, pg_id=pg_id, read_group=read_group, sort_order=sort, num_threads=args.threads) return [agd_to_bam], []
def make_central_pipeline(self, args, input_gen, pass_around_gen): self.write_columns.append('results') for i in range(args.max_secondary): self.write_columns.append('secondary{}'.format(i)) self.write_columns = [{ "type": "structured", "extension": a } for a in self.write_columns] joiner = tuple( tuple(a) + tuple(b) for a, b in zip(input_gen, pass_around_gen)) ready_to_process = pipeline.join( upstream_tensors=joiner, parallel=args.parallel, capacity=args.parallel, # multiplied by some factor? multi=True, name="ready_to_process") # need to unpack better here to_agd_reader, pass_around_agd_reader = zip( *((a[:2], a[2:]) for a in ready_to_process)) multi_column_gen = pipeline.agd_reader_multi_column_pipeline( upstream_tensorz=to_agd_reader) def process_processed_bufs(): for processed_column, pass_around in zip(multi_column_gen, pass_around_agd_reader): if isinstance(pass_around, tf.Tensor): pass_around = (pass_around, ) yield tuple( a for a in itertools.chain(processed_column, pass_around)) processed_bufs = tuple(a for a in process_processed_bufs()) ready_to_assemble = pipeline.join( upstream_tensors=processed_bufs, parallel=args.assemblers, capacity=args.assemblers * 2, multi=True, name="ready_to_assemble" ) # TODO these params are kinda arbitrary :/ # ready_to_assemble: [output_buffers, num_records, first_ordinal, record_id, pass_around {flattened}) x N] to_assembler, pass_around_assembler = zip( *((a[:2], a[1:]) for a in ready_to_assemble)) # each item out of this is a handle to AGDReads agd_read_assembler_gen = tuple( pipeline.agd_read_assembler(upstream_tensors=to_assembler, include_meta=False)) # assembled_records, ready_to_align: [(agd_reads_handle, (num_records, first_ordinal, record_id), (pass_around)) x N] assembled_records_gen = tuple( zip(agd_read_assembler_gen, pass_around_assembler)) assembled_records = tuple( (a, ) + tuple(b) for a, b in assembled_records_gen) ready_to_align = pipeline.join( upstream_tensors=assembled_records, parallel=args.aligners, capacity=int(args.aligners * 1.5), multi=True, name="ready_to_align") # TODO still have default capacity here :/ if args.paired: aligner_type = persona_ops.snap_align_paired aligner_options = persona_ops.paired_aligner_options( cmd_line=args.snap_args.split(), name="paired_aligner_options") executor_type = persona_ops.snap_paired_executor else: aligner_type = persona_ops.snap_align_single aligner_options = persona_ops.aligner_options( cmd_line=args.snap_args.split(), name="aligner_options" ) # -o output.sam will not actually do anything executor_type = persona_ops.snap_single_executor first_assembled_result = ready_to_align[0][1:] sink_queue_shapes = [a.get_shape() for a in first_assembled_result] sink_queue_dtypes = [a.dtype for a in first_assembled_result] aligner_dtype = tf.string aligner_shape = (args.max_secondary + 1, 2) sink_queue_shapes.append(aligner_shape) sink_queue_dtypes.append(aligner_dtype) pass_around_aligners = tuple( a[1:] for a in ready_to_align ) # type: [(num_records, first_ordinal, record_id, pass_around x N) x N] pass_to_aligners = tuple(a[0] for a in ready_to_align) buffer_list_pool = persona_ops.buffer_list_pool( **pipeline.pool_default_args) genome = persona_ops.genome_index(genome_location=args.index_path, name="genome_loader") def make_aligners(): single_executor = executor_type(num_threads=args.aligner_threads, work_queue_size=args.aligners + 1, options_handle=aligner_options, genome_handle=genome) for read_handle, pass_around in zip(pass_to_aligners, pass_around_aligners): aligner_results = aligner_type( read=read_handle, buffer_list_pool=buffer_list_pool, subchunk_size=args.subchunking, executor_handle=single_executor, max_secondary=args.max_secondary) yield (aligner_results, ) + tuple(pass_around) aligners = tuple(make_aligners()) # aligners: [(buffer_list_handle, num_records, first_ordinal, record_id, pass_around X N) x N], that is COMPLETELY FLAT if args.compress_parallel > 0: aligner_results_to_compress = pipeline.join( upstream_tensors=aligners, parallel=args.compress_parallel, multi=True, capacity=4, name="ready_to_compress") to_compressors = (a[0] for a in aligner_results_to_compress) around_compressors = (a[1:] for a in aligner_results_to_compress) compressed_buffers = pipeline.aligner_compress_pipeline( upstream_tensors=to_compressors) after_compression = ( (a, ) + tuple(b) for a, b in zip(compressed_buffers, around_compressors)) aligners = tuple(after_compression) aligned_results = pipeline.join(upstream_tensors=aligners, parallel=args.writers, multi=True, capacity=4, name="aligned_results") ref_seqs, lens = persona_ops.snap_index_reference_sequences( genome_handle=genome) # Taking this out because it currently breaks distributed runtime return aligned_results, ( genome, ref_seqs, lens ) # returns [(buffer_list_handle, num_records, first_ordinal, record_id, pass_around X N) x N], that is COMPLETELY FLAT
def make_central_pipeline(self, args, input_gen, pass_around_gen): self.write_columns.append('results') for i in range(args.max_secondary): self.write_columns.append('secondary{}'.format(i)) self.write_columns = [{ "type": "structured", "extension": a } for a in self.write_columns] joiner = tuple( tuple(a) + tuple(b) for a, b in zip(input_gen, pass_around_gen)) ready_to_process = pipeline.join(upstream_tensors=joiner, parallel=args.parallel, capacity=args.mmap_queue, multi=True) # need to unpack better here to_agd_reader, pass_around_agd_reader = zip( *((a[:2], a[2:]) for a in ready_to_process)) multi_column_gen = pipeline.agd_reader_multi_column_pipeline( upstream_tensorz=to_agd_reader) def process_processed_bufs(): for processed_column, pass_around in zip(multi_column_gen, pass_around_agd_reader): if isinstance(pass_around, tf.Tensor): pass_around = (pass_around, ) yield tuple( a for a in itertools.chain(processed_column, pass_around)) processed_bufs = tuple(a for a in process_processed_bufs()) ready_to_assemble = pipeline.join( upstream_tensors=processed_bufs, parallel=1, capacity=8, multi=True) # TODO these params are kinda arbitrary :/ # ready_to_assemble: [output_buffers, num_records, first_ordinal, record_id, pass_around {flattened}) x N] to_assembler, pass_around_assembler = zip( *((a[:2], a[1:]) for a in ready_to_assemble)) #print("reads {}".format(to_assembler)) base, qual = tf.unstack(to_assembler[0][0]) num_recs = to_assembler[0][1] two_bit_base = persona_ops.two_bit_converter(num_recs, base) to_assembler_converted = [[tf.stack([two_bit_base, qual]), num_recs]] # each item out of this is a handle to AGDReads agd_read_assembler_gen = tuple( pipeline.agd_bwa_read_assembler( upstream_tensors=to_assembler_converted, include_meta=False)) # assembled_records, ready_to_align: [(agd_reads_handle, (num_records, first_ordinal, record_id), (pass_around)) x N] assembled_records_gen = tuple( zip(agd_read_assembler_gen, pass_around_assembler)) assembled_records = tuple( ((a, ) + tuple(b)) for a, b in assembled_records_gen) ready_to_align = pipeline.join( upstream_tensors=assembled_records, parallel=args.aligners, capacity=8, multi=True) # TODO still have default capacity here :/ options = persona_ops.bwa_options(options=args.bwa_args.split(), name="bwa_aligner_options") bwa_index = persona_ops.bwa_index(index_location=args.index_path, ignore_alt=False, name="index_loader") first_assembled_result = ready_to_align[0][1:] sink_queue_shapes = [a.get_shape() for a in first_assembled_result] sink_queue_dtypes = [a.dtype for a in first_assembled_result] aligner_dtype = tf.string aligner_shape = (args.max_secondary + 1, 2) sink_queue_shapes.append(aligner_shape) sink_queue_dtypes.append(aligner_dtype) pass_around_aligners = tuple( a[1:] for a in ready_to_align ) # type: [(num_records, first_ordinal, record_id, pass_around x N) x N] pass_to_aligners = tuple(a[0] for a in ready_to_align) buffer_list_pool = persona_ops.buffer_list_pool( **pipeline.pool_default_args) if args.paired: executor = persona_ops.bwa_paired_executor( max_secondary=args.max_secondary, num_threads=args.aligner_threads, work_queue_size=args.aligners + 1, options_handle=options, index_handle=bwa_index, thread_ratio=0.66) else: executor = persona_ops.bwa_single_executor( max_secondary=args.max_secondary, num_threads=args.aligner_threads, work_queue_size=args.aligners + 1, options_handle=options, index_handle=bwa_index) def make_aligners(): aligner_op = persona_ops.bwa_align_single if not args.paired else persona_ops.bwa_align_paired for read_handle, pass_around in zip(pass_to_aligners, pass_around_aligners): aligner_results = aligner_op(read=read_handle, buffer_list_pool=buffer_list_pool, subchunk_size=args.subchunking, executor_handle=executor, max_secondary=args.max_secondary) #print(aligner_results) yield (aligner_results, ) + tuple(pass_around) aligners = tuple(make_aligners()) if args.compress_parallel > 0: aligner_results_to_compress = pipeline.join( upstream_tensors=aligners, parallel=args.compress_parallel, multi=True, capacity=4, name="ready_to_compress") to_compressors = (a[0] for a in aligner_results_to_compress) around_compressors = (a[1:] for a in aligner_results_to_compress) compressed_buffers = pipeline.aligner_compress_pipeline( upstream_tensors=to_compressors) after_compression = ( (a, ) + tuple(b) for a, b in zip(compressed_buffers, around_compressors)) aligners = tuple(after_compression) aligned_results = pipeline.join(upstream_tensors=aligners, parallel=args.writers, multi=True, capacity=32) ref_seqs, lens = persona_ops.bwa_index_reference_sequences( index_handle=bwa_index) return aligned_results, ( bwa_index, ref_seqs, lens ) # returns [(buffer_list_handle, num_records, first_ordinal, record_id, pass_around X N) x N], that is COMPLETELY FLAT