Skip to content

Create pull request #26

New issue

Have a question about this project? Sign up for a free GitHub account to open an issue and contact its maintainers and the community.

Sign up for GitHub

By clicking “Sign up for GitHub”, you agree to our terms of service and privacy statement. We’ll occasionally send you account related emails.

Already on GitHub? Sign in to your account

Merged
vadanamu merged 2 commits into from
Apr 9, 2026
Merged

Create pull request #26

Show file tree
Hide file tree
Changes from all commits
Commits
Show all changes
2 commits
Select commit Hold shift + click to select a range
c02b49c
[preprocess] Fix concurrent Pod5FileReader access crash
donghyuka Apr 7, 2026
3121020
[preprocess] Optimize dispatch and convert-later pattern
donghyuka Apr 7, 2026
File filter

Filter by extension

Filter by extension
Conversations
Failed to load comments.
Loading
Jump to
Jump to file
Failed to load files.
Loading
Diff view
Diff view
12 changes: 11 additions & 1 deletion cpp/src/args/ArgumentParser.cpp
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -57,6 +57,8 @@ void ArgumentParser::print_help(const char* program_name)
<< " -w, --sig-window NUM Signal window size (default: " << args.sig_window << ")\n"
<< " -g, --filter-flag NUM BAM flag bits to filter (default: " <<
args.filter_flag << ")\n"
<< " -Q, --max-queue NUM Max BAM queue size, 0=auto (default: " <<
args.max_queue << ")\n"
<< " -d, --label-div NUM Label division factor (default: " << args.label_div << ")\n"
<< " -h, --help Show this help message\n"
<< " -v, --version Show version information\n";
Expand Down Expand Up @@ -89,6 +91,7 @@ Arguments ArgumentParser::parse(int argc, char* argv[])
{"dwell-shift", required_argument, 0, 'f'},
{"sig-window", required_argument, 0, 'w'},
{"filter-flag", required_argument, 0, 'g'},
{"max-queue", required_argument, 0, 'Q'},
{"label-div", required_argument, 0, 'd'},
{"consistency", no_argument, 0, 'C'},
{"help", no_argument, 0, 'h'},
Expand All @@ -97,7 +100,7 @@ Arguments ArgumentParser::parse(int argc, char* argv[])
};

