Skip to content

Unet3d Object: s3dlio per-process GET throughput #701

Description

@bissantz

The issue

MLPerf Storage unet3d on S3 is client-bound, not storage-bound. On a 32-vCPU / 128 GB client node against an endpoint that serves ~10 GB/s to that single node (independently proven by warp), stock s3dlio 0.9.106 delivers:

  • ~1.1 GB/s per process through the loader path the benchmark uses — flat, regardless of prefetch depth;
  • In the benchmark with 8 DataLoader worker processes this contends down to ~700 MB/s each resulting in ~5.3 GB/s node total, where it caps with increasing the worker number, which is below the 6,046 MB/s one emulated B200 accelerator demands.

Root causes (both in s3dlio, verified by code audit + isolated A/B):

  1. The loader drives all fetches from one async task (buffer_unordered inside a single spawned producer) — all request
    driving and body accumulation share ~one core.
  2. The HTTP connector fully buffers every response body before the SDK caller sees byte one, and the range engine then copies everything again into a second assembly buffer.

Reproduction environment

Nodes: 32 vCPU, 128 GB RAM (the datagen sizing derives from this).
Stack: unmodified mlpstorage tree + s3dlio 0.9.106 venv; Credentials/endpoint in .env.

./01_datagen.sh          # 4,687 x ~140 MiB (5 x 128 GiB rule) -> s3://<bucket>/s3dlio-repro
./02_read_s3dlio.sh      # 1 / 4 / 8 s3dlio processes -> 1,147 / 3,371 / 5,338 MB/s
./03_read_mlpstorage.sh  # full benchmark, 1 node x 1 accelerator -> 65% AU / 3,929 MB/s (fails)
./04_read_s5cmd.sh       # 1x32 vs 4x8 procs, same 32 streams     -> 6,204 vs 8,111 MB/s
./05_read_warp.sh        # single warp process, same objects      -> ~9-10 GiB/s

01_datagen.sh

#!/bin/bash
# Generate the unet3d dataset for ONE client node with 128 GB RAM.
# MLPerf sizing rule: dataset >= 5 x client RAM
#   -> 5 x 128 GiB / 146,600,628 B per record = 4,687 files (~642 GiB).
set -e
set -a; . "$(dirname "$0")/.env"; set +a

PREFIX="s3://$BUCKET/s3dlio-repro"
STACK=<path_to_mlpstorage_repo> 
NUM_FILES=4687
SUBFOLDERS=14
WRITERS=16

export PATH="$STACK/.venv/bin:$PATH"
[ -f ./results/mlperf-results.yaml ] || mlpstorage init Repro ./results

mlpstorage whatif training unet3d datagen object \
    --hosts localhost --num-processes $WRITERS --data-dir "$PREFIX" \
    --results-dir ./results --systemname s3dlio-repro \
    --allow-run-as-root --oversubscribe --mpi-params="-x TMPDIR=/tmp" \
    --params dataset.num_files_train=$NUM_FILES dataset.num_subfolders_train=$SUBFOLDERS

02_read_s3dlio.sh

#!/bin/bash
# Read the ENTIRE dataset via the exact loader path the MLPerf benchmark
# uses (PyDataset -> PyBytesAsyncDataLoader), three ways (measured):
#   A) 1 process,  prefetch 16      -> 1,147 MB/s (the per-process wall)
#   B) 4 processes, prefetch 4 each -> 3,371 MB/s (SAME 16 total streams)
#   C) 8 processes, prefetch 4 each -> 5,338 MB/s (the benchmark's own
#      worker shape: per-process degrades to ~700 MB/s under contention —
#      NOT 8 x 1.1; this IS the benchmark's observed node plateau)
# Each pid also reports batch-arrival stalls (p50/max gap): 0.9 s -> 1.5 s
# p50 and up to 4 s maxima at 8 processes. The single producer task makes
# completions slow and bursty; the benchmark consumes its 8 workers
# round-robin, so one worker's stall idles the whole rotation — which is
# why the benchmark realizes even less (~3.9 GB/s) than isolated C.
set -e
set -a; . "$(dirname "$0")/.env"; set +a

PREFIX="s3://$BUCKET/s3dlio-repro"
PY=<path_to_venv_python>

read_uris() { # $1 = uri file, $2 = prefetch depth
    URIS=$1 PREFETCH=$2 "$PY" - <<'EOF'
import os, statistics, time, s3dlio
uris = open(os.environ["URIS"]).read().split()
ds = s3dlio.PyDataset.from_uris(uris)
it = s3dlio.PyBytesAsyncDataLoader(ds, {"prefetch": int(os.environ["PREFETCH"])}).items()
total, gaps, last, t0 = 0, [], time.time(), time.time()
while batch := it.collect_batch(7):
    now = time.time(); gaps.append(now - last); last = now
    total += sum(len(b) for b in batch)
dt = time.time() - t0
print(f"   pid {os.getpid()}: {total/dt/1e6:.0f} MB/s ({len(uris)} objects, {dt:.0f} s)"
      f"  batch gaps p50 {statistics.median(gaps)*1e3:.0f} ms / max {max(gaps)*1e3:.0f} ms")
EOF
}

aws --endpoint-url "$AWS_ENDPOINT_URL" s3 ls "$PREFIX/" --recursive \
    | awk -v b="$BUCKET" '$NF ~ /npz$/ {print "s3://" b "/" $NF}' > /tmp/repro_uris.txt
BYTES=$(aws --endpoint-url "$AWS_ENDPOINT_URL" s3 ls "$PREFIX/" --recursive | awk '$NF ~ /npz$/ {s += $3} END {print s}')

run_procs() { # $1 = number of processes, $2 = prefetch per process
    split -n l/$1 -d /tmp/repro_uris.txt /tmp/repro_part.
    t0=$(date +%s.%N)
    for p in /tmp/repro_part.*; do read_uris "$p" "$2" & done
    wait
    t1=$(date +%s.%N)
    awk -v b=$BYTES -v t0=$t0 -v t1=$t1 'BEGIN {printf "   aggregate: %.0f MB/s\n", b/(t1-t0)/1e6}'
    rm -f /tmp/repro_part.*
}

