Dorado polish low GPU utilisation for bam regions of variable coverage

[lorenzwidmer]

New issue checks

• I have read the Dorado Documentation.
• I did not find an existing issue.

Dorado version

1.1.1+e72f1492

Dorado subcommand

Polish

The issue

Dorado polish can sometimes show inconsistent GPU utilisation when processing BAM files containing many short regions (e.g. in my case, >260,000 regions of ~1.5 kB for amplicon polishing) with inconsistent region coverage, or BAMs with inconsistent region lengths. I have mostly seen this on our cluster with a very fast graphics card but sometimes low I/O speeds. In this case, the infer threads seem to be faster than the encoder threads, with the trace log showing a queue size of 0: [consumer x] Popped data: item.samples.size() = 1, queue size: 0. Increasing the number of threads for encoding doesn't seem to help either.

I've noticed that the CPU usage rises after the debug message [debug] Starting to encode regions for X windows using X threads and then gradually drops until the next batch of windows is encoded. This indicates that the computational load is not evenly distributed among the encoders and that all encoders wait for an entire batch of windows to finish before pushing it into the infer queue and moving on to the next one. This cannot be fixed by providing more threads for encoding, as they all wait for the slowest thread to complete.

System specifications

Operating System: Ubuntu 22.04.4 LTS
CPU: AMD EPYC 7742 64-Core Processor (64 cores / 64 threads)
RAM: 120 GB
GPU: NVIDIA A100-SXM4-40GB / NVIDIA H100-80GB
Disk: Sometimes very slow

[lorenzwidmer]

From my understanding of the code, the issue stems from how batches are created in dorado::polish::sample_producer() and split in dorado::polish::encode_windows_in_parallel().
Currently, BAM windows are split into batches with a maximum size of num_threads BAM regions:

// Divide draft sequences into groups of specified size, as sort of a barrier.
const std::vector<secondary::Interval> bam_region_batches = secondary::create_batches(
        bam_region_intervals, num_threads,
        [](const secondary::Interval& val) { return val.end - val.start; });

For BAMs with regions shorter than the specified window size, the code produces batches of up to num_threads windows. For BAMs with longer regions, the batches contain more windows.

Later, in dorado::polish::encode_windows_in_parallel(), these batches are split again into num_threads chunks.

const std::vector<secondary::Interval> thread_chunks = secondary::compute_partitions(
        static_cast<int32_t>(std::size(windows)),
        std::max(num_threads, static_cast<int32_t>(std::size(encoders))));
        
spdlog::debug("Starting to encode regions for {} windows using {} threads.", 
              std::size(windows), std::size(thread_chunks));

This results in a fixed number of windows being allocated to each thread, which can lead to an uneven distribution of workloads and prevent a thread that has been blocked for a long time from catching up.

One possible improvement would be to specify the number of windows to be encoded together (with num_threads being the default) and replace the fixed partitioning of windows with a thread-safe queue. This would allow encoders to dynamically pull work until all windows have been processed.

I have implemented these improvements and they work reasonably well, achieving much higher speeds on our cluster. I can submit a pull request if desired. However, this change may introduce other issues that I did not consider.

[malton-ont]

Hi @lorenzwidmer,

Thanks for taking the time to investigate this. Please do submit a pull request and we can test it internally.