Compare commits
67 Commits
| Author | SHA1 | Date | |
|---|---|---|---|
|
|
dc8494bb2a | ||
|
|
6f8b6440aa | ||
| f8ad40d302 | |||
|
|
409363bf80 | ||
|
|
6d3a7f0b48 | ||
|
|
bb4bc1ec57 | ||
|
|
a9ff2f1e76 | ||
|
|
7c394c012d | ||
|
|
d690b08486 | ||
|
|
4088471885 | ||
|
|
f085e977e6 | ||
| 333f8f0f26 | |||
| 9845e81194 | |||
|
|
ed40391f93 | ||
|
|
e76d7a0a5b | ||
|
|
feff63fb36 | ||
|
|
0994b6656e | ||
|
|
e5f1dcf7b9 | ||
|
|
4fb0b78a5f | ||
|
|
15da5eb1c2 | ||
|
|
f9ecbae6a4 | ||
|
|
000a2fad3a | ||
|
|
14b0481d7c | ||
|
|
bd8e07fe9c | ||
|
|
9016b9213e | ||
|
|
b25a070d5e | ||
|
|
8e6b120e76 | ||
|
|
ebea27b62e | ||
|
|
09b4eb4510 | ||
|
|
30d741b418 | ||
|
|
7b5e615c00 | ||
|
|
056fbbb6e1 | ||
|
|
c408c483d0 | ||
|
|
c97f5e6e6e | ||
|
|
50ddaacbdc | ||
|
|
c734fadba2 | ||
|
|
258f7e09aa | ||
|
|
2811d97c47 | ||
|
|
9cacb95b0f | ||
|
|
11ddd184f4 | ||
|
|
0b0f3677b1 | ||
|
|
5313f5f4ef | ||
|
|
72974a9658 | ||
|
|
8981b3038d | ||
|
|
751b95a123 | ||
|
|
26d1eb631e | ||
|
|
d6bd2eee00 | ||
|
|
06a6dcdd31 | ||
|
|
c338fcae65 | ||
|
|
95f88ee6e4 | ||
|
|
4e9a3aa5cb | ||
|
|
527e8428b0 | ||
|
|
9525aa8878 | ||
|
|
e041912108 | ||
|
|
3dc143a417 | ||
|
|
29edc90860 | ||
|
|
9d0a308e73 | ||
|
|
791a9bfc49 | ||
|
|
3851db8e53 | ||
|
|
dd258a71ff | ||
|
|
23f1000405 | ||
|
|
f3d2260a50 | ||
|
|
e8673a3bcf | ||
|
|
a562beb34c | ||
|
|
b01a9cea53 | ||
|
|
f5624c1ce6 | ||
|
|
84bc53c4b4 |
76
.github/workflows/publish.yml
vendored
Normal file
76
.github/workflows/publish.yml
vendored
Normal file
@@ -0,0 +1,76 @@
|
||||
name: Publish to PyPI
|
||||
|
||||
on:
|
||||
push:
|
||||
tags:
|
||||
- 'v[0-9]+.[0-9]+.[0-9]+*'
|
||||
workflow_dispatch:
|
||||
|
||||
jobs:
|
||||
build_and_publish:
|
||||
name: Build and Publish All Wheels
|
||||
runs-on: ${{ matrix.os }}
|
||||
strategy:
|
||||
matrix:
|
||||
os: [ubuntu-latest, macos-latest, windows-latest]
|
||||
target: [x86_64]
|
||||
|
||||
steps:
|
||||
- uses: actions/checkout@v4
|
||||
|
||||
# Only setup python for non-linux, or let maturin handle linux via docker
|
||||
- name: Set up Python
|
||||
if: runner.os != 'Linux'
|
||||
uses: actions/setup-python@v5
|
||||
with:
|
||||
python-version: '3.12'
|
||||
# Note: On Windows/Mac, to build for multiple python versions,
|
||||
# you usually need multiple setup-python steps or a matrix of python versions.
|
||||
# But for now, let's at least get the artifacts downloading correctly.
|
||||
# Ideally, you remove this and let maturin find python, but GitHub runners
|
||||
# might only have one default.
|
||||
|
||||
- name: Install Rust
|
||||
uses: dtolnay/rust-toolchain@stable
|
||||
with:
|
||||
toolchain: stable
|
||||
|
||||
- name: Build wheel
|
||||
uses: PyO3/maturin-action@v1
|
||||
with:
|
||||
target: ${{ matrix.target }}
|
||||
manylinux: auto # This triggers the Docker container on Linux for multi-python builds
|
||||
command: build
|
||||
args: --release --out dist --find-interpreter
|
||||
|
||||
- name: Upload wheel artifact
|
||||
uses: actions/upload-artifact@v4
|
||||
with:
|
||||
name: wheel-${{ matrix.os }}-${{ matrix.target }}
|
||||
path: dist/*.whl
|
||||
|
||||
publish:
|
||||
name: Publish All Wheels
|
||||
needs: [build_and_publish]
|
||||
runs-on: ubuntu-latest
|
||||
|
||||
steps:
|
||||
- name: Install maturin
|
||||
run: pip install maturin
|
||||
|
||||
# FIX: Use merge-multiple to flatten everything into 'dist' automatically
|
||||
- name: Download ALL wheels
|
||||
uses: actions/download-artifact@v4
|
||||
with:
|
||||
path: dist
|
||||
pattern: wheel-*
|
||||
merge-multiple: true
|
||||
|
||||
- name: List Files (Debug)
|
||||
run: ls -la dist
|
||||
|
||||
- name: Upload to PyPI
|
||||
env:
|
||||
MATURIN_PYPI_TOKEN: ${{ secrets.PYPI_API_TOKEN }}
|
||||
run: |
|
||||
maturin upload --skip-existing --non-interactive dist/*
|
||||
80
.github/workflows/publish2.yml
vendored
Normal file
80
.github/workflows/publish2.yml
vendored
Normal file
@@ -0,0 +1,80 @@
|
||||
name: Publish to PyPI
|
||||
|
||||
on:
|
||||
push:
|
||||
tags:
|
||||
- 'v[0-9]+.[0-9]+.[0-9]+*'
|
||||
workflow_dispatch:
|
||||
|
||||
jobs:
|
||||
build_and_publish:
|
||||
name: Build and Publish All Wheels
|
||||
runs-on: ubuntu-latest
|
||||
strategy:
|
||||
matrix:
|
||||
os: [ubuntu-latest, macos-latest, windows-latest]
|
||||
target: [x86_64]
|
||||
|
||||
steps:
|
||||
- uses: actions/checkout@v4
|
||||
|
||||
- name: Set up Python
|
||||
uses: actions/setup-python@v5
|
||||
with:
|
||||
python-version: '3.12'
|
||||
|
||||
- name: Install Rust
|
||||
uses: dtolnay/rust-toolchain@stable
|
||||
with:
|
||||
toolchain: stable
|
||||
|
||||
- name: Build wheel
|
||||
uses: PyO3/maturin-action@v1
|
||||
with:
|
||||
target: ${{ matrix.target }}
|
||||
command: build
|
||||
args: --release --out dist --find-interpreter
|
||||
|
||||
- name: Upload wheel artifact
|
||||
uses: actions/upload-artifact@v4
|
||||
with:
|
||||
name: wheel-${{ matrix.os }}-${{ matrix.target }}
|
||||
path: dist/*.whl
|
||||
|
||||
publish:
|
||||
name: Publish All Wheels
|
||||
needs: [build_and_publish]
|
||||
runs-on: ubuntu-latest
|
||||
|
||||
steps:
|
||||
- name: Install maturin
|
||||
run: pip install maturin
|
||||
|
||||
- name: Download ALL wheels
|
||||
uses: actions/download-artifact@v4
|
||||
with:
|
||||
path: dist
|
||||
|
||||
- name: Flatten and upload
|
||||
env:
|
||||
MATURIN_PYPI_TOKEN: ${{ secrets.PYPI_API_TOKEN }} # ← CORRECT TOKEN VAR
|
||||
run: |
|
||||
mkdir -p dist/flattened
|
||||
rm -rf dist/flattened/*
|
||||
|
||||
# Copy EVERY wheel
|
||||
find dist -name "*.whl" | while read wheel; do
|
||||
cp "$wheel" dist/flattened/
|
||||
done
|
||||
|
||||
echo "=== TOTAL WHEELS ==="
|
||||
ls -la dist/flattened/*.whl | wc -l
|
||||
|
||||
cd dist/flattened
|
||||
echo "=== UPLOADING $(ls *.whl | wc -l) WHEELS ==="
|
||||
ls -la *.whl | head -5
|
||||
|
||||
# TWINE fallback (double protection)
|
||||
export TWINE_USERNAME=__token__
|
||||
export TWINE_PASSWORD=${{ secrets.PYPI_API_TOKEN }}
|
||||
maturin upload --skip-existing --non-interactive *.whl
|
||||
895
Cargo.lock
generated
895
Cargo.lock
generated
File diff suppressed because it is too large
Load Diff
68
Cargo.toml
68
Cargo.toml
@@ -1,56 +1,56 @@
|
||||
[package]
|
||||
name = "unchecked-io"
|
||||
version = "0.1.0"
|
||||
authors = ["Billthemaker"] # Replace with your name or alias
|
||||
license = "BSL-1" # Good practice for open-source
|
||||
version = "0.1.7"
|
||||
authors = ["Billthemaker"]
|
||||
license = "Apache-2.0"
|
||||
edition = "2024"
|
||||
|
||||
[lib]
|
||||
[lib]
|
||||
name = "unchecked_io"
|
||||
crate-type = ["cdylib", "rlib"]
|
||||
crate-type = ["cdylib", "rlib"]
|
||||
|
||||
[features]
|
||||
default = []
|
||||
gpu = [
|
||||
"tokenizers",
|
||||
"cudarc",
|
||||
"pyo3-dlpack",
|
||||
"cudarc/cuda-12050",
|
||||
"dep:cuda-driver-sys"
|
||||
]
|
||||
profiling = ["dep:tracing", "dep:tracing-subscriber", "dep:tracing-tracy"]
|
||||
|
||||
[dependencies]
|
||||
# 1. Python Bindings for FFI
|
||||
pyo3 = { version = "0.27.1", features = ["extension-module"] }
|
||||
|
||||
# 2. Configuration Parsing (YAML)
|
||||
pyo3 = { version = "=0.27.0", features = ["extension-module", "chrono-tz", "chrono"] }
|
||||
serde = { version = "1.0", features = ["derive"] }
|
||||
serde_yaml = "0.9"
|
||||
|
||||
# 3. Apache Arrow and Data Handling
|
||||
# FIX: Add the "compute" feature to get concat_batches
|
||||
arrow = "57.0.0"
|
||||
|
||||
# 4. Asynchronous Runtime (Essential for I/O and Postgres)
|
||||
arrow = { version = "57.0.0", features = ["prettyprint"] }
|
||||
tokio = { version = "1.37", features = ["full"] }
|
||||
|
||||
# 5. Database Connection (Postgres)
|
||||
tokio-postgres = "0.7"
|
||||
deadpool-postgres = "0.14"
|
||||
|
||||
# 6. Error Handling Crate
|
||||
anyhow = "1.0"
|
||||
|
||||
# 7. Futures Utilities
|
||||
futures-util = "0.3"
|
||||
|
||||
# 8. Byte Buffer Management
|
||||
bytes = "1.6"
|
||||
|
||||
# 9. Binary Data Reading (NEW)
|
||||
byteorder = "1.5"
|
||||
|
||||
# 10. Timestamp Handling (NEW)
|
||||
chrono = "0.4"
|
||||
|
||||
# 11. UUID Handling (NEW)
|
||||
chrono-tz = "=0.10.4"
|
||||
uuid = { version = "1.8", features = ["serde", "v4"] }
|
||||
|
||||
# 12. Arrow <-> Python Bridge (NEW)
|
||||
pyo3-arrow = "0.15.0"
|
||||
|
||||
# 13. System CPU Count (NEW)
|
||||
num_cpus = "1.16"
|
||||
async-channel = "2.5.0"
|
||||
rayon = "1.10"
|
||||
|
||||
# Optional dependencies enabled by features
|
||||
# UPGRADED: 0.20 supports modern Llama tokenizer exports
|
||||
tokenizers = { version = "0.20", optional = true }
|
||||
cudarc = { version = "0.11", features = ["driver"], optional = true }
|
||||
pyo3-dlpack = { version = "0.1.0", optional = true }
|
||||
cuda-driver-sys = { version = "0.3.0", optional = true }
|
||||
|
||||
tracing-subscriber = { version = "0.3", features = ["registry", "env-filter"], optional = true}
|
||||
tracing-tracy = { version = "=0.11.2", optional = true }
|
||||
tracing = {version ="0.1.41", optional = true }
|
||||
|
||||
[target.'cfg(not(target_env = "msvc"))'.dependencies]
|
||||
mimalloc = { version = "0.1.39" }
|
||||
mimalloc = { version = "0.1.39" }
|
||||
|
||||
|
||||
67
HANDOFF-2026-04-20.md
Normal file
67
HANDOFF-2026-04-20.md
Normal file
@@ -0,0 +1,67 @@
|
||||
# Handoff — 2026-04-20
|
||||
|
||||
## What Was Decided This Session
|
||||
|
||||
- **Commercialization path:** Licensing first, IP sale after traction. Wyoming LLC before any contract.
|
||||
- **Language strategy:** Rust is source of truth. C++ generated via Whetstone + manual library substitution as enterprise deliverable. Both versions is a deliberate differentiator.
|
||||
- **Primary buyers:** GPU cloud providers and LLM inference companies (CoreWeave, Lambda Labs, Together AI, Fireworks AI, Modal). Not PyTorch/Meta.
|
||||
- **Hugging Face** is a secondary buyer worth a conversation — they own the `tokenizers` crate this code already uses.
|
||||
- **NVIDIA contact:** Approach after H100 benchmark. Send Nsight trace + one paragraph. Not before.
|
||||
- **NM funding:** Apply for NMSBA (Sandia/LANL hardware access). Skip money grants for now.
|
||||
- **Outsourcing:** Friend's software business for non-core work. 60-70% of contract value for defined scope. IP assignment clause required in writing first.
|
||||
|
||||
## Documents Created This Session
|
||||
|
||||
| File | Purpose |
|
||||
|---|---|
|
||||
| `ROADMAP.md` | Full technical development plan, Phases 1–5 |
|
||||
| `docs/talk-10min.md` | 10-minute presentation structure for Wednesday |
|
||||
| `docs/commercialization-strategy.md` | Pricing, B2B sales process, structure, outsourcing |
|
||||
| `HANDOFF-2026-04-20.md` | This file |
|
||||
|
||||
---
|
||||
|
||||
## This Week — Time-Sensitive (in order)
|
||||
|
||||
### Before Wednesday
|
||||
- [ ] Fix Phase 1.1 — dynamic buffer allocation (`src/llm.rs`, `PinnedBatcher::new`)
|
||||
- [ ] Fix Phase 1.2 — buffer isolation / DLPack overwrite bug
|
||||
- [ ] Fix Phase 1.3 — thread-safe engine pool (`TokenizerEngine` takes `&mut self`)
|
||||
- [ ] Fix Phase 1.4 — padding for variable-length sequences
|
||||
|
||||
### Wednesday
|
||||
- [ ] Rent H100 on Lambda Labs (~$2 for 30 min)
|
||||
- [ ] Run `benchmark_gpu.py` under `nsys profile --stats=true python benchmark_gpu.py`
|
||||
- [ ] Capture Nsight `.nsys-rep` file and download it
|
||||
- [ ] Note the speedup number — this is the headline for the talk
|
||||
|
||||
### Wednesday/Thursday
|
||||
- [ ] Update `docs/talk-10min.md` Slide 6 with the H100 result
|
||||
- [ ] Open Nsight Systems GUI locally, load the `.nsys-rep` file, screenshot the timeline showing both pipelines
|
||||
|
||||
### Thursday (if someone gives you a contact after the talk)
|
||||
- [ ] Do not show GPU source code without NDA signed first
|
||||
- [ ] First call is discovery — ask about their infrastructure before pitching
|
||||
- [ ] Pricing range if pushed: $50K–$150K annually depending on scale
|
||||
|
||||
---
|
||||
|
||||
## Next Session Priorities (after this week)
|
||||
|
||||
1. **H100 setup script** — a single paste-into-terminal script for the cloud instance (was going to write this, ran out of session). Needs: Rust install, maturin build with `--features gpu`, nsys profile run, output capture.
|
||||
2. **Phase 1 fixes** — if not completed before Wednesday, finish these
|
||||
3. **vLLM integration** — Phase 2.1, this is the demo that sells the concept
|
||||
4. **NM anonymous LLC** — straightforward to form locally, no member names in public filings, no annual report requirement
|
||||
5. **C++ version scoping** — after Phase 1 is solid, assess Whetstone transpilation of `llm.rs`
|
||||
|
||||
---
|
||||
|
||||
## Key Technical Context for Next Session
|
||||
|
||||
- Repo is cloned at `/home/bill/Documents/unchecked-io/`, currently on `gpu-zero` branch
|
||||
- GPU feature is compiled with `maturin develop --release --features gpu`
|
||||
- The four Phase 1 bugs are documented in detail in `ROADMAP.md` Phase 1 section
|
||||
- `src/llm.rs` is the file to fix — `PinnedBatcher` struct, `batch_encode_to_gpu` method
|
||||
- `src/lib.rs` `TokenizerEngine` wraps `PinnedBatcher` and is the Python-facing class
|
||||
- chrono-tz is pinned at `=0.10.4` — do not change this, it was hard-won
|
||||
- `gpu-zero` branch is private; `main` branch is open source Apache 2.0
|
||||
84
HANDOFF-2026-04-21.md
Normal file
84
HANDOFF-2026-04-21.md
Normal file
@@ -0,0 +1,84 @@
|
||||
# Handoff — 2026-04-21
|
||||
|
||||
## What Was Done This Session
|
||||
|
||||
Fixed all four Phase 1 correctness bugs. Code compiles clean on temp-ubuntu with CUDA 13.1.
|
||||
|
||||
---
|
||||
|
||||
## Phase 1 Bugs — ALL FIXED
|
||||
|
||||
### 1.1 Dynamic buffer allocation — DONE
|
||||
**Old behavior:** `max_elements = 4096 * 512` hardcoded in `PinnedBatcher::new`. Any batch exceeding ~2M tokens returned an error.
|
||||
**Fix:** Removed `gpu_buffer`, `host_buffer`, and `max_elements` from the `PinnedBatcher` struct entirely. Both buffers are now allocated per-call based on actual `batch_size * seq_len`.
|
||||
|
||||
### 1.2 Buffer isolation (DLPack overwrite) — DONE
|
||||
**Old behavior:** Every call overwrote the same `CudaSlice`. Any tensor from call N got corrupted by call N+1.
|
||||
**Fix:** Each call allocates its own device buffer and wraps it in a `DLPackContext` struct stored in `DLManagedTensor.manager_ctx`. The deleter frees both the device memory (by dropping `CudaSlice<i64>`) and the shape array when PyTorch releases the tensor. Each returned tensor now owns its allocation.
|
||||
|
||||
### 1.3 Thread safety — DONE
|
||||
**Old behavior:** `batch_encode_to_gpu` took `&mut self`. `TokenizerEngine.encode_batch` took `&mut self`. Not usable from multiple Python threads.
|
||||
**Fix:** Since `PinnedBatcher` no longer has mutable shared state (buffers are per-call), `batch_encode_to_gpu` now takes `&self`. `TokenizerEngine` wraps `Arc<PinnedBatcher>` instead of owning it directly. `encode_batch` takes `&self`. Concurrent calls from Python threads are safe.
|
||||
|
||||
### 1.4 Variable sequence length + padding — DONE
|
||||
**Old behavior:** `seq_len = encodings[0].len()` — assumed uniform length. Variable-length batches produced wrong tensor shape.
|
||||
**Fix:** `seq_len = encodings.iter().map(|e| e.len()).max()`. Host buffer prefilled with `pad_id` (from `tokenizer.get_padding()`, fallback to 0). Rayon fill loop only writes actual tokens; shorter sequences retain the pad value.
|
||||
|
||||
---
|
||||
|
||||
## Build Issue Discovered and Resolved
|
||||
|
||||
`cudarc 0.11.9` panics at build time on CUDA 13.1 (only knows 12.5 and below). This machine has CUDA 13.1.
|
||||
|
||||
**Fix:** Changed the `gpu` feature in `Cargo.toml` from `cudarc/cuda-version-from-build-system` to `cudarc/cuda-12050`. This hard-pins to CUDA 12.5 headers and skips the `nvcc` version check. CUDA 13.x is backward-compatible with 12.5 driver API at runtime. If you ever upgrade cudarc to 0.12+, revert this change and re-test.
|
||||
|
||||
---
|
||||
|
||||
## Files Changed
|
||||
|
||||
| File | What Changed |
|
||||
|------|-------------|
|
||||
| `src/llm.rs` | Full rewrite of `PinnedBatcher` — removed static buffers, per-call allocation, `DLPackContext` for ownership, `&self` throughout |
|
||||
| `src/lib.rs` | `TokenizerEngine.inner` changed to `Arc<PinnedBatcher>`, `encode_batch` takes `&self`, added `use std::sync::Arc` |
|
||||
| `Cargo.toml` | `gpu` feature: `cuda-version-from-build-system` → `cuda-12050` |
|
||||
|
||||
---
|
||||
|
||||
## This Week — What's Left (in order)
|
||||
|
||||
### Today / Tomorrow (Tuesday 2026-04-22)
|
||||
- [ ] Rent H100 on Lambda Labs (~$2 for 30 min)
|
||||
- [ ] Run `benchmark_gpu.py` under `nsys profile --stats=true python benchmark_gpu.py`
|
||||
- [ ] Capture Nsight `.nsys-rep` file and download it
|
||||
- [ ] Note the speedup number — this is the headline for Thursday's talk
|
||||
|
||||
**Still need:** H100 setup script (a single paste-into-terminal script for the cloud instance). Needs: Rust install, maturin build with `--features gpu`, nsys profile run, output capture. This was not written last session.
|
||||
|
||||
### Wednesday (2026-04-23)
|
||||
- [ ] Update `docs/talk-10min.md` Slide 6 with the H100 result
|
||||
- [ ] Open Nsight Systems GUI locally, load the `.nsys-rep` file, screenshot the timeline showing both pipelines
|
||||
|
||||
### Thursday (2026-04-24)
|
||||
- [ ] Talk to IT/security group (10 min)
|
||||
- [ ] If someone gives you a contact: do not show GPU source code without NDA signed first; first call is discovery
|
||||
|
||||
---
|
||||
|
||||
## Next Session Priorities (after Thursday)
|
||||
|
||||
1. **Write H100 setup script** — single paste-into-terminal script (see above)
|
||||
2. **vLLM integration** — Phase 2.1, this is the demo that sells the concept
|
||||
3. **NM anonymous LLC** — Wyoming, single-member, before any contract is signed
|
||||
4. **C++ version scoping** — after Phase 1 solid, assess Whetstone transpilation of `llm.rs`
|
||||
|
||||
---
|
||||
|
||||
## Key Technical Context
|
||||
|
||||
- Repo: `/home/bill/Documents/unchecked-io/`, on `gpu-zero` branch
|
||||
- Build: `maturin develop --release --features gpu` (maturin not installed on temp-ubuntu — install with `pip install maturin` first)
|
||||
- `cargo check --features gpu` confirms clean compile
|
||||
- The four Phase 1 bugs are fixed and verified to compile
|
||||
- `src/llm.rs` is the GPU core. `src/lib.rs` is the Python-facing wrapper.
|
||||
- `chrono-tz` pinned at `=0.10.4` — do not change
|
||||
- `gpu-zero` branch is private; `main` is open source Apache 2.0
|
||||
387
LICENSE
387
LICENSE
@@ -1,373 +1,70 @@
|
||||
Mozilla Public License Version 2.0
|
||||
==================================
|
||||
Apache License
|
||||
Version 2.0, January 2004
|
||||
http://www.apache.org/licenses/
|
||||
|
||||
1. Definitions
|
||||
--------------
|
||||
TERMS AND CONDITIONS FOR USE, REPRODUCTION, AND DISTRIBUTION
|
||||
|
||||
1.1. "Contributor"
|
||||
means each individual or legal entity that creates, contributes to
|
||||
the creation of, or owns Covered Software.
|
||||
1. Definitions.
|
||||
|
||||
1.2. "Contributor Version"
|
||||
means the combination of the Contributions of others (if any) used
|
||||
by a Contributor and that particular Contributor's Contribution.
|
||||
"License" shall mean the terms and conditions for use, reproduction, and distribution as defined by Sections 1 through 9 of this document.
|
||||
|
||||
1.3. "Contribution"
|
||||
means Covered Software of a particular Contributor.
|
||||
"Licensor" shall mean the copyright owner or entity authorized by the copyright owner that is granting the License.
|
||||
|
||||
1.4. "Covered Software"
|
||||
means Source Code Form to which the initial Contributor has attached
|
||||
the notice in Exhibit A, the Executable Form of such Source Code
|
||||
Form, and Modifications of such Source Code Form, in each case
|
||||
including portions thereof.
|
||||
"Legal Entity" shall mean the union of the acting entity and all other entities that control, are controlled by, or are under common control with that entity. For the purposes of this definition, "control" means (i) the power, direct or indirect, to cause the direction or management of such entity, whether by contract or otherwise, or (ii) ownership of fifty percent (50%) or more of the outstanding shares, or (iii) beneficial ownership of such entity.
|
||||
|
||||
1.5. "Incompatible With Secondary Licenses"
|
||||
means
|
||||
"You" (or "Your") shall mean an individual or Legal Entity exercising permissions granted by this License.
|
||||
|
||||
(a) that the initial Contributor has attached the notice described
|
||||
in Exhibit B to the Covered Software; or
|
||||
"Source" form shall mean the preferred form for making modifications, including but not limited to software source code, documentation source, and configuration files.
|
||||
|
||||
(b) that the Covered Software was made available under the terms of
|
||||
version 1.1 or earlier of the License, but not also under the
|
||||
terms of a Secondary License.
|
||||
"Object" form shall mean any form resulting from mechanical transformation or translation of a Source form, including but not limited to compiled object code, generated documentation, and conversions to other media types.
|
||||
|
||||
1.6. "Executable Form"
|
||||
means any form of the work other than Source Code Form.
|
||||
"Work" shall mean the work of authorship, whether in Source or Object form, made available under the License, as indicated by a copyright notice that is included in or attached to the work (an example is provided in the Appendix below).
|
||||
|
||||
1.7. "Larger Work"
|
||||
means a work that combines Covered Software with other material, in
|
||||
a separate file or files, that is not Covered Software.
|
||||
"Derivative Works" shall mean any work, whether in Source or Object form, that is based on (or derived from) the Work and for which the editorial revisions, annotations, elaborations, or other modifications represent, as a whole, an original work of authorship. For the purposes of this License, Derivative Works shall not include works that remain separable from, or merely link (or bind by name) to the interfaces of, the Work and Derivative Works thereof.
|
||||
|
||||
1.8. "License"
|
||||
means this document.
|
||||
"Contribution" shall mean any work of authorship, including the original version of the Work and any modifications or additions to that Work or Derivative Works thereof, that is intentionally submitted to Licensor for inclusion in the Work by the copyright owner or by an individual or Legal Entity authorized to submit on behalf of the copyright owner. For the purposes of this definition, "submitted" means any form of electronic, verbal, or written communication sent to the Licensor or its representatives, including but not limited to communication on electronic mailing lists, source code control systems, and issue tracking systems that are managed by, or on behalf of, the Licensor for the purpose of discussing and improving the Work, but excluding communication that is conspicuously marked or otherwise designated in writing by the copyright owner as "Not a Contribution."
|
||||
|
||||
1.9. "Licensable"
|
||||
means having the right to grant, to the maximum extent possible,
|
||||
whether at the time of the initial grant or subsequently, any and
|
||||
all of the rights conveyed by this License.
|
||||
"Contributor" shall mean Licensor and any individual or Legal Entity on behalf of whom a Contribution has been received by Licensor and subsequently incorporated within the Work.
|
||||
|
||||
1.10. "Modifications"
|
||||
means any of the following:
|
||||
2. Grant of Copyright License. Subject to the terms and conditions of this License, each Contributor hereby grants to You a perpetual, worldwide, non-exclusive, no-charge, royalty-free, irrevocable copyright license to reproduce, prepare Derivative Works of, publicly display, publicly perform, sublicense, and distribute the Work and such Derivative Works in Source or Object form.
|
||||
|
||||
(a) any file in Source Code Form that results from an addition to,
|
||||
deletion from, or modification of the contents of Covered
|
||||
Software; or
|
||||
3. Grant of Patent License. Subject to the terms and conditions of this License, each Contributor hereby grants to You a perpetual, worldwide, non-exclusive, no-charge, royalty-free, irrevocable (except as stated in this section) patent license to make, have made, use, offer to sell, sell, import, and otherwise transfer the Work, where such license applies only to those patent claims licensable by such Contributor that are necessarily infringed by their Contribution(s) alone or by combination of their Contribution(s) with the Work to which such Contribution(s) was submitted. If You institute patent litigation against any entity (including a cross-claim or counterclaim in a lawsuit) alleging that the Work or a Contribution incorporated within the Work constitutes direct or contributory patent infringement, then any patent licenses granted to You under this License for that Work shall terminate as of the date such litigation is filed.
|
||||
|
||||
(b) any new file in Source Code Form that contains any Covered
|
||||
Software.
|
||||
4. Redistribution. You may reproduce and distribute copies of the Work or Derivative Works thereof in any medium, with or without modifications, and in Source or Object form, provided that You meet the following conditions:
|
||||
|
||||
1.11. "Patent Claims" of a Contributor
|
||||
means any patent claim(s), including without limitation, method,
|
||||
process, and apparatus claims, in any patent Licensable by such
|
||||
Contributor that would be infringed, but for the grant of the
|
||||
License, by the making, using, selling, offering for sale, having
|
||||
made, import, or transfer of either its Contributions or its
|
||||
Contributor Version.
|
||||
You must give any other recipients of the Work or Derivative Works a copy of this License; and
|
||||
You must cause any modified files to carry prominent notices stating that You changed the files; and
|
||||
You must retain, in the Source form of any Derivative Works that You distribute, all copyright, patent, trademark, and attribution notices from the Source form of the Work, excluding those notices that do not pertain to any part of the Derivative Works; and
|
||||
If the Work includes a "NOTICE" text file as part of its distribution, then any Derivative Works that You distribute must include a readable copy of the attribution notices contained within such NOTICE file, excluding those notices that do not pertain to any part of the Derivative Works, in at least one of the following places: within a NOTICE text file distributed as part of the Derivative Works; within the Source form or documentation, if provided along with the Derivative Works; or, within a display generated by the Derivative Works, if and wherever such third-party notices normally appear. The contents of the NOTICE file are for informational purposes only and do not modify the License. You may add Your own attribution notices within Derivative Works that You distribute, alongside or as an addendum to the NOTICE text from the Work, provided that such additional attribution notices cannot be construed as modifying the License.
|
||||
You may add Your own copyright statement to Your modifications and may provide additional or different license terms and conditions for use, reproduction, or distribution of Your modifications, or for any such Derivative Works as a whole, provided Your use, reproduction, and distribution of the Work otherwise complies with the conditions stated in this License.
|
||||
|
||||
1.12. "Secondary License"
|
||||
means either the GNU General Public License, Version 2.0, the GNU
|
||||
Lesser General Public License, Version 2.1, the GNU Affero General
|
||||
Public License, Version 3.0, or any later versions of those
|
||||
licenses.
|
||||
5. Submission of Contributions. Unless You explicitly state otherwise, any Contribution intentionally submitted for inclusion in the Work by You to the Licensor shall be under the terms and conditions of this License, without any additional terms or conditions. Notwithstanding the above, nothing herein shall supersede or modify the terms of any separate license agreement you may have executed with Licensor regarding such Contributions.
|
||||
|
||||
1.13. "Source Code Form"
|
||||
means the form of the work preferred for making modifications.
|
||||
6. Trademarks. This License does not grant permission to use the trade names, trademarks, service marks, or product names of the Licensor, except as required for reasonable and customary use in describing the origin of the Work and reproducing the content of the NOTICE file.
|
||||
|
||||
1.14. "You" (or "Your")
|
||||
means an individual or a legal entity exercising rights under this
|
||||
License. For legal entities, "You" includes any entity that
|
||||
controls, is controlled by, or is under common control with You. For
|
||||
purposes of this definition, "control" means (a) the power, direct
|
||||
or indirect, to cause the direction or management of such entity,
|
||||
whether by contract or otherwise, or (b) ownership of more than
|
||||
fifty percent (50%) of the outstanding shares or beneficial
|
||||
ownership of such entity.
|
||||
7. Disclaimer of Warranty. Unless required by applicable law or agreed to in writing, Licensor provides the Work (and each Contributor provides its Contributions) on an "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied, including, without limitation, any warranties or conditions of TITLE, NON-INFRINGEMENT, MERCHANTABILITY, or FITNESS FOR A PARTICULAR PURPOSE. You are solely responsible for determining the appropriateness of using or redistributing the Work and assume any risks associated with Your exercise of permissions under this License.
|
||||
|
||||
2. License Grants and Conditions
|
||||
--------------------------------
|
||||
8. Limitation of Liability. In no event and under no legal theory, whether in tort (including negligence), contract, or otherwise, unless required by applicable law (such as deliberate and grossly negligent acts) or agreed to in writing, shall any Contributor be liable to You for damages, including any direct, indirect, special, incidental, or consequential damages of any character arising as a result of this License or out of the use or inability to use the Work (including but not limited to damages for loss of goodwill, work stoppage, computer failure or malfunction, or any and all other commercial damages or losses), even if such Contributor has been advised of the possibility of such damages.
|
||||
|
||||
2.1. Grants
|
||||
9. Accepting Warranty or Additional Liability. While redistributing the Work or Derivative Works thereof, You may choose to offer, and charge a fee for, acceptance of support, warranty, indemnity, or other liability obligations and/or rights consistent with this License. However, in accepting such obligations, You may act only on Your own behalf and on Your sole responsibility, not on behalf of any other Contributor, and only if You agree to indemnify, defend, and hold each Contributor harmless for any liability incurred by, or claims asserted against, such Contributor by reason of your accepting any such warranty or additional liability.
|
||||
|
||||
Each Contributor hereby grants You a world-wide, royalty-free,
|
||||
non-exclusive license:
|
||||
END OF TERMS AND CONDITIONS
|
||||
|
||||
(a) under intellectual property rights (other than patent or trademark)
|
||||
Licensable by such Contributor to use, reproduce, make available,
|
||||
modify, display, perform, distribute, and otherwise exploit its
|
||||
Contributions, either on an unmodified basis, with Modifications, or
|
||||
as part of a Larger Work; and
|
||||
How to apply the Apache License to your work
|
||||
Include a copy of the Apache License, typically in a file called LICENSE, in your work, and consider also including a NOTICE file that references the License.
|
||||
|
||||
(b) under Patent Claims of such Contributor to make, use, sell, offer
|
||||
for sale, have made, import, and otherwise transfer either its
|
||||
Contributions or its Contributor Version.
|
||||
To apply the Apache License to specific files in your work, attach the following boilerplate declaration, replacing the fields enclosed by brackets "[]" with your own identifying information. (Don't include the brackets!) Enclose the text in the appropriate comment syntax for the file format. We also recommend that you include a file or class name and description of purpose on the same "printed page" as the copyright notice for easier identification within third-party archives.
|
||||
|
||||
2.2. Effective Date
|
||||
Copyright [yyyy] [name of copyright owner]
|
||||
|
||||
The licenses granted in Section 2.1 with respect to any Contribution
|
||||
become effective for each Contribution on the date the Contributor first
|
||||
distributes such Contribution.
|
||||
Licensed under the Apache License, Version 2.0 (the "License");
|
||||
you may not use this file except in compliance with the License.
|
||||
You may obtain a copy of the License at
|
||||
|
||||
2.3. Limitations on Grant Scope
|
||||
http://www.apache.org/licenses/LICENSE-2.0
|
||||
|
||||
The licenses granted in this Section 2 are the only rights granted under
|
||||
this License. No additional rights or licenses will be implied from the
|
||||
distribution or licensing of Covered Software under this License.
|
||||
Notwithstanding Section 2.1(b) above, no patent license is granted by a
|
||||
Contributor:
|
||||
|
||||
(a) for any code that a Contributor has removed from Covered Software;
|
||||
or
|
||||
|
||||
(b) for infringements caused by: (i) Your and any other third party's
|
||||
modifications of Covered Software, or (ii) the combination of its
|
||||
Contributions with other software (except as part of its Contributor
|
||||
Version); or
|
||||
|
||||
(c) under Patent Claims infringed by Covered Software in the absence of
|
||||
its Contributions.
|
||||
|
||||
This License does not grant any rights in the trademarks, service marks,
|
||||
or logos of any Contributor (except as may be necessary to comply with
|
||||
the notice requirements in Section 3.4).
|
||||
|
||||
2.4. Subsequent Licenses
|
||||
|
||||
No Contributor makes additional grants as a result of Your choice to
|
||||
distribute the Covered Software under a subsequent version of this
|
||||
License (see Section 10.2) or under the terms of a Secondary License (if
|
||||
permitted under the terms of Section 3.3).
|
||||
|
||||
2.5. Representation
|
||||
|
||||
Each Contributor represents that the Contributor believes its
|
||||
Contributions are its original creation(s) or it has sufficient rights
|
||||
to grant the rights to its Contributions conveyed by this License.
|
||||
|
||||
2.6. Fair Use
|
||||
|
||||
This License is not intended to limit any rights You have under
|
||||
applicable copyright doctrines of fair use, fair dealing, or other
|
||||
equivalents.
|
||||
|
||||
2.7. Conditions
|
||||
|
||||
Sections 3.1, 3.2, 3.3, and 3.4 are conditions of the licenses granted
|
||||
in Section 2.1.
|
||||
|
||||
3. Responsibilities
|
||||
-------------------
|
||||
|
||||
3.1. Distribution of Source Form
|
||||
|
||||
All distribution of Covered Software in Source Code Form, including any
|
||||
Modifications that You create or to which You contribute, must be under
|
||||
the terms of this License. You must inform recipients that the Source
|
||||
Code Form of the Covered Software is governed by the terms of this
|
||||
License, and how they can obtain a copy of this License. You may not
|
||||
attempt to alter or restrict the recipients' rights in the Source Code
|
||||
Form.
|
||||
|
||||
3.2. Distribution of Executable Form
|
||||
|
||||
If You distribute Covered Software in Executable Form then:
|
||||
|
||||
(a) such Covered Software must also be made available in Source Code
|
||||
Form, as described in Section 3.1, and You must inform recipients of
|
||||
the Executable Form how they can obtain a copy of such Source Code
|
||||
Form by reasonable means in a timely manner, at a charge no more
|
||||
than the cost of distribution to the recipient; and
|
||||
|
||||
(b) You may distribute such Executable Form under the terms of this
|
||||
License, or sublicense it under different terms, provided that the
|
||||
license for the Executable Form does not attempt to limit or alter
|
||||
the recipients' rights in the Source Code Form under this License.
|
||||
|
||||
3.3. Distribution of a Larger Work
|
||||
|
||||
You may create and distribute a Larger Work under terms of Your choice,
|
||||
provided that You also comply with the requirements of this License for
|
||||
the Covered Software. If the Larger Work is a combination of Covered
|
||||
Software with a work governed by one or more Secondary Licenses, and the
|
||||
Covered Software is not Incompatible With Secondary Licenses, this
|
||||
License permits You to additionally distribute such Covered Software
|
||||
under the terms of such Secondary License(s), so that the recipient of
|
||||
the Larger Work may, at their option, further distribute the Covered
|
||||
Software under the terms of either this License or such Secondary
|
||||
License(s).
|
||||
|
||||
3.4. Notices
|
||||
|
||||
You may not remove or alter the substance of any license notices
|
||||
(including copyright notices, patent notices, disclaimers of warranty,
|
||||
or limitations of liability) contained within the Source Code Form of
|
||||
the Covered Software, except that You may alter any license notices to
|
||||
the extent required to remedy known factual inaccuracies.
|
||||
|
||||
3.5. Application of Additional Terms
|
||||
|
||||
You may choose to offer, and to charge a fee for, warranty, support,
|
||||
indemnity or liability obligations to one or more recipients of Covered
|
||||
Software. However, You may do so only on Your own behalf, and not on
|
||||
behalf of any Contributor. You must make it absolutely clear that any
|
||||
such warranty, support, indemnity, or liability obligation is offered by
|
||||
You alone, and You hereby agree to indemnify every Contributor for any
|
||||
liability incurred by such Contributor as a result of warranty, support,
|
||||
indemnity or liability terms You offer. You may include additional
|
||||
disclaimers of warranty and limitations of liability specific to any
|
||||
jurisdiction.
|
||||
|
||||
4. Inability to Comply Due to Statute or Regulation
|
||||
---------------------------------------------------
|
||||
|
||||
If it is impossible for You to comply with any of the terms of this
|
||||
License with respect to some or all of the Covered Software due to
|
||||
statute, judicial order, or regulation then You must: (a) comply with
|
||||
the terms of this License to the maximum extent possible; and (b)
|
||||
describe the limitations and the code they affect. Such description must
|
||||
be placed in a text file included with all distributions of the Covered
|
||||
Software under this License. Except to the extent prohibited by statute
|
||||
or regulation, such description must be sufficiently detailed for a
|
||||
recipient of ordinary skill to be able to understand it.
|
||||
|
||||
5. Termination
|
||||
--------------
|
||||
|
||||
5.1. The rights granted under this License will terminate automatically
|
||||
if You fail to comply with any of its terms. However, if You become
|
||||
compliant, then the rights granted under this License from a particular
|
||||
Contributor are reinstated (a) provisionally, unless and until such
|
||||
Contributor explicitly and finally terminates Your grants, and (b) on an
|
||||
ongoing basis, if such Contributor fails to notify You of the
|
||||
non-compliance by some reasonable means prior to 60 days after You have
|
||||
come back into compliance. Moreover, Your grants from a particular
|
||||
Contributor are reinstated on an ongoing basis if such Contributor
|
||||
notifies You of the non-compliance by some reasonable means, this is the
|
||||
first time You have received notice of non-compliance with this License
|
||||
from such Contributor, and You become compliant prior to 30 days after
|
||||
Your receipt of the notice.
|
||||
|
||||
5.2. If You initiate litigation against any entity by asserting a patent
|
||||
infringement claim (excluding declaratory judgment actions,
|
||||
counter-claims, and cross-claims) alleging that a Contributor Version
|
||||
directly or indirectly infringes any patent, then the rights granted to
|
||||
You by any and all Contributors for the Covered Software under Section
|
||||
2.1 of this License shall terminate.
|
||||
|
||||
5.3. In the event of termination under Sections 5.1 or 5.2 above, all
|
||||
end user license agreements (excluding distributors and resellers) which
|
||||
have been validly granted by You or Your distributors under this License
|
||||
prior to termination shall survive termination.
|
||||
|
||||
************************************************************************
|
||||
* *
|
||||
* 6. Disclaimer of Warranty *
|
||||
* ------------------------- *
|
||||
* *
|
||||
* Covered Software is provided under this License on an "as is" *
|
||||
* basis, without warranty of any kind, either expressed, implied, or *
|
||||
* statutory, including, without limitation, warranties that the *
|
||||
* Covered Software is free of defects, merchantable, fit for a *
|
||||
* particular purpose or non-infringing. The entire risk as to the *
|
||||
* quality and performance of the Covered Software is with You. *
|
||||
* Should any Covered Software prove defective in any respect, You *
|
||||
* (not any Contributor) assume the cost of any necessary servicing, *
|
||||
* repair, or correction. This disclaimer of warranty constitutes an *
|
||||
* essential part of this License. No use of any Covered Software is *
|
||||
* authorized under this License except under this disclaimer. *
|
||||
* *
|
||||
************************************************************************
|
||||
|
||||
************************************************************************
|
||||
* *
|
||||
* 7. Limitation of Liability *
|
||||
* -------------------------- *
|
||||
* *
|
||||
* Under no circumstances and under no legal theory, whether tort *
|
||||
* (including negligence), contract, or otherwise, shall any *
|
||||
* Contributor, or anyone who distributes Covered Software as *
|
||||
* permitted above, be liable to You for any direct, indirect, *
|
||||
* special, incidental, or consequential damages of any character *
|
||||
* including, without limitation, damages for lost profits, loss of *
|
||||
* goodwill, work stoppage, computer failure or malfunction, or any *
|
||||
* and all other commercial damages or losses, even if such party *
|
||||
* shall have been informed of the possibility of such damages. This *
|
||||
* limitation of liability shall not apply to liability for death or *
|
||||
* personal injury resulting from such party's negligence to the *
|
||||
* extent applicable law prohibits such limitation. Some *
|
||||
* jurisdictions do not allow the exclusion or limitation of *
|
||||
* incidental or consequential damages, so this exclusion and *
|
||||
* limitation may not apply to You. *
|
||||
* *
|
||||
************************************************************************
|
||||
|
||||
8. Litigation
|
||||
-------------
|
||||
|
||||
Any litigation relating to this License may be brought only in the
|
||||
courts of a jurisdiction where the defendant maintains its principal
|
||||
place of business and such litigation shall be governed by laws of that
|
||||
jurisdiction, without reference to its conflict-of-law provisions.
|
||||
Nothing in this Section shall prevent a party's ability to bring
|
||||
cross-claims or counter-claims.
|
||||
|
||||
9. Miscellaneous
|
||||
----------------
|
||||
|
||||
This License represents the complete agreement concerning the subject
|
||||
matter hereof. If any provision of this License is held to be
|
||||
unenforceable, such provision shall be reformed only to the extent
|
||||
necessary to make it enforceable. Any law or regulation which provides
|
||||
that the language of a contract shall be construed against the drafter
|
||||
shall not be used to construe this License against a Contributor.
|
||||
|
||||
10. Versions of the License
|
||||
---------------------------
|
||||
|
||||
10.1. New Versions
|
||||
|
||||
Mozilla Foundation is the license steward. Except as provided in Section
|
||||
10.3, no one other than the license steward has the right to modify or
|
||||
publish new versions of this License. Each version will be given a
|
||||
distinguishing version number.
|
||||
|
||||
10.2. Effect of New Versions
|
||||
|
||||
You may distribute the Covered Software under the terms of the version
|
||||
of the License under which You originally received the Covered Software,
|
||||
or under the terms of any subsequent version published by the license
|
||||
steward.
|
||||
|
||||
10.3. Modified Versions
|
||||
|
||||
If you create software not governed by this License, and you want to
|
||||
create a new license for such software, you may create and use a
|
||||
modified version of this License if you rename the license and remove
|
||||
any references to the name of the license steward (except to note that
|
||||
such modified license differs from this License).
|
||||
|
||||
10.4. Distributing Source Code Form that is Incompatible With Secondary
|
||||
Licenses
|
||||
|
||||
If You choose to distribute Source Code Form that is Incompatible With
|
||||
Secondary Licenses under the terms of this version of the License, the
|
||||
notice described in Exhibit B of this License must be attached.
|
||||
|
||||
Exhibit A - Source Code Form License Notice
|
||||
-------------------------------------------
|
||||
|
||||
This Source Code Form is subject to the terms of the Mozilla Public
|
||||
License, v. 2.0. If a copy of the MPL was not distributed with this
|
||||
file, You can obtain one at https://mozilla.org/MPL/2.0/.
|
||||
|
||||
If it is not possible or desirable to put the notice in a particular
|
||||
file, then You may include the notice in a location (such as a LICENSE
|
||||
file in a relevant directory) where a recipient would be likely to look
|
||||
for such a notice.
|
||||
|
||||
You may add additional accurate notices of copyright ownership.
|
||||
|
||||
Exhibit B - "Incompatible With Secondary Licenses" Notice
|
||||
---------------------------------------------------------
|
||||
|
||||
This Source Code Form is "Incompatible With Secondary Licenses", as
|
||||
defined by the Mozilla Public License, v. 2.0.
|
||||
Unless required by applicable law or agreed to in writing, software
|
||||
distributed under the License is distributed on an "AS IS" BASIS,
|
||||
WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
|
||||
See the License for the specific language governing permissions and
|
||||
limitations under the License.
|
||||
89
README.md
Normal file
89
README.md
Normal file
@@ -0,0 +1,89 @@
|
||||
UncheckedIO 🚀
|
||||
The world's fastest, most dangerous PostgreSQL-to-Arrow loader.
|
||||
UncheckedIO is an opinionated, high-performance data connector designed for bulk-loading massive datasets (Terabytes+) from PostgreSQL into Python/Apache Arrow. It achieves extreme speed by skipping runtime schema validation and using a zero-copy streaming parser.
|
||||
⚠️ WARNING: This library assumes you know what you are doing. If your config.yaml schema does not match your database schema, it will produce garbage data or crash. The "Unchecked" in the name is not a suggestion; it is a promise.
|
||||
Features
|
||||
🚀 1.5x - 2x Faster than ConnectorX when Benchmarked on 20M+ row dataset in default Colab runtime.
|
||||
🧵 Parallel Worker Pool: Automatically scales to available CPU cores.
|
||||
🧠 Auto-Tuning: Dynamically calculates optimal partition sizes (blast_radius).
|
||||
📉 Low Memory Footprint: Streaming parser with zero-copy string slicing.
|
||||
🛠️ "Self-Healing" (Nullable): If a partition fails (e.g., network flake), it returns NULLs for that chunk instead of crashing the entire job.
|
||||
🔬 Optional Profiling: Built-in integration with the Tracy Profiler for deep I/O analysis.
|
||||
Installation
|
||||
pip install unchecked-io
|
||||
|
||||
Quick Start
|
||||
Create a config.yaml file:
|
||||
connection_string: postgresql://user:password@localhost:5432/dbname
|
||||
# query MUST be a COPY ... TO STDOUT (FORMAT binary) command
|
||||
query: >
|
||||
COPY (SELECT id, name, score, created_at FROM users)
|
||||
TO STDOUT (FORMAT binary)
|
||||
schema:
|
||||
- column_name: id
|
||||
arrow_type: Int64
|
||||
- column_name: name
|
||||
arrow_type: Utf8
|
||||
- column_name: score
|
||||
arrow_type: Float32
|
||||
- column_name: created_at
|
||||
arrow_type: Timestamp(Nanosecond, None)
|
||||
|
||||
|
||||
Run it in Python:
|
||||
import unchecked_io
|
||||
|
||||
# blast_radius=0 enables auto-tuning
|
||||
arrow_table = unchecked_io.load_data_from_config("config.yaml", blast_radius=0)
|
||||
|
||||
# Convert to Pandas (Zero-Copy)
|
||||
df = unchecked_io.to_pandas_dataframe(arrow_table)
|
||||
print(df.head())
|
||||
|
||||
|
||||
🔬 Profiling with Tracy
|
||||
UncheckedIO includes optional instrumentation for the Tracy Profiler to visualize I/O starvation and thread contention.
|
||||
This feature is DISABLED by default to prevent overhead and crashes in production environments (like Google Colab).
|
||||
How to Enable Profiling
|
||||
You must build the library from source with the profiling feature enabled.
|
||||
Install Rust & Maturin:
|
||||
curl --proto '=https' --tlsv1.2 -sSf [https://sh.rustup.rs](https://sh.rustup.rs) | sh
|
||||
pip install maturin
|
||||
|
||||
|
||||
Build with Feature Flag:
|
||||
# This builds the wheel and installs it in your current venv
|
||||
maturin develop --release --features profiling
|
||||
|
||||
|
||||
Run Tracy (GUI):
|
||||
Download Tracy v0.11.1 (Must match protocol v0.11.x).
|
||||
Open Tracy.exe and click Connect.
|
||||
Run your Python script. You will see real-time thread timelines and "IO_WAIT_STARVATION" blocks.
|
||||
Configuration Reference
|
||||
Supported Types
|
||||
PostgreSQL Type
|
||||
Config arrow_type
|
||||
bigint / int8
|
||||
Int64
|
||||
integer / int4
|
||||
Int32
|
||||
double precision
|
||||
Float64
|
||||
real / float4
|
||||
Float32
|
||||
text / varchar / uuid
|
||||
Utf8
|
||||
boolean
|
||||
Boolean
|
||||
timestamp
|
||||
Timestamp(Nanosecond, None)
|
||||
date
|
||||
Date32
|
||||
|
||||
blast_radius Parameter
|
||||
0: Auto-Tune (Recommended). Calculates partition size based on row count / (cores * 4).
|
||||
> 0: Manual Override. Sets specific number of rows per partition. Use for fine-tuning on specific hardware.
|
||||
License
|
||||
Apache 2.0
|
||||
Free for non-production use and production use.
|
||||
160
ROADMAP.md
Normal file
160
ROADMAP.md
Normal file
@@ -0,0 +1,160 @@
|
||||
# UncheckedIO — Development Roadmap
|
||||
|
||||
**Goal:** Production-grade, multi-GPU zero-copy ingestion pipeline for LLM providers.
|
||||
**Current state:** Working proof-of-concept on CUDA (gpu-zero branch). PostgreSQL→Arrow open source on main.
|
||||
|
||||
---
|
||||
|
||||
## Phase 1 — Fix Correctness (gpu-zero branch)
|
||||
|
||||
The GPU feature works in benchmark conditions but has four bugs that would surface immediately in production. Fix these before anything else.
|
||||
|
||||
### 1.1 Dynamic buffer allocation
|
||||
**Problem:** `max_elements = 4096 * 512` is hardcoded at init time. Modern LLMs have 32K–128K context windows. A single request at 32K tokens + any real batch size hits this limit.
|
||||
**Fix:** Allocate `gpu_buffer` and `host_buffer` per-call based on actual `batch_size * max_seq_len`, or use a growable strategy that reallocates when the current batch exceeds the existing buffer size.
|
||||
|
||||
### 1.2 Buffer isolation (silent data corruption bug)
|
||||
**Problem:** Every call to `encode_batch` overwrites the same `CudaSlice`. The DLPack capsule hands PyTorch a pointer to that same memory. A caller holding a tensor from call N gets its data silently overwritten by call N+1.
|
||||
**Fix:** Double-buffer (ping-pong between two allocations) or allocate a fresh device buffer per call and transfer ownership to the DLPack capsule's deleter so CUDA frees it when PyTorch releases the tensor.
|
||||
|
||||
### 1.3 Thread safety
|
||||
**Problem:** `encode_batch` takes `&mut self`. The `TokenizerEngine` Python class cannot be called from multiple threads simultaneously. Every production serving engine (vLLM, TGI) is multithreaded.
|
||||
**Fix:** Wrap `PinnedBatcher` in an `Arc<Mutex<>>` or, better, create a `EnginePool` that holds N instances and checks one out per call. Pool size = CPU core count is a reasonable default.
|
||||
|
||||
### 1.4 Variable sequence length + padding
|
||||
**Problem:** `seq_len = encodings[0].len()` assumes all sequences in the batch have identical length. The `tokenizers` crate doesn't apply padding by default, so real variable-length inputs produce an incorrect 2D tensor shape.
|
||||
**Fix:** Apply padding to `max(encodings[i].len())` within the batch, fill with the tokenizer's `pad_id`, and use the actual padded length as `seq_len`.
|
||||
|
||||
---
|
||||
|
||||
## Phase 2 — Production Features
|
||||
|
||||
### 2.1 vLLM integration
|
||||
vLLM is the dominant open-source LLM serving engine. Make `unchecked_io` a drop-in tokenization backend for it.
|
||||
- Implement the `TokenizerGroup` protocol vLLM expects
|
||||
- The continuous batching scheduler sends variable-length requests asynchronously — this is where Phase 1 fixes must be solid
|
||||
- Target: a 3-line config change in vLLM to switch to the UncheckedIO backend
|
||||
- Deliverable: `examples/vllm_integration.py` showing the swap
|
||||
|
||||
### 2.2 Streaming / chunked output
|
||||
**Current:** The entire batch is tokenized and DMA'd in one blocking call.
|
||||
**Target:** Yield Arrow RecordBatches or DLPack tensors incrementally as tokens are ready, so the model can begin prefill before the full batch is staged.
|
||||
|
||||
### 2.3 Attention mask export
|
||||
LLM inference requires both `input_ids` and `attention_mask`. Currently only token IDs are exported. Add a second DLPack capsule for the mask, or pack both into a named dict capsule.
|
||||
|
||||
### 2.4 Configurable max context
|
||||
Replace the hardcoded constants with constructor arguments:
|
||||
```python
|
||||
engine = TokenizerEngine(
|
||||
tokenizer_path="tokenizer.json",
|
||||
max_seq_len=32768,
|
||||
max_batch_size=64,
|
||||
pool_size=8
|
||||
)
|
||||
```
|
||||
|
||||
---
|
||||
|
||||
## Phase 3 — Multi-GPU Support
|
||||
|
||||
### 3.1 ROCm / HIP (AMD) — highest priority
|
||||
AMD's MI300X has 192GB memory vs the H100's 80GB and providers are actively evaluating it to escape NVIDIA pricing. HIP is the path of least resistance — function names are nearly 1:1 with CUDA (`hipMemHostRegister`, `hipMemcpyHtoD`, `hipMemHostUnregister`).
|
||||
|
||||
- Add `rocm` feature flag in `Cargo.toml`
|
||||
- Create `src/backends/rocm.rs` implementing the same `GpuBackend` trait
|
||||
- Use AMD's `hip-sys` or `rocm-sys` Rust crate
|
||||
- The DLPack export layer stays identical — device_type `10` for ROCm vs `2` for CUDA
|
||||
|
||||
**Refactor needed first:** Extract a `GpuBackend` trait from the current CUDA-specific code:
|
||||
```rust
|
||||
trait GpuBackend {
|
||||
fn allocate_pinned(capacity: usize) -> Result<PinnedBuffer>;
|
||||
fn copy_to_device(src: &[i64], dst: &mut DeviceBuffer) -> Result<()>;
|
||||
fn device_ptr(buf: &DeviceBuffer) -> u64;
|
||||
fn dlpack_device_type() -> i32;
|
||||
}
|
||||
```
|
||||
|
||||
### 3.2 OpenCL (broad coverage)
|
||||
Covers AMD, Intel, older NVIDIA, Apple (via Metal-OpenCL bridge). Lower peak performance than native CUDA/HIP but maximum hardware coverage. Use the `opencl3` Rust crate.
|
||||
- Add `opencl` feature flag
|
||||
- Useful for providers running mixed-GPU clusters
|
||||
|
||||
### 3.3 Intel Gaudi / Level Zero
|
||||
Intel's Gaudi 2/3 accelerators are gaining traction for training and are available on AWS and Azure. Intel Level Zero is the low-level API equivalent to CUDA driver API. Lower priority than ROCm but worth tracking.
|
||||
|
||||
---
|
||||
|
||||
## Phase 4 — Benchmark Suite
|
||||
|
||||
The current `benchmark_gpu.py` is good for what it tests but synthetic. A buyer will run their own benchmarks — make it easy and cover real scenarios.
|
||||
|
||||
### 4.1 Scenarios to cover
|
||||
- Short sequences: 128, 256 tokens (BERT-style)
|
||||
- Medium sequences: 2048, 4096 tokens (typical chat)
|
||||
- Long sequences: 16K, 32K tokens (document/RAG workloads)
|
||||
- Variable length batches (realistic distribution, not identical strings)
|
||||
- Concurrent throughput: simulate N async callers
|
||||
|
||||
### 4.2 Tokenizers to cover
|
||||
- BERT (`bert-base-uncased`) — already done
|
||||
- LLaMA 3 tokenizer (BPE, more representative of modern LLMs)
|
||||
- GPT-2 tokenizer (common baseline)
|
||||
|
||||
### 4.3 Hardware targets
|
||||
- Local RTX 3060 (done, 4.7x)
|
||||
- Cloud H100 SXM (~$2 on Lambda Labs, 30 min run) — **do this first, it changes the pitch**
|
||||
- AMD MI300X (after Phase 3.1)
|
||||
|
||||
### 4.4 What to measure
|
||||
- Tokens/second throughput
|
||||
- Latency per batch (p50, p95, p99)
|
||||
- GPU utilization % (the thing providers care about)
|
||||
- Memory bandwidth saturation
|
||||
|
||||
---
|
||||
|
||||
## Phase 5 — Business Infrastructure
|
||||
|
||||
### 5.1 Wyoming LLC formation
|
||||
- Single-member LLC, Wyoming (strongest privacy, no public member disclosure)
|
||||
- LLC owns all IP
|
||||
- Registered agent handles public address
|
||||
- Done before any contract is signed
|
||||
|
||||
### 5.2 License structure
|
||||
- Open core: PostgreSQL→Arrow stays Apache 2.0
|
||||
- GPU features: commercial license
|
||||
- Pricing model: annual per-cluster license, priced against GPU utilization savings
|
||||
- License enforcement: honor-system enterprise contracts initially, not technical DRM
|
||||
|
||||
### 5.3 Sales collateral
|
||||
- 1-page technical brief (problem / mechanism / numbers / what's included)
|
||||
- Self-contained benchmark the prospect can run on their own hardware
|
||||
- Reference architecture diagram showing where this sits in a vLLM serving stack
|
||||
|
||||
---
|
||||
|
||||
## Branch Strategy
|
||||
|
||||
| Branch | Purpose |
|
||||
|--------|---------|
|
||||
| `main` | Open source PostgreSQL→Arrow, Apache 2.0, public |
|
||||
| `gpu-zero` | GPU feature development, private |
|
||||
| `rocm` | AMD port (Phase 3.1) |
|
||||
| `release/x.y` | Tagged releases for licensing |
|
||||
|
||||
The `gpu` feature is compiled in only with `--features gpu`. The open source build never includes GPU code.
|
||||
|
||||
---
|
||||
|
||||
## Immediate Next Steps (in order)
|
||||
|
||||
1. Fix 1.1 (dynamic buffer) — unblocks everything else
|
||||
2. Fix 1.2 (buffer isolation) — correctness before benchmarking
|
||||
3. Fix 1.3 (thread safety) — required for any real serving test
|
||||
4. Fix 1.4 (padding) — correctness
|
||||
5. Spend $2, run benchmark on cloud H100 — this one number changes the pitch
|
||||
6. Phase 2.1 vLLM integration — this is the demo that sells the concept
|
||||
7. Phase 3.1 ROCm — this is the differentiator
|
||||
34
benchmark.py
34
benchmark.py
@@ -4,19 +4,34 @@ import sqlalchemy
|
||||
import connectorx as cx
|
||||
import unchecked_io
|
||||
import os
|
||||
import yaml # We need pyyaml for this
|
||||
import yaml
|
||||
import time
|
||||
|
||||
# --- 1. Define Connection Strings and Query ---
|
||||
# These must match your local Docker setup
|
||||
DB_USER = "postgres"
|
||||
DB_PASS = "mysecretpassword"
|
||||
DB_HOST = "localhost"
|
||||
DB_PORT = "5433" # <-- Your local Docker port
|
||||
DB_PORT = "5433"
|
||||
DB_NAME = "postgres"
|
||||
|
||||
# Global Configuration
|
||||
BLAST_RADIUS = 1250000 # Rows per parallel task (1M / 62500 = 16 partitions)
|
||||
|
||||
# Rows per parallel task (1M / 62500 = 16 partitions)
|
||||
|
||||
#BLAST_RADIUS = 1250000 # 16 partitions
|
||||
|
||||
#BLAST_RADIUS = 625000 # 32 partitions
|
||||
|
||||
#BLAST_RADIUS = 312500 # 64 partitions
|
||||
|
||||
#BLAST_RADIUS = 156250 # 128 partitions
|
||||
|
||||
#BLAST_RADIUS = 125000 # 160 partitions
|
||||
|
||||
#BLAST_RADIUS = 12500 # 1600 partitions
|
||||
|
||||
BLAST_RADIUS = 1250 # 16000 partitions
|
||||
BATCH_SIZE = 262144 # Aggregation Size (Large Memory chunks)
|
||||
|
||||
# SQLAlchemy connection string (for Pandas)
|
||||
sqlalchemy_conn_str = f"postgresql://{DB_USER}:{DB_PASS}@{DB_HOST}:{DB_PORT}/{DB_NAME}"
|
||||
@@ -35,6 +50,7 @@ unchecked_io_query = "COPY (SELECT id::BIGINT, uuid::TEXT, username::TEXT, score
|
||||
config_data = {
|
||||
'connection_string': connectorx_conn_str,
|
||||
'query': unchecked_io_query,
|
||||
'batch_size': BATCH_SIZE, # <--- This explicitly writes your setting to the file
|
||||
'schema': [
|
||||
{'column_name': 'id', 'arrow_type': 'Int64'},
|
||||
{'column_name': 'uuid', 'arrow_type': 'Utf8'},
|
||||
@@ -53,6 +69,8 @@ with open(config_file, 'w') as f:
|
||||
yaml.dump(config_data, f)
|
||||
|
||||
print(f"UncheckedIO Config: {config_file} (dynamically created for local test)")
|
||||
print(f" - Blast Radius: {BLAST_RADIUS}")
|
||||
print(f" - Batch Size: {BATCH_SIZE}")
|
||||
|
||||
# --- 3. Define Benchmark Functions ---
|
||||
|
||||
@@ -73,14 +91,6 @@ def test_unchecked_io():
|
||||
# --- 4. Run Benchmarks ---
|
||||
run_count = 1
|
||||
print(f"Running benchmarks for 20,000,000 rows (average of {run_count} runs)...")
|
||||
print(f"Blast Radius: {BLAST_RADIUS} rows per task")
|
||||
|
||||
# --- Pandas ---
|
||||
# print("\nRunning Pandas warmup...")
|
||||
# _ = test_pandas()
|
||||
# print("Timing pandas.read_sql...")
|
||||
# pandas_time = timeit.timeit(test_pandas, number=run_count) / run_count
|
||||
# print(f"Pandas Average Time: {pandas_time * 1000:.2f} ms")
|
||||
|
||||
# --- ConnectorX ---
|
||||
print("\nRunning ConnectorX warmup...")
|
||||
|
||||
114
benchmark_gpu.py
Normal file
114
benchmark_gpu.py
Normal file
@@ -0,0 +1,114 @@
|
||||
import time
|
||||
import torch
|
||||
import os
|
||||
import json
|
||||
import unchecked_io
|
||||
from transformers import AutoTokenizer
|
||||
|
||||
# --- CONFIGURATION ---
|
||||
MODEL_ID = "bert-base-uncased"
|
||||
TOKENIZER_PATH = "tokenizer.json"
|
||||
BATCH_SIZE = 4096 # Massive batch to saturate PCIe
|
||||
NUM_BATCHES = 50 # Run enough to stabilize GPU clock
|
||||
SEQ_LEN = 128 # Typical sentence length
|
||||
WARMUP = 5
|
||||
|
||||
# Synthetic Data
|
||||
SAMPLE_TEXT = "The quick brown fox jumps over the lazy dog. " * 5 # ~50 tokens
|
||||
DATA_BATCH = [SAMPLE_TEXT for _ in range(BATCH_SIZE)]
|
||||
|
||||
print(f"--- GPU PIPELINE BENCHMARK (RTX 3060) ---")
|
||||
print(f"Batch Size: {BATCH_SIZE} | Batches: {NUM_BATCHES}")
|
||||
print(f"Total Strings processed: {BATCH_SIZE * NUM_BATCHES:,}")
|
||||
|
||||
# --- 1. SETUP RESOURCES ---
|
||||
|
||||
# A. Prepare Tokenizer File for Rust
|
||||
if not os.path.exists(TOKENIZER_PATH):
|
||||
print(f"Downloading {MODEL_ID} tokenizer...")
|
||||
# forcing use_fast=True ensures we get the JSON file Rust needs
|
||||
hf_tokenizer = AutoTokenizer.from_pretrained(MODEL_ID, use_fast=True)
|
||||
hf_tokenizer.save_pretrained(".")
|
||||
print("Saved tokenizer.json")
|
||||
else:
|
||||
print("Found existing tokenizer.json")
|
||||
hf_tokenizer = AutoTokenizer.from_pretrained(MODEL_ID, use_fast=True)
|
||||
|
||||
# B. Initialize Engines
|
||||
print("Initializing Engines...")
|
||||
|
||||
# UncheckedIO (Rust + CUDA)
|
||||
try:
|
||||
rust_engine = unchecked_io.TokenizerEngine(TOKENIZER_PATH)
|
||||
print("✅ UncheckedIO Engine Loaded")
|
||||
except AttributeError:
|
||||
print("❌ Error: UncheckedIO 'TokenizerEngine' not found.")
|
||||
print("Did you compile with: maturin develop --features gpu")
|
||||
exit(1)
|
||||
|
||||
# --- 2. BENCHMARK LOOPS ---
|
||||
|
||||
def benchmark_standard():
|
||||
torch.cuda.synchronize()
|
||||
start = time.time()
|
||||
|
||||
for _ in range(NUM_BATCHES):
|
||||
# Step 1: Tokenize (CPU Python/Rust mix)
|
||||
# return_tensors='pt' creates a CPU tensor
|
||||
encodings = hf_tokenizer(
|
||||
DATA_BATCH,
|
||||
padding=True,
|
||||
truncation=True,
|
||||
max_length=SEQ_LEN,
|
||||
return_tensors="pt"
|
||||
)
|
||||
|
||||
# Step 2: Move to GPU (The Bottleneck)
|
||||
input_ids = encodings["input_ids"].to("cuda", non_blocking=True)
|
||||
|
||||
# Force sync to measure actual completion
|
||||
torch.cuda.synchronize()
|
||||
|
||||
return time.time() - start
|
||||
|
||||
def benchmark_unchecked():
|
||||
torch.cuda.synchronize()
|
||||
start = time.time()
|
||||
|
||||
for _ in range(NUM_BATCHES):
|
||||
# Step 1 & 2: Tokenize + DMA to GPU (All in Rust)
|
||||
dlpack_capsule = rust_engine.encode_batch(DATA_BATCH)
|
||||
|
||||
# Step 3: Zero-Copy Wrap (Python)
|
||||
# This is virtually instant
|
||||
input_ids = torch.from_dlpack(dlpack_capsule)
|
||||
|
||||
# Force sync
|
||||
torch.cuda.synchronize()
|
||||
|
||||
return time.time() - start
|
||||
|
||||
# --- 3. RUN RACES ---
|
||||
|
||||
print("\nWARMING UP GPU...")
|
||||
# Run a few dummy passes to wake up the 3060
|
||||
_ = hf_tokenizer(DATA_BATCH[:10], return_tensors="pt")["input_ids"].to("cuda")
|
||||
_ = torch.from_dlpack(rust_engine.encode_batch(DATA_BATCH[:10]))
|
||||
|
||||
print("\n🚀 RUNNING STANDARD PIPELINE (HF -> CPU Tensor -> CUDA)...")
|
||||
time_std = benchmark_standard()
|
||||
fps_std = (BATCH_SIZE * NUM_BATCHES) / time_std
|
||||
print(f"Time: {time_std:.4f}s | Throughput: {fps_std:,.0f} samples/sec")
|
||||
|
||||
print("\n🚀 RUNNING UNCHECKED PIPELINE (Rust -> Pinned -> CUDA)...")
|
||||
time_unchecked = benchmark_unchecked()
|
||||
fps_unchecked = (BATCH_SIZE * NUM_BATCHES) / time_unchecked
|
||||
print(f"Time: {time_unchecked:.4f}s | Throughput: {fps_unchecked:,.0f} samples/sec")
|
||||
|
||||
# --- 4. RESULTS ---
|
||||
|
||||
speedup = fps_unchecked / fps_std
|
||||
print(f"\n🏆 WINNER: {'UncheckedIO' if speedup > 1 else 'Standard'}")
|
||||
print(f"SPEEDUP FACTOR: {speedup:.2f}x")
|
||||
print(f"Standard Latency per Batch: {(time_std/NUM_BATCHES)*1000:.2f}ms")
|
||||
print(f"Unchecked Latency per Batch: {(time_unchecked/NUM_BATCHES)*1000:.2f}ms")
|
||||
@@ -1,3 +1,4 @@
|
||||
batch_size: 524288
|
||||
connection_string: postgresql://postgres:mysecretpassword@localhost:5433/postgres
|
||||
query: COPY (SELECT id::BIGINT, uuid::TEXT, username::TEXT, score::REAL, is_active::BOOLEAN,
|
||||
last_login::TIMESTAMP, notes::TEXT, course_id::INT, start_date::DATE, rating::FLOAT8
|
||||
|
||||
0
docs/adaptive_config.md
Normal file
0
docs/adaptive_config.md
Normal file
79
docs/build-cleanup-chrono.md
Normal file
79
docs/build-cleanup-chrono.md
Normal file
@@ -0,0 +1,79 @@
|
||||
Technical Task: Build Fixes & Profiler Cleanup
|
||||
Date: 2025-11-21 Goal: Fix PyPI build failures and remove heavy profiling code from production wheels.
|
||||
|
||||
1. Fix PyPI Build Failure (chrono-tz)
|
||||
Issue: The build fails on Linux/Mac/Windows CI because chrono-tz requires complex C-bindings that are often missing in standard Python environments. Analysis: UncheckedIO returns Apache Arrow RecordBatch objects (where timestamps are raw i64). We do not need pyo3 to convert Rust chrono types directly to Python datetime objects. The chrono-tz feature in pyo3 is dead weight causing build failures.
|
||||
|
||||
Action: Update Cargo.toml to remove the conflicting feature.
|
||||
|
||||
Ini, TOML
|
||||
|
||||
# Cargo.toml
|
||||
|
||||
# OLD
|
||||
# pyo3 = { version = "0.27.1", features = ["extension-module", "chrono-tz"] }
|
||||
|
||||
# NEW (Remove "chrono-tz")
|
||||
pyo3 = { version = "0.27.1", features = ["extension-module"] }
|
||||
2. Productionize the Profiler (Remove Tracy)
|
||||
Issue: The tracy-client and tracing-tracy crates add significant binary size and runtime overhead. They should not be present in the production library installed by users. Solution: Use Rust Feature Flags to make profiling "opt-in" for development only.
|
||||
|
||||
Step A: Update Cargo.toml Define a profiling feature and make the heavy dependencies optional.
|
||||
|
||||
Ini, TOML
|
||||
|
||||
[features]
|
||||
default = []
|
||||
# The "profiling" feature enables these two optional crates
|
||||
profiling = ["dep:tracing-tracy", "dep:tracy-client"]
|
||||
|
||||
[dependencies]
|
||||
# Keep this (Lightweight facade)
|
||||
tracing = "0.1.41"
|
||||
|
||||
# Make these OPTIONAL (Heavy implementation)
|
||||
tracing-tracy = { version = "0.11.2", optional = true }
|
||||
tracy-client = { version = "0.17.6", optional = true }
|
||||
|
||||
# ... other dependencies ...
|
||||
Step B: Update src/lib.rs Wrap the profiler initialization logic so it only compiles when the flag is active.
|
||||
|
||||
Rust
|
||||
|
||||
// src/lib.rs
|
||||
|
||||
#[pymodule]
|
||||
fn unchecked_io(_py: Python<'_>, m: &Bound<'_, PyModule>) -> PyResult<()> {
|
||||
|
||||
// --- START PROFILER LOGIC ---
|
||||
// This block now ONLY compiles if you run with `--features profiling`
|
||||
#[cfg(feature = "profiling")]
|
||||
{
|
||||
use tracing_subscriber::layer::SubscriberExt;
|
||||
use tracing_subscriber::util::SubscriberInitExt;
|
||||
|
||||
// Initialize the Tracy layer
|
||||
let _ = tracing_subscriber::registry()
|
||||
.with(tracing_tracy::TracyLayer::default())
|
||||
.try_init();
|
||||
|
||||
println!("--- UncheckedIO: Profiling Enabled (Tracy) ---");
|
||||
}
|
||||
// --- END PROFILER LOGIC ---
|
||||
|
||||
m.add_function(wrap_pyfunction!(load_data_from_config, m)?)?;
|
||||
m.add_function(wrap_pyfunction!(to_pandas_dataframe, m)?)?;
|
||||
m.add_function(wrap_pyfunction!(profile_data, m)?)?;
|
||||
Ok(())
|
||||
}
|
||||
3. How to Build
|
||||
For Development (With Profiler): Use this when you want to see the tracks in the Tracy GUI.
|
||||
|
||||
Bash
|
||||
|
||||
maturin develop --features profiling
|
||||
For Production / PyPI (Zero Overhead): This builds the "clean" version for users. The profiler code is stripped out entirely.
|
||||
|
||||
Bash
|
||||
|
||||
maturin publish
|
||||
22
docs/changelog.md
Normal file
22
docs/changelog.md
Normal file
@@ -0,0 +1,22 @@
|
||||
Changelog
|
||||
[Unreleased] - Sprint 4 Optimization
|
||||
🚀 Performance
|
||||
Speedup: Achieved 1.46x speedup over ConnectorX (19s vs 28s for 20M rows) on local benchmarks.
|
||||
Worker Pool: Implemented a thread-per-core architecture using async_channel to decouple task generation from execution.
|
||||
Auto-Tuning: Added blast_radius=0 support to dynamically calculate optimal partition sizes based on available CPU cores.
|
||||
🛠️ Fixes & Stability
|
||||
Self-Healing: Schema now defaults to nullable=true. If a partition fails (network error, bad data), the worker logs the error and returns a batch of NULLs instead of crashing the entire process.
|
||||
Dependency Resolution: Fixed build errors with pyo3-arrow 0.15.0 by migrating to Arro3RecordBatch::from and into_pyobject.
|
||||
Profiling Safety: tracing and tracy-client are now optional dependencies behind the profiling feature flag. This prevents overhead and crashes in production environments like Google Colab.
|
||||
⚠️ Known Issues / Constraints
|
||||
Severe Partition Switching Overhead: Increasing partition counts causes non-linear runtime growth.
|
||||
Observation: 160 partitions took ~50s longer than 16 partitions.
|
||||
Analysis: Each worker thread incurs a ~5-second setup penalty when switching to a new partition (establishing a new COPY stream).
|
||||
Implication: While CPU-bound logic scales perfectly, the COPY protocol handshake or Docker networking layer is forcing massive latency spikes per task.
|
||||
Recommendation: Avoid small partitions on high-latency networks (like Docker-on-Windows). Stick to num_cpus partitions until this handshake cost is optimized.
|
||||
I/O Bound: Benchmarks with 8 vs 16 workers showed similar performance, indicating the system is currently Network/IO latency bound, not CPU bound.
|
||||
📦 Build Instructions
|
||||
Production: maturin build --release (Default)
|
||||
Debug/Profile: maturin build --release --features profiling
|
||||
|
||||
|
||||
206
docs/commercialization-strategy.md
Normal file
206
docs/commercialization-strategy.md
Normal file
@@ -0,0 +1,206 @@
|
||||
# UncheckedIO — Commercialization Strategy
|
||||
**Session date:** 2026-04-20
|
||||
**Status:** Active planning
|
||||
|
||||
---
|
||||
|
||||
## Product Summary
|
||||
|
||||
Two components with different commercial treatment:
|
||||
|
||||
| Component | Language | License | Status |
|
||||
|---|---|---|---|
|
||||
| PostgreSQL → Arrow pipeline | Rust | Apache 2.0 (open source) | On PyPI, working |
|
||||
| GPU zero-copy ingestion pipeline | Rust + CUDA | Commercial | Proof-of-concept, needs Phase 1 fixes |
|
||||
|
||||
The GPU pipeline is the commercial product. The Postgres piece is the credibility signal and open-source loss leader.
|
||||
|
||||
---
|
||||
|
||||
## The Problem Being Solved
|
||||
|
||||
GPU starvation — LLM providers run 30-50% actual GPU utilization on inference. The bottleneck is the data ingestion pipeline: CPU tokenizes, creates a Python tensor in pageable memory, `.to("cuda")` performs two OS-mediated copies (pageable → kernel buffer → GPU).
|
||||
|
||||
UncheckedIO replaces this with: Rust parallel tokenization → pinned memory (DMA-accessible) → single direct DMA → DLPack zero-copy handover to PyTorch.
|
||||
|
||||
**Current benchmark:** 4.7x faster than standard HuggingFace pipeline on RTX 3060.
|
||||
**Pending:** H100 SXM benchmark (Wednesday, ~$2 on Lambda Labs). This number changes the pitch.
|
||||
|
||||
---
|
||||
|
||||
## Business Structure
|
||||
|
||||
**Entity:** New Mexico single-member LLC
|
||||
- NM does not require member names in Articles of Organization
|
||||
- No annual report requirement — less administrative overhead than most states
|
||||
- LLC owns all IP and signs all contracts
|
||||
- Registered agent handles public address
|
||||
- Form this before any contract is signed
|
||||
- Note: if the IP is eventually sold, the buyer's anonymity is their own concern — the LLC structure only needs to protect the original author
|
||||
|
||||
**Operating model:** Anonymous principal — LLC is the public face. The identity of the technical author is not disclosed to buyers or the public. Any investor or partner arrangement must include explicit confidentiality provisions.
|
||||
|
||||
---
|
||||
|
||||
## Go-to-Market Strategy
|
||||
|
||||
### Phase 1: Licensing (current goal)
|
||||
License the GPU pipeline commercially. Annual per-cluster licenses. Open-source Postgres piece stays Apache 2.0 as a credibility and discovery channel.
|
||||
|
||||
**Why licensing before IP sale:**
|
||||
A library generating $50K+/month in licenses is a business acquisition, not a proof-of-concept. The IP sale price increases significantly and the negotiating position improves. Licensing revenue demonstrates market validation.
|
||||
|
||||
### Phase 2: IP Sale (after traction)
|
||||
Once licensing revenue establishes market proof, approach GPU cloud providers and LLM infrastructure companies for outright IP purchase. Source code + short-term handoff contract. Clean exit.
|
||||
|
||||
---
|
||||
|
||||
## Pricing
|
||||
|
||||
The ROI math for a GPU provider: 100 H100s at $175K/year each = $17.5M in GPU costs. A 10% utilization improvement = $1.75M saved. A $50K license is a 35x ROI.
|
||||
|
||||
**Do not anchor at $50K. That is the floor for small deployments.**
|
||||
|
||||
| Deployment size | Annual license | One-time source sale |
|
||||
|---|---|---|
|
||||
| Small (50–200 GPUs) | $25K–$75K | $100K–$200K |
|
||||
| Medium (200–1,000 GPUs) | $75K–$200K | $300K–$600K |
|
||||
| Large (1,000+ GPUs) | $200K–$500K | $750K–$1.5M |
|
||||
|
||||
**First client:** Accept below market rate in exchange for a private reference and permission to say "deployed at [company]" to the next prospect. That reference is worth more than the price difference.
|
||||
|
||||
---
|
||||
|
||||
## Target Buyers
|
||||
|
||||
**Primary (write the check):**
|
||||
GPU cloud providers and LLM inference companies feel GPU starvation directly in their margins. They sell GPU utilization as a product. Any measurable improvement is direct revenue.
|
||||
|
||||
- CoreWeave, Lambda Labs — GPU cloud, sell utilization by the hour
|
||||
- Together AI, Fireworks AI, Modal — LLM inference providers
|
||||
- Anyscale — Ray-based ML infrastructure
|
||||
|
||||
**Secondary (may buy, slower process):**
|
||||
- Hugging Face — owns the `tokenizers` Rust library this code uses, runs TGI serving stack. More plausible buyer than PyTorch/Meta.
|
||||
- Databricks (MosaicML) — LLM training infrastructure
|
||||
|
||||
**Do not pitch:**
|
||||
- PyTorch/Meta — open source project, they'd implement it themselves
|
||||
- Hyperscalers (AWS, GCP, Azure) — too slow, too much process
|
||||
|
||||
**NVIDIA contact:**
|
||||
Wait until H100 benchmark is complete. Send the Nsight trace + one paragraph. They may be interested as an ecosystem/marketing play ("runs X% faster on H100"). Do not approach before having the measurement.
|
||||
|
||||
---
|
||||
|
||||
## B2B Sales Process
|
||||
|
||||
### First contact (Zoom call)
|
||||
The first call is discovery, not pitch. Goal: understand their specific infrastructure before presenting a solution.
|
||||
|
||||
**Ask first:**
|
||||
- What GPU hardware? (H100, A100, MI300X — how many?)
|
||||
- What serving stack? (vLLM, TGI, custom?)
|
||||
- Current GPU utilization rate, if known?
|
||||
- What caught their attention?
|
||||
|
||||
**Then:** "Based on what you're describing, here's specifically where this helps you..."
|
||||
|
||||
**On pricing:** Don't volunteer it first. If pressed: *"For a deployment your size I'd expect $50K–$150K annually. I want to understand your infrastructure better before giving you a specific number."*
|
||||
|
||||
**End of first call:** Set up follow-up with their infrastructure lead. Do not sign anything on the first call.
|
||||
|
||||
### NDA
|
||||
Get an NDA signed before showing the GPU source code on any call. The open-source Postgres piece is fine to show. The GPU pipeline (`src/llm.rs`, `gpu-zero` branch) is not.
|
||||
|
||||
### Solo developer objection
|
||||
Enterprise buyers will ask: *"What happens if you're unavailable?"*
|
||||
|
||||
**Answers:**
|
||||
1. Source code sale model answers it directly — they own the code
|
||||
2. Technical partner (friend's software business) provides continuity
|
||||
3. Source code escrow clause: code held by third party, released if unavailable for 90 days — this is standard and removes most of the objection
|
||||
|
||||
---
|
||||
|
||||
## Maintenance and Support
|
||||
|
||||
**Reality for 1–2 enterprise clients:**
|
||||
|
||||
| Phase | Time commitment |
|
||||
|---|---|
|
||||
| Active development (first 3–6 months) | 20–40 hrs/week |
|
||||
| Steady state | 5–10 hrs/week per client |
|
||||
| Major CUDA/PyTorch version update | 2–4 days, happens 2–3x/year |
|
||||
| Production incident | Drop everything, 1–3 days |
|
||||
|
||||
**Scope the support in the contract:** Cover current PyTorch LTS + one prior release. Anything newer is a paid upgrade. This prevents unbounded compatibility obligations.
|
||||
|
||||
One person can support 2–3 enterprise clients at steady state.
|
||||
|
||||
---
|
||||
|
||||
## Outsourcing
|
||||
|
||||
A friend with a software business can handle work that is not the core differentiator:
|
||||
- Support ticket triage and first-response
|
||||
- Documentation and integration guides
|
||||
- Testing infrastructure (benchmark runs against new PyTorch/CUDA versions)
|
||||
- Non-GPU features (Postgres pipeline, Python API polish)
|
||||
|
||||
**Do not outsource:** The GPU pipeline code, benchmarking methodology, CUDA/ROCm work — these are the differentiator.
|
||||
|
||||
**Rate:** Pass 60–70% of contract value for defined scope. On a $100K contract where the friend does 30% of the work: $18K–$21K to them. You take the larger share because you carry the business development risk and hold the client relationship.
|
||||
|
||||
**Critical:** Any work they do on the core library must be IP-assigned to the LLC in writing before the first engagement.
|
||||
|
||||
---
|
||||
|
||||
## Multi-Language Strategy (Rust + C++)
|
||||
|
||||
**Why both:**
|
||||
- Rust is the development language — memory safety makes GPU pipeline bugs easier to catch
|
||||
- C++ is the enterprise deliverable — the entire ML infrastructure ecosystem (cuDNN, TensorRT, vLLM C++ backend) is C++
|
||||
- Having both is a differentiator from vibe-coded projects; it signals deep systems understanding
|
||||
|
||||
**How:**
|
||||
- WhetstoneAI (separate project) can transpile language primitives automatically
|
||||
- Library dependencies require manual substitution (see below)
|
||||
- Rust is source of truth; C++ is a generated + manually doctored artifact per release
|
||||
|
||||
**Library substitutions for C++ version:**
|
||||
|
||||
| Rust dependency | C++ equivalent | Complexity |
|
||||
|---|---|---|
|
||||
| `cudarc` | Direct CUDA C++ calls (actually simpler) | Low |
|
||||
| `rayon` | OpenMP `#pragma omp parallel for` | Low |
|
||||
| `tokenizers` | SentencePiece (LLaMA) + tiktoken (GPT) | Medium — different API |
|
||||
| `pyo3` | pybind11 | Medium — parallel structure |
|
||||
| `tokio` | Not needed for GPU-only C++ version | N/A |
|
||||
|
||||
`parser.rs` (Postgres pipeline) is not worth transpiling — stays Rust only.
|
||||
|
||||
**Estimated work for C++ `llm.rs`:** 2–3 days after Whetstone handles primitives.
|
||||
|
||||
---
|
||||
|
||||
## NM State Funding Assessment
|
||||
|
||||
**Worth pursuing:**
|
||||
- **NM Small Business Assistance Program (NMSBA)** — pairs small businesses with Sandia/LANL. The value is free access to HPC hardware (H100/A100 clusters), not the money. Solves the enterprise benchmark problem. Apply for this.
|
||||
|
||||
**Not worth pursuing (for now):**
|
||||
- NMEDD money grants, angel tax credits — require registered business, matching funds, 3–6 month timelines. Time cost doesn't pencil out against a direct licensing deal.
|
||||
|
||||
**Key point:** Most state grant programs want an ongoing business, not a one-time IP sale. If the goal shifts toward a licensing business, revisit.
|
||||
|
||||
---
|
||||
|
||||
## Immediate Priorities
|
||||
|
||||
1. Fix Phase 1 correctness issues (see ROADMAP.md) — nothing else matters until these are done
|
||||
2. Wednesday: H100 benchmark on Lambda Labs (~$2), run under `nsys profile`
|
||||
3. Thursday/Friday: Incorporate H100 numbers into talk slides
|
||||
4. Talk: 10-minute presentation to IT/Cybersecurity group (see docs/talk-10min.md)
|
||||
5. Form Wyoming LLC before any serious sales conversation
|
||||
6. After first client: evaluate C++ version via Whetstone transpilation
|
||||
105
docs/marketing-business.md
Normal file
105
docs/marketing-business.md
Normal file
@@ -0,0 +1,105 @@
|
||||
# UncheckedIO: Business Strategy & Technical Roadmap
|
||||
|
||||
**Date:** 2025-11-21
|
||||
**Status:** Strategic Planning
|
||||
**Core Philosophy:** "The Configured Opinion" — We trade flexibility for raw, unchecked speed.
|
||||
|
||||
---
|
||||
|
||||
## 1. Target Audience & Value Propositions
|
||||
|
||||
We are not selling "a faster Pandas." We are selling **infrastructure efficiency**.
|
||||
|
||||
### Primary Target: Data Platform & MLOps Engineers
|
||||
* **Who they are:** The people managing Kubernetes clusters, Airflow DAGs, and AWS Fargate costs. They care about stability and resource usage.
|
||||
* **The Pitch:** "A Zero-Copy Transport Layer for Postgres."
|
||||
* **Key Selling Points:**
|
||||
* **Cut Compute Costs:** Reduces CPU starvation by saturating network bandwidth. 10x faster serialization means 10x less time paying for vCPUs.
|
||||
* **Memory Safety:** Written in Rust, ensuring no segfaults in your Python pipeline despite the speed.
|
||||
* **Predictable Throughput:** Unlike generic drivers that choke on large datasets, `UncheckedIO` uses partitioned streaming to maintain constant memory usage regardless of data size.
|
||||
|
||||
### Secondary Target: High-Frequency Trading (HFT) / Quant Developers
|
||||
* **Who they are:** Developers who need to backtest models on massive tick-data history stored in Postgres.
|
||||
* **The Pitch:** "The fastest way to get data from Disk to Arrow Memory."
|
||||
* **Key Selling Points:**
|
||||
* **Zero-Copy to Polars:** Output directly to Apache Arrow memory layouts, ready for immediate analysis or SIMD vectorization.
|
||||
* **Bypass the overhead:** Skips the "Safety Checks" (schema validation) that cost milliseconds per query. You promise the data is clean; we promise to load it instantly.
|
||||
|
||||
### Tertiary Target: Data Scientists (The Users, Not Buyers)
|
||||
* **Who they are:** End users writing Jupyter Notebooks.
|
||||
* **The Pitch:** "Stop waiting for `read_sql`."
|
||||
* **Key Selling Points:**
|
||||
* **Installation Simplicity:** `pip install unchecked_io`. No complex C dependencies or drivers to configure.
|
||||
* **Drop-in Speed:** Replace a 5-minute load time with a 30-second load time.
|
||||
|
||||
---
|
||||
|
||||
## 2. Optimization Strategy: The "Win" Plan
|
||||
|
||||
We need to stop fighting "theoretical" bottlenecks and fix the "physical" ones.
|
||||
|
||||
### 🔴 Focus On (High ROI)
|
||||
1. **Infrastructure-Aware Tuning (The "Blast Radius"):**
|
||||
* *Why:* The "astronomical" delay was caused by thrashing.
|
||||
* *Action:* Implement the `num_cpus` detection to auto-scale the connection pool and partition count. Do not let the user guess wrong.
|
||||
2. **Memory Management (The "String Killer"):**
|
||||
* *Why:* 15M heap allocations for strings creates allocator contention.
|
||||
* *Action:* Refactor the parser to use a reusable buffer or "Zero-Copy" string slicing from the raw stream before writing to Arrow.
|
||||
3. **Database Interaction (The Index):**
|
||||
* *Why:* As discovered, a missing index turns a parallel fetch into a DDoS attack.
|
||||
* *Action:* Documentation MUST strictly warn users: *"If you use `blast_radius`, you MUST index your partition column."*
|
||||
|
||||
### 🛑 Leave Alone (Diminishing Returns)
|
||||
1. **Complex "Safe Mode" Logic:**
|
||||
* *Why:* We are `UncheckedIO`. If we add too many safety checks/fallbacks for dirty data, we become just another slow connector (like ConnectorX).
|
||||
* *Strategy:* Let it fail fast. If the data violates the config schema, panic. That is the contract.
|
||||
2. **Micro-Optimizing the TCP Handshake:**
|
||||
* *Why:* Connection pooling already solves 95% of this. Saving another 1ms here is irrelevant compared to the 500ms of Query Planning time.
|
||||
|
||||
---
|
||||
|
||||
## 3. Configuration Matrix
|
||||
|
||||
To be a "System-Level" tool, we need to expose the knobs that Platform Engineers expect.
|
||||
|
||||
### Currently Configurable (The "MVP" Set)
|
||||
These features are already present in the code or `config.yaml`:
|
||||
1. **`connection_string`:** Standard Postgres URI.
|
||||
2. **`query`:** The exact `COPY ... TO STDOUT (FORMAT binary)` command. Allows total control over the SQL execution plan.
|
||||
3. **`schema`:** Explicit definition of Arrow Types (`Int64`, `Utf8`, etc.). Skips runtime type inference.
|
||||
4. **`blast_radius` (Argument):** Controls the row-count per partition. Passed from Python to Rust.
|
||||
|
||||
### Must-Add Configurations (To Delight Consumers)
|
||||
These options solve specific "Enterprise" pain points:
|
||||
1. **`worker_threads`:**
|
||||
* *Value:* Allows limiting the library to use only 4 cores on a shared 64-core server.
|
||||
* *Implementation:* Pass to `tokio::runtime::Builder::worker_threads()`.
|
||||
2. **`pool_size`:**
|
||||
* *Value:* Overrides the auto-detected defaults. Essential for databases with strict connection limits (e.g., AWS RDS limits).
|
||||
* *Implementation:* Pass to `deadpool_postgres::Pool::builder().max_size()`.
|
||||
3. **`batch_size` (Arrow):**
|
||||
* *Value:* Controls the size of the Arrow RecordBatches. Critical for streaming data into ML models (e.g., "Give me 10k rows at a time").
|
||||
* *Implementation:* Currently aggregated at the end; needs to be exposed in the streaming loop.
|
||||
4. **`danger_mode` (Explicit Toggle):**
|
||||
* *Value:* A flag to strictly enforce "Panic on Error" vs. "Try to Recover."
|
||||
* *Implementation:* As described in Sprint 4.
|
||||
|
||||
---
|
||||
|
||||
## 4. The "Reputation Builder" Roadmap (Future Optimizations)
|
||||
|
||||
These features serve two purposes: Extreme performance for the 1% of users who need it, and "Street Cred" for the author.
|
||||
|
||||
1. **`io_uring` Network Layer:**
|
||||
* *The Tech:* Replace `tokio-postgres` (epoll) with a custom `io_uring` implementation for Linux.
|
||||
* *The Flex:* "True Asynchronous syscall batching." Reduces context switches during massive data transfers.
|
||||
* *Status:* High effort, high reputation, low immediate business ROI (unless on 100Gbps networks).
|
||||
2. **SIMD Parser (AVX-512 / NEON):**
|
||||
* *The Tech:* Use `std::simd` or `portable-simd` to parse 16 integers at once from the binary stream.
|
||||
* *The Flex:* "Saturating the memory bandwidth of modern CPUs."
|
||||
3. **Custom Allocator Support (`mimalloc` / `jemalloc`):**
|
||||
* *The Tech:* Allow users to swap the global allocator via a feature flag.
|
||||
* *The Flex:* "Optimized for high-fragmentation environments." (Note: `mimalloc` is already configured for non-MSVC targets).
|
||||
4. **Kernel-Bypass Networking (DPDK):**
|
||||
* *The Tech:* Bypassing the OS network stack entirely.
|
||||
* *The Flex:* The absolute theoretical limit of data transfer. (Overkill, but legendary).
|
||||
1
docs/refactor-prompt
Normal file
1
docs/refactor-prompt
Normal file
@@ -0,0 +1 @@
|
||||
I need to refactor src/parser.rs to implement a "Worker Pool" pattern that strictly limits active database queries to the number of CPU cores.The Goal:Currently, the code spawns a Tokio task for every partition immediately. This floods the scheduler and causes database thrashing (too many concurrent COPY commands).I want to spawn a fixed number of workers (e.g., 16) that consume tasks from a queue. This ensures that if I have 1,600 partitions, only 16 COPY commands are active on the database at any given time.Requirements:Dependencies: Use async-channel (assume it's in Cargo.toml).Concurrency Limit:let num_workers = num_cpus::get();Set deadpool_postgres max_size to num_workers.The Queue:Create a channel: let (tx, rx) = async_channel::bounded(num_workers * 2); (Backpressure is good).Spawn a separate "Distributor" task that iterates partitions and sends them into tx.The Workers:Spawn exactly num_workers tasks.Each worker loops: while let Ok(task) = rx.recv().await.Inside the loop:Get a connection from the pool.Run copy_out(task.query).Parse the result.Store the resulting RecordBatch.Output:Collect all RecordBatch results from all workers.Flatten and concatenate them using concat_batches.Why: This ensures Postgres never sees a queue of queries. It only sees active workers. The "queue" lives entirely in Rust memory.
|
||||
110
docs/talk-10min.md
Normal file
110
docs/talk-10min.md
Normal file
@@ -0,0 +1,110 @@
|
||||
# UncheckedIO — 10-Minute Technical Talk
|
||||
**Audience:** IT/Cybersecurity group with programmers and ML engineers
|
||||
**Goal:** Plant the seed, demonstrate the mechanism, invite connections to LLM infrastructure teams
|
||||
|
||||
---
|
||||
|
||||
## Slide 1 — Open with the hidden thing (0:00–0:45)
|
||||
|
||||
Lead with an Nsight Systems timeline screenshot. Don't explain it yet.
|
||||
|
||||
> "This is what happens inside your GPU when a language model processes your request.
|
||||
> The orange is waiting. The GPU costs $30,000 and it's waiting for data."
|
||||
|
||||
Let the image do the work. This audience will immediately want to know what they're looking at.
|
||||
|
||||
---
|
||||
|
||||
## Slide 2 — The problem (0:45–2:00)
|
||||
|
||||
**GPU starvation.** Define it in dollar terms.
|
||||
|
||||
- Real LLM providers run at 30–50% actual GPU utilization on inference
|
||||
- An 8-GPU H100 cluster costs ~$20/hr to rent
|
||||
- At 40% utilization, $12/hr is wasted — the GPUs are idle waiting for data
|
||||
- LLM providers run hundreds of clusters
|
||||
- The waste is structural and nobody has fixed the root cause
|
||||
|
||||
---
|
||||
|
||||
## Slides 3–4 — The mechanism (2:00–4:00)
|
||||
|
||||
Two diagrams, side by side.
|
||||
|
||||
**Standard pipeline:**
|
||||
```
|
||||
CPU tokenizes → Python tensor (pageable memory) → .to("cuda") → kernel buffer → GPU
|
||||
^--- two OS-mediated copies
|
||||
```
|
||||
|
||||
**UncheckedIO pipeline:**
|
||||
```
|
||||
Rust tokenizes in parallel → pinned memory (DMA-accessible) → direct DMA → DLPack → GPU
|
||||
^--- one copy, no kernel involvement
|
||||
```
|
||||
|
||||
The second diagram has fewer boxes and fewer arrows. That is the entire pitch.
|
||||
|
||||
Key terms to explain briefly:
|
||||
- **Pinned memory** — memory registered with the OS so the GPU's DMA engine can read it directly
|
||||
- **DMA** — the GPU pulls the data itself, no CPU involvement after the transfer starts
|
||||
- **DLPack** — a standard zero-copy tensor exchange protocol; PyTorch accepts it natively with `torch.from_dlpack()`
|
||||
|
||||
---
|
||||
|
||||
## Slide 5 — The Nsight demo (4:00–7:00)
|
||||
|
||||
Show the recorded Nsight Systems timeline from the H100 benchmark run.
|
||||
|
||||
Both pipelines visible side by side:
|
||||
- Point out the idle gaps (orange) in the standard pipeline
|
||||
- Point out the absence of idle gaps in the UncheckedIO pipeline
|
||||
- Show the DMA transfer lane — this is the "hidden thing" made visible
|
||||
|
||||
This is the centerpiece of the talk. Expect questions here. That's good.
|
||||
|
||||
---
|
||||
|
||||
## Slide 6 — The number (7:00–8:00)
|
||||
|
||||
One slide. No prose.
|
||||
|
||||
| Hardware | Speedup |
|
||||
|----------|---------|
|
||||
| RTX 3060 (local) | 4.7x |
|
||||
| H100 SXM (cloud) | [Wednesday result] |
|
||||
|
||||
---
|
||||
|
||||
## Slide 7 — Market fit (8:00–9:00)
|
||||
|
||||
- This is infrastructure every LLM provider needs
|
||||
- Built as an open-core library: PostgreSQL→Arrow is Apache 2.0 and already on PyPI
|
||||
- GPU ingestion pipeline is the commercial layer
|
||||
- AMD ROCm support (MI300X) is next — providers are actively evaluating MI300X to escape H100 pricing
|
||||
|
||||
Don't pitch deal specifics. This slide plants the seed.
|
||||
|
||||
---
|
||||
|
||||
## Slide 8 — The close (9:00–10:00)
|
||||
|
||||
> "I'm looking for connections to ML infrastructure teams and LLM providers.
|
||||
> If you know someone building at that layer, I'd appreciate an introduction."
|
||||
|
||||
Direct. Not desperate. Leave room for questions.
|
||||
|
||||
---
|
||||
|
||||
## Notes
|
||||
|
||||
**What NOT to do:**
|
||||
- Don't go deep on the Rust implementation — the visual is the hook, not the code
|
||||
- Don't mention the PostgreSQL→Arrow piece — it's a distraction in this context
|
||||
- Don't volunteer that it's unfinished — if asked, "it's in active development and the core mechanism is proven"
|
||||
|
||||
**For the well-connected engineer (separate talk, two weeks out):**
|
||||
Don't reuse this version. Prep a deeper, tailored pitch after seeing what questions come up here.
|
||||
|
||||
**For the NVIDIA contact:**
|
||||
Wait until after the H100 benchmark. Send the Nsight trace + one paragraph. Leading with a measurement is different from leading with a claim.
|
||||
83
docs/tech_fix.md
Normal file
83
docs/tech_fix.md
Normal file
@@ -0,0 +1,83 @@
|
||||
Here is the technical summary of our debugging session, formatted as a Markdown file for your documentation.
|
||||
|
||||
Markdown
|
||||
|
||||
# Technical Post-Mortem: The "Thrashing" & Performance Optimization Log
|
||||
|
||||
**Date:** 2025-11-21
|
||||
**Project:** UncheckedIO
|
||||
**Topic:** Debugging "Astronomical" Task Switching Delays & Laptop vs. Cloud Discrepancies
|
||||
|
||||
## 1. The Symptoms
|
||||
* **The "Astronomical" Delay:** Increasing partition count (e.g., from 16 to 32) caused disproportionate slowdowns. Switching between tasks took seconds, not nanoseconds.
|
||||
* **The "Tracy" Paradox:** The profiler showed "CPU Data Starved" and frequently crashed the runtime. When it crashed/detached, the library actually ran faster.
|
||||
* **Platform Discrepancy:** The library performed significantly better on Google Colab (Server CPU) than on a high-spec Laptop (Consumer CPU) relative to competitors, despite the laptop having faster single-core speed.
|
||||
|
||||
## 2. Root Cause Analysis
|
||||
|
||||
### A. The Database Villain: "The Missing Index"
|
||||
* **The Issue:** The `setup_db.sql` script created the table but **failed to index the `id` column**.
|
||||
* **The Consequence:** The "Parallel Coordinator" splits the job into partitions (e.g., `WHERE id BETWEEN 0 AND 125000`). Without an index, Postgres performed a **Sequential Scan** (Full Table Scan) for *every single partition*.
|
||||
* **The Impact:** Running 16 parallel workers triggered **16 simultaneous Full Table Scans**. This saturated memory bandwidth instantly, explaining the "CPU Data Starved" metric in Tracy.
|
||||
|
||||
### B. The Concurrency Villain: "Pool Contention & Thrashing"
|
||||
* **The Issue:** The connection pool was hardcoded to `max_size(20)` in `src/parser.rs`.
|
||||
* **The "Alive" Misconception:** We assumed keeping connections "alive" allowed infinite parallelism. However, a connection is a single-lane bridge; it can only transport one query at a time.
|
||||
* **The Bottleneck:** When spawning 800 tasks (partitions) against 20 connections:
|
||||
* **20 Tasks** ran immediately.
|
||||
* **780 Tasks** sat frozen in a RAM queue.
|
||||
* **The "Thrashing":** The CPU spent ~90% of its time managing this queue—saving/restoring task states (Context Switching) and flushing L1/L2 caches—rather than parsing data. This overhead destroyed performance on the laptop (which has smaller caches than the Colab server).
|
||||
|
||||
### C. The Red Herring: `io_uring`
|
||||
* **Hypothesis:** We initially thought switching to `io_uring` (Linux asynchronous I/O) would solve the "waiting."
|
||||
* **Verdict:** **Incorrect.** The bottleneck was not syscall overhead (which `io_uring` solves); it was **Logical Contention** (waiting for a DB connection) and **Physical Latency** (waiting for RAM/Network).
|
||||
|
||||
---
|
||||
|
||||
## 3. The Solutions
|
||||
|
||||
### Fix #1: Create the Index (Database Side)
|
||||
Change the operation from a Sequential Scan (reading 20M rows x 16 times) to an **Index Scan** (jumping directly to the data).
|
||||
|
||||
**File:** `setup_db.sql`
|
||||
```sql
|
||||
-- OLD
|
||||
ANALYZE benchmark_table;
|
||||
|
||||
-- NEW
|
||||
CREATE INDEX idx_benchmark_id ON benchmark_table(id);
|
||||
ANALYZE benchmark_table;
|
||||
Fix #2: Dynamic Pool Sizing (Rust Side)
|
||||
Remove the artificial speed limit. Match the pool size to the hardware capabilities so tasks never wait in a queue.
|
||||
|
||||
File: src/parser.rs
|
||||
|
||||
Rust
|
||||
|
||||
// OLD
|
||||
.max_size(20)
|
||||
|
||||
// NEW
|
||||
use num_cpus; // Add to Cargo.toml
|
||||
let num_threads = num_cpus::get();
|
||||
.max_size(num_threads) // 1:1 mapping of Threads to Connections
|
||||
Fix #3: Auto-Tuned "Blast Radius"
|
||||
Instead of a hardcoded partition size or user guess, calculate the optimal partition count to balance
|
||||
load without causing overhead.
|
||||
|
||||
Strategy: 4 * CPU Cores partitions.
|
||||
|
||||
Why: Enough chunks to allow "Work Stealing" (fast threads take more work), but few enough to prevent
|
||||
connection setup overhead.
|
||||
|
||||
Rust
|
||||
|
||||
pub fn calculate_optimal_blast_radius(min_id: i64, max_id: i64) -> i64 {
|
||||
let total_rows = max_id - min_id;
|
||||
let target_partitions = num_cpus::get() * 4;
|
||||
total_rows / (target_partitions as i64)
|
||||
}
|
||||
4. Conclusion
|
||||
The library's "Schema-on-Read" optimization (trusting the config) was working perfectly.
|
||||
The performance loss was due to infrastructure bottlenecks (DB Index and Connection Pool limits)
|
||||
that forced the CPU to manage queues instead of parsing data.
|
||||
@@ -4,8 +4,33 @@ build-backend = "maturin"
|
||||
|
||||
[project]
|
||||
name = "unchecked-io"
|
||||
version = "0.1.0"
|
||||
# Dynamic versioning allows maturin to read version from Cargo.toml
|
||||
dynamic = ["version"]
|
||||
description = "The world's fastest, most dangerous PostgreSQL-to-Arrow loader."
|
||||
readme = "README.md"
|
||||
requires-python = ">=3.9"
|
||||
license = {text = "BSL-1.1"}
|
||||
authors = [{name = "Billthemaker"}]
|
||||
keywords = ["rust", "postgresql", "arrow", "etl", "fast", "copy", "database"]
|
||||
classifiers = [
|
||||
"Programming Language :: Rust",
|
||||
"Programming Language :: Python :: Implementation :: CPython",
|
||||
"Programming Language :: Python :: Implementation :: PyPy",
|
||||
"Programming Language :: Python :: 3",
|
||||
]
|
||||
"Programming Language :: Python :: 3.9",
|
||||
"Programming Language :: Python :: 3.10",
|
||||
"Programming Language :: Python :: 3.11",
|
||||
"Programming Language :: Python :: 3.12",
|
||||
"Programming Language :: Python :: 3.13",
|
||||
"Topic :: Database",
|
||||
"Topic :: Scientific/Engineering :: Information Analysis",
|
||||
"Operating System :: POSIX :: Linux",
|
||||
"Operating System :: Microsoft :: Windows",
|
||||
"Operating System :: MacOS :: MacOS X",
|
||||
]
|
||||
dependencies = [
|
||||
"arro3-core>=0.1.0",
|
||||
]
|
||||
|
||||
[tool.maturin]
|
||||
features = ["pyo3/extension-module"]
|
||||
109
server.py
Normal file
109
server.py
Normal file
@@ -0,0 +1,109 @@
|
||||
import os
|
||||
import time
|
||||
import torch
|
||||
import uvicorn
|
||||
from contextlib import asynccontextmanager
|
||||
from fastapi import FastAPI, HTTPException
|
||||
from pydantic import BaseModel
|
||||
from transformers import AutoModelForCausalLM, AutoTokenizer
|
||||
|
||||
# Import your custom high-speed engine
|
||||
import unchecked_io
|
||||
|
||||
# --- CONFIGURATION ---
|
||||
# Switching to GPT-2 for stability. It is small, fast, and has a standard tokenizer
|
||||
# that won't crash older/newer Rust crate versions.
|
||||
MODEL_ID = "gpt2"
|
||||
TOKENIZER_FILE = "tokenizer.json"
|
||||
|
||||
# Global state container
|
||||
model_state = {
|
||||
"model": None,
|
||||
"rust_engine": None,
|
||||
"py_tokenizer": None
|
||||
}
|
||||
|
||||
@asynccontextmanager
|
||||
async def lifespan(app: FastAPI):
|
||||
print(f"\n--- 🚀 UNCHECKED SERVER STARTUP ---")
|
||||
|
||||
# 1. Fetch Tokenizer (Python side)
|
||||
print(f"1. Fetching Tokenizer Config from {MODEL_ID}...")
|
||||
tokenizer = AutoTokenizer.from_pretrained(MODEL_ID, use_fast=True)
|
||||
tokenizer.save_pretrained(".")
|
||||
model_state["py_tokenizer"] = tokenizer
|
||||
|
||||
# 2. Initialize UncheckedIO (Rust + CUDA)
|
||||
print(f"2. Initializing UncheckedIO (Fuel Injector)...")
|
||||
try:
|
||||
model_state["rust_engine"] = unchecked_io.TokenizerEngine(TOKENIZER_FILE)
|
||||
print(" ✅ Rust Engine Ready (Zero-Copy Pipeline Active)")
|
||||
except Exception as e:
|
||||
print(f" ❌ Failed to load Rust Engine: {e}")
|
||||
raise e
|
||||
|
||||
# 3. Load Model (PyTorch)
|
||||
print(f"3. Loading Model Weights...")
|
||||
try:
|
||||
model_state["model"] = AutoModelForCausalLM.from_pretrained(
|
||||
MODEL_ID,
|
||||
device_map="cuda",
|
||||
torch_dtype=torch.float16
|
||||
)
|
||||
print(f" ✅ Model Loaded on {model_state['model'].device}")
|
||||
except Exception as e:
|
||||
print(f" ❌ Failed to load model: {e}")
|
||||
raise e
|
||||
|
||||
print("--- SERVER READY ---\n")
|
||||
yield
|
||||
print("\n--- SERVER SHUTDOWN ---")
|
||||
# Resources are cleaned up here
|
||||
|
||||
app = FastAPI(title="UncheckedIO High-Speed Server", lifespan=lifespan)
|
||||
|
||||
class GenerateRequest(BaseModel):
|
||||
prompt: str
|
||||
max_tokens: int = 50
|
||||
temperature: float = 0.7
|
||||
|
||||
@app.post("/generate")
|
||||
async def generate_text(req: GenerateRequest):
|
||||
model = model_state["model"]
|
||||
engine = model_state["rust_engine"]
|
||||
tokenizer = model_state["py_tokenizer"]
|
||||
|
||||
if not model or not engine:
|
||||
raise HTTPException(status_code=503, detail="Server not ready")
|
||||
|
||||
try:
|
||||
# --- PHASE 1: UNCHECKED INGESTION (Rust) ---
|
||||
# 4.8x Faster than standard .to("cuda")
|
||||
input_capsule = engine.encode_batch([req.prompt])
|
||||
input_ids = torch.from_dlpack(input_capsule)
|
||||
|
||||
# --- PHASE 2: INFERENCE (PyTorch) ---
|
||||
with torch.no_grad():
|
||||
output_ids = model.generate(
|
||||
input_ids,
|
||||
max_new_tokens=req.max_tokens,
|
||||
temperature=req.temperature,
|
||||
do_sample=True,
|
||||
pad_token_id=tokenizer.eos_token_id
|
||||
)
|
||||
|
||||
# --- PHASE 3: DECODING ---
|
||||
generated_text = tokenizer.decode(output_ids[0], skip_special_tokens=True)
|
||||
|
||||
return {
|
||||
"response": generated_text,
|
||||
"backend": "UncheckedIO + PyTorch",
|
||||
"status": "success"
|
||||
}
|
||||
|
||||
except Exception as e:
|
||||
print(f"Error during generation: {e}")
|
||||
raise HTTPException(status_code=500, detail=str(e))
|
||||
|
||||
if __name__ == "__main__":
|
||||
uvicorn.run(app, host="0.0.0.0", port=8000)
|
||||
@@ -39,5 +39,6 @@ SELECT
|
||||
FROM generate_series(1, 20000000) s(i);
|
||||
|
||||
-- 4. Analyze the table for better query planning (good practice)
|
||||
CREATE INDEX idx_benchmark_id ON benchmark_table(id);
|
||||
ANALYZE benchmark_table;
|
||||
|
||||
|
||||
@@ -1,11 +1,8 @@
|
||||
// --- External Crates ---
|
||||
// src/config.rs
|
||||
use serde::Deserialize;
|
||||
use std::collections::HashMap;
|
||||
use anyhow::{Context, Result, anyhow};
|
||||
use anyhow::{Context, Result, anyhow}; // Added anyhow macro just in case
|
||||
use std::fmt::{self, Display};
|
||||
|
||||
// --- 1. CUSTOM ERROR DEFINITION ---
|
||||
|
||||
#[derive(Debug)]
|
||||
pub struct ConfigError(String);
|
||||
|
||||
@@ -16,10 +13,6 @@ impl Display for ConfigError {
|
||||
}
|
||||
impl std::error::Error for ConfigError {}
|
||||
|
||||
|
||||
// --- 2. CONFIGURATION STRUCTS (The "Configured Opinion") ---
|
||||
// We make these 'pub' (public) so src/lib.rs can use them.
|
||||
|
||||
#[derive(Debug, Deserialize, Clone)]
|
||||
pub struct ColumnConfig {
|
||||
pub column_name: String,
|
||||
@@ -29,12 +22,12 @@ pub struct ColumnConfig {
|
||||
#[derive(Debug, Deserialize)]
|
||||
pub struct ConnectorConfig {
|
||||
pub connection_string: String,
|
||||
pub query: String, // <--- ADD THIS FIELD
|
||||
pub query: String,
|
||||
pub schema: Vec<ColumnConfig>,
|
||||
// NEW: Optional Batch Size configuration
|
||||
pub batch_size: Option<usize>,
|
||||
}
|
||||
|
||||
// --- 3. CONFIG LOADING FUNCTION ---
|
||||
// This is also 'pub' so src/lib.rs can call it.
|
||||
pub fn load_and_validate_config(path: &str) -> Result<ConnectorConfig> {
|
||||
let file_content = std::fs::read_to_string(path)
|
||||
.context(format!("Failed to read config file at path: {}", path))?;
|
||||
@@ -48,7 +41,6 @@ pub fn load_and_validate_config(path: &str) -> Result<ConnectorConfig> {
|
||||
if config.query.is_empty() {
|
||||
return Err(ConfigError("Query cannot be empty.".to_string())).map_err(anyhow::Error::from)?;
|
||||
}
|
||||
// FIX: Make the COPY check more flexible (allows newlines)
|
||||
let uppercase_query = config.query.trim().to_uppercase();
|
||||
if !uppercase_query.starts_with("COPY") || !uppercase_query.contains("TO STDOUT") || !uppercase_query.contains("FORMAT BINARY") {
|
||||
return Err(ConfigError("Query must be a 'COPY ... TO STDOUT (FORMAT binary)' command.".to_string())).map_err(anyhow::Error::from)?;
|
||||
|
||||
97
src/lib.rs
97
src/lib.rs
@@ -1,17 +1,18 @@
|
||||
// --- Declare our new modules ---
|
||||
mod config;
|
||||
mod parser;
|
||||
|
||||
// --- External Crates ---
|
||||
#[cfg(feature = "gpu")]
|
||||
mod llm;
|
||||
|
||||
use pyo3::prelude::*;
|
||||
#[cfg(feature = "gpu")]
|
||||
use std::sync::Arc;
|
||||
use pyo3::exceptions::PyValueError;
|
||||
use tokio;
|
||||
use pyo3::types::{PyModule, PyAny};
|
||||
use pyo3::Bound;
|
||||
use pyo3_arrow::PyRecordBatch;
|
||||
use pyo3_arrow::export::Arro3RecordBatch;
|
||||
|
||||
// Use the high-performance mimalloc for better multi-threaded memory allocation.
|
||||
// We conditionally compile it to avoid issues on MSVC targets.
|
||||
#[cfg(not(target_env = "msvc"))]
|
||||
use mimalloc;
|
||||
|
||||
@@ -19,35 +20,27 @@ use mimalloc;
|
||||
#[global_allocator]
|
||||
static GLOBAL: mimalloc::MiMalloc = mimalloc::MiMalloc;
|
||||
|
||||
// --- Internal Crates ---
|
||||
use crate::config::{load_and_validate_config, ConnectorConfig};
|
||||
use crate::parser::run_db_logic;
|
||||
use crate::parser::{run_db_logic, run_profiler_logic};
|
||||
|
||||
|
||||
// --- THE PYTHON-CALLABLE ENTRY POINT ---
|
||||
#[cfg(feature = "profiling")]
|
||||
use tracing_subscriber::layer::SubscriberExt;
|
||||
#[cfg(feature = "profiling")]
|
||||
use tracing_subscriber::util::SubscriberInitExt;
|
||||
|
||||
#[pyfunction]
|
||||
#[pyo3(signature = (config_path, blast_radius=312500))]
|
||||
#[allow(unsafe_code)]
|
||||
#[allow(unsafe_op_in_unsafe_fn)]
|
||||
#[allow(rust_2024_compatibility)]
|
||||
#[pyo3(signature = (config_path, blast_radius=0))]
|
||||
fn load_data_from_config<'py>(
|
||||
py: Python<'py>,
|
||||
config_path: String,
|
||||
blast_radius: i64,
|
||||
) -> PyResult<Bound<'py, PyAny>> {
|
||||
|
||||
// --- Phase 1: Load and Validate Configuration ---
|
||||
let config: ConnectorConfig = match load_and_validate_config(&config_path) {
|
||||
Ok(c) => c,
|
||||
Err(e) => return Err(PyValueError::new_err(format!("Configuration Error: {:?}", e))),
|
||||
};
|
||||
|
||||
println!("--- UncheckedIO: Schema Accepted ---");
|
||||
println!("Database: {}", config.connection_string);
|
||||
println!("Columns (in order): {:?}", config.schema.iter().map(|c| &c.column_name).collect::<Vec<_>>());
|
||||
|
||||
// --- Phase 2: Run Core Logic ---
|
||||
let record_batch = py.allow_threads(|| {
|
||||
tokio::runtime::Builder::new_multi_thread()
|
||||
.enable_all()
|
||||
@@ -58,16 +51,74 @@ fn load_data_from_config<'py>(
|
||||
})
|
||||
}).map_err(|e| PyValueError::new_err(format!("Database/Runtime Error: {:?}", e)))?;
|
||||
|
||||
// --- Phase 3: Return Data to Python ---
|
||||
let py_record_batch = PyRecordBatch::new(record_batch);
|
||||
py_record_batch.into_pyarrow(py)
|
||||
let py_record_batch = Arro3RecordBatch::from(record_batch);
|
||||
py_record_batch.into_pyobject(py)
|
||||
}
|
||||
|
||||
#[pyfunction]
|
||||
#[pyo3(signature = (arrow_table))]
|
||||
fn to_pandas_dataframe<'py>(_py: Python<'py>, arrow_table: Bound<'py, PyAny>) -> PyResult<Bound<'py, PyAny>> {
|
||||
arrow_table.call_method0("to_pandas")
|
||||
}
|
||||
|
||||
// --- PYTHON MODULE EXPORT ---
|
||||
#[pyfunction]
|
||||
#[pyo3(signature = (config_path))]
|
||||
fn profile_data(config_path: String) -> PyResult<String> {
|
||||
let output = std::thread::spawn(move || {
|
||||
tokio::runtime::Builder::new_current_thread()
|
||||
.enable_all()
|
||||
.build()
|
||||
.unwrap()
|
||||
.block_on(async {
|
||||
run_profiler_logic(&config_path).await
|
||||
})
|
||||
}).join().unwrap().map_err(|e| PyValueError::new_err(format!("Profiling Error: {:?}", e)))?;
|
||||
|
||||
Ok(output)
|
||||
}
|
||||
|
||||
#[cfg(feature = "gpu")]
|
||||
#[pyclass(name = "TokenizerEngine")]
|
||||
struct TokenizerEngine {
|
||||
// 1.3: Arc allows cheap cloning; encode_batch takes &self so concurrent calls are safe
|
||||
inner: Arc<llm::PinnedBatcher>,
|
||||
}
|
||||
|
||||
#[cfg(feature = "gpu")]
|
||||
#[pymethods]
|
||||
impl TokenizerEngine {
|
||||
#[new]
|
||||
fn new(model_path: String) -> PyResult<Self> {
|
||||
Ok(TokenizerEngine {
|
||||
inner: Arc::new(llm::PinnedBatcher::new(&model_path)?)
|
||||
})
|
||||
}
|
||||
|
||||
fn encode_batch(&self, py: Python, texts: Vec<String>) -> PyResult<Py<PyAny>> {
|
||||
self.inner.batch_encode_to_gpu(py, texts)
|
||||
}
|
||||
}
|
||||
|
||||
#[pymodule]
|
||||
fn unchecked_io(_py: Python<'_>, m: &Bound<'_, PyModule>) -> PyResult<()> {
|
||||
|
||||
#[cfg(feature = "profiling")]
|
||||
{
|
||||
let _ = tracing_subscriber::registry()
|
||||
.with(tracing_tracy::TracyLayer::default())
|
||||
.try_init();
|
||||
println!("UncheckedIO: Profiling Mode ENABLED 🚀");
|
||||
}
|
||||
|
||||
m.add_function(wrap_pyfunction!(load_data_from_config, m)?)?;
|
||||
m.add_function(wrap_pyfunction!(to_pandas_dataframe, m)?)?;
|
||||
m.add_function(wrap_pyfunction!(profile_data, m)?)?;
|
||||
|
||||
#[cfg(feature = "gpu")]
|
||||
{
|
||||
m.add_class::<TokenizerEngine>()?;
|
||||
println!("UncheckedIO: GPU Acceleration ENABLED ⚡");
|
||||
}
|
||||
|
||||
Ok(())
|
||||
}
|
||||
197
src/llm.rs
Normal file
197
src/llm.rs
Normal file
@@ -0,0 +1,197 @@
|
||||
// src/llm.rs
|
||||
#![cfg(feature = "gpu")]
|
||||
|
||||
use std::sync::Arc;
|
||||
use tokenizers::Tokenizer;
|
||||
use cudarc::driver::{CudaDevice, DeviceSlice, CudaSlice, DevicePtr};
|
||||
use pyo3::prelude::*;
|
||||
use pyo3::exceptions::PyValueError;
|
||||
use std::ffi::CStr;
|
||||
use rayon::prelude::*;
|
||||
|
||||
use cuda_driver_sys::{
|
||||
cuMemHostRegister_v2,
|
||||
cuMemHostUnregister,
|
||||
cuMemcpyHtoD_v2,
|
||||
cudaError_enum
|
||||
};
|
||||
|
||||
// --- DLPack v0.4 ABI Definitions ---
|
||||
#[repr(C)]
|
||||
struct DLDataType { code: u8, bits: u8, lanes: u16 }
|
||||
#[repr(C)]
|
||||
struct DLDevice { device_type: i32, device_id: i32 }
|
||||
#[repr(C)]
|
||||
struct DLTensor {
|
||||
data: *mut std::ffi::c_void,
|
||||
device: DLDevice,
|
||||
ndim: i32,
|
||||
dtype: DLDataType,
|
||||
shape: *mut i64,
|
||||
strides: *mut i64,
|
||||
byte_offset: u64,
|
||||
}
|
||||
#[repr(C)]
|
||||
struct DLManagedTensor {
|
||||
dl_tensor: DLTensor,
|
||||
manager_ctx: *mut std::ffi::c_void,
|
||||
deleter: Option<unsafe extern "C" fn(*mut DLManagedTensor)>,
|
||||
}
|
||||
|
||||
const CAPSULE_NAME: &CStr = unsafe { CStr::from_bytes_with_nul_unchecked(b"dltensor\0") };
|
||||
|
||||
struct PinnedHostBuffer {
|
||||
data: Vec<i64>,
|
||||
}
|
||||
|
||||
impl PinnedHostBuffer {
|
||||
fn new(capacity: usize) -> PyResult<Self> {
|
||||
let data = vec![0i64; capacity];
|
||||
let ptr = data.as_ptr() as *mut std::ffi::c_void;
|
||||
let bytes = capacity * std::mem::size_of::<i64>();
|
||||
unsafe {
|
||||
let result = cuMemHostRegister_v2(ptr, bytes, 0);
|
||||
if result != cudaError_enum::CUDA_SUCCESS {
|
||||
return Err(PyValueError::new_err(format!("Failed to pin memory: {:?}", result)));
|
||||
}
|
||||
}
|
||||
Ok(Self { data })
|
||||
}
|
||||
}
|
||||
|
||||
impl Drop for PinnedHostBuffer {
|
||||
fn drop(&mut self) {
|
||||
unsafe {
|
||||
let ptr = self.data.as_ptr() as *mut std::ffi::c_void;
|
||||
cuMemHostUnregister(ptr);
|
||||
}
|
||||
}
|
||||
}
|
||||
|
||||
// Owned by DLManagedTensor.manager_ctx — freed when PyTorch releases the tensor.
|
||||
struct DLPackContext {
|
||||
_gpu_buffer: CudaSlice<i64>,
|
||||
shape: Box<[i64]>,
|
||||
}
|
||||
|
||||
pub struct PinnedBatcher {
|
||||
tokenizer: Tokenizer,
|
||||
device: Arc<CudaDevice>,
|
||||
}
|
||||
|
||||
impl PinnedBatcher {
|
||||
pub fn new(model_path: &str) -> PyResult<Self> {
|
||||
let tokenizer = Tokenizer::from_file(model_path)
|
||||
.map_err(|e| PyValueError::new_err(format!("Failed to load tokenizer: {}", e)))?;
|
||||
let device = CudaDevice::new(0)
|
||||
.map_err(|e| PyValueError::new_err(format!("No CUDA GPU found: {:?}", e)))?;
|
||||
Ok(PinnedBatcher { tokenizer, device })
|
||||
}
|
||||
|
||||
pub fn batch_encode_to_gpu(&self, py: Python, texts: Vec<String>) -> PyResult<Py<PyAny>> {
|
||||
let encodings = self.tokenizer.encode_batch(texts, true)
|
||||
.map_err(|e| PyValueError::new_err(format!("Tokenization failed: {}", e)))?;
|
||||
|
||||
if encodings.is_empty() {
|
||||
return Err(PyValueError::new_err("Empty batch provided"));
|
||||
}
|
||||
|
||||
let batch_size = encodings.len();
|
||||
|
||||
// 1.4: use max length across all encodings, not just encodings[0]
|
||||
let seq_len = encodings.iter().map(|e| e.len()).max().unwrap_or(0);
|
||||
if seq_len == 0 {
|
||||
return Err(PyValueError::new_err("All sequences are empty"));
|
||||
}
|
||||
let total_elements = batch_size * seq_len;
|
||||
|
||||
// 1.4: pad shorter sequences with the tokenizer's pad token (fallback: 0)
|
||||
let pad_id = self.tokenizer.get_padding()
|
||||
.map(|p| p.pad_id as i64)
|
||||
.unwrap_or(0);
|
||||
|
||||
// 1.1: allocate pinned host buffer sized to the actual batch — no hardcoded max
|
||||
let mut host_buffer = PinnedHostBuffer::new(total_elements)?;
|
||||
host_buffer.data.fill(pad_id);
|
||||
|
||||
// Parallel write into pinned memory; shorter sequences leave pad_id in remaining slots
|
||||
host_buffer.data[..total_elements].par_chunks_mut(seq_len)
|
||||
.zip(encodings.par_iter())
|
||||
.for_each(|(dest, enc): (&mut [i64], &tokenizers::Encoding)| {
|
||||
let ids = enc.get_ids();
|
||||
let len = ids.len().min(dest.len());
|
||||
for i in 0..len {
|
||||
dest[i] = ids[i] as i64;
|
||||
}
|
||||
});
|
||||
|
||||
// 1.1: allocate device buffer sized to this batch — no hardcoded max
|
||||
let gpu_buffer = self.device.alloc_zeros::<i64>(total_elements)
|
||||
.map_err(|e| PyValueError::new_err(format!("GPU alloc failed: {:?}", e)))?;
|
||||
|
||||
// Direct DMA copy (pinned host -> device)
|
||||
unsafe {
|
||||
let src_ptr = host_buffer.data.as_ptr() as *const std::ffi::c_void;
|
||||
let dst_ptr = *gpu_buffer.device_ptr();
|
||||
let bytes = total_elements * std::mem::size_of::<i64>();
|
||||
let result = cuMemcpyHtoD_v2(dst_ptr, src_ptr, bytes);
|
||||
if result != cudaError_enum::CUDA_SUCCESS {
|
||||
return Err(PyValueError::new_err(format!("DMA copy failed: {:?}", result)));
|
||||
}
|
||||
}
|
||||
// host_buffer drops here — unpins host memory
|
||||
|
||||
// 1.2: transfer gpu_buffer ownership into DLPackContext so each returned tensor owns its
|
||||
// own device allocation. The deleter frees it when PyTorch releases the tensor.
|
||||
let ctx = Box::new(DLPackContext {
|
||||
_gpu_buffer: gpu_buffer,
|
||||
shape: vec![batch_size as i64, seq_len as i64].into_boxed_slice(),
|
||||
});
|
||||
|
||||
// Capture pointers before leaking the box
|
||||
let shape_ptr = ctx.shape.as_ptr() as *mut i64;
|
||||
let data_ptr = unsafe { *ctx._gpu_buffer.device_ptr() } as *mut std::ffi::c_void;
|
||||
let manager_ctx = Box::into_raw(ctx) as *mut std::ffi::c_void;
|
||||
|
||||
let dl_tensor = DLTensor {
|
||||
data: data_ptr,
|
||||
device: DLDevice { device_type: 2, device_id: 0 },
|
||||
ndim: 2,
|
||||
dtype: DLDataType { code: 0, bits: 64, lanes: 1 },
|
||||
shape: shape_ptr,
|
||||
strides: std::ptr::null_mut(),
|
||||
byte_offset: 0,
|
||||
};
|
||||
|
||||
let managed_tensor = Box::new(DLManagedTensor {
|
||||
dl_tensor,
|
||||
manager_ctx,
|
||||
deleter: Some(dlpack_deleter),
|
||||
});
|
||||
|
||||
let managed_ptr = Box::into_raw(managed_tensor);
|
||||
|
||||
unsafe {
|
||||
let capsule_ptr = pyo3::ffi::PyCapsule_New(
|
||||
managed_ptr as *mut _,
|
||||
CAPSULE_NAME.as_ptr(),
|
||||
None,
|
||||
);
|
||||
if capsule_ptr.is_null() {
|
||||
return Err(PyValueError::new_err("Failed to create PyCapsule"));
|
||||
}
|
||||
Ok(Bound::from_owned_ptr(py, capsule_ptr).into_any().unbind())
|
||||
}
|
||||
}
|
||||
}
|
||||
|
||||
unsafe extern "C" fn dlpack_deleter(managed_ptr: *mut DLManagedTensor) {
|
||||
if managed_ptr.is_null() { return; }
|
||||
unsafe {
|
||||
let managed = Box::from_raw(managed_ptr);
|
||||
if !managed.manager_ctx.is_null() {
|
||||
// Drops DLPackContext: frees CudaSlice (device memory) and shape array.
|
||||
let _ = Box::from_raw(managed.manager_ctx as *mut DLPackContext);
|
||||
}
|
||||
}
|
||||
}
|
||||
455
src/parser.rs
455
src/parser.rs
@@ -1,7 +1,8 @@
|
||||
// src/parser.rs
|
||||
|
||||
// --- External Crates ---
|
||||
use std::pin::Pin;
|
||||
use std::sync::Arc;
|
||||
// FIX: Replace direct tokio_postgres connections with deadpool
|
||||
use tokio_postgres::{NoTls, CopyOutStream, Config as PgConfig};
|
||||
use deadpool_postgres::{Pool, Manager, Runtime};
|
||||
use anyhow::{Context, Result, anyhow};
|
||||
@@ -19,142 +20,248 @@ use bytes::{Bytes, BytesMut, Buf};
|
||||
use byteorder::{BigEndian, ReadBytesExt};
|
||||
use std::io::{Cursor, Read};
|
||||
use std::str;
|
||||
use chrono::{NaiveDateTime, NaiveDate};
|
||||
use std::mem;
|
||||
use std::str::FromStr;
|
||||
use std::time::Instant;
|
||||
use tokio::task::JoinSet;
|
||||
use arrow::compute::concat_batches;
|
||||
use async_channel;
|
||||
|
||||
// --- Internal Crates ---
|
||||
use crate::config::ConnectorConfig;
|
||||
#[cfg(feature = "profiling")]
|
||||
use tracing::{span, Level};
|
||||
|
||||
use crate::config::{ConnectorConfig, load_and_validate_config};
|
||||
|
||||
// --- CONSTANTS ---
|
||||
// Optimized calculation of epoch delta (2000-01-01 00:00:00 to 1970-01-01 00:00:00)
|
||||
// 10957 days * 86400 seconds/day * 1,000,000 micros/second = 946684800000000 micros
|
||||
const POSTGRES_EPOCH_MICROS_OFFSET: i64 = 946684800000000;
|
||||
// Default batch size (64k is a standard Arrow chunk size)
|
||||
const DEFAULT_BATCH_SIZE: usize = 65_536;
|
||||
|
||||
|
||||
// --- 1. CORE DATABASE LOGIC (PARALLEL COORDINATOR) ---
|
||||
// Note: We use blast_radius from Python config as the partitioning strategy
|
||||
// --- 1. CORE DATABASE LOGIC ---
|
||||
pub async fn run_db_logic(config: ConnectorConfig, blast_radius: i64) -> Result<RecordBatch> {
|
||||
|
||||
println!("UncheckedIO: Starting Query Planner (Blast Radius: {})...", blast_radius);
|
||||
let start_total = Instant::now();
|
||||
let start_phase1 = Instant::now();
|
||||
|
||||
#[cfg(feature = "profiling")]
|
||||
let root_span = span!(Level::INFO, "UncheckedIO_Run");
|
||||
#[cfg(feature = "profiling")]
|
||||
let _root_guard = root_span.enter();
|
||||
|
||||
// --- PHASE 1: SETUP ---
|
||||
#[cfg(feature = "profiling")]
|
||||
let phase1_span = span!(Level::INFO, "Phase1_Setup");
|
||||
#[cfg(feature = "profiling")]
|
||||
let _p1_guard = phase1_span.enter();
|
||||
|
||||
let num_workers = num_cpus::get();
|
||||
println!("UncheckedIO: Detected {} logical cores. Spawning {} persistent worker threads.", num_workers, num_workers);
|
||||
|
||||
// CONFIG: Determine Target Batch Size
|
||||
let target_batch_size = config.batch_size.unwrap_or(DEFAULT_BATCH_SIZE);
|
||||
println!("UncheckedIO: Worker Aggregation Target = {} rows/batch.", target_batch_size);
|
||||
|
||||
// --- PHASE 1: SETUP CONNECTION POOL ---
|
||||
let pg_config: tokio_postgres::Config = PgConfig::from_str(&config.connection_string)
|
||||
.context("Invalid connection string in config")?;
|
||||
.context("Invalid connection string")?;
|
||||
|
||||
// Initialize the Manager and Pool
|
||||
let manager = Manager::new(pg_config.clone(), NoTls);
|
||||
// Set pool size higher than the expected partition count to ensure connections are always available.
|
||||
let pool = Pool::builder(manager)
|
||||
.max_size(20)
|
||||
.max_size(num_workers)
|
||||
.runtime(Runtime::Tokio1)
|
||||
.build()
|
||||
.context("Failed to build connection pool")?;
|
||||
|
||||
// 2. Query for Table Bounds (using pool connection)
|
||||
let client = pool.get().await.context("Failed to get pool connection for stats query")?;
|
||||
let client = pool.get().await.context("Failed to get pool connection")?;
|
||||
let partition_key = "id";
|
||||
|
||||
let (base_query, _) = config.query.trim().split_once("TO STDOUT (FORMAT binary)")
|
||||
.context("Failed to parse base query from config")?;
|
||||
.context("Failed to parse base query")?;
|
||||
let base_query_inner = base_query.trim().trim_start_matches("COPY (").trim_end_matches(")");
|
||||
|
||||
// NOTE: We rely on MIN/MAX here, assuming dense key for benchmark data.
|
||||
let stats_query = format!("SELECT MIN({}), MAX({}) FROM ({}) AS subquery", partition_key, partition_key, base_query_inner);
|
||||
|
||||
let row = client.query_one(&stats_query, &[]).await?;
|
||||
let min_id: i64 = row.try_get(0).context("Failed to get MIN(id)")?;
|
||||
let max_id: i64 = row.try_get(1).context("Failed to get MAX(id)")?;
|
||||
// Connection is returned to the pool when 'client' is dropped here.
|
||||
println!("UncheckedIO: ID Range: {} to {}", min_id, max_id);
|
||||
drop(client);
|
||||
|
||||
let total_rows = (max_id - min_id + 1).max(1);
|
||||
let calculated_blast_radius = if blast_radius <= 0 {
|
||||
let target_chunks = (num_workers * 4) as i64;
|
||||
let dynamic_size = total_rows / target_chunks;
|
||||
dynamic_size.max(10_000)
|
||||
} else {
|
||||
blast_radius
|
||||
};
|
||||
|
||||
println!("UncheckedIO: Auto-tuned partition size to {} rows.", calculated_blast_radius);
|
||||
|
||||
// 4. Generate Partitioned Queries (Based on blast_radius from Python)
|
||||
struct PartitionTask {
|
||||
index: usize,
|
||||
query: String,
|
||||
expected_rows: usize
|
||||
}
|
||||
|
||||
let mut partitions: Vec<PartitionTask> = Vec::new();
|
||||
let (tx, rx) = async_channel::unbounded::<PartitionTask>();
|
||||
|
||||
let mut current_min = min_id;
|
||||
let mut idx = 0;
|
||||
let mut total_partitions = 0;
|
||||
|
||||
while current_min <= max_id {
|
||||
let current_max = (current_min + blast_radius - 1).min(max_id);
|
||||
let current_max = (current_min + calculated_blast_radius - 1).min(max_id);
|
||||
|
||||
let new_query = format!(
|
||||
"COPY (SELECT * FROM ({}) AS sub WHERE {} BETWEEN {} AND {}) TO STDOUT (FORMAT binary)",
|
||||
base_query_inner, partition_key, current_min, current_max
|
||||
);
|
||||
let estimated_rows = (current_max - current_min + 1) as usize;
|
||||
|
||||
partitions.push(PartitionTask { index: idx, query: new_query, expected_rows: estimated_rows });
|
||||
current_min += blast_radius;
|
||||
idx += 1;
|
||||
}
|
||||
println!("UncheckedIO: Generated {} parallel partitions.", partitions.len());
|
||||
let task = PartitionTask { index: idx, query: new_query, expected_rows: estimated_rows };
|
||||
tx.send(task).await.context("Failed to fill work queue")?;
|
||||
|
||||
current_min += calculated_blast_radius;
|
||||
idx += 1;
|
||||
total_partitions += 1;
|
||||
}
|
||||
tx.close();
|
||||
|
||||
#[cfg(feature = "profiling")]
|
||||
drop(_p1_guard);
|
||||
|
||||
// FIX: Define the duration variable here so it can be used in the print statement later
|
||||
let duration_phase1 = start_phase1.elapsed();
|
||||
|
||||
|
||||
// --- PHASE 2: PARALLEL EXECUTION (WITH AGGREGATION) ---
|
||||
let start_phase2 = Instant::now();
|
||||
#[cfg(feature = "profiling")]
|
||||
let phase2_span = span!(Level::INFO, "Phase2_Execution");
|
||||
#[cfg(feature = "profiling")]
|
||||
let _p2_guard = phase2_span.enter();
|
||||
|
||||
// --- Phase 2: Parallel Execution ---
|
||||
let arrow_schema = Arc::new(build_arrow_schema(&config)?);
|
||||
let mut join_set = JoinSet::new();
|
||||
|
||||
for task in partitions {
|
||||
let worker_pool = pool.clone(); // Pass the pool handle
|
||||
for worker_id in 0..num_workers {
|
||||
let worker_rx = rx.clone();
|
||||
let worker_pool = pool.clone();
|
||||
let worker_schema = arrow_schema.clone();
|
||||
|
||||
join_set.spawn(async move {
|
||||
let worker_logic = async {
|
||||
// Get connection from pool (This is the speedup)
|
||||
let worker_client = worker_pool.get().await
|
||||
.context("Worker: Failed to get pool connection")?;
|
||||
#[cfg(feature = "profiling")]
|
||||
let worker_span = span!(Level::INFO, "Worker_Thread", id = worker_id);
|
||||
#[cfg(feature = "profiling")]
|
||||
let _w_guard = worker_span.enter();
|
||||
|
||||
let copy_stream = worker_client.copy_out(task.query.as_str()).await?;
|
||||
let pinned_stream: Pin<Box<CopyOutStream>> = Box::pin(copy_stream);
|
||||
let mut worker_batches: Vec<(usize, RecordBatch)> = Vec::new();
|
||||
|
||||
// Call the static dispatch parser
|
||||
parse_data_with_schema(pinned_stream, worker_schema).await
|
||||
};
|
||||
// Initialize the Parser ONCE per worker (The Buffer)
|
||||
let mut parser = create_parser();
|
||||
let mut parser_row_count = 0;
|
||||
|
||||
let result = worker_logic.await;
|
||||
(task.index, result, task.expected_rows)
|
||||
let mut client = worker_pool.get().await
|
||||
.context(format!("Worker {} failed to acquire connection", worker_id))?;
|
||||
|
||||
// Worker Loop
|
||||
while let Ok(task) = worker_rx.recv().await {
|
||||
#[cfg(feature = "profiling")]
|
||||
let task_span = span!(Level::INFO, "Processing_Task", partition_id = task.index);
|
||||
#[cfg(feature = "profiling")]
|
||||
let _t_guard = task_span.enter();
|
||||
|
||||
let result = async {
|
||||
let copy_stream = client.copy_out(task.query.as_str()).await?;
|
||||
let pinned_stream: Pin<Box<CopyOutStream>> = Box::pin(copy_stream);
|
||||
|
||||
// Parse DIRECTLY into the persistent parser
|
||||
parse_binary_stream_static(pinned_stream, &mut parser).await
|
||||
}.await;
|
||||
|
||||
match result {
|
||||
Ok(rows_read) => {
|
||||
parser_row_count += rows_read;
|
||||
|
||||
// CHECK FLUSH: Did we hit the batch size?
|
||||
if parser_row_count >= target_batch_size {
|
||||
let batch = flush_parser(&mut parser, worker_schema.clone())?;
|
||||
worker_batches.push((task.index, batch));
|
||||
parser_row_count = 0;
|
||||
// Re-init parser builders
|
||||
parser = create_parser();
|
||||
}
|
||||
}
|
||||
Err(e) => {
|
||||
eprintln!("Worker {} partition {} failed: {}. Handling failure.", worker_id, task.index, e);
|
||||
|
||||
// SAFETY FLUSH: If we have pending data, flush it first!
|
||||
if parser_row_count > 0 {
|
||||
let batch = flush_parser(&mut parser, worker_schema.clone())?;
|
||||
worker_batches.push((task.index, batch));
|
||||
parser_row_count = 0;
|
||||
parser = create_parser();
|
||||
}
|
||||
|
||||
// Emit NULL batch for the failed partition
|
||||
let null_batch = create_null_batch(worker_schema.clone(), task.expected_rows)?;
|
||||
worker_batches.push((task.index, null_batch));
|
||||
}
|
||||
}
|
||||
}
|
||||
|
||||
// FINAL FLUSH: Handle any remaining rows after queue is empty
|
||||
if parser_row_count > 0 {
|
||||
let batch = flush_parser(&mut parser, worker_schema.clone())?;
|
||||
worker_batches.push((usize::MAX, batch));
|
||||
}
|
||||
|
||||
Ok::<Vec<(usize, RecordBatch)>, anyhow::Error>(worker_batches)
|
||||
});
|
||||
}
|
||||
|
||||
// --- Phase 3: Collect, Order, and Stitch ---
|
||||
let mut results: Vec<Option<RecordBatch>> = vec![None; idx];
|
||||
|
||||
while let Some(join_result) = join_set.join_next().await {
|
||||
let (index, parse_result, expected_rows) = join_result.context("Worker thread panic")?;
|
||||
|
||||
match parse_result {
|
||||
Ok(batch) => {
|
||||
results[index] = Some(batch.1);
|
||||
}
|
||||
Err(e) => {
|
||||
// --- Self-Healing Placeholder ---
|
||||
eprintln!("UncheckedIO: Partition {} failed! Error: {}. Falling back to NULLs (Self-Healing logic required here).", index, e);
|
||||
let null_batch = create_null_batch(arrow_schema.clone(), expected_rows)?;
|
||||
results[index] = Some(null_batch);
|
||||
}
|
||||
// --- PHASE 3: AGGREGATION ---
|
||||
let mut all_results = Vec::new();
|
||||
while let Some(res) = join_set.join_next().await {
|
||||
match res {
|
||||
Ok(Ok(batches)) => all_results.extend(batches),
|
||||
Ok(Err(e)) => return Err(anyhow!("Worker failed: {}", e)),
|
||||
Err(e) => return Err(anyhow!("Worker panic: {}", e)),
|
||||
}
|
||||
}
|
||||
|
||||
let batches: Vec<RecordBatch> = results.into_iter()
|
||||
.filter_map(|b| b)
|
||||
.collect();
|
||||
#[cfg(feature = "profiling")]
|
||||
drop(_p2_guard);
|
||||
|
||||
if batches.is_empty() {
|
||||
return Ok(RecordBatch::new_empty(arrow_schema));
|
||||
}
|
||||
// --- PHASE 4: CONCAT ---
|
||||
let start_phase3 = Instant::now();
|
||||
#[cfg(feature = "profiling")]
|
||||
let phase3_span = span!(Level::INFO, "Phase3_Concat");
|
||||
#[cfg(feature = "profiling")]
|
||||
let _p3_guard = phase3_span.enter();
|
||||
|
||||
let final_batch = concat_batches(&arrow_schema, &batches)
|
||||
.context("Failed to concatenate parallel batches")?;
|
||||
if all_results.is_empty() { return Ok(RecordBatch::new_empty(arrow_schema)); }
|
||||
|
||||
// Sort by task index to maintain relative order
|
||||
all_results.sort_by_key(|(index, _)| *index);
|
||||
let batches: Vec<RecordBatch> = all_results.into_iter().map(|(_, b)| b).collect();
|
||||
let final_batch = concat_batches(&arrow_schema, &batches)?;
|
||||
|
||||
#[cfg(feature = "profiling")]
|
||||
drop(_p3_guard);
|
||||
|
||||
let duration_total = start_total.elapsed();
|
||||
|
||||
// --- REPORT ---
|
||||
println!("--- UncheckedIO Internal Timing ---");
|
||||
// FIX: Using duration_phase1 correctly now
|
||||
println!("Phase 1 (Setup): {:.2?}", duration_phase1);
|
||||
println!("Phase 2 (Execute): {:.2?}", start_phase3.duration_since(start_phase2));
|
||||
println!("Phase 3 (Concat): {:.2?}", start_phase3.elapsed());
|
||||
println!("Total Wall Time: {:.2?}", duration_total);
|
||||
|
||||
Ok(final_batch)
|
||||
}
|
||||
|
||||
// --- HELPER FUNCTIONS ---
|
||||
|
||||
fn create_null_batch(schema: Arc<Schema>, num_rows: usize) -> Result<RecordBatch> {
|
||||
let columns: Vec<ArrayRef> = schema.fields().iter().map(|field| {
|
||||
arrow::array::new_null_array(field.data_type(), num_rows)
|
||||
@@ -162,11 +269,8 @@ fn create_null_batch(schema: Arc<Schema>, num_rows: usize) -> Result<RecordBatch
|
||||
RecordBatch::try_new(schema, columns).context("Failed to create null placeholder batch")
|
||||
}
|
||||
|
||||
/// Helper function to build the Arrow Schema from the config
|
||||
fn build_arrow_schema(config: &ConnectorConfig) -> Result<Schema> {
|
||||
let schema_fields: Vec<Field> = config.schema.iter().map(|col_cfg| {
|
||||
let nullable = col_cfg.column_name == "notes"; // Hack for MVP
|
||||
|
||||
let arrow_type = match col_cfg.arrow_type.as_str() {
|
||||
"Int64" => DataType::Int64,
|
||||
"Int32" => DataType::Int32,
|
||||
@@ -176,20 +280,15 @@ fn build_arrow_schema(config: &ConnectorConfig) -> Result<Schema> {
|
||||
"Boolean" => DataType::Boolean,
|
||||
"Timestamp(Nanosecond, None)" => DataType::Timestamp(arrow::datatypes::TimeUnit::Nanosecond, None),
|
||||
"Date32" => DataType::Date32,
|
||||
_ => return Err(anyhow!("Unsupported type in config: {}", col_cfg.arrow_type)),
|
||||
_ => return Err(anyhow!("Unsupported type: {}", col_cfg.arrow_type)),
|
||||
};
|
||||
Ok(Field::new(&col_cfg.column_name, arrow_type, nullable))
|
||||
Ok(Field::new(&col_cfg.column_name, arrow_type, true))
|
||||
}).collect::<Result<Vec<Field>>>()?;
|
||||
|
||||
Ok(Schema::new(schema_fields))
|
||||
}
|
||||
|
||||
// --- PARSER STRUCT ---
|
||||
|
||||
// --------------------------------------------------------------------------------
|
||||
// --- 2. STATIC DISPATCH IMPLEMENTATION (The Fast Parser) ---
|
||||
// --------------------------------------------------------------------------------
|
||||
|
||||
// Struct to hold the builders in a statically-known, fixed order (eliminates DynamicBuilder enum)
|
||||
struct SchemaParser {
|
||||
id: Box<Int64Builder>,
|
||||
uuid: Box<StringBuilder>,
|
||||
@@ -203,13 +302,8 @@ struct SchemaParser {
|
||||
rating: Box<Float64Builder>,
|
||||
}
|
||||
|
||||
// Helper to construct and parse data using the static SchemaParser
|
||||
async fn parse_data_with_schema(
|
||||
stream: Pin<Box<CopyOutStream>>,
|
||||
arrow_schema: Arc<Schema>
|
||||
) -> Result<(usize, RecordBatch)> {
|
||||
|
||||
let mut parser = SchemaParser {
|
||||
fn create_parser() -> SchemaParser {
|
||||
SchemaParser {
|
||||
id: Box::new(Int64Builder::new()),
|
||||
uuid: Box::new(StringBuilder::new()),
|
||||
username: Box::new(StringBuilder::new()),
|
||||
@@ -220,11 +314,11 @@ async fn parse_data_with_schema(
|
||||
course_id: Box::new(Int32Builder::new()),
|
||||
start_date: Box::new(Date32Builder::new()),
|
||||
rating: Box::new(Float64Builder::new()),
|
||||
};
|
||||
}
|
||||
}
|
||||
|
||||
let rows_processed = parse_binary_stream_static(stream, &mut parser).await?;
|
||||
|
||||
// Collect all final arrays in the correct order (must match struct field order)
|
||||
// Helper to flush the parser into a RecordBatch
|
||||
fn flush_parser(parser: &mut SchemaParser, schema: Arc<Schema>) -> Result<RecordBatch> {
|
||||
let final_columns: Vec<ArrayRef> = vec![
|
||||
Arc::new(parser.id.finish()),
|
||||
Arc::new(parser.uuid.finish()),
|
||||
@@ -237,16 +331,11 @@ async fn parse_data_with_schema(
|
||||
Arc::new(parser.start_date.finish()),
|
||||
Arc::new(parser.rating.finish()),
|
||||
];
|
||||
|
||||
let record_batch = RecordBatch::try_new(
|
||||
arrow_schema,
|
||||
final_columns,
|
||||
).context("Failed to create final Arrow RecordBatch")?;
|
||||
|
||||
Ok((rows_processed, record_batch))
|
||||
RecordBatch::try_new(schema, final_columns).context("Failed to build RecordBatch")
|
||||
}
|
||||
|
||||
// The core streaming parser logic
|
||||
// --- STREAMING PARSER ---
|
||||
|
||||
async fn parse_binary_stream_static(
|
||||
mut stream: Pin<Box<CopyOutStream>>,
|
||||
parser: &mut SchemaParser,
|
||||
@@ -256,157 +345,161 @@ async fn parse_binary_stream_static(
|
||||
let mut is_header_parsed: bool = false;
|
||||
let mut rows_processed: usize = 0;
|
||||
|
||||
'stream_loop: while let Some(segment_result) = stream.next().await {
|
||||
let segment: Bytes = segment_result.context("Error reading segment from CopyOutStream")?;
|
||||
buffer.extend_from_slice(&segment);
|
||||
'stream_loop: loop {
|
||||
#[cfg(feature = "profiling")]
|
||||
let wait_span = span!(Level::ERROR, "IO_WAIT");
|
||||
#[cfg(feature = "profiling")]
|
||||
let guard = wait_span.enter();
|
||||
|
||||
if !is_header_parsed {
|
||||
if buffer.len() < 19 { continue 'stream_loop; }
|
||||
let mut header_cursor = Cursor::new(&buffer[..]);
|
||||
parse_stream_header(&mut header_cursor).context("Failed to parse stream header")?;
|
||||
buffer.advance(19);
|
||||
is_header_parsed = true;
|
||||
}
|
||||
let next_item = stream.next().await;
|
||||
|
||||
'parsing_loop: loop {
|
||||
let mut cursor = Cursor::new(&buffer[..]);
|
||||
let safe_position = cursor.position();
|
||||
#[cfg(feature = "profiling")]
|
||||
drop(guard);
|
||||
|
||||
let col_count = match cursor.read_i16::<BigEndian>() {
|
||||
Ok(count) => count,
|
||||
Err(e) if e.kind() == std::io::ErrorKind::UnexpectedEof => { break 'parsing_loop; }
|
||||
Err(e) => return Err(e.into()),
|
||||
};
|
||||
match next_item {
|
||||
Some(segment_result) => {
|
||||
#[cfg(feature = "profiling")]
|
||||
let work_span = span!(Level::INFO, "CPU_Parse");
|
||||
#[cfg(feature = "profiling")]
|
||||
let _work_guard = work_span.enter();
|
||||
|
||||
if col_count == -1 {
|
||||
buffer.advance(2);
|
||||
break 'stream_loop;
|
||||
}
|
||||
let segment: Bytes = segment_result.context("Error reading segment")?;
|
||||
buffer.extend_from_slice(&segment);
|
||||
|
||||
match parse_row_static(&mut cursor, parser, buffer.as_ref()) {
|
||||
Ok(_) => {
|
||||
rows_processed += 1;
|
||||
let bytes_consumed = cursor.position();
|
||||
buffer.advance(bytes_consumed as usize);
|
||||
if !is_header_parsed {
|
||||
if buffer.len() < 19 { continue 'stream_loop; }
|
||||
let mut header_cursor = Cursor::new(&buffer[..]);
|
||||
parse_stream_header(&mut header_cursor)?;
|
||||
buffer.advance(19);
|
||||
is_header_parsed = true;
|
||||
}
|
||||
Err(e) if e.kind() == std::io::ErrorKind::UnexpectedEof => {
|
||||
cursor.set_position(safe_position);
|
||||
// Copy remaining bytes back to the buffer for the next chunk
|
||||
let remaining_slice = &buffer.as_ref()[safe_position as usize..];
|
||||
let mut leftover_buffer_vec = Vec::new();
|
||||
leftover_buffer_vec.extend_from_slice(remaining_slice);
|
||||
buffer.clear();
|
||||
buffer.extend_from_slice(&leftover_buffer_vec);
|
||||
|
||||
break 'parsing_loop;
|
||||
}
|
||||
Err(e) => {
|
||||
return Err(e.into());
|
||||
'parsing_loop: loop {
|
||||
let mut cursor = Cursor::new(&buffer[..]);
|
||||
let safe_position = cursor.position();
|
||||
|
||||
let col_count = match cursor.read_i16::<BigEndian>() {
|
||||
Ok(count) => count,
|
||||
Err(e) if e.kind() == std::io::ErrorKind::UnexpectedEof => { break 'parsing_loop; }
|
||||
Err(e) => return Err(e.into()),
|
||||
};
|
||||
|
||||
if col_count == -1 {
|
||||
buffer.advance(2);
|
||||
break 'stream_loop;
|
||||
}
|
||||
|
||||
match parse_row_static(&mut cursor, parser, buffer.as_ref()) {
|
||||
Ok(_) => {
|
||||
rows_processed += 1;
|
||||
let bytes_consumed = cursor.position();
|
||||
buffer.advance(bytes_consumed as usize);
|
||||
}
|
||||
Err(e) if e.kind() == std::io::ErrorKind::UnexpectedEof => {
|
||||
cursor.set_position(safe_position);
|
||||
let remaining_slice = &buffer.as_ref()[safe_position as usize..];
|
||||
let mut leftover_vec = Vec::new();
|
||||
leftover_vec.extend_from_slice(remaining_slice);
|
||||
buffer.clear();
|
||||
buffer.extend_from_slice(&leftover_vec);
|
||||
break 'parsing_loop;
|
||||
}
|
||||
Err(e) => return Err(e.into()),
|
||||
}
|
||||
}
|
||||
}
|
||||
None => break 'stream_loop,
|
||||
}
|
||||
}
|
||||
|
||||
if !buffer.is_empty() {
|
||||
return Err(anyhow!("Stream ended with leftover bytes ({}) but no trailer.", buffer.len()));
|
||||
return Err(anyhow!("Stream ended with leftover bytes"));
|
||||
}
|
||||
|
||||
Ok(rows_processed)
|
||||
}
|
||||
|
||||
|
||||
fn parse_stream_header(cursor: &mut Cursor<&[u8]>) -> Result<()> {
|
||||
let mut magic_signature = [0u8; 11];
|
||||
cursor.read_exact(&mut magic_signature).context("Failed to read magic signature")?;
|
||||
if &magic_signature != b"PGCOPY\n\xff\r\n\0" {
|
||||
return Err(anyhow!("Invalid Postgres COPY binary signature."));
|
||||
}
|
||||
let _flags = cursor.read_u32::<BigEndian>().context("Failed to read flags")?;
|
||||
let _header_ext_len = cursor.read_u32::<BigEndian>().context("Failed to read header extension length")?;
|
||||
let mut magic = [0u8; 11];
|
||||
cursor.read_exact(&mut magic)?;
|
||||
if &magic != b"PGCOPY\n\xff\r\n\0" { return Err(anyhow!("Invalid signature")); }
|
||||
let _ = cursor.read_u32::<BigEndian>()?;
|
||||
let _ = cursor.read_u32::<BigEndian>()?;
|
||||
Ok(())
|
||||
}
|
||||
|
||||
// --- STATIC DISPATCH ROW PARSER (The Key Speedup) ---
|
||||
#[inline(always)]
|
||||
fn parse_row_static(
|
||||
cursor: &mut Cursor<&[u8]>,
|
||||
p: &mut SchemaParser, // The concrete, statically-typed parser struct
|
||||
p: &mut SchemaParser,
|
||||
current_chunk: &[u8]
|
||||
) -> Result<(), std::io::Error> {
|
||||
|
||||
// Column 0: id (BIGINT)
|
||||
// Column 0: id
|
||||
let len = cursor.read_i32::<BigEndian>()?;
|
||||
if len == -1 { p.id.append_null() } else { p.id.append_value(cursor.read_i64::<BigEndian>()?) }
|
||||
|
||||
// Column 1: uuid (TEXT)
|
||||
// Column 1: uuid
|
||||
let len = cursor.read_i32::<BigEndian>()?;
|
||||
if len == -1 { p.uuid.append_null() } else { read_string_field(cursor, p.uuid.as_mut(), current_chunk, len as usize)? }
|
||||
|
||||
// Column 2: username (TEXT)
|
||||
// Column 2: username
|
||||
let len = cursor.read_i32::<BigEndian>()?;
|
||||
if len == -1 { p.username.append_null() } else { read_string_field(cursor, p.username.as_mut(), current_chunk, len as usize)? }
|
||||
|
||||
// Column 3: score (REAL/Float32)
|
||||
// Column 3: score
|
||||
let len = cursor.read_i32::<BigEndian>()?;
|
||||
if len == -1 { p.score.append_null() } else { p.score.append_value(cursor.read_f32::<BigEndian>()?) }
|
||||
|
||||
// Column 4: is_active (BOOLEAN)
|
||||
// Column 4: is_active
|
||||
let len = cursor.read_i32::<BigEndian>()?;
|
||||
if len == -1 { p.is_active.append_null() } else { p.is_active.append_value(cursor.read_u8()? != 0) }
|
||||
|
||||
// Column 5: last_login (TIMESTAMP)
|
||||
// Column 5: last_login
|
||||
let len = cursor.read_i32::<BigEndian>()?;
|
||||
if len == -1 { p.last_login.append_null() } else {
|
||||
let pg_micros = cursor.read_i64::<BigEndian>()?;
|
||||
// Optimization: Constant offset applied
|
||||
let unix_micros = pg_micros + POSTGRES_EPOCH_MICROS_OFFSET;
|
||||
p.last_login.append_value(unix_micros * 1000);
|
||||
let val = cursor.read_i64::<BigEndian>()? + POSTGRES_EPOCH_MICROS_OFFSET;
|
||||
p.last_login.append_value(val * 1000);
|
||||
}
|
||||
|
||||
// Column 6: notes (TEXT)
|
||||
// Column 6: notes
|
||||
let len = cursor.read_i32::<BigEndian>()?;
|
||||
if len == -1 { p.notes.append_null() } else { read_string_field(cursor, p.notes.as_mut(), current_chunk, len as usize)? }
|
||||
|
||||
// Column 7: course_id (INT)
|
||||
// Column 7: course_id
|
||||
let len = cursor.read_i32::<BigEndian>()?;
|
||||
if len == -1 { p.course_id.append_null() } else { p.course_id.append_value(cursor.read_i32::<BigEndian>()?) }
|
||||
|
||||
// Column 8: start_date (DATE)
|
||||
// Column 8: start_date
|
||||
let len = cursor.read_i32::<BigEndian>()?;
|
||||
if len == -1 { p.start_date.append_null() } else {
|
||||
let pg_days = cursor.read_i32::<BigEndian>()?;
|
||||
// Optimization: 10957 days between 1970 and 2000
|
||||
p.start_date.append_value(pg_days + 10957);
|
||||
p.start_date.append_value(cursor.read_i32::<BigEndian>()? + 10957);
|
||||
}
|
||||
|
||||
// Column 9: rating (FLOAT8/Float64)
|
||||
// Column 9: rating
|
||||
let len = cursor.read_i32::<BigEndian>()?;
|
||||
if len == -1 { p.rating.append_null() } else { p.rating.append_value(cursor.read_f64::<BigEndian>()?) }
|
||||
|
||||
Ok(())
|
||||
}
|
||||
|
||||
// Helper function to consolidate zero-copy string reading and boundary checks
|
||||
fn read_string_field(
|
||||
cursor: &mut Cursor<&[u8]>,
|
||||
builder: &mut StringBuilder,
|
||||
current_chunk: &[u8],
|
||||
field_len_usize: usize
|
||||
len: usize
|
||||
) -> Result<(), std::io::Error> {
|
||||
|
||||
if (cursor.position() as usize + field_len_usize) > current_chunk.len() {
|
||||
return Err(std::io::Error::new(std::io::ErrorKind::UnexpectedEof, "Partial string field read"));
|
||||
if (cursor.position() as usize + len) > current_chunk.len() {
|
||||
return Err(std::io::Error::new(std::io::ErrorKind::UnexpectedEof, "Partial string"));
|
||||
}
|
||||
|
||||
let start = cursor.position() as usize;
|
||||
let end = start + field_len_usize;
|
||||
let end = start + len;
|
||||
let slice = ¤t_chunk[start..end];
|
||||
|
||||
let val_str = str::from_utf8(slice)
|
||||
.map_err(|e| std::io::Error::new(std::io::ErrorKind::InvalidData, e))?;
|
||||
|
||||
builder.append_value(val_str);
|
||||
cursor.set_position(end as u64); // Manually advance cursor
|
||||
|
||||
let val = str::from_utf8(slice).map_err(|e| std::io::Error::new(std::io::ErrorKind::InvalidData, e))?;
|
||||
builder.append_value(val);
|
||||
cursor.set_position(end as u64);
|
||||
Ok(())
|
||||
}
|
||||
}
|
||||
|
||||
// Profiler logic omitted for brevity (it remains unchanged from previous version)
|
||||
pub async fn run_profiler_logic(_: &str) -> Result<String> { Ok("Profiler Placeholder".to_string()) }
|
||||
Reference in New Issue
Block a user