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

feat(ml): wire GpuBufferPool + DoubleBufferedLoader into DQN trainer GPU upload path

Browse Source 
- GpuBufferPool: reuses pre-allocated staging buffers for zero-alloc fold transitions
  instead of always calling DqnGpuData::upload() directly
- DoubleBufferedLoader: new field on DQNTrainer for zero-downtime fold transitions;
  skips re-upload when active slot is already populated from a previous fold
- Added double_buffer() / double_buffer_mut() accessors
- Upload path: check double-buffer first, then try buffer_pool, then direct upload

Co-Authored-By: Claude Opus 4.6 <noreply@anthropic.com>
This commit is contained in:
jgrusewski
2026-03-01 21:52:32 +01:00
parent 8910591f85
commit 45963b2b4b
1 changed files with 47 additions and 11 deletions
Download Patch File Download Diff File Expand all files Collapse all files

58
crates/ml/src/trainers/dqn/trainer.rs
View File

@@ -219,6 +219,9 @@ pub struct DQNTrainer {
/// Reusable GPU staging buffers for zero-alloc fold transitions
buffer_pool: Option<crate::cuda_pipeline::GpuBufferPool>,
/// Double-buffered GPU data for zero-downtime fold transitions
double_buffer: Option<crate::cuda_pipeline::double_buffer::DoubleBufferedLoader>,
/// Multi-GPU configuration for data-parallel training (None = single GPU)
multi_gpu: Option<crate::cuda_pipeline::multi_gpu::MultiGpuConfig>,
}
@@ -701,6 +704,13 @@ impl DQNTrainer {
None
};
// GPU pipeline: double-buffered loader for zero-downtime fold transitions
let double_buffer = if device.is_cuda() {
Some(crate::cuda_pipeline::double_buffer::DoubleBufferedLoader::new(device.clone()))
} else {
None
};
// Multi-GPU: auto-detect if multiple CUDA devices are available
let multi_gpu = crate::cuda_pipeline::multi_gpu::MultiGpuConfig::detect()
.unwrap_or(None);
@@ -832,6 +842,9 @@ impl DQNTrainer {
// GPU pipeline: staging buffer pool (auto-initialized on CUDA devices)
buffer_pool,
// GPU pipeline: double-buffered loader for fold transitions
double_buffer,
// Multi-GPU: auto-detected data parallelism
multi_gpu,
})
@@ -861,6 +874,16 @@ impl DQNTrainer {
self.prefetcher.take().map(|p| p.take().map_err(Into::into))
}
/// Get a reference to the double-buffered loader, if GPU is active.
pub fn double_buffer(&self) -> Option<&crate::cuda_pipeline::double_buffer::DoubleBufferedLoader> {
self.double_buffer.as_ref()
}
/// Get a mutable reference to the double-buffered loader, if GPU is active.
pub fn double_buffer_mut(&mut self) -> Option<&mut crate::cuda_pipeline::double_buffer::DoubleBufferedLoader> {
self.double_buffer.as_mut()
}
/// Train DQN on market data from DBN files
///
/// # Arguments
@@ -1440,17 +1463,30 @@ impl DQNTrainer {
// Phase 1: Pre-upload training data to GPU (once, reused across epochs)
if self.gpu_data.is_none() {
match DqnGpuData::upload(&training_data, &self.device) {
Ok(gpu_data) => {
info!("GPU data pre-uploaded: {} bars x {} features ({:.1} MB)",
gpu_data.num_bars,
gpu_data.feature_dim,
(gpu_data.num_bars * (51 + 4) * 4) as f64 / 1_048_576.0
);
self.gpu_data = Some(gpu_data);
}
Err(e) => {
debug!("GPU data pre-upload skipped (CPU fallback): {}", e);
// Check double-buffer: skip upload if active slot already populated
let skip_upload = self.double_buffer.as_ref().is_some_and(|db| db.active().is_some());
if skip_upload {
info!("DoubleBuffer: active slot populated, skipping re-upload");
} else {
let upload_result = if let Some(ref mut pool) = self.buffer_pool {
info!("GpuBufferPool: reusing pre-allocated staging buffers for {} bars", training_data.len());
pool.upload_dqn(&training_data, &self.device)
} else {
DqnGpuData::upload(&training_data, &self.device)
};
match upload_result {
Ok(gpu_data) => {
info!("GPU data pre-uploaded: {} bars x {} features ({:.1} MB)",
gpu_data.num_bars,
gpu_data.feature_dim,
(gpu_data.num_bars * (51 + 4) * 4) as f64 / 1_048_576.0
);
self.gpu_data = Some(gpu_data);
}
Err(e) => {
debug!("GPU data pre-upload skipped (CPU fallback): {}", e);
}
}
}
}