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