gitadmin/foxhunt
1
0
Fork 0
Code Issues Pull Requests Packages Projects Releases Wiki Activity

fix(rl): streaming per-file GPU upload — prevent host OOM at 45M snaps

Browse Source 
Convert + upload one file at a time instead of all 45M snapshots at
once. Peak host memory drops from ~10GB (all converted snapshots) to
~2GB (one file). Fixes OOMKilled on 40Gi L40S pods.

Co-Authored-By: Claude Opus 4.7 <noreply@anthropic.com>
This commit is contained in:
jgrusewski
2026-05-26 21:04:30 +02:00
parent 1f672c3b09
commit 3552d08501
1 changed files with 27 additions and 22 deletions
Download Patch File Download Diff File Expand all files Collapse all files

49
crates/ml-alpha/src/data/loader.rs
View File

@@ -714,43 +714,48 @@ impl MultiHorizonLoader {
"upload_to_gpu: pre-converting snapshots",
);
let mut all_raws: Vec<Mbp10RawInput> = Vec::with_capacity(total_snaps);
let mut file_offsets: Vec<i32> = Vec::with_capacity(n_files);
let mut file_sizes: Vec<i32> = Vec::with_capacity(n_files);
let mut offset: usize = 0;
for lf in &self.files_loaded {
file_offsets.push(offset as i32);
file_sizes.push(lf.snapshots.len() as i32);
for (idx, snap) in lf.snapshots.iter().enumerate() {
let prev = if idx == 0 { snap } else { &lf.snapshots[idx - 1] };
all_raws.push(convert(snap, prev, lf.regime_full[idx]));
}
offset += lf.snapshots.len();
}
debug_assert_eq!(all_raws.len(), total_snaps);
// ── 2. Upload snapshots as raw bytes ─────────────────────────────
// Reinterpret Vec<Mbp10RawInput> as &[u8] for the upload. The GPU
// kernel reads the same struct layout via its `Mbp10Raw` definition.
// ── 2. Upload snapshots per-file (streaming) ────────────────────
// Allocate the full GPU buffer upfront, then convert + upload one
// file at a time. Peak host memory = one file's worth of converted
// snapshots (~1-2GB) instead of all files (~10GB).
let total_bytes = total_snaps * MBP10_RAW_INPUT_BYTES;
let snap_bytes: &[u8] = unsafe {
std::slice::from_raw_parts(
all_raws.as_ptr() as *const u8,
total_bytes,
)
};
let snapshots_d = stream
.alloc_zeros::<u8>(total_bytes)
.context("upload_to_gpu: snapshots alloc")?;
unsafe {
cudarc::driver::sys::cuMemcpyHtoD_v2(
snapshots_d.raw_ptr(),
snap_bytes.as_ptr() as *const std::ffi::c_void,
total_bytes,
let mut gpu_offset_bytes: usize = 0;
for (f_idx, lf) in self.files_loaded.iter().enumerate() {
let n = lf.snapshots.len();
let mut file_raws: Vec<Mbp10RawInput> = Vec::with_capacity(n);
for (idx, snap) in lf.snapshots.iter().enumerate() {
let prev = if idx == 0 { snap } else { &lf.snapshots[idx - 1] };
file_raws.push(convert(snap, prev, lf.regime_full[idx]));
}
let file_bytes = n * MBP10_RAW_INPUT_BYTES;
unsafe {
cudarc::driver::sys::cuMemcpyHtoD_v2(
snapshots_d.raw_ptr() + gpu_offset_bytes as u64,
file_raws.as_ptr() as *const std::ffi::c_void,
file_bytes,
);
}
gpu_offset_bytes += file_bytes;
tracing::info!(
file = f_idx,
n_snapshots = n,
gpu_offset_mb = gpu_offset_bytes / (1024 * 1024),
"upload_to_gpu: file uploaded",
);
drop(file_raws);
}
// Drop the host-side Vec now — it can be multi-GB.
drop(all_raws);
tracing::info!(
gpu_bytes = total_bytes,
"upload_to_gpu: snapshots uploaded",