use anyhow::Result;
use clap::Args;
use serde::{Deserialize, Serialize};
use std::collections::HashMap;
use crate::config::Config;
use crate::utils::{display_banner, output, system};
#[derive(Debug, Args)]
pub struct InfoCommand {
#[arg(long)]
pub detailed: bool,
#[arg(long)]
pub installation: bool,
#[arg(long)]
pub devices: bool,
#[arg(long)]
pub features: bool,
#[arg(long)]
pub diagnostics: bool,
#[arg(long)]
pub show_config: bool,
}
#[derive(Debug, Serialize, Deserialize)]
pub struct SystemInformation {
pub torsh: TorshInfo,
pub system: system::SystemInfo,
pub devices: HashMap<String, serde_json::Value>,
pub features: FeatureInfo,
pub installation: InstallationInfo,
}
#[derive(Debug, Serialize, Deserialize)]
pub struct TorshInfo {
pub version: String,
pub build_type: String,
pub build_date: String,
pub git_commit: String,
pub rust_version: String,
pub target_triple: String,
}
#[derive(Debug, Serialize, Deserialize)]
pub struct FeatureInfo {
pub enabled_features: Vec<String>,
pub disabled_features: Vec<String>,
pub experimental_features: Vec<String>,
}
#[derive(Debug, Serialize, Deserialize)]
pub struct InstallationInfo {
pub install_path: String,
pub config_path: String,
pub cache_path: String,
pub models_path: String,
pub size_on_disk: String,
}
#[derive(Debug, Serialize, Deserialize)]
pub struct DiagnosticResult {
pub name: String,
pub status: DiagnosticStatus,
pub message: String,
pub details: Option<serde_json::Value>,
}
#[derive(Debug, Serialize, Deserialize)]
pub enum DiagnosticStatus {
Pass,
Warning,
Fail,
Info,
}
pub async fn execute(args: InfoCommand, config: &Config, output_format: &str) -> Result<()> {
display_banner();
if !args.detailed
&& !args.installation
&& !args.devices
&& !args.features
&& !args.diagnostics
&& !args.show_config
{
show_basic_info(output_format).await?;
} else {
if args.detailed || args.installation {
show_detailed_info(output_format).await?;
}
if args.devices {
show_device_info(output_format).await?;
}
if args.features {
show_feature_info(output_format).await?;
}
if args.show_config {
show_config_info(config, output_format).await?;
}
if args.diagnostics {
run_diagnostics(config, output_format).await?;
}
}
Ok(())
}
async fn show_basic_info(output_format: &str) -> Result<()> {
let torsh_info = get_torsh_info();
let system_info = system::get_system_info();
let basic_info = serde_json::json!({
"torsh_version": torsh_info.version,
"os": system_info.os,
"total_memory": system_info.total_memory,
"cpu_count": system_info.cpu_count,
"available_devices": get_available_devices_summary(),
});
output::print_table("ToRSh System Information", &basic_info, output_format)?;
Ok(())
}
async fn show_detailed_info(output_format: &str) -> Result<()> {
let system_info = SystemInformation {
torsh: get_torsh_info(),
system: system::get_system_info(),
devices: system::get_device_info(),
features: get_feature_info(),
installation: get_installation_info().await?,
};
output::print_table("Detailed System Information", &system_info, output_format)?;
Ok(())
}
async fn show_device_info(output_format: &str) -> Result<()> {
let device_info = system::get_device_info();
output::print_table("Available Devices", &device_info, output_format)?;
for (device_name, info) in &device_info {
if let Some(available) = info.get("available").and_then(|v| v.as_bool()) {
if available {
output::print_success(&format!(
"✓ {} device is available",
device_name.to_uppercase()
));
} else {
output::print_warning(&format!(
"⚠ {} device is not available",
device_name.to_uppercase()
));
}
}
}
Ok(())
}
async fn show_feature_info(output_format: &str) -> Result<()> {
let feature_info = get_feature_info();
output::print_table("Feature Information", &feature_info, output_format)?;
output::print_info(&format!(
"Enabled features: {}",
feature_info.enabled_features.len()
));
output::print_info(&format!(
"Disabled features: {}",
feature_info.disabled_features.len()
));
if !feature_info.experimental_features.is_empty() {
output::print_warning(&format!(
"Experimental features: {}",
feature_info.experimental_features.len()
));
}
Ok(())
}
async fn show_config_info(config: &Config, output_format: &str) -> Result<()> {
let config_summary = serde_json::json!({
"output_dir": config.general.output_dir,
"cache_dir": config.general.cache_dir,
"default_device": config.general.default_device,
"num_workers": config.general.num_workers,
"default_dtype": config.general.default_dtype,
"hub_endpoint": config.hub.api_endpoint,
"mixed_precision": config.training.mixed_precision,
});
output::print_table("Configuration", &config_summary, output_format)?;
Ok(())
}
async fn run_diagnostics(config: &Config, output_format: &str) -> Result<()> {
output::print_info("Running system diagnostics...");
let mut diagnostics = Vec::new();
diagnostics.push(check_torsh_installation().await);
diagnostics.push(check_dependencies().await);
diagnostics.extend(check_device_availability().await);
diagnostics.push(check_configuration(config).await);
diagnostics.push(check_permissions(config).await);
diagnostics.push(check_disk_space(config).await);
output::print_table("Diagnostic Results", &diagnostics, output_format)?;
let pass_count = diagnostics
.iter()
.filter(|d| matches!(d.status, DiagnosticStatus::Pass))
.count();
let warning_count = diagnostics
.iter()
.filter(|d| matches!(d.status, DiagnosticStatus::Warning))
.count();
let fail_count = diagnostics
.iter()
.filter(|d| matches!(d.status, DiagnosticStatus::Fail))
.count();
println!();
output::print_info(&format!(
"Diagnostic Summary: {} passed, {} warnings, {} failed",
pass_count, warning_count, fail_count
));
if fail_count > 0 {
output::print_error("Some diagnostics failed. Please check the results above.");
} else if warning_count > 0 {
output::print_warning("Some diagnostics have warnings. Please review the results above.");
} else {
output::print_success("All diagnostics passed!");
}
Ok(())
}
fn get_torsh_info() -> TorshInfo {
TorshInfo {
version: env!("CARGO_PKG_VERSION").to_string(),
build_type: if cfg!(debug_assertions) {
"debug"
} else {
"release"
}
.to_string(),
build_date: "2024-01-15".to_string(), git_commit: "abc123def".to_string(), rust_version: std::env::var("RUST_VERSION").unwrap_or_else(|_| "unknown".to_string()),
target_triple: format!("{}-{}", std::env::consts::ARCH, std::env::consts::OS),
}
}
fn get_feature_info() -> FeatureInfo {
#[allow(unused_mut)]
let mut enabled_features = Vec::new();
#[allow(unused_mut)]
let mut disabled_features = Vec::new();
#[allow(unused_mut)]
let mut experimental_features = Vec::new();
#[cfg(feature = "nn")]
enabled_features.push("nn".to_string());
#[cfg(not(feature = "nn"))]
disabled_features.push("nn".to_string());
#[cfg(feature = "optim")]
enabled_features.push("optim".to_string());
#[cfg(not(feature = "optim"))]
disabled_features.push("optim".to_string());
#[cfg(feature = "data")]
enabled_features.push("data".to_string());
#[cfg(not(feature = "data"))]
disabled_features.push("data".to_string());
#[cfg(feature = "vision")]
enabled_features.push("vision".to_string());
#[cfg(not(feature = "vision"))]
disabled_features.push("vision".to_string());
#[cfg(feature = "text")]
enabled_features.push("text".to_string());
#[cfg(not(feature = "text"))]
disabled_features.push("text".to_string());
#[cfg(feature = "quantization")]
enabled_features.push("quantization".to_string());
#[cfg(not(feature = "quantization"))]
disabled_features.push("quantization".to_string());
#[cfg(feature = "jit")]
experimental_features.push("jit".to_string());
#[cfg(feature = "hub")]
enabled_features.push("hub".to_string());
#[cfg(not(feature = "hub"))]
disabled_features.push("hub".to_string());
FeatureInfo {
enabled_features,
disabled_features,
experimental_features,
}
}
async fn get_installation_info() -> Result<InstallationInfo> {
let current_exe = std::env::current_exe().unwrap_or_else(|_| "unknown".into());
let install_path = current_exe
.parent()
.unwrap_or_else(|| std::path::Path::new("unknown"));
let home_dir = dirs::home_dir().unwrap_or_else(|| std::path::PathBuf::from("."));
let config_dir = dirs::config_dir().unwrap_or_else(|| home_dir.join(".config"));
let cache_dir = dirs::cache_dir().unwrap_or_else(|| home_dir.join(".cache"));
let torsh_config = config_dir.join("torsh");
let torsh_cache = cache_dir.join("torsh");
let torsh_models = torsh_cache.join("models");
let mut total_size = 0u64;
if let Ok(metadata) = tokio::fs::metadata(¤t_exe).await {
total_size += metadata.len();
}
Ok(InstallationInfo {
install_path: install_path.display().to_string(),
config_path: torsh_config.display().to_string(),
cache_path: torsh_cache.display().to_string(),
models_path: torsh_models.display().to_string(),
size_on_disk: crate::utils::fs::format_file_size(total_size),
})
}
fn get_available_devices_summary() -> HashMap<String, bool> {
let device_info = system::get_device_info();
let mut summary = HashMap::new();
for (device_name, info) in device_info {
if let Some(available) = info.get("available").and_then(|v| v.as_bool()) {
summary.insert(device_name, available);
}
}
summary
}
async fn check_torsh_installation() -> DiagnosticResult {
let current_exe = std::env::current_exe();
match current_exe {
Ok(exe_path) => {
if exe_path.exists() {
DiagnosticResult {
name: "ToRSh Installation".to_string(),
status: DiagnosticStatus::Pass,
message: "ToRSh CLI is properly installed".to_string(),
details: Some(serde_json::json!({
"executable_path": exe_path.display().to_string()
})),
}
} else {
DiagnosticResult {
name: "ToRSh Installation".to_string(),
status: DiagnosticStatus::Fail,
message: "ToRSh executable not found".to_string(),
details: None,
}
}
}
Err(e) => DiagnosticResult {
name: "ToRSh Installation".to_string(),
status: DiagnosticStatus::Fail,
message: format!("Cannot determine executable path: {}", e),
details: None,
},
}
}
async fn check_dependencies() -> DiagnosticResult {
let mut dependency_status = HashMap::new();
let mut issues = Vec::new();
if let Ok(output) = std::process::Command::new("python3")
.arg("--version")
.output()
{
if output.status.success() {
let version = String::from_utf8_lossy(&output.stdout);
dependency_status.insert("python3", version.trim().to_string());
} else {
issues.push("Python3 not found");
dependency_status.insert("python3", "Not Available".to_string());
}
} else {
issues.push("Python3 not found");
dependency_status.insert("python3", "Not Available".to_string());
}
if let Ok(output) = std::process::Command::new("git").arg("--version").output() {
if output.status.success() {
let version = String::from_utf8_lossy(&output.stdout);
dependency_status.insert("git", version.trim().to_string());
} else {
issues.push("Git not found");
dependency_status.insert("git", "Not Available".to_string());
}
} else {
issues.push("Git not found");
dependency_status.insert("git", "Not Available".to_string());
}
if let Ok(output) = std::process::Command::new("curl").arg("--version").output() {
if output.status.success() {
dependency_status.insert("curl", "Available".to_string());
} else {
dependency_status.insert("curl", "Not Available".to_string());
}
} else {
dependency_status.insert("curl", "Not Available".to_string());
}
let status = if issues.is_empty() {
DiagnosticStatus::Pass
} else if issues.len() <= 2 {
DiagnosticStatus::Warning
} else {
DiagnosticStatus::Fail
};
let message = if issues.is_empty() {
"All external dependencies are available".to_string()
} else {
format!("Some dependencies missing: {}", issues.join(", "))
};
DiagnosticResult {
name: "External Dependencies".to_string(),
status,
message,
details: Some(serde_json::json!(dependency_status)),
}
}
async fn check_device_availability() -> Vec<DiagnosticResult> {
let mut results = Vec::new();
let device_info = system::get_device_info();
for (device_name, info) in device_info {
let available = info
.get("available")
.and_then(|v| v.as_bool())
.unwrap_or(false);
let status = if available {
DiagnosticStatus::Pass
} else {
DiagnosticStatus::Warning
};
let message = if available {
format!("{} device is available", device_name.to_uppercase())
} else {
format!("{} device is not available", device_name.to_uppercase())
};
results.push(DiagnosticResult {
name: format!("{} Device", device_name.to_uppercase()),
status,
message,
details: Some(info),
});
}
results
}
async fn check_configuration(config: &Config) -> DiagnosticResult {
let mut issues = Vec::new();
if !config.general.cache_dir.exists() {
issues.push("Cache directory does not exist");
}
if config.general.num_workers == 0 {
issues.push("Number of workers is set to 0");
}
if issues.is_empty() {
DiagnosticResult {
name: "Configuration".to_string(),
status: DiagnosticStatus::Pass,
message: "Configuration is valid".to_string(),
details: None,
}
} else {
DiagnosticResult {
name: "Configuration".to_string(),
status: DiagnosticStatus::Warning,
message: format!("Configuration has {} issues", issues.len()),
details: Some(serde_json::json!({
"issues": issues
})),
}
}
}
async fn check_permissions(config: &Config) -> DiagnosticResult {
let test_file = config.general.cache_dir.join(".torsh_test");
match tokio::fs::write(&test_file, "test").await {
Ok(_) => {
let _ = tokio::fs::remove_file(&test_file).await;
DiagnosticResult {
name: "Permissions".to_string(),
status: DiagnosticStatus::Pass,
message: "Write permissions are available".to_string(),
details: None,
}
}
Err(e) => DiagnosticResult {
name: "Permissions".to_string(),
status: DiagnosticStatus::Fail,
message: format!("Cannot write to cache directory: {}", e),
details: Some(serde_json::json!({
"cache_dir": config.general.cache_dir.display().to_string(),
"error": e.to_string(),
})),
},
}
}
async fn check_disk_space(config: &Config) -> DiagnosticResult {
use byte_unit::Byte;
use sysinfo::Disks;
let cache_dir = &config.general.cache_dir;
let disks = Disks::new_with_refreshed_list();
let cache_path = cache_dir
.canonicalize()
.unwrap_or_else(|_| cache_dir.clone());
let mut target_disk = None;
let mut longest_mount_len = 0;
for disk in disks.list() {
let mount_point = disk.mount_point();
if cache_path.starts_with(mount_point) {
let mount_len = mount_point.as_os_str().len();
if mount_len > longest_mount_len {
target_disk = Some((
disk.total_space(),
disk.available_space(),
mount_point.to_path_buf(),
));
longest_mount_len = mount_len;
}
}
}
if let Some((total_bytes, available_bytes, mount_point)) = target_disk {
let used_bytes = total_bytes.saturating_sub(available_bytes);
let usage_percent = if total_bytes > 0 {
(used_bytes as f64 / total_bytes as f64) * 100.0
} else {
0.0
};
let status = if usage_percent > 90.0 {
DiagnosticStatus::Fail
} else if usage_percent > 80.0 {
DiagnosticStatus::Warning
} else {
DiagnosticStatus::Pass
};
let message = if usage_percent > 90.0 {
"Very low disk space available (>90% used)".to_string()
} else if usage_percent > 80.0 {
"Low disk space warning (>80% used)".to_string()
} else {
"Sufficient disk space available".to_string()
};
return DiagnosticResult {
name: "Disk Space".to_string(),
status,
message,
details: Some(serde_json::json!({
"cache_dir": cache_dir.display().to_string(),
"mount_point": mount_point.display().to_string(),
"total_space": Byte::from_u128(total_bytes as u128).unwrap_or_else(|| Byte::from_u128(0).expect("zero bytes should always be valid")).get_appropriate_unit(byte_unit::UnitType::Binary).to_string(),
"available_space": Byte::from_u128(available_bytes as u128).unwrap_or_else(|| Byte::from_u128(0).expect("zero bytes should always be valid")).get_appropriate_unit(byte_unit::UnitType::Binary).to_string(),
"used_space": Byte::from_u128(used_bytes as u128).unwrap_or_else(|| Byte::from_u128(0).expect("zero bytes should always be valid")).get_appropriate_unit(byte_unit::UnitType::Binary).to_string(),
"usage_percent": format!("{:.1}%", usage_percent),
})),
};
}
DiagnosticResult {
name: "Disk Space".to_string(),
status: DiagnosticStatus::Warning,
message: "Could not determine disk space usage".to_string(),
details: Some(serde_json::json!({
"cache_dir": cache_dir.display().to_string(),
"note": "Could not find disk containing cache directory"
})),
}
}
#[cfg(test)]
mod tests {
use super::*;
#[test]
fn test_torsh_info() {
let info = get_torsh_info();
assert!(!info.version.is_empty());
assert!(!info.target_triple.is_empty());
}
#[test]
fn test_feature_info() {
let features = get_feature_info();
assert!(!features.enabled_features.is_empty() || !features.disabled_features.is_empty());
}
#[tokio::test]
async fn test_installation_info() {
let info = get_installation_info()
.await
.expect("operation should succeed");
assert!(!info.install_path.is_empty());
}
}