int option_index = 0;
while ((opt = getopt_long(argc, argv, "p:b:o:t:q:k:z:s:y:e:l:a:n:f:w:g:d:Chv", long_options,
while ((opt = getopt_long(argc, argv, "p:b:o:t:q:k:z:s:y:e:l:a:n:f:w:g:Q:d:Chv", long_options,
&option_index)) != -1) {
switch (opt) {
case 'p':
Expand Down Expand Up @@ -211,6 +214,13 @@ Arguments ArgumentParser::parse(int argc, char* argv[])
exit(1);
}
break;
case 'Q':
args.max_queue = atoi(optarg);
if (args.max_queue < 0) {
cerr << "Error: Max queue size must be non-negative\n";
exit(1);
}
break;
case 'C':
args.consistency = true;
break;
Expand Down
2 changes: 2 additions & 0 deletions cpp/src/args/ArgumentParser.h
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -37,6 +37,7 @@ struct Arguments {
int dwell_shift;
int sig_window;
int filter_flag;
int max_queue;
uint64_t label_div;
bool consistency;

Expand All @@ -56,6 +57,7 @@ struct Arguments {
dwell_shift = 10;
sig_window = 5;
filter_flag = 276; // unmapped 4 + reversed 16 non-primary 256
max_queue = 0; // 0 = auto (process_once * 4 * worker_count)
label_div = 1000000000ULL;
consistency = false;
}
Expand Down
38 changes: 35 additions & 3 deletions cpp/src/main.cpp
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -117,6 +117,9 @@ void process_merged_data_meta_worker(int worker_id, const Arguments& args,
log_info() << "Starting record merger worker " << worker_id
<< " with " << pod5_meta_records.size() << " POD5 records" << endl;

// Per-worker POD5 readers, opened lazily for thread-safe access
unordered_map<string, Pod5FileReader_t*> worker_readers;

// Set normalization factors same as Python code
NormalizationFactors norm_factors;

Expand Down Expand Up @@ -155,7 +158,21 @@ void process_merged_data_meta_worker(int worker_id, const Arguments& args,
auto bam_it = bam_index.find(pod5_meta.read_id);
if (bam_it != bam_index.end()) {
for (auto& bam_rec : bam_it->second) {
batch_pod5_meta.push_back(pod5_meta);
// Lazy-open per-worker reader for thread-safe POD5 access
Pod5RecordMeta meta_copy = pod5_meta;
auto [rit, inserted] = worker_readers.try_emplace(meta_copy.file_path, nullptr);
if (inserted) {
rit->second = pod5_open_file(meta_copy.file_path.c_str());
if (!rit->second) {
log_err() << "Record merger worker " << worker_id
<< " failed to open POD5: "
<< meta_copy.file_path << endl;
}
}
if (!rit->second)
continue;
meta_copy.reader = rit->second;
batch_pod5_meta.push_back(move(meta_copy));
batch_bam.push_back(move(bam_rec));
}
keys_to_erase.push_back(pod5_meta.read_id);
Expand All @@ -179,7 +196,7 @@ void process_merged_data_meta_worker(int worker_id, const Arguments& args,

if (!processed_records.empty()) {
// Write to NPZ
writer.add_records(processed_records);
writer.add_records(move(processed_records));
}

writer.increment_processing_unit();
Expand All @@ -196,13 +213,20 @@ void process_merged_data_meta_worker(int worker_id, const Arguments& args,

// Flush remaining records (same as Python code)
writer.flush();

// Close per-worker POD5 readers
for (auto& [path, reader] : worker_readers)
if (reader)
pod5_close_and_free_reader(reader);

log_info() << "Record merger worker " << worker_id << " completed" << endl;
}


int main(int argc, char* argv[])
{
auto start_time = chrono::high_resolution_clock::now();
log_start_time();

// Parse arguments
ArgumentParser parser;
Expand Down Expand Up @@ -358,6 +382,8 @@ int main(int argc, char* argv[])
pod5_reader_records.clear();
pod5_reader_records.shrink_to_fit();

log_info() << "[MEM] POD5 redistributed: " << get_rss_mb() << " MB" << endl;

// Create POD5 index (read_id -> worker_id mapping)
log_info() << "Creating POD5 index for merged data workers" << endl;
unordered_map<string, int> pod5_index;
Expand All @@ -368,6 +394,8 @@ int main(int argc, char* argv[])
}
}

log_info() << "[MEM] pod5_index created: " << get_rss_mb() << " MB" << endl;

// Check if consistency mode is enabled
if (args.consistency) {
// Consistency mode: Read all BAM records first with multiple workers, and distribute to workers
Expand Down Expand Up @@ -434,7 +462,7 @@ int main(int argc, char* argv[])
log_info() << "Starting merged data processing with "
<< num_workers << " workers" << endl;

// Index BAM records by read_id - multiple BAM records can have same read_id
// Index BAM records by parent_id - multiple BAM records can have same read_id
unordered_map<string, vector<BamRecord>> bam_index;
for (auto& bam_rec : entire_bam_records) {
bam_index[bam_rec.parent_id].push_back(move(bam_rec));
Expand Down Expand Up @@ -469,12 +497,16 @@ int main(int argc, char* argv[])
}

log_info() << "Started " << num_workers << " merged data workers" << endl;
log_info() << "[MEM] workers created: " << get_rss_mb() << " MB" << endl;

// Set workers for SAM dispatcher
sam_dispatcher.set_workers(&merged_workers, &pod5_index);

log_info() << "[MEM] dispatch started: " << get_rss_mb() << " MB" << endl;

// Wait for SAM dispatcher to finish
sam_dispatcher.stop();
log_info() << "[MEM] dispatch completed: " << get_rss_mb() << " MB" << endl;

// Signal all workers to stop (notify them to process final batches)
for (auto worker : merged_workers) {
Expand Down
121 changes: 69 additions & 52 deletions cpp/src/merger/MergedDataWorker.cpp
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -83,6 +83,12 @@ namespace deeprm {
void MergedDataWorker::process_loop()
{
NormalizationFactors norm_factors;
RecordMerger merger(norm_factors, args.cb_len, args.kmer_len,
args.max_token_len, args.sampling, args.dwell_shift,
args.sig_window, args.label_div);

// Per-worker POD5 readers, opened lazily for thread-safe access
unordered_map<string, Pod5FileReader_t*> worker_readers;

unordered_map<string, Pod5RecordMeta> pod5_index;
// Build POD5 index for this worker
Expand All @@ -94,7 +100,8 @@ namespace deeprm {
<< " indexed " << pod5_meta_records.size()
<< " POD5 records" << endl;

vector<pair<Pod5RecordMeta, BamRecord>> batch_merged;
vector<pair<const Pod5RecordMeta*, bam1_t*>> batch_queued;
vector<ProcessedRecord> processed_records;

while (is_running || !bam_queue.empty()) {
unique_lock<mutex> lock(queue_mutex);
Expand All @@ -110,82 +117,86 @@ namespace deeprm {
bam_queue.pop();
lock.unlock();

// Convert bam1_t to BamRecord
BamRecord bam_record;
if (!convert_bam1_to_record(read, bam_header, bam_record))
continue;
// Extract parent_id from bam1_t for POD5 lookup
uint8_t* pi_tag = bam_aux_get(read, "pi");
string parent_id = pi_tag ? string(bam_aux2Z(pi_tag))
: string(bam_get_qname(read));

// Find matching POD5 record
auto pod5_it = pod5_index.find(bam_record.parent_id);
auto pod5_it = pod5_index.find(parent_id);
if (pod5_it != pod5_index.end()) {
batch_merged.emplace_back(pod5_it->second, bam_record);
// Lazy-open per-worker reader for thread-safe POD5 access
auto& meta = pod5_it->second;
auto [rit, inserted] = worker_readers.try_emplace(meta.file_path, nullptr);
if (inserted) {
rit->second = pod5_open_file(meta.file_path.c_str());
if (!rit->second) {
log_err() << "merged data worker " << worker_id
<< " failed to open POD5: " << meta.file_path << endl;
}
}
if (!rit->second) {
bam_destroy1(read);
continue;
}
meta.reader = rit->second;
batch_queued.emplace_back(&meta, read);

// Process batch when it reaches process_once size
if (batch_merged.size() >= static_cast<size_t>(args.process_once)) {
if (batch_queued.size() >= static_cast<size_t>(args.process_once)) {
output_index++;
log_info() << "merged data worker " << worker_id
<< " processing batch " << output_index
<< " with " << batch_merged.size() << " records" << endl;

// Create merger and process
RecordMerger merger(norm_factors, args.cb_len, args.kmer_len, args.max_token_len,
args.sampling, args.dwell_shift, args.sig_window, args.label_div);

vector<Pod5RecordMeta> batch_pod5_meta;
vector<BamRecord> batch_bam;
batch_pod5_meta.reserve(batch_merged.size());
batch_bam.reserve(batch_merged.size());

for (auto& pair : batch_merged) {
batch_pod5_meta.push_back(move(pair.first));
batch_bam.push_back(move(pair.second));
if ((output_index * args.process_once) % 1000 < args.process_once) {
if (worker_id == 0) {
log_info() << "[MEM] worker 0 pre-batch " << output_index
<< " (queue=" << bam_queue.size()
<< ", readers=" << worker_readers.size()
<< "): " << get_rss_mb() << " MB" << endl;
}
log_info() << "merged data worker " << worker_id
<< " processing batch " << output_index
<< " with " << batch_queued.size() << " records" << endl;
}

merger.add_bam_records(move(batch_bam));
merger.add_pod5_meta_records(move(batch_pod5_meta));

vector<ProcessedRecord> processed_records = merger.merge_and_process_with_meta();
for (auto& [meta_ptr, bam_read] : batch_queued) {
BamRecord bam_record;
if (convert_bam1_to_record(bam_read, bam_header, bam_record))
merger.process_merged_record_with_meta(bam_record, *meta_ptr, processed_records);
bam_destroy1(bam_read);
}

if (!processed_records.empty()) {
writer->add_records(processed_records);
writer->add_records(move(processed_records));
processed_records.clear();
}

writer->increment_processing_unit();
batch_merged.clear();
batch_queued.clear();
if (worker_id == 0 && (output_index * args.process_once) % 1000 < args.process_once) {
log_info() << "[MEM] worker 0 post-batch " << output_index
<< ": " << get_rss_mb() << " MB" << endl;
}
}
} else {
bam_destroy1(read);
}

// Clean up
bam_destroy1(read);
}

// Process remaining records
if (!batch_merged.empty()) {
if (!batch_queued.empty()) {
output_index++;
log_info() << "merged data worker " << worker_id
<< " processing final batch " << output_index
<< " with " << batch_merged.size() << " records" << endl;

RecordMerger merger(norm_factors, args.cb_len, args.kmer_len, args.max_token_len,
args.sampling, args.dwell_shift, args.sig_window, args.label_div);

vector<Pod5RecordMeta> batch_pod5_meta;
vector<BamRecord> batch_bam;
batch_pod5_meta.reserve(batch_merged.size());
batch_bam.reserve(batch_merged.size());
<< " with " << batch_queued.size() << " records" << endl;

for (auto& pair : batch_merged) {
batch_pod5_meta.push_back(move(pair.first));
batch_bam.push_back(move(pair.second));
for (auto& [meta_ptr, bam_read] : batch_queued) {
BamRecord bam_record;
if (convert_bam1_to_record(bam_read, bam_header, bam_record))
merger.process_merged_record_with_meta(bam_record, *meta_ptr, processed_records);
bam_destroy1(bam_read);
}

merger.add_bam_records(move(batch_bam));
merger.add_pod5_meta_records(move(batch_pod5_meta));

vector<ProcessedRecord> processed_records = merger.merge_and_process_with_meta();

if (!processed_records.empty()) {
writer->add_records(processed_records);
writer->add_records(move(processed_records));
}

writer->increment_processing_unit();
Expand All @@ -197,6 +208,12 @@ namespace deeprm {
writer->flush();
}

// Close per-worker POD5 readers
for (auto& [path, reader] : worker_readers) {
if (reader)
pod5_close_and_free_reader(reader);
}

log_info() << "merged data worker " << worker_id
<< " completed processing. Total batches: "
<< output_index << endl;
Expand Down
8 changes: 4 additions & 4 deletions cpp/src/merger/RecordMerger.h
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -50,15 +50,15 @@ namespace deeprm {
vector<BamRecord> bam_records;
vector<Pod5RecordMeta> pod5_meta_records;

bool process_merged_record_with_meta(const BamRecord& bam_rec,
const Pod5RecordMeta& pod5_meta,
vector<ProcessedRecord>& output) const;

public:
RecordMerger(const NormalizationFactors& nf, int cb_len, int kmer_len,
int max_token_len, int sampling, int dwell_shift, int sig_window,
uint64_t label_div);

bool process_merged_record_with_meta(const BamRecord& bam_rec,
const Pod5RecordMeta& pod5_meta,
vector<ProcessedRecord>& output) const;

void add_bam_records(vector<BamRecord>&& records);
void add_pod5_meta_records(vector<Pod5RecordMeta>&& meta_records);
vector<ProcessedRecord> merge_and_process_with_meta();
Expand Down
Loading
Loading