Skip to main content

rust_ai_core/
lib.rs

1// SPDX-License-Identifier: MIT
2// Copyright 2026 Tyler Zervas
3
4//! # rust-ai-core
5//!
6//! Shared core utilities for the rust-ai ecosystem, providing unified abstractions
7//! for device selection, error handling, configuration validation, and `CubeCL` interop.
8//!
9//! ## Why This Crate Exists
10//!
11//! The rust-ai ecosystem consists of multiple specialized crates (peft-rs, qlora-rs,
12//! unsloth-rs, etc.) that share common patterns. Without a foundation layer, each crate
13//! would independently implement device selection, error types, and GPU dispatch logic,
14//! leading to inconsistency and code duplication.
15//!
16//! rust-ai-core consolidates these patterns into a single, well-tested foundation that
17//! ensures consistency across the ecosystem and provides a unified user experience.
18//!
19//! ## Design Philosophy
20//!
21//! **CUDA-first**: All operations prefer GPU execution. CPU is a fallback that emits
22//! warnings, not a silent alternative. This design ensures users are immediately aware
23//! when they're not getting optimal performance, rather than silently running slower.
24//!
25//! **Zero-cost abstractions**: Traits compile to static dispatch with no vtable overhead.
26//! The abstraction layer adds no runtime cost compared to direct implementations.
27//!
28//! **Fail-fast validation**: Configuration errors are caught at construction time, not
29//! deep in a training loop. This saves users from wasted compute time.
30//!
31//! ## Modules
32//!
33//! - [`device`] - CUDA-first device selection with environment variable overrides
34//! - [`error`] - Unified error types across all rust-ai crates
35//! - [`traits`] - Common traits for configs, quantization, and GPU dispatch
36//! - [`memory`] - Memory estimation and tracking utilities
37//! - [`dtype`] - Data type utilities and precision helpers
38//! - [`logging`] - Unified logging and observability
39//! - `cubecl` - `CubeCL` ↔ Candle tensor interoperability (requires `cuda` feature)
40//!
41//! ## Quick Start
42//!
43//! ```rust
44//! use rust_ai_core::{get_device, DeviceConfig, CoreError, Result};
45//!
46//! fn main() -> Result<()> {
47//!     // Get CUDA device with automatic fallback + warning
48//!     // Why: Users should know when they're not getting GPU acceleration
49//!     let device = get_device(&DeviceConfig::default())?;
50//!
51//!     // Or with explicit configuration from environment
52//!     // Why: Production deployments need runtime configuration without recompilation
53//!     let config = DeviceConfig::from_env();
54//!     let device = get_device(&config)?;
55//!
56//!     Ok(())
57//! }
58//! ```
59//!
60//! ## Feature Flags
61//!
62//! | Feature | Description | Dependencies |
63//! |---------|-------------|--------------|
64//! | `cuda` | Enable CUDA support via Candle and `CubeCL` kernels | cubecl, cubecl-cuda |
65//! | `python` | Enable Python bindings via `PyO3` | pyo3, numpy |
66//!
67//! ## Crate Integration
68//!
69//! All rust-ai crates should depend on rust-ai-core and use its shared types.
70//! This ensures consistent error handling, device selection, and trait implementations
71//! across the ecosystem.
72//!
73//! ```toml
74//! [dependencies]
75//! rust-ai-core = { version = "0.1", features = ["cuda"] }
76//! ```
77//!
78//! ## Environment Variables
79//!
80//! | Variable | Description | Example |
81//! |----------|-------------|---------|
82//! | `RUST_AI_FORCE_CPU` | Force CPU execution | `1` or `true` |
83//! | `RUST_AI_CUDA_DEVICE` | Select CUDA device ordinal | `0`, `1` |
84//! | `RUST_LOG` | Control logging level | `rust_ai_core=debug` |
85
86#![warn(missing_docs)]
87#![warn(clippy::all)]
88#![warn(clippy::pedantic)]
89#![allow(clippy::module_name_repetitions)]
90
91pub mod device;
92pub mod dtype;
93pub mod error;
94pub mod logging;
95pub mod memory;
96pub mod traits;
97
98#[cfg(feature = "cuda")]
99pub mod cubecl;
100
101#[cfg(feature = "python")]
102pub mod python;
103
104// Re-exports for convenience.
105//
106// Why: Users shouldn't need to navigate the module hierarchy for common types.
107// Flat re-exports provide a clean, discoverable API surface.
108pub use device::{get_device, warn_if_cpu, DeviceConfig};
109pub use dtype::{bytes_per_element, is_floating_point, DTypeExt};
110pub use error::{CoreError, Result};
111pub use logging::{init_logging, LogConfig};
112pub use memory::{estimate_tensor_bytes, MemoryTracker};
113pub use traits::{Dequantize, GpuDispatchable, Quantize, ValidatableConfig};
114
115#[cfg(feature = "cuda")]
116pub use cubecl::{
117    allocate_output_buffer, candle_to_cubecl_handle, cubecl_to_candle_tensor,
118    has_cubecl_cuda_support, TensorBuffer,
119};
120
121/// Crate version for runtime version checking.
122///
123/// Why: Dependent crates may need to verify compatibility at runtime,
124/// especially when loading serialized data that includes version info.
125pub const VERSION: &str = env!("CARGO_PKG_VERSION");