perf(ai): make variable-length text ingestion more efficient
- Dominant language
- Rust
- Stars
- 244
- Forks
- 29
- Avg merge
- 2d 41m
- Merged PRs (30d)
- 19
Description
## Summary
Code-indexing workloads contain text inputs with extreme length variation: short functions, medium methods, long traits, and generated service blocks. A DevMind run against Ahnlich indexed 2,217 chunks across 190 files and took more than 40 minutes on CPU. DevMind's request shape contributes, but Ahnlich's text path should also avoid doing model work for padding that carries no information.
Related client-side findings: https://github.com/uche09/dev_mind/issues/1
## Primary issue: padding scope is larger than the model inference batch
The model manager preprocesses every input in an incoming request before model inference:
```rust
let processed_inputs = self.preprocess_store_input(process_action, inputs)?;
let mut store_key = self
.model
.model_ndarray(processed_inputs, &action_type, execution_provider, model_params)
.await?;
```
For text, preprocessing tokenizes the input collection together:
```rust
let outputs = provider.preprocess_texts(inputs, truncate)?;
let token_size = outputs
.first()
.ok_or(/* ... */)?
.len();
```
Only later does the ONNX model divide the already-prepared encodings into configured model batches:
```rust
for batch_encoding in encodings
.into_iter()
.chunks(self.model_batch_size)
.into_iter()
{
let embeddings =
self.batch_inference_text(batch_encoding.collect(), &session)?;
}
```
`batch_inference_text()` constructs dense tensors using one encoding length for the whole inference batch:
```rust
let batch_size = encodings.len();
let encoding_length = encodings[0].len();
let max_size = encoding_length * batch_size;
for encoding in &encodings {
ids_array.extend(encoding.get_ids().iter().map(|x| *x as i64));
}
Array::from_shape_vec((batch_size, encoding_length), ids_array)
```
For variable-length code, a long chunk can cause many shorter chunks to be padded and run at its sequence length. Transformer inference becomes materially more expensive as sequence length grows, so this wastes CPU and allocation pressure.
## Proposed implementation
1. Tokenize text inputs without padding the entire request together.
2. Retain each encoding's original input position and unpadded token length.
3. Bucket/sort similarly sized encodings into batches up to `model_batch_size`.
4. Pad only within an individual model-inference batch.
5. Run inference and restore `ModelResponse`s to the original request order.
6. Validate the token limit for every individual input before inference, rather than relying on the first processed encoding's length.
This should be internal to Ahnlich: callers continue to submit `Set.inputs` in their preferred order and receive responses in that same order.
## Additional Ahnlich work exposed by the workload
### Apply adaptive parallelism to AI preparation
AI ingestion uses Rayon both when assembling request data and when cloning raw strings for preprocessing:
```rust
params.inputs.into_par_iter().flat_map(/* ... */).collect()
```
```rust
inputs.par_iter().filter_map(/* clone RawString */).collect()
```
ONNX Runtime also uses CPU threads. For sustained indexing, those pools can oversubscribe the machine. Use the existing parallelism policy or equivalent request-load-aware thresholds for preparation work.
### Add ingestion phase timings
Expose timings/metrics for:
- validation and request preparation;
- tokenization;
- ONNX inference;
- response/vector construction;
- DB pipeline persistence.
Without this, a client sees only end-to-end latency and cannot establish whether model execution or storage is the limiting factor.
### Make execution resources explicit
When callers omit an execution provider, text inference uses CPU. Expose documented configuration for appropriate execution providers and ONNX Runtime CPU thread settings so operators can tune ingestion to available hardware.
### Future: bounded dynamic batching across requests
The per-model worker processes one request at a time. A bounded micro-batcher could coalesce concurrently arriving small requests, split responses back to callers, and improve general small-write throughput. This is a follow-up: it does not help a strictly sequential client as much as fixing padding scope and client request batching.
## Acceptance criteria
- Mixed-length text inputs preserve result count and caller input order.
- Token-limit validation identifies the offending input deterministically.
- Benchmarks cover representative short, medium, and long code-like inputs, not only uniform short texts.
- The revised path demonstrates lower runtime and allocation for mixed-length batches than the current implementation.
- Existing text-model behavior and public API remain compatible.
Contributor guide
Research direction
Start by tracing model manager preprocessing through preprocess_store_input, model_ndarray, provider.preprocess_texts, and batch_inference_text; then inspect the params.inputs preparation path and existing benchmark coverage. Done means mixed-length inputs preserve count and order, validate each token limit deterministically, and demonstrate lower runtime and allocation without changing the public API.
Written by the indexing model from the issue text.
Assessment
- Tech stack
- rust
- Domain
- backend, machine-learning, performance
- Issue type
- Feature
- Difficulty
- 5/5
- Estimated time
- Over a week
- Activity status
- Active
- Clarity
- Mostly clear
- Newbie friendliness
- 42/100