0lai0 opened a new pull request, #1466: URL: https://github.com/apache/mahout/pull/1466
### Related Issues Closes #1430 ### Changes - [ ] Bug fix - [x] New feature - [x] Refactoring - [x] Documentation - [x] Test - [x] CI/CD pipeline - [ ] Other ### Why An oversized QDP pipeline currently runs until `encode_batch` allocates GPU buffers before it fails — slow, and the error does not say what to change. This adds a fail-fast check against free VRAM so a configuration that cannot fit is rejected in milliseconds, before any input file is read, with a message naming the values to lower. ### How - After `normalize()`, the three `PipelineIterator` constructors (`new_synthetic`, `new_from_file`, `new_from_file_streaming`) compare `estimate_memory().gpu_state_bytes` to free device memory. The check is skipped when the CUDA runtime reports no usable device — stub / no-toolkit builds, and also an empty `CUDA_VISIBLE_DEVICES`. - **Refactor (needed to test the rule on CI, which has no GPU):** `ensure_device_memory_available` now queries `cudaMemGetInfo` and delegates the comparison to `ensure_fits_in_free_memory`. Tests inject `(free, total)` through `ensure_config_fits_device_with`. Production path, error type, and message are unchanged; both helpers are `pub(crate)` / private. - The budget is **two** concurrent batch state buffers. That is deliberate: a `for qt in loader:` loop keeps the previous DLPack tensor alive while the next batch is allocated. A config that fits one buffer but not two is rejected. There is no override. **This budget is a floor, not the true peak:** the encoders also upload the input batch to the device (`htod_sync_copy`) and hold it alongside the state buffers — for amplitude, half a buffer, so the real steady-state peak is nearer 2.5. A configuration sitting just under free memory can therefore still OOM mid-run; this narrows that window rather than closing it. Modelling it properly means changing the memory model, which #1430 puts out of scope. - Precision is the wider of the loader dtype and the engine precision. File loaders budget basis as float64 (indices are read as f64); synthetic basis is budgeted at the requested dtype. - The comparison is against free memory *sampled* at that moment — nothing is reserved. Another process or a second loader can take it before the first batch allocates, so passing the check is a fast sanity check, not a guarantee. - On the Python surface the check runs when **iteration starts** (`for qt in loader:`), not when `QuantumDataLoader(...)` returns. The rejection is a `RuntimeError` wrapping the Rust `MemoryAllocation` message (the existing conversion for every backend error). A dedicated exception type is left as follow-up. - `qumat_qdp.estimate_memory()` exposes the same arithmetic without opening a device, so a caller can size `qubits` / `batch_size` before iterating. - CI covers the accept/reject rule with injected free-memory figures (no GPU required). Integration tests still exercise the live-GPU path and self-skip on stub builds. Out of scope, per #1430: changing the underlying memory model (see the 2-buffer note above), automatic config downsizing, and multi-GPU accounting — the guard queries the current CUDA device and does not call `cudaSetDevice`. Also not covered: host-side streaming chunk limits, and the benchmark helpers (`run_throughput_pipeline` / `run_latency_pipeline`), which build their own engine and bypass `PipelineIterator`. ### Verification Run on a host with CUDA 12.5 and a live GPU: - `make pre-commit` — clean. - `make test_rust` — 379 passed, 0 failed, 0 ignored. - `make test_python` — 888 passed, 1 skipped (pre-existing, unrelated), 0 failed. - `QDP_NO_CUDA=1 cargo check --workspace --tests` — clean; both new CI steps pass on the stub build (15 and 47 tests). - `cd website && npm run typecheck && npm run build` — clean; the two broken-link warnings are pre-existing (verified by building `upstream/main`). ## Checklist - [x] Added or updated unit tests for all changes - [x] Added or updated documentation for all changes -- This is an automated message from the Apache Git Service. To respond to the message, please log on to GitHub and use the URL above to go to the specific comment. To unsubscribe, e-mail: [email protected] For queries about this service, please contact Infrastructure at: [email protected]
