Improve load time (mpsc channel, post page-lock, directory buckets)

[eaypek-tfh]

Change

This PR improves load time via:

1. Switching from (returning vector of streams followed by select_all and try_next) to (spawning a separate tokio task for each get_object and sending items to a tokio mpsc channel). Instead of polling multiple dynamically changed list of streams, we now wait for a single channel.
2. Having page locking later than import: Working on memory at the same time with page-lock and import writes is affecting the performance. When we move the page-lock after the import, we saw that page-lock time decreases by 70% while import time remains almost the same
3. Infra change: Directory buckets (s3express) bring more stable s3 throughput which used to be varying a lot for general purpose buckets. Overall, it improved the performance in stage by around 20-25%.

Background & More Details

• We need page-lock for better batch performance and this is why we do not start processing them before page lock finishes. It takes around 7,75s to page-lock 1M records in memory. Parallelizing it does not help as it's probably serialized on OS level.
• We used to start page-lock in the background and immediately start s3 and aurora streams to write into the memory. This used to affect memory write and stream performance indirectly.
• Since we move page-locked after importing everything, the real s3 throughput became obvious where we exceeded aws defined network throughput by around 20-25% in our stage instances.
• In order not to waste the whole time for page-lock, we tried to divide it into chunks and start writing into the memory chunk which is already page-locked. However, chunked page locks resulted in CUDA_ERROR_INVALID_VALUE errors. So, we abandoned them.
• With these changes, current bottleneck seems to be memory writing
• Please note that we can further avoid the time to send items to tokio mpsc channel by directly writing into the memory in importer.

[danielle-tfh]

out of interest - why do we divide by the count?

[philsippl]

Because we actually don't allocate one blob of memory, but device_count many. Elements are allocated to devices by serial_id % device_index. This is mostly a leftover from when the memory was loaded onto the devices, we could also do it differently now.

[philsippl]

@eaypek-tfh what's the reason for the manual drop?

[eaypek-tfh]

Yeah, I don't need to drop it expicitly as it'd be implicitly dropped after all clones are dropped, right? I think I just wanted to make it obvious but let me remove it 👍

[philsippl]

Yeah should be all fine, if anything this is a bit awkward since you drop it before the handles are awaited.

[philsippl]

lgtm!