echo "A) 1 s3dlio process, prefetch 16, whole dataset:"
read_uris /tmp/repro_uris.txt 16
echo "B) 4 s3dlio processes, prefetch 4 each (same 16 total streams):"
run_procs 4 4
echo "C) 8 s3dlio processes, prefetch 4 each (the benchmark's worker shape):"
run_procs 8 4
echo "D) 16 s3dlio processes, prefetch 4 each (the benchmark's worker shape):"
run_procs 16 4

03_read_mlpstorage.sh

#!/bin/bash
# The full benchmark: mlpstorage -> DLIO -> s3dlio, unet3d, 1 node x 1
# emulated b200 accelerator, 2 epochs over the entire dataset.
# Stock stack: ~70% AU (fails the 90% bar). With the three patches
# applied to s3dlio: ~95% AU at identical settings.
set -e
set -a; . "$(dirname "$0")/.env"; set +a

PREFIX="s3://$BUCKET/s3dlio-repro"
STACK=<path_to_mlpstorage_repo>        
NUM_FILES=4687
SUBFOLDERS=14
RT=8                                        # reader.read_threads (CLOSED-allowed)
PW=4                                       # storage_options.prefetch_window (CLOSED-allowed, #666)

export PATH="$STACK/.venv/bin:$PATH"
[ -f ./results/mlperf-results.yaml ] || mlpstorage init Repro ./results

mlpstorage whatif training unet3d run object \
    --hosts localhost --num-accelerators 1 --num-client-hosts 1 --accelerator-type b200 \
    --client-host-memory-in-gb 128 --data-dir "$PREFIX" --results-dir ./results \
    --systemname s3dlio-repro --allow-run-as-root --oversubscribe --mpi-params="-x TMPDIR=/tmp" \
    --params dataset.num_files_train=$NUM_FILES dataset.num_subfolders_train=$SUBFOLDERS \
             dataset.skip_listing=True reader.read_threads=$RT train.epochs=2 \
             storage.storage_options.prefetch_window=$PW

python3 - <<'EOF'
import glob, json
s = sorted(glob.glob("results/whatif/*/results/s3dlio-repro/training/unet3d/run/*/summary.json"))[-1]
m = json.load(open(s))["metric"]
print(f"mlpstorage unet3d: AU {m['train_au_mean_percentage']:.2f}%  "
      f"train {m['train_io_mean_MB_per_second']:.0f} MB/s   (pass bar: AU >= 90%)")
EOF

04_read_s5cmd.sh

#!/bin/bash
# Read the ENTIRE dataset with s5cmd (independent Go client), two ways:
#   A) 1 process  x 32 workers  -> the per-process ceiling
#   B) 4 processes x 8 workers  -> SAME 32 total streams   
# Throughput scales with PROCESS count at constant stream count: the
# ceiling is per-process client behavior, not the endpoint or connections.
set -e
set -a; . "$(dirname "$0")/.env"; set +a

PREFIX="s3://$BUCKET/s3dlio-repro"
S5CMD=<path_to_s5cmd>

aws --endpoint-url "$AWS_ENDPOINT_URL" s3 ls "$PREFIX/" --recursive \
    | awk -v b="$BUCKET" '$NF ~ /npz$/ {print $3 " cat s3://" b "/" $NF}' > /tmp/repro_sized.txt
BYTES=$(awk '{s += $1} END {print s}' /tmp/repro_sized.txt)
cut -d' ' -f2- /tmp/repro_sized.txt > /tmp/repro_all.cmds

echo "A) 1 process x 32 workers, $BYTES bytes:"
t0=$(date +%s.%N)
"$S5CMD" --endpoint-url "$AWS_ENDPOINT_URL" --numworkers 32 run /tmp/repro_all.cmds > /dev/null
t1=$(date +%s.%N)
awk -v b=$BYTES -v t0=$t0 -v t1=$t1 'BEGIN {printf "   %.0f MB/s (%.0f s)\n", b/(t1-t0)/1e6, t1-t0}'

echo "B) 4 processes x 8 workers (same 32 total streams):"
split -n l/4 -d /tmp/repro_all.cmds /tmp/repro_part.
t0=$(date +%s.%N)
for p in /tmp/repro_part.0*; do
    "$S5CMD" --endpoint-url "$AWS_ENDPOINT_URL" --numworkers 8 run "$p" > /dev/null &
done
wait
t1=$(date +%s.%N)
awk -v b=$BYTES -v t0=$t0 -v t1=$t1 'BEGIN {printf "   %.0f MB/s (%.0f s)\n", b/(t1-t0)/1e6, t1-t0}'

05_read_warp.sh

#!/bin/bash
# Read the same objects with MinIO warp (Go, minio-go): 
# NOTE: --noclear is mandatory — warp's DEFAULT DELETES the bucket contents.
set -e
set -a; . "$(dirname "$0")/.env"; set +a

WARP=<path_to_warp>
HOST="${AWS_ENDPOINT_URL#https://}"

"$WARP" get --host "$HOST" --tls \
    --access-key "$AWS_ACCESS_KEY_ID" --secret-key "$AWS_SECRET_ACCESS_KEY" \
    --bucket "$BUCKET" --prefix s3dlio-repro --list-existing --noclear \
    --concurrent 32 --duration 2m 2>&1 | grep -E "Report:|Average:|TTFB"

