use crate::core::types::{BackendMode, ExecutorKind};
use serde::{Deserialize, Serialize};
use std::collections::HashMap;
#[derive(Debug, Clone)]
pub struct BackendRouter {
mode: BackendMode,
overrides: HashMap<String, ExecutorKind>,
}
impl BackendRouter {
pub fn new(mode: BackendMode) -> Self {
Self {
mode,
overrides: HashMap::new(),
}
}
pub fn set_override(&mut self, component_id: &str, kind: ExecutorKind) {
self.overrides.insert(component_id.to_string(), kind);
}
pub fn executor_for(&self, component_id: &str) -> ExecutorKind {
if let Some(kind) = self.overrides.get(component_id) {
return *kind;
}
if !is_known_component(component_id) {
return ExecutorKind::Unsupported;
}
match self.mode {
BackendMode::LinuxNative => self.route_linux_native(component_id),
BackendMode::WslBackedLinux => self.route_wsl_backed(component_id),
BackendMode::WindowsNative => self.route_windows_native(component_id),
}
}
pub fn mode(&self) -> BackendMode {
self.mode
}
pub fn unsupported_message(&self, component_id: &str) -> String {
let kind = self.executor_for(component_id);
if kind != ExecutorKind::Unsupported {
return String::new();
}
if !is_known_component(component_id) {
return format!(
"Component '{}' is not recognized by Rusty Stack.",
component_id
);
}
match self.mode {
BackendMode::WindowsNative => format!(
"Component '{}' requires a Linux environment with ROCm support. \
It cannot be installed natively on Windows. \
Consider using WSL2 with ROCm to enable Linux GPU computing on Windows.",
display_name(component_id)
),
BackendMode::WslBackedLinux => format!(
"Component '{}' is not supported in the current WSL2 configuration.",
display_name(component_id)
),
BackendMode::LinuxNative => format!(
"Component '{}' is not supported on this Linux configuration.",
display_name(component_id)
),
}
}
fn route_linux_native(&self, component_id: &str) -> ExecutorKind {
match component_id {
"wandb" => ExecutorKind::ExternalPackageManager,
_ => ExecutorKind::LegacyScript,
}
}
fn route_wsl_backed(&self, component_id: &str) -> ExecutorKind {
match component_id {
"wandb" => ExecutorKind::ExternalPackageManager,
_ => ExecutorKind::LegacyScript,
}
}
fn route_windows_native(&self, component_id: &str) -> ExecutorKind {
match component_id {
"wandb" => ExecutorKind::ExternalPackageManager,
_ => ExecutorKind::Unsupported,
}
}
}
pub fn detect_backend_mode() -> BackendMode {
#[cfg(target_os = "windows")]
{
BackendMode::WindowsNative
}
#[cfg(not(target_os = "windows"))]
{
if is_wsl_environment() {
BackendMode::WslBackedLinux
} else {
BackendMode::LinuxNative
}
}
}
#[cfg(not(target_os = "windows"))]
fn is_wsl_environment() -> bool {
if let Ok(version) = std::fs::read_to_string("/proc/version") {
let version_lower = version.to_lowercase();
return version_lower.contains("microsoft") || version_lower.contains("wsl");
}
false
}
fn is_known_component(id: &str) -> bool {
matches!(
id,
"rocm"
| "pytorch"
| "triton"
| "deepspeed"
| "vllm"
| "aiter"
| "onnx"
| "bitsandbytes"
| "migraphx"
| "flash-attn"
| "mpi4py"
| "wandb"
| "comfyui"
| "vllm-studio"
| "textgen"
| "rocm-smi"
| "permanent-env"
)
}
fn display_name(id: &str) -> &str {
match id {
"rocm" => "ROCm",
"pytorch" => "PyTorch",
"triton" => "Triton",
"deepspeed" => "DeepSpeed",
"vllm" => "vLLM",
"aiter" => "AITER",
"onnx" => "ONNX Runtime",
"bitsandbytes" => "bitsandbytes",
"migraphx" => "MIGraphX",
"flash-attn" => "Flash Attention",
"mpi4py" => "MPI4Py",
"wandb" => "Weights & Biases",
"comfyui" => "ComfyUI",
"vllm-studio" => "vLLM Studio",
"textgen" => "Text Generation WebUI",
"rocm-smi" => "ROCm SMI",
"permanent-env" => "Permanent Environment",
_ => id,
}
}
#[derive(Debug, Clone, Default, Serialize, Deserialize, PartialEq)]
pub struct WindowsInfo {
pub version: String,
pub build: u32,
pub wsl2_available: bool,
pub has_amd_gpu: bool,
}
#[cfg(test)]
mod tests {
use super::*;
#[test]
fn test_router_linux_native_rocm() {
let router = BackendRouter::new(BackendMode::LinuxNative);
assert_eq!(router.executor_for("rocm"), ExecutorKind::LegacyScript);
}
#[test]
fn test_router_linux_native_all_components() {
let router = BackendRouter::new(BackendMode::LinuxNative);
let components = [
"rocm",
"pytorch",
"triton",
"deepspeed",
"vllm",
"aiter",
"onnx",
"bitsandbytes",
"migraphx",
"flash-attn",
"mpi4py",
"comfyui",
"vllm-studio",
"textgen",
"rocm-smi",
"permanent-env",
];
for c in &components {
assert_ne!(
router.executor_for(c),
ExecutorKind::Unsupported,
"{} should not be unsupported on Linux",
c
);
}
}
#[test]
fn test_router_linux_native_wandb() {
let router = BackendRouter::new(BackendMode::LinuxNative);
assert_eq!(
router.executor_for("wandb"),
ExecutorKind::ExternalPackageManager
);
}
#[test]
fn test_router_windows_native_rocm_unsupported() {
let router = BackendRouter::new(BackendMode::WindowsNative);
assert_eq!(router.executor_for("rocm"), ExecutorKind::Unsupported);
}
#[test]
fn test_router_windows_native_all_gpu_unsupported() {
let router = BackendRouter::new(BackendMode::WindowsNative);
let gpu_components = [
"rocm",
"pytorch",
"triton",
"deepspeed",
"vllm",
"aiter",
"onnx",
"bitsandbytes",
"migraphx",
"flash-attn",
"mpi4py",
"comfyui",
"vllm-studio",
"textgen",
"rocm-smi",
"permanent-env",
];
for c in &gpu_components {
assert_eq!(
router.executor_for(c),
ExecutorKind::Unsupported,
"{} should be unsupported on Windows native",
c
);
}
}
#[test]
fn test_router_windows_native_wandb() {
let router = BackendRouter::new(BackendMode::WindowsNative);
assert_eq!(
router.executor_for("wandb"),
ExecutorKind::ExternalPackageManager
);
}
#[test]
fn test_router_wsl_backed_rocm() {
let router = BackendRouter::new(BackendMode::WslBackedLinux);
assert_eq!(router.executor_for("rocm"), ExecutorKind::LegacyScript);
}
#[test]
fn test_router_unknown_component() {
let router = BackendRouter::new(BackendMode::LinuxNative);
assert_eq!(
router.executor_for("nonexistent"),
ExecutorKind::Unsupported
);
}
#[test]
fn test_router_override() {
let mut router = BackendRouter::new(BackendMode::LinuxNative);
router.set_override("rocm", ExecutorKind::Rust);
assert_eq!(router.executor_for("rocm"), ExecutorKind::Rust);
}
#[test]
fn test_unsupported_message_windows_rocm() {
let router = BackendRouter::new(BackendMode::WindowsNative);
let msg = router.unsupported_message("rocm");
assert!(!msg.is_empty());
assert!(msg.contains("ROCm"));
assert!(msg.contains("Windows") || msg.contains("WSL"));
}
#[test]
fn test_unsupported_message_unknown_component() {
let router = BackendRouter::new(BackendMode::LinuxNative);
let msg = router.unsupported_message("nonexistent");
assert!(msg.contains("not recognized"));
}
#[test]
fn test_unsupported_message_supported_component_empty() {
let router = BackendRouter::new(BackendMode::LinuxNative);
let msg = router.unsupported_message("rocm");
assert!(
msg.is_empty(),
"Supported component should return empty message"
);
}
#[test]
fn test_detect_backend_mode() {
let mode = detect_backend_mode();
match mode {
BackendMode::LinuxNative | BackendMode::WslBackedLinux => {}
BackendMode::WindowsNative => {
panic!("Should not detect WindowsNative on Linux")
}
}
}
#[test]
fn test_is_known_component_all_17() {
let known = [
"rocm",
"pytorch",
"triton",
"deepspeed",
"vllm",
"aiter",
"onnx",
"bitsandbytes",
"migraphx",
"flash-attn",
"mpi4py",
"wandb",
"comfyui",
"vllm-studio",
"textgen",
"rocm-smi",
"permanent-env",
];
for c in &known {
assert!(is_known_component(c), "{} should be known", c);
}
}
#[test]
fn test_is_known_component_unknowns() {
let unknowns = ["foo", "bar", "cuda", "nvidia", ""];
for c in &unknowns {
assert!(!is_known_component(c), "{} should be unknown", c);
}
}
#[test]
fn test_display_name_known() {
assert_eq!(display_name("rocm"), "ROCm");
assert_eq!(display_name("pytorch"), "PyTorch");
assert_eq!(display_name("wandb"), "Weights & Biases");
}
#[test]
fn test_display_name_unknown_passthrough() {
assert_eq!(display_name("custom-component"), "custom-component");
}
#[test]
fn test_windows_info_serde_roundtrip() {
let info = WindowsInfo {
version: "10.0.19045".to_string(),
build: 19045,
wsl2_available: true,
has_amd_gpu: true,
};
let json = serde_json::to_string(&info).unwrap();
let back: WindowsInfo = serde_json::from_str(&json).unwrap();
assert_eq!(info, back);
}
#[test]
fn test_windows_info_default() {
let info = WindowsInfo::default();
assert!(info.version.is_empty());
assert_eq!(info.build, 0);
assert!(!info.wsl2_available);
assert!(!info.has_amd_gpu);
}
#[test]
fn test_router_mode() {
let router = BackendRouter::new(BackendMode::WindowsNative);
assert_eq!(router.mode(), BackendMode::WindowsNative);
}
}