Scripts 02–05 read the identical objects (numbers above are measured on
this environment). The chain of evidence:

  • 02-A: one s3dlio process is walled at 1,147 MB/s, flat vs prefetc depth. 02-B/C: aggregate scales only with process count at constant
    stream count, and sub-linearly — at the benchmark's own shape (8 processes x prefetch 4) per-process throughput degrades in. my caese to ~700 MB/s and the node totals ~5GB/s. That, not 8 x 1.1, is the benchmark's node plateau; one emulated B200 accelerator demands 6,046 MB/s, so the stock client cannot pass regardless of worker count.
  • 02's stall stats: batch-arrival gaps grow from p50 0.9 s (solo) to p50 1.5 s / max 4 s at 8 processes. DLIO's consumer drains its workers
    round-robin, so one worker's stall idles the rotation — which is why the benchmark (03) realizes only ~4GB/s of the ~5GB/s available.
  • 04: an independent Go client shows the same per-process scaling law at a higher level; 05: a single warp process reaches ~10 GiB/s on the same node and objects — the wall is s3dlio's, not the endpoint's.

Why it likely went unnoticed so far: the defect only binds when the storage outruns ~1 GB/s per client process AND the emulated accelerator
demands more. Slower endpoints hide it (storage is the bottleneck), and older accelerator profiles hide it (an a100/h100 unet3d profile demands a fraction of the B200's 6,046 MB/s, so 8 walled workers still keep up). B200-class demand against competent object storage — the target regime of MLPerf Storage v3 — is exactly where it becomes the limiting factor, for every submitter.

Proposed fixes (one per issue, apply in order)

With patches 0001+0003+0004 applied, one s3dlio process moves ~4× the data of stock (1.1 → 4.6 GB/s, matching what s5cmd achieves per process), with warp's 10 GB/s showing the remaining headroom in the HTTP layer.

0001 — loader: spawn each fetch as its own runtime task

All fetch futures were polled inside ONE spawned producer task via buffer_unordered — request driving and body accumulation cooperatively shared a single worker thread (the ~1.1 GB/s per-process wall, flat vs prefetch depth). One tokio::spawn per fetch parallelizes across the runtime; buffer_unordered(prefetch) still bounds in-flight tasks and the bounded channel still bounds delivery. Same change at all three sites (__iter__, items(), Parquet loader). Measured: loader 1,147 → ~4,600 MB/s per process.

--- a/src/python_api/python_aiml_api.rs
+++ b/src/python_api/python_aiml_api.rs
                 let mut stream = stream::iter(0..len)
                     .map(|idx| {
                         let ds = dataset.clone();
-                        async move { ds.inner.get(idx).await }
+                        tokio::spawn(async move { ds.inner.get(idx).await })
                     })
                     .buffer_unordered(prefetch);

-                while let Some(result) = stream.next().await {
+                while let Some(joined) = stream.next().await {
+                    let result = match joined {
+                        Ok(r) => r,
+                        Err(e) => Err(DatasetError::from(format!(
+                            "loader fetch task panicked: {e}"
+                        ))),
+                    };
full patch: 0001-perf-loader-spawn-each-fetch-as-its-own-runtime-task.patch (apply with git am)
---
 src/python_api/python_aiml_api.rs | 46 ++++++++++++++++++++++++++-----
 1 file changed, 39 insertions(+), 7 deletions(-)

diff --git a/src/python_api/python_aiml_api.rs b/src/python_api/python_aiml_api.rs
index f3fb2af..419c0cb 100644
--- a/src/python_api/python_aiml_api.rs
+++ b/src/python_api/python_aiml_api.rs
@@ -233,14 +233,30 @@ impl PyBytesAsyncDataLoader {
                     prefetch
                 );

                 let mut stream = stream::iter(0..len)
                     .map(|idx| {
                         let ds = dataset.clone();
-                        async move { ds.inner.get(idx).await }
+                        tokio::spawn(async move { ds.inner.get(idx).await })
                     })
                     .buffer_unordered(prefetch);
 
-                while let Some(result) = stream.next().await {
+                while let Some(joined) = stream.next().await {
+                    let result = match joined {
+                        Ok(r) => r,
+                        Err(e) => Err(DatasetError::from(format!(
+                            "loader fetch task panicked: {e}"
+                        ))),
+                    };
                     if let Ok(ref bytes) = result {
                         total_bytes += bytes.len() as u64;
                         count += 1;
@@ -356,22 +372,30 @@ impl PyBytesAsyncDataLoader {
                     prefetch,
                 );
 
+                // PERF: per-fetch tokio::spawn — see __iter__ above for the
+                // full rationale (single-producer-task ceiling).
                 let mut stream = stream::iter(0..len)
                     .map(|idx| {
                         let ds = dataset.clone();
                         let k = Arc::clone(&keys);
-                        async move {
+                        tokio::spawn(async move {
                             let uri = k
                                 .as_ref()
                                 .as_ref()
                                 .and_then(|v| v.get(idx).cloned())
                                 .unwrap_or_else(|| format!("item:{}", idx));
                             ds.inner.get(idx).await.map(|b| (uri, b))
-                        }
+                        })
                     })
                     .buffer_unordered(prefetch);
 
-                while let Some(result) = stream.next().await {
+                while let Some(joined) = stream.next().await {
+                    let result = match joined {
+                        Ok(r) => r,
+                        Err(e) => Err(DatasetError::from(format!(
+                            "loader fetch task panicked: {e}"
+                        ))),
+                    };
                     if let Ok((_, ref bytes)) = result {
                         total_bytes += bytes.len() as u64;
                         count += 1;
@@ -2205,14 +2229,22 @@ impl PyParquetStreamLoader {
             let mut total_bytes: u64 = 0;
             let mut t_last_log = Instant::now();
 
+            // PERF: per-fetch tokio::spawn — see PyBytesAsyncDataLoader::__iter__
+            // for the full rationale (single-producer-task ceiling).
             let mut stream = stream::iter(0..total)
                 .map(|idx| {
                     let ds = Arc::clone(&dataset);
-                    async move { ds.get(idx).await.map(|b| (idx, b)) }
+                    tokio::spawn(async move { ds.get(idx).await.map(|b| (idx, b)) })
                 })
                 .buffer_unordered(concurrency);
 
-            while let Some(result) = stream.next().await {
+            while let Some(joined) = stream.next().await {
+                let result = match joined {
+                    Ok(r) => r,
+                    Err(e) => Err(DatasetError::from(format!(
+                        "loader fetch task panicked: {e}"
+                    ))),
+                };
                 if let Ok((_, ref b)) = result {
                     total_bytes += b.len() as u64;
                     count += 1;
-- 
2.43.0

0002 — http: pin the HTTPS client to HTTP/1.1

Hardening, not required for the numbers here. select_client() routes all https:// requests to a client that is neither .http1_only() nor
H2-tuned (all S3DLIO_H2_* env vars only configure the h2c client): on an h2-capable endpoint, ALPN silently negotiates HTTP/2 and hyper multiplexes ONE connection with a default 5 MiB window — an uncappable ~5 MiB/RTT per-process limit.

--- a/src/reqwest_client.rs
+++ b/src/reqwest_client.rs
@@ fn build_reqwest_client_raw()
                 win_cfg.conn_window_mb,
             );
         }
+    } else {
+        builder = builder.http1_only();
     }

     Ok(builder
full patch: 0002-fix-http-pin-the-HTTPS-client-to-HTTP-1.1.patch (apply with git am)
---
 src/reqwest_client.rs | 10 ++++++++++
 1 file changed, 10 insertions(+)

diff --git a/src/reqwest_client.rs b/src/reqwest_client.rs
index f2eafe9..07e784c 100644
--- a/src/reqwest_client.rs
+++ b/src/reqwest_client.rs
@@ -483,6 +483,16 @@ fn build_reqwest_client_raw(
                 win_cfg.conn_window_mb,
             );
         }
+    } else {
+        // This client serves ALL https:// requests (see select_client). The
+        // reqwest `http2` cargo feature adds h2 to the TLS ALPN list, so an
+        // h2-capable endpoint would silently negotiate HTTP/2 here and
+        // hyper-util would multiplex every request over ONE connection with
+        // hyper's default 5 MiB connection window — a hard ~5 MiB/RTT
+        // per-process cap that none of the S3DLIO_H2_* env vars can reach
+        // (they only configure the h2c client above). Pin to HTTP/1.1 so
+        // HTTPS traffic fans out across the pooled connections instead.
+        builder = builder.http1_only();
     }
 
     Ok(builder
-- 
2.43.0

0003 — range engine: stream chunks into pre-split segments, O(1) unsplit

concurrent_range_get_impl collected every chunk into its own Bytes and then re-copied all of them into a second BytesMut — an extra full copy of every byte and a transient 2× memory peak per object. Pre-split ONE zeroed allocation into per-chunk segments, write each chunk directly into its segment, rejoin in offset order with BytesMut::unsplit (O(1) for segments of one contiguous allocation).

--- a/src/s3_utils.rs  (concurrent_range_get_impl)
+++ b/src/s3_utils.rs
+    let mut backing = BytesMut::zeroed(total_bytes);
+    for (range_start, range_end, _) in &ranges {
+        segments.push(backing.split_to((*range_end - *range_start) as usize));
+    }
@@ per-chunk task
-            let body = resp.body.collect().await.context("collect chunk body")?;
-            let chunk_data = body.into_bytes();
-            Ok((buffer_offset, chunk_data))
+            let mut body = resp.body;
+            let mut written = 0usize;
+            while let Some(frag) = body.try_next().await.context("read chunk body")? {
+                seg[written..written + frag.len()].copy_from_slice(&frag);
+                written += frag.len();
+            }
+            Ok((idx, seg))
@@ assembly
-    let mut output = BytesMut::with_capacity(total_bytes);
-    for (_, chunk) in chunks {
-        output.extend_from_slice(&chunk);
-    }
+    parts.sort_unstable_by_key(|(idx, _)| *idx);
+    let (_, mut output) = iter.next().unwrap();
+    for (_, seg) in iter {
+        output.unsplit(seg);
+    }
     Ok(output.freeze())
full patch: 0003-perf-range-stream-chunks-into-pre-split-segments-O-1.patch (apply with git am)
---
 src/s3_utils.rs | 71 ++++++++++++++++++++++++++++++++++++++-----------
 1 file changed, 56 insertions(+), 15 deletions(-)

diff --git a/src/s3_utils.rs b/src/s3_utils.rs
index 1049351..d6460ca 100644
--- a/src/s3_utils.rs
+++ b/src/s3_utils.rs
@@ -1172,10 +1172,30 @@ async fn concurrent_range_get_impl(
         current_offset = chunk_end;
     }
 
+    if ranges.is_empty() {
+        return Ok(Bytes::new());
+    }
+
+    // PERF (issue 9.3): pre-split ONE zeroed allocation into per-chunk owned
+    // segments; each chunk task streams its response body directly into its
+    // segment (exactly one copy per byte), then contiguity is restored with
+    // O(1) BytesMut::unsplit in offset order. The old code collected every
+    // chunk into its own Bytes and then re-copied all of them into a second
+    // BytesMut — an extra full-object copy and a transient 2x memory peak.
+    let n_chunks = ranges.len();
+    let mut backing = BytesMut::zeroed(total_bytes);
+    let mut segments: Vec<BytesMut> = Vec::with_capacity(n_chunks);
+    for (range_start, range_end, _) in &ranges {
+        segments.push(backing.split_to((*range_end - *range_start) as usize));
+    }
+    debug_assert!(backing.is_empty());
+
     let semaphore = Arc::new(Semaphore::new(max_concurrency));
     let mut futures = FuturesUnordered::new();
 
-    for (range_start, range_end, buffer_offset) in ranges {
+    for (idx, ((range_start, range_end, _buffer_offset), mut seg)) in
+        ranges.into_iter().zip(segments).enumerate()
+    {
         let client = client.clone();
         let bucket = bucket.to_string();
         let key = key.to_string();
@@ -1192,7 +1212,7 @@ async fn concurrent_range_get_impl(
                 .get_object()
                 .bucket(&bucket)
                 .key(key.trim_start_matches('/'))
-                .range(range_header)
+                .range(&range_header)
                 .send()
                 .await
                 .sdk_context(format!(
@@ -1203,27 +1223,48 @@ async fn concurrent_range_get_impl(
                     range_end - 1
                 ))?;
 
-            let body = resp.body.collect().await.context("collect chunk body")?;
-            let chunk_data = body.into_bytes();
+            // Stream fragments straight into this chunk's segment —
+            // no intermediate whole-chunk buffer.
+            let mut body = resp.body;
+            let mut written = 0usize;
+            while let Some(frag) = body.try_next().await.context("read chunk body")? {
+                let end = written + frag.len();
+                if end > seg.len() {
+                    anyhow::bail!(
+                        "range GET s3://{}/{} bytes={}-{}: server returned more than requested ({} > {})",
+                        bucket, key, range_start, range_end - 1, end, seg.len()
+                    );
+                }
+                seg[written..end].copy_from_slice(&frag);
+                written = end;
+            }
+            if written != seg.len() {
+                anyhow::bail!(
+                    "range GET s3://{}/{} bytes={}-{}: short read ({} of {} bytes)",
+                    bucket, key, range_start, range_end - 1, written, seg.len()
+                );
+            }
 
-            // Return (buffer_offset, data) — no shared mutex, no contention
-            Ok::<(usize, Bytes), anyhow::Error>((buffer_offset, chunk_data))
+            Ok::<(usize, BytesMut), anyhow::Error>((idx, seg))
         });
     }
 
-    // Collect all (offset, chunk) pairs — order is non-deterministic (FuturesUnordered)
-    let mut chunks: Vec<(usize, Bytes)> = Vec::new();
+    // Collect filled segments — completion order is non-deterministic
+    let mut parts: Vec<(usize, BytesMut)> = Vec::with_capacity(n_chunks);
     while let Some(result) = futures.next().await {
-        chunks.push(result.context("concurrent range chunk failed")?);
+        parts.push(result.context("concurrent range chunk failed")?);
     }
 
-    // Sort by buffer offset, then assemble with a single sequential pass
-    chunks.sort_unstable_by_key(|(offset, _)| *offset);
-
-    let mut output = BytesMut::with_capacity(total_bytes);
-    for (_, chunk) in chunks {
-        output.extend_from_slice(&chunk);
+    // Restore original order, then rejoin adjacent segments zero-copy.
+    parts.sort_unstable_by_key(|(idx, _)| *idx);
+    let mut iter = parts.into_iter();
+    let (_, mut output) = iter
+        .next()
+        .ok_or_else(|| anyhow::anyhow!("range GET produced no chunks"))?;
+    for (_, seg) in iter {
+        output.unsplit(seg); // O(1): segments came from one contiguous allocation
     }
+    debug_assert_eq!(output.len(), total_bytes);
 
     Ok(output.freeze())
 }
-- 
2.43.0

0004 — connector: stream response bodies; bounded body-read retries

The connector fully buffered every response (resp.bytes().await) — the SDK caller saw byte one only after the last byte arrived: no
receive/consume overlap, an extra whole-body buffer per GET, and first_byte_time metrics that actually measured last-byte. Hand the SDK a live stream instead. Consequence handled in the same patch: with streaming, send() resolves at headers, so the SDK's transient-error retry no longer covers mid-body failures — every streaming-GET consumer (range chunks, S3Ops::get_object/get_object_range,
S3ObjectStore::get) re-issues the request on body-read errors, bounded by the existing S3DLIO_MAX_RETRY_ATTEMPTS budget.

--- a/src/reqwest_client.rs
+++ b/src/reqwest_client.rs
-            let resp_body = resp
-                .bytes()
-                .await
-                .map_err(|e| ConnectorError::io(e.into()))?;
+            use futures_util::TryStreamExt as _;
+            let frame_stream = resp.bytes_stream().map_ok(http_body::Frame::data);
+            let stream_body = http_body_util::StreamBody::new(
+                sync_wrapper::SyncStream::new(frame_stream),
+            );

             let mut response = HttpResponse::new(
                 http::StatusCode::from_u16(status)
                     .map_err(|e| ConnectorError::other(e.into(), None))?
                     .into(),
-                SdkBody::from(resp_body),
+                SdkBody::from_body_1_x(stream_body),
             );
--- a/src/s3_ops.rs  (same pattern at get_object_range and S3ObjectStore::get)
+++ b/src/s3_ops.rs
+        let max_attempts = crate::constants::max_retry_attempts().max(1);
+        let mut attempt: u32 = 0;
+        loop {
+            attempt += 1;
             ...send()...
-                let body = output.body.collect().await?.into_bytes();
+                Ok(output) => match output.body.collect().await {
+                    Ok(body) => { ... return Ok(body.into_bytes()); }
+                    Err(e) if attempt < max_attempts => {
+                        tokio::time::sleep(Duration::from_millis(100 * attempt as u64)).await;
+                    }
+                    Err(e) => { ... return Err(...); }
+                },
full patch: 0004-perf-connector-stream-response-bodies-to-the-SDK-bou.patch (apply with git am)
---
 Cargo.lock            |   2 +
 Cargo.toml            |   6 +-
 src/object_store.rs   |  45 +++++++----
 src/reqwest_client.rs |  20 +++--
 src/s3_ops.rs         | 173 ++++++++++++++++++++++++++++--------------
 src/s3_utils.rs       | 100 +++++++++++++++---------
 6 files changed, 232 insertions(+), 114 deletions(-)

diff --git a/Cargo.lock b/Cargo.lock
index 9c2fe74..eded635 100644
--- a/Cargo.lock
+++ b/Cargo.lock
@@ -5208,6 +5208,7 @@ dependencies = [
  "hdf5-metno-sys",
  "hdrhistogram",
  "http 1.4.0",
+ "http-body 1.0.1",
  "http-body-util",
  "humantime",
  "hwlocality",
@@ -5240,6 +5241,7 @@ dependencies = [
  "rustls-native-certs",
  "serde",
  "serde_json",
+ "sync_wrapper",
  "tempfile",
  "tfrecord",
  "thiserror 1.0.69",
diff --git a/Cargo.toml b/Cargo.toml
index c25f260..0017583 100644
--- a/Cargo.toml
+++ b/Cargo.toml
@@ -57,10 +57,12 @@ tokio-stream = { version = "^0.1", features = ["sync"] }
 tokio-util = "^0.7"  # For CancellationToken and other utilities
 async-stream = "^0.3"  # For stream! macro in list_stream implementations
 # HTTP stack
-reqwest = { version = "0.13", default-features = false, features = ["rustls", "http2"] }
+reqwest = { version = "0.13", default-features = false, features = ["rustls", "http2", "stream"] }
 http    = "^1"             # Re-exported by reqwest, but needed as direct dep for response building
 http-body-util  = "^0.1"   # BodyExt::collect() for SdkBody async collection
-aws-smithy-types = "^1"   # SdkBody type used by reqwest_client.rs bridge
+http-body       = "^1"     # Frame type for the streaming response bridge (issue 9.2)
+sync_wrapper    = { version = "^1", features = ["futures"] }  # SyncStream: reqwest byte streams are Send but not Sync; SdkBody::from_body_1_x requires both
+aws-smithy-types = { version = "^1", features = ["http-body-1-x"] }   # SdkBody + from_body_1_x for streaming responses
 aws-smithy-http-client = { version = "^1", features = ["rustls-aws-lc"] }
 aws-smithy-runtime-api = "^1"                         # ← Added for HTTP client types
 
diff --git a/src/object_store.rs b/src/object_store.rs
index 892b5bb..7c1d02a 100644
--- a/src/object_store.rs
+++ b/src/object_store.rs
@@ -1180,19 +1180,38 @@ impl ObjectStore for S3ObjectStore {
 
         if let Some(client) = &self.client {
             let (bucket, key) = parse_s3_uri(uri)?;
-            let resp = client
-                .get_object()
-                .bucket(&bucket)
-                .key(&key)
-                .send()
-                .await
-                .sdk_context(format!("S3 GET failed for '{}'", uri))?;
-            let data = resp
-                .body
-                .collect()
-                .await
-                .context("Failed to read S3 GET response body")?;
-            return Ok(data.into_bytes());
+            // Bounded body-read retry: with the streaming connector
+            // (issue 9.2) the SDK retry no longer covers mid-body transport
+            // errors, so re-issue the whole GET on a body-read failure.
+            let max_attempts = crate::constants::max_retry_attempts().max(1);
+            let mut attempt: u32 = 0;
+            loop {
+                attempt += 1;
+                let resp = client
+                    .get_object()
+                    .bucket(&bucket)
+                    .key(&key)
+                    .send()
+                    .await
+                    .sdk_context(format!("S3 GET failed for '{}'", uri))?;
+                match resp.body.collect().await {
+                    Ok(data) => return Ok(data.into_bytes()),
+                    Err(e) if attempt < max_attempts => {
+                        tracing::warn!(
+                            "S3 GET '{}' body read failed (attempt {}/{}): {}; retrying",
+                            uri, attempt, max_attempts, e
+                        );
+                        tokio::time::sleep(std::time::Duration::from_millis(
+                            100 * attempt as u64,
+                        ))
+                        .await;
+                    }
+                    Err(e) => {
+                        return Err(anyhow::Error::new(e)
+                            .context("Failed to read S3 GET response body"));
+                    }
+                }
+            }
         }
 
         // Range optimization is ENABLED by default (v0.9.60+).
diff --git a/src/reqwest_client.rs b/src/reqwest_client.rs
index 07e784c..ef4ed0f 100644
--- a/src/reqwest_client.rs
+++ b/src/reqwest_client.rs
@@ -315,17 +315,25 @@ impl HttpConnector for ReqwestHttpConnector {
                 tracing::debug!("HTTP protocol: {:?}", version);
             }
 
-            let resp_body = resp
-                .bytes()
-                .await
-                .map_err(|e| ConnectorError::io(e.into()))?;
+            // ── Build Smithy response (STREAMING body — issue 9.2) ───────
+            // Previously this did `resp.bytes().await`, fully buffering the
+            // body inside the connector: the SDK caller saw byte one only
+            // after the LAST byte arrived (no receive/consume overlap, an
+            // extra whole-body buffer per GET, and `first_byte_time` metrics
+            // that actually measured last-byte). Hand the SDK the live byte
+            // stream instead; consumers (collect(), the range engine, the
+            // loader) pull fragments as they arrive.
+            use futures_util::TryStreamExt as _;
+            let frame_stream = resp.bytes_stream().map_ok(http_body::Frame::data);
+            let stream_body = http_body_util::StreamBody::new(
+                sync_wrapper::SyncStream::new(frame_stream),
+            );
 
-            // ── Build Smithy response ─────────────────────────────────────
             let mut response = HttpResponse::new(
                 http::StatusCode::from_u16(status)
                     .map_err(|e| ConnectorError::other(e.into(), None))?
                     .into(),
-                SdkBody::from(resp_body),
+                SdkBody::from_body_1_x(stream_body),
             );
 
             for (name, value) in &headers {
diff --git a/src/s3_ops.rs b/src/s3_ops.rs
index cebf926..6e06428 100644
--- a/src/s3_ops.rs
+++ b/src/s3_ops.rs
@@ -88,36 +88,69 @@ impl S3Ops {
             start_time: SystemTime::now(),
         };
 
-        let result = self
-            .client
-            .get_object()
-            .bucket(bucket)
-            .key(key)
-            .send()
-            .await;
-        let first_byte_time = SystemTime::now();
+        // Streaming bodies (issue 9.2): `send()` resolves at headers, so the
+        // SDK retry policy no longer covers transport failures during body
+        // transfer. Compensate with a bounded whole-request retry on
+        // body-read errors (budget: S3DLIO_MAX_RETRY_ATTEMPTS). The retry
+        // arm does not consume `ctx`; every exit arm returns, so `ctx` is
+        // moved at most once.
+        let max_attempts = crate::constants::max_retry_attempts().max(1);
+        let mut attempt: u32 = 0;
+        loop {
+            attempt += 1;
+            let result = self
+                .client
+                .get_object()
+                .bucket(bucket)
+                .key(key)
+                .send()
+                .await;
+            let first_byte_time = SystemTime::now();
 
-        match result {
-            Ok(output) => {
-                let body = output.body.collect().await?.into_bytes();
-                let bytes = body.len() as u64;
-                self.log_op(ctx, &Ok::<(), &str>(()), bytes, Some(first_byte_time));
-                // Zero-copy: return Bytes directly instead of converting to Vec
-                Ok(body)
-            }
-            Err(e) => {
-                // format_sdk_error preserves the connector chain (I/O, timeout,
-                // TLS, DNS, refused) plus a hint string.  Without this the bare
-                // `SdkError::DispatchFailure` Display impl is the literal text
-                // "dispatch failure" — the opaque message reported in
-                // mlcommons/storage#506.  Log uses the same detail so the
-                // s3dlio op-log file matches what the caller sees.
-                let detail = format_sdk_error(&e);
-                self.log_op(ctx, &Err::<(), _>(detail), 0, None);
-                Err(sdk_anyhow(
-                    format!("S3 GET for s3://{}/{} failed", bucket, key),
-                    e,
-                ))
+            match result {
+                Ok(output) => match output.body.collect().await {
+                    Ok(body) => {
+                        let body = body.into_bytes();
+                        let bytes = body.len() as u64;
+                        self.log_op(ctx, &Ok::<(), &str>(()), bytes, Some(first_byte_time));
+                        // Zero-copy: return Bytes directly instead of converting to Vec
+                        return Ok(body);
+                    }
+                    Err(e) if attempt < max_attempts => {
+                        tracing::warn!(
+                            "S3 GET s3://{}/{} body read failed (attempt {}/{}): {}; retrying",
+                            bucket, key, attempt, max_attempts, e
+                        );
+                        tokio::time::sleep(std::time::Duration::from_millis(
+                            100 * attempt as u64,
+                        ))
+                        .await;
+                    }
+                    Err(e) => {
+                        // Previously a body-read failure `?`-escaped before
+                        // log_op ran; now the op-log records it too.
+                        let detail = format!("body read failed: {e}");
+                        self.log_op(ctx, &Err::<(), _>(detail), 0, Some(first_byte_time));
+                        return Err(anyhow::Error::new(e).context(format!(
+                            "S3 GET for s3://{}/{} failed reading response body",
+                            bucket, key
+                        )));
+                    }
+                },
+                Err(e) => {
+                    // format_sdk_error preserves the connector chain (I/O, timeout,
+                    // TLS, DNS, refused) plus a hint string.  Without this the bare
+                    // `SdkError::DispatchFailure` Display impl is the literal text
+                    // "dispatch failure" — the opaque message reported in
+                    // mlcommons/storage#506.  Log uses the same detail so the
+                    // s3dlio op-log file matches what the caller sees.
+                    let detail = format_sdk_error(&e);
+                    self.log_op(ctx, &Err::<(), _>(detail), 0, None);
+                    return Err(sdk_anyhow(
+                        format!("S3 GET for s3://{}/{} failed", bucket, key),
+                        e,
+                    ));
+                }
             }
         }
     }
@@ -138,35 +171,61 @@ impl S3Ops {
         };
 
         let range = format!("bytes={}-{}", start, end);
-        let result = self
-            .client
-            .get_object()
-            .bucket(bucket)
-            .key(key)
-            .range(range)
-            .send()
-            .await;
-        let first_byte_time = SystemTime::now();
+        // Bounded body-read retry — see the GET path above (issue 9.2).
+        let max_attempts = crate::constants::max_retry_attempts().max(1);
+        let mut attempt: u32 = 0;
+        loop {
+            attempt += 1;
+            let result = self
+                .client
+                .get_object()
+                .bucket(bucket)
+                .key(key)
+                .range(&range)
+                .send()
+                .await;
+            let first_byte_time = SystemTime::now();
 
-        match result {
-            Ok(output) => {
-                let body = output.body.collect().await?.into_bytes();
-                let bytes = body.len() as u64;
-                self.log_op(ctx, &Ok::<(), &str>(()), bytes, Some(first_byte_time));
-                // Zero-copy: return Bytes directly
-                Ok(body)
-            }
-            Err(e) => {
-                // See note on the GET path above re: format_sdk_error and #506.
-                let detail = format_sdk_error(&e);
-                self.log_op(ctx, &Err::<(), _>(detail), 0, None);
-                Err(sdk_anyhow(
-                    format!(
-                        "S3 GET_RANGE for s3://{}/{} bytes={}-{} failed",
-                        bucket, key, start, end
-                    ),
-                    e,
-                ))
+            match result {
+                Ok(output) => match output.body.collect().await {
+                    Ok(body) => {
+                        let body = body.into_bytes();
+                        let bytes = body.len() as u64;
+                        self.log_op(ctx, &Ok::<(), &str>(()), bytes, Some(first_byte_time));
+                        // Zero-copy: return Bytes directly
+                        return Ok(body);
+                    }
+                    Err(e) if attempt < max_attempts => {
+                        tracing::warn!(
+                            "S3 GET_RANGE s3://{}/{} bytes={}-{} body read failed (attempt {}/{}): {}; retrying",
+                            bucket, key, start, end, attempt, max_attempts, e
+                        );
+                        tokio::time::sleep(std::time::Duration::from_millis(
+                            100 * attempt as u64,
+                        ))
+                        .await;
+                    }
+                    Err(e) => {
+                        let detail = format!("body read failed: {e}");
+                        self.log_op(ctx, &Err::<(), _>(detail), 0, Some(first_byte_time));
+                        return Err(anyhow::Error::new(e).context(format!(
+                            "S3 GET_RANGE for s3://{}/{} bytes={}-{} failed reading response body",
+                            bucket, key, start, end
+                        )));
+                    }
+                },
+                Err(e) => {
+                    // See note on the GET path above re: format_sdk_error and #506.
+                    let detail = format_sdk_error(&e);
+                    self.log_op(ctx, &Err::<(), _>(detail), 0, None);
+                    return Err(sdk_anyhow(
+                        format!(
+                            "S3 GET_RANGE for s3://{}/{} bytes={}-{} failed",
+                            bucket, key, start, end
+                        ),
+                        e,
+                    ));
+                }
             }
         }
     }
diff --git a/src/s3_utils.rs b/src/s3_utils.rs
index d6460ca..82546be 100644
--- a/src/s3_utils.rs
+++ b/src/s3_utils.rs
@@ -1208,44 +1208,72 @@ async fn concurrent_range_get_impl(
                 .map_err(|e| anyhow::anyhow!("Semaphore error: {}", e))?;
 
             let range_header = format!("bytes={}-{}", range_start, range_end - 1);
-            let resp = client
-                .get_object()
-                .bucket(&bucket)
-                .key(key.trim_start_matches('/'))
-                .range(&range_header)
-                .send()
-                .await
-                .sdk_context(format!(
-                    "concurrent range GET for s3://{}/{} bytes={}-{} failed",
-                    bucket,
-                    key,
-                    range_start,
-                    range_end - 1
-                ))?;
-
-            // Stream fragments straight into this chunk's segment —
-            // no intermediate whole-chunk buffer.
-            let mut body = resp.body;
-            let mut written = 0usize;
-            while let Some(frag) = body.try_next().await.context("read chunk body")? {
-                let end = written + frag.len();
-                if end > seg.len() {
-                    anyhow::bail!(
-                        "range GET s3://{}/{} bytes={}-{}: server returned more than requested ({} > {})",
-                        bucket, key, range_start, range_end - 1, end, seg.len()
-                    );
+
+            // With streaming bodies (issue 9.2) `send()` resolves at response
+            // headers, so the SDK's transient-error retry no longer covers
+            // failures DURING body transfer. Compensate with a bounded
+            // whole-chunk retry: the segment never escapes until fully
+            // written, so re-issuing the ranged GET and rewriting from
+            // offset 0 is idempotent. Budget matches the SDK request-phase
+            // policy (S3DLIO_MAX_RETRY_ATTEMPTS, default 3).
+            let max_attempts = crate::constants::max_retry_attempts().max(1);
+            let mut attempt: u32 = 0;
+            loop {
+                attempt += 1;
+                let outcome: anyhow::Result<()> = async {
+                    let resp = client
+                        .get_object()
+                        .bucket(&bucket)
+                        .key(key.trim_start_matches('/'))
+                        .range(&range_header)
+                        .send()
+                        .await
+                        .sdk_context(format!(
+                            "concurrent range GET for s3://{}/{} bytes={}-{} failed",
+                            bucket,
+                            key,
+                            range_start,
+                            range_end - 1
+                        ))?;
+
+                    // Stream fragments straight into this chunk's segment —
+                    // no intermediate whole-chunk buffer.
+                    let mut body = resp.body;
+                    let mut written = 0usize;
+                    while let Some(frag) = body.try_next().await.context("read chunk body")? {
+                        let end = written + frag.len();
+                        if end > seg.len() {
+                            anyhow::bail!(
+                                "range GET s3://{}/{} bytes={}-{}: server returned more than requested ({} > {})",
+                                bucket, key, range_start, range_end - 1, end, seg.len()
+                            );
+                        }
+                        seg[written..end].copy_from_slice(&frag);
+                        written = end;
+                    }
+                    if written != seg.len() {
+                        anyhow::bail!(
+                            "range GET s3://{}/{} bytes={}-{}: short read ({} of {} bytes)",
+                            bucket, key, range_start, range_end - 1, written, seg.len()
+                        );
+                    }
+                    Ok(())
+                }
+                .await;
+
+                match outcome {
+                    Ok(()) => return Ok::<(usize, BytesMut), anyhow::Error>((idx, seg)),
+                    Err(e) if attempt < max_attempts => {
+                        tracing::warn!(
+                            "range chunk s3://{}/{} bytes={}-{} attempt {}/{} failed: {:#}; retrying",
+                            bucket, key, range_start, range_end - 1, attempt, max_attempts, e
+                        );
+                        tokio::time::sleep(std::time::Duration::from_millis(100 * attempt as u64))
+                            .await;
+                    }
+                    Err(e) => return Err(e),
                 }
-                seg[written..end].copy_from_slice(&frag);
-                written = end;
-            }
-            if written != seg.len() {
-                anyhow::bail!(
-                    "range GET s3://{}/{} bytes={}-{}: short read ({} of {} bytes)",
-                    bucket, key, range_start, range_end - 1, written, seg.len()
-                );
             }
-
-            Ok::<(usize, BytesMut), anyhow::Error>((idx, seg))
         });
     }
 
-- 
2.43.0

Metadata

Metadata

Assignees

Labels

enhancementNew feature or request

Fields

No fields configured for Feature.

Projects

No projects

Milestone

No milestone

Relationships

None yet

Development

No branches or pull requests

Issue actions