1use anyhow::Result;
10use std::path::Path;
11use std::sync::OnceLock;
12
13pub fn machine_name() -> String {
14 let name = hostname::get()
15 .ok()
16 .and_then(|s| s.into_string().ok())
17 .unwrap_or_else(|| "unknown-host".to_string());
18 tracing::debug!(
19 target: "studio_worker::sys",
20 op = "machine_name",
21 value = %name,
22 "resolved host machine name"
23 );
24 name
25}
26
27pub fn username() -> String {
28 username_from_probe(whoami::username())
29}
30
31fn username_from_probe<E: std::fmt::Display>(probe: std::result::Result<String, E>) -> String {
37 let user = match probe {
38 Ok(user) => user,
39 Err(e) => {
40 tracing::warn!(
41 target: "studio_worker::sys",
42 op = "username",
43 error = %e,
44 "failed to resolve OS user; falling back to unknown-user"
45 );
46 "unknown-user".to_string()
47 }
48 };
49 tracing::debug!(
50 target: "studio_worker::sys",
51 op = "username",
52 value = %user,
53 "resolved OS user"
54 );
55 user
56}
57
58static VRAM_GB: OnceLock<f32> = OnceLock::new();
63
64pub fn detect_vram_gb() -> Result<f32> {
78 Ok(*VRAM_GB.get_or_init(probe_vram_gb))
79}
80
81fn probe_vram_gb() -> f32 {
82 #[cfg(target_os = "linux")]
85 {
86 let from_sysfs = detect_vram_gb_from_sysfs(Path::new("/proc/driver/nvidia/gpus"));
87 if from_sysfs > 0.0 {
88 return from_sysfs;
89 }
90 let from_amd = detect_vram_gb_from_amd_sysfs(Path::new("/sys/class/drm"));
92 if from_amd > 0.0 {
93 return from_amd;
94 }
95 }
96 #[cfg(target_os = "macos")]
98 {
99 if let Some(gb) = detect_vram_gb_via_sysctl() {
100 return gb;
101 }
102 }
103 if let Some(gb) = detect_vram_gb_via_nvidia_smi() {
106 return gb;
107 }
108 #[cfg(target_os = "windows")]
110 {
111 if let Some(gb) = detect_vram_gb_via_wmic() {
112 return gb;
113 }
114 }
115 0.0
116}
117
118fn bytes_to_gib(bytes: u64) -> f32 {
121 (bytes as f64 / (1024.0 * 1024.0 * 1024.0)) as f32
122}
123
124fn is_drm_card_dir(name: &str) -> bool {
127 name.strip_prefix("card")
128 .is_some_and(|rest| !rest.is_empty() && rest.chars().all(|c| c.is_ascii_digit()))
129}
130
131pub fn detect_vram_gb_from_amd_sysfs(root: &Path) -> f32 {
137 let entries = match std::fs::read_dir(root) {
138 Ok(e) => e,
139 Err(_) => return 0.0,
140 };
141 let mut total_bytes: u64 = 0;
142 let mut cards: u32 = 0;
143 for entry in entries.flatten() {
144 let name = entry.file_name().to_string_lossy().into_owned();
145 if !is_drm_card_dir(&name) {
146 continue;
147 }
148 let vram_file = entry.path().join("device").join("mem_info_vram_total");
149 if let Ok(content) = std::fs::read_to_string(&vram_file) {
150 if let Some(bytes) = parse_amd_vram_total_bytes(&content) {
151 total_bytes = total_bytes.saturating_add(bytes);
152 cards += 1;
153 }
154 }
155 }
156 let gib = bytes_to_gib(total_bytes);
157 if cards > 0 {
158 tracing::info!(
159 target: "studio_worker::sys",
160 op = "probe_vram",
161 source = "amd_drm_sysfs",
162 vram_gb = gib,
163 cards,
164 "detected VRAM via AMD/DRM sysfs"
165 );
166 }
167 gib
168}
169
170pub fn parse_amd_vram_total_bytes(content: &str) -> Option<u64> {
173 content.trim().parse::<u64>().ok().filter(|b| *b > 0)
174}
175
176pub fn parse_sysctl_memsize(stdout: &str, fraction: f64) -> Option<f32> {
181 let bytes = stdout.trim().parse::<u64>().ok().filter(|b| *b > 0)?;
182 Some((bytes_to_gib(bytes) as f64 * fraction) as f32)
183}
184
185pub fn parse_wmic_adapter_ram(stdout: &str) -> Option<f32> {
190 let mut total: u64 = 0;
191 let mut found = false;
192 for line in stdout.lines() {
193 if let Some(bytes) = line.trim().parse::<u64>().ok().filter(|b| *b > 0) {
194 total = total.saturating_add(bytes);
195 found = true;
196 }
197 }
198 found.then(|| bytes_to_gib(total))
199}
200
201#[cfg_attr(coverage_nightly, coverage(off))]
211fn detect_vram_gb_via_nvidia_smi() -> Option<f32> {
212 let output = std::process::Command::new("nvidia-smi")
213 .args(["--query-gpu=memory.total", "--format=csv,noheader,nounits"])
214 .output();
215 match output {
216 Ok(o) if o.status.success() => vram_gb_from_smi_stdout(&String::from_utf8_lossy(&o.stdout)),
217 Ok(o) => {
218 tracing::warn!(
219 target: "studio_worker::sys",
220 op = "probe_vram",
221 source = "nvidia_smi_failed",
222 code = ?o.status.code(),
223 "nvidia-smi exited non-zero while probing VRAM — defaulting to 0 GB"
224 );
225 None
226 }
227 Err(e) => {
228 tracing::info!(
229 target: "studio_worker::sys",
230 op = "probe_vram",
231 source = "nvidia_smi_absent",
232 error = %e,
233 "nvidia-smi not available — cannot probe VRAM; defaulting to 0 GB"
234 );
235 None
236 }
237 }
238}
239
240#[cfg(target_os = "macos")]
244#[cfg_attr(coverage_nightly, coverage(off))]
245fn detect_vram_gb_via_sysctl() -> Option<f32> {
246 let output = std::process::Command::new("sysctl")
247 .args(["-n", "hw.memsize"])
248 .output()
249 .ok()?;
250 if !output.status.success() {
251 return None;
252 }
253 parse_sysctl_memsize(&String::from_utf8_lossy(&output.stdout), 0.75)
254}
255
256#[cfg(target_os = "windows")]
262#[cfg_attr(coverage_nightly, coverage(off))]
263fn detect_vram_gb_via_wmic() -> Option<f32> {
264 let output = std::process::Command::new("powershell")
265 .args([
266 "-NoProfile",
267 "-Command",
268 "(Get-CimInstance Win32_VideoController).AdapterRAM",
269 ])
270 .output()
271 .ok()?;
272 if !output.status.success() {
273 return None;
274 }
275 parse_wmic_adapter_ram(&String::from_utf8_lossy(&output.stdout))
276}
277
278struct SmiMemTotal {
289 mib: f64,
290 dropped: u32,
291}
292
293fn vram_gb_from_smi_stdout(stdout: &str) -> Option<f32> {
298 let SmiMemTotal { mib, dropped } = parse_nvidia_smi_mib(stdout)?;
299 let vram_gb = (mib / 1024.0) as f32;
300 tracing::info!(
301 target: "studio_worker::sys",
302 op = "probe_vram",
303 source = "nvidia_smi",
304 vram_gb = vram_gb,
305 dropped = dropped,
306 "detected NVIDIA VRAM via nvidia-smi fallback"
307 );
308 Some(vram_gb)
309}
310
311fn parse_nvidia_smi_mib(stdout: &str) -> Option<SmiMemTotal> {
319 let mut total: f64 = 0.0;
320 let mut any = false;
321 let mut dropped: u32 = 0;
322 for (idx, line) in stdout.lines().enumerate() {
323 let trimmed = line.trim();
324 if trimmed.is_empty() {
325 continue;
326 }
327 match trimmed
328 .split_whitespace()
329 .next()
330 .and_then(|tok| tok.parse::<f64>().ok())
331 {
332 Some(mib) => {
333 total += mib;
334 any = true;
335 }
336 None => {
337 dropped += 1;
338 tracing::warn!(
339 target: "studio_worker::sys",
340 op = "probe_vram",
341 source = "nvidia_smi",
342 line = idx,
343 content = trimmed,
344 "nvidia-smi VRAM line did not parse as MiB — dropping this GPU from the total"
345 );
346 }
347 }
348 }
349 any.then_some(SmiMemTotal {
350 mib: total,
351 dropped,
352 })
353}
354
355pub fn detect_vram_gb_from_sysfs(root: &Path) -> f32 {
379 let entries = match std::fs::read_dir(root) {
380 Ok(e) => e,
381 Err(_) => {
382 tracing::info!(
383 target: "studio_worker::sys",
384 op = "probe_vram",
385 source = "no_nvidia_sysfs",
386 vram_gb = 0.0,
387 root = %root.display(),
388 "no NVIDIA sysfs tree at probe root — defaulting to 0 GB VRAM"
389 );
390 return 0.0;
391 }
392 };
393
394 let mut total_mib: f64 = 0.0;
395 let mut gpu_count: u32 = 0;
396 let mut parseable: u32 = 0;
397 for entry in entries.flatten() {
398 gpu_count += 1;
399 let gpu_path = entry.path();
400 let info_path = gpu_path.join("information");
401 match std::fs::read_to_string(&info_path) {
402 Ok(content) => {
403 let mut found = false;
404 let mut unparseable: Option<String> = None;
411 for line in content.lines() {
412 if let Some(rest) = line.trim().strip_prefix("Video Memory:") {
413 if let Some(mib) = parse_mib(rest) {
414 total_mib += mib;
415 found = true;
416 } else if unparseable.is_none() {
417 unparseable = Some(rest.trim().to_string());
418 }
419 }
420 }
421 if found {
422 parseable += 1;
423 } else if let Some(content) = unparseable {
424 tracing::warn!(
425 target: "studio_worker::sys",
426 op = "probe_vram",
427 source = "nvidia_sysfs",
428 reason = "video_memory_unparseable",
429 gpu = %gpu_path.display(),
430 content = content.as_str(),
431 "sysfs GPU Video Memory line did not parse as MiB — dropping it from the total"
432 );
433 } else {
434 tracing::warn!(
435 target: "studio_worker::sys",
436 op = "probe_vram",
437 source = "nvidia_sysfs",
438 reason = "no_video_memory_line",
439 gpu = %gpu_path.display(),
440 "sysfs GPU has no parseable Video Memory line — dropping it from the total"
441 );
442 }
443 }
444 Err(e) => {
445 tracing::warn!(
446 target: "studio_worker::sys",
447 op = "probe_vram",
448 source = "nvidia_sysfs",
449 reason = "info_unreadable",
450 gpu = %gpu_path.display(),
451 error = %e,
452 "could not read a sysfs GPU information file — dropping it from the total"
453 );
454 }
455 }
456 }
457
458 let vram_gb = (total_mib / 1024.0) as f32;
459 let dropped = gpu_count.saturating_sub(parseable);
460 if parseable > 0 {
461 tracing::info!(
462 target: "studio_worker::sys",
463 op = "probe_vram",
464 source = "nvidia_sysfs",
465 vram_gb = vram_gb,
466 gpu_count = parseable,
467 dropped = dropped,
468 "detected NVIDIA VRAM via sysfs"
469 );
470 } else {
471 tracing::warn!(
472 target: "studio_worker::sys",
473 op = "probe_vram",
474 source = "sysfs_unparseable",
475 vram_gb = 0.0,
476 gpu_count = gpu_count,
477 root = %root.display(),
478 "NVIDIA sysfs entries present but no Video Memory line (current 5xx drivers dropped it) — falling back to nvidia-smi"
479 );
480 }
481 vram_gb
482}
483
484fn parse_mib(s: &str) -> Option<f64> {
485 let trimmed = s.trim();
487 let mut parts = trimmed.split_whitespace();
488 let value = parts.next()?.parse::<f64>().ok()?;
489 let unit = parts.next().unwrap_or("MiB");
490 match unit.to_ascii_lowercase().as_str() {
491 "mib" | "mb" => Some(value),
492 "gib" | "gb" => Some(value * 1024.0),
493 _ => Some(value),
494 }
495}
496
497#[cfg(test)]
498mod tests {
499 use super::*;
500
501 #[test]
502 fn parse_mib_handles_mib() {
503 assert_eq!(parse_mib(" 24576 MiB"), Some(24576.0));
504 assert_eq!(parse_mib("12288 MB"), Some(12288.0));
505 assert_eq!(parse_mib("24 GiB"), Some(24576.0));
506 assert_eq!(parse_mib("8 GB"), Some(8192.0));
507 }
508
509 #[test]
510 fn parse_mib_defaults_to_mib_when_the_unit_is_omitted() {
511 assert_eq!(parse_mib("4096"), Some(4096.0));
515 }
516
517 #[test]
518 fn parse_mib_treats_an_unknown_unit_as_raw_mib() {
519 assert_eq!(parse_mib("2048 KiB"), Some(2048.0));
525 assert_eq!(parse_mib("4 TB"), Some(4.0));
526 }
527
528 #[test]
529 fn parse_mib_rejects_unparseable_or_empty_values() {
530 assert_eq!(parse_mib("N/A MiB"), None);
534 assert_eq!(parse_mib(""), None);
535 assert_eq!(parse_mib(" "), None);
536 }
537
538 #[test]
539 fn machine_name_returns_non_empty() {
540 assert!(!machine_name().is_empty());
541 }
542
543 #[test]
544 fn username_returns_non_empty() {
545 assert!(!username().is_empty());
546 }
547
548 #[test]
549 fn username_from_probe_returns_the_resolved_value() {
550 let user = username_from_probe(Ok::<_, std::io::Error>("alice".to_string()));
551 assert_eq!(user, "alice");
552 }
553
554 #[test]
555 fn username_from_probe_falls_back_to_unknown_user_on_error() {
556 let user =
557 username_from_probe(Err::<String, _>(std::io::Error::other("no entropy source")));
558 assert_eq!(user, "unknown-user");
559 }
560
561 #[test]
562 fn username_from_probe_warns_with_the_error_on_failure() {
563 let logs = crate::test_support::capture(|| {
567 let _ =
568 username_from_probe(Err::<String, _>(std::io::Error::other("permission denied")));
569 });
570 assert!(logs.contains("WARN"), "expected WARN level, got: {logs}");
571 assert!(
572 logs.contains("op=\"username\""),
573 "expected username op, got: {logs}"
574 );
575 assert!(
576 logs.contains("permission denied"),
577 "expected underlying error, got: {logs}"
578 );
579 }
580
581 #[test]
582 fn username_from_probe_emits_debug_value_on_success() {
583 let logs = crate::test_support::capture(|| {
584 let _ = username_from_probe(Ok::<_, std::io::Error>("bob".to_string()));
585 });
586 assert!(logs.contains("DEBUG"), "expected DEBUG event, got: {logs}");
587 assert!(
588 logs.contains("value=bob"),
589 "expected resolved value, got: {logs}"
590 );
591 }
592
593 #[test]
594 fn detect_vram_gb_from_sysfs_returns_zero_when_root_missing() {
595 let dir = tempfile::tempdir().unwrap();
596 let missing = dir.path().join("nope");
597 assert_eq!(detect_vram_gb_from_sysfs(&missing), 0.0);
598 }
599
600 #[cfg(target_os = "linux")]
606 #[test]
607 fn detect_vram_gb_from_sysfs_sums_parseable_gpus() {
608 let dir = tempfile::tempdir().unwrap();
609 for (bus, mib) in [("0000:01:00.0", "12288"), ("0000:02:00.0", "24576")] {
610 let gpu = dir.path().join(bus);
611 std::fs::create_dir_all(&gpu).unwrap();
612 std::fs::write(
613 gpu.join("information"),
614 format!("Model: x\nVideo Memory: {mib} MiB\n"),
615 )
616 .unwrap();
617 }
618 let gb = detect_vram_gb_from_sysfs(dir.path());
620 assert!((gb - 36.0).abs() < 1e-3, "got {gb}");
621 }
622
623 #[cfg(target_os = "linux")]
624 #[test]
625 fn detect_vram_gb_from_sysfs_sums_only_survivors_when_one_gpu_is_unreadable() {
626 let dir = tempfile::tempdir().unwrap();
631 let good = dir.path().join("0000:01:00.0");
632 std::fs::create_dir_all(&good).unwrap();
633 std::fs::write(good.join("information"), "Video Memory: 12288 MiB\n").unwrap();
634 let bad = dir.path().join("0000:02:00.0");
635 std::fs::create_dir_all(bad.join("information")).unwrap();
636 let gb = detect_vram_gb_from_sysfs(dir.path());
638 assert!((gb - 12.0).abs() < 1e-3, "got {gb}");
639 }
640
641 #[test]
649 fn is_drm_card_dir_matches_only_primary_nodes() {
650 assert!(is_drm_card_dir("card0"));
651 assert!(is_drm_card_dir("card12"));
652 assert!(!is_drm_card_dir("card0-DP-1"));
653 assert!(!is_drm_card_dir("renderD128"));
654 assert!(!is_drm_card_dir("card"));
655 assert!(!is_drm_card_dir("controlD64"));
656 }
657
658 #[test]
659 fn parse_amd_vram_total_bytes_reads_a_byte_count() {
660 assert_eq!(
661 parse_amd_vram_total_bytes("17163091968\n"),
662 Some(17163091968)
663 );
664 assert_eq!(
665 parse_amd_vram_total_bytes("0"),
666 None,
667 "zero = no VRAM file value"
668 );
669 assert_eq!(parse_amd_vram_total_bytes("N/A"), None);
670 assert_eq!(parse_amd_vram_total_bytes(""), None);
671 }
672
673 #[test]
674 fn detect_vram_gb_from_amd_sysfs_sums_cards_and_ignores_siblings() {
675 let dir = tempfile::tempdir().unwrap();
676 for (card, bytes) in [
678 ("card0", 16u64 * 1024 * 1024 * 1024),
679 ("card1", 8 * 1024 * 1024 * 1024),
680 ] {
681 let dev = dir.path().join(card).join("device");
682 std::fs::create_dir_all(&dev).unwrap();
683 std::fs::write(dev.join("mem_info_vram_total"), bytes.to_string()).unwrap();
684 }
685 std::fs::create_dir_all(dir.path().join("card0-DP-1")).unwrap();
687 std::fs::create_dir_all(dir.path().join("renderD128")).unwrap();
688 std::fs::create_dir_all(dir.path().join("card2").join("device")).unwrap();
690
691 let gb = detect_vram_gb_from_amd_sysfs(dir.path());
692 assert!((gb - 24.0).abs() < 1e-3, "expected 24 GiB, got {gb}");
693 }
694
695 #[test]
696 fn detect_vram_gb_from_amd_sysfs_returns_zero_without_a_tree() {
697 let missing = std::path::Path::new("/definitely/no/drm/here");
698 assert_eq!(detect_vram_gb_from_amd_sysfs(missing), 0.0);
699 }
700
701 #[test]
702 fn parse_sysctl_memsize_scales_unified_memory() {
703 let bytes = (32u64 * 1024 * 1024 * 1024).to_string();
705 let gb = parse_sysctl_memsize(&bytes, 0.75).unwrap();
706 assert!((gb - 24.0).abs() < 1e-2, "got {gb}");
707 assert_eq!(parse_sysctl_memsize("0", 0.75), None);
708 assert_eq!(parse_sysctl_memsize("garbage", 0.75), None);
709 }
710
711 #[test]
712 fn parse_wmic_adapter_ram_sums_adapters_and_ignores_noise() {
713 let out = format!(
715 "AdapterRAM\n\n{}\n{}\n",
716 8u64 * 1024 * 1024 * 1024,
717 4u64 * 1024 * 1024 * 1024
718 );
719 let gb = parse_wmic_adapter_ram(&out).unwrap();
720 assert!((gb - 12.0).abs() < 1e-2, "got {gb}");
721 assert_eq!(parse_wmic_adapter_ram("AdapterRAM\n\n"), None, "no numbers");
722 assert_eq!(parse_wmic_adapter_ram(""), None);
723 }
724
725 #[test]
734 fn parse_nvidia_smi_mib_reads_a_single_bare_value() {
735 let total = parse_nvidia_smi_mib("24564\n").unwrap();
736 assert_eq!(total.mib, 24564.0);
737 assert_eq!(total.dropped, 0);
738 }
739
740 #[test]
741 fn parse_nvidia_smi_mib_sums_multiple_gpus() {
742 let total = parse_nvidia_smi_mib("24564\n24564\n").unwrap();
743 assert_eq!(total.mib, 49128.0);
744 assert_eq!(total.dropped, 0);
745 }
746
747 #[test]
748 fn parse_nvidia_smi_mib_tolerates_units_and_crlf_whitespace() {
749 let total = parse_nvidia_smi_mib(" 24564 MiB \r\n").unwrap();
751 assert_eq!(total.mib, 24564.0);
752 assert_eq!(total.dropped, 0);
753 }
754
755 #[test]
756 fn parse_nvidia_smi_mib_returns_none_on_empty_or_na() {
757 assert!(parse_nvidia_smi_mib("").is_none());
758 assert!(parse_nvidia_smi_mib("\n[N/A]\n").is_none());
759 }
760
761 #[test]
762 fn parse_nvidia_smi_mib_sums_survivors_and_counts_a_dropped_gpu() {
763 let total = parse_nvidia_smi_mib("24564\n[N/A]\n24564\n").unwrap();
767 assert_eq!(total.mib, 49128.0);
768 assert_eq!(total.dropped, 1);
769 }
770
771 #[test]
772 fn parse_nvidia_smi_mib_warns_on_each_dropped_gpu_line() {
773 let logs = crate::test_support::capture(|| {
777 let _ = parse_nvidia_smi_mib("24564\n[N/A]\n");
778 });
779 assert!(logs.contains("WARN"), "expected WARN level, got: {logs}");
780 assert!(
781 logs.contains("op=\"probe_vram\""),
782 "expected probe_vram op, got: {logs}"
783 );
784 assert!(
785 logs.contains("source=\"nvidia_smi\""),
786 "expected source=nvidia_smi, got: {logs}"
787 );
788 assert!(
789 logs.contains("[N/A]"),
790 "the warning must name the unparseable value, got: {logs}"
791 );
792 assert!(
793 logs.contains("dropping this GPU"),
794 "the warning must explain the drop, got: {logs}"
795 );
796 }
797
798 #[test]
799 fn vram_gb_from_smi_stdout_reports_dropped_count_in_breadcrumb() {
800 let logs = crate::test_support::capture(|| {
803 let gb = vram_gb_from_smi_stdout("24564\n[N/A]\n").unwrap();
804 assert!((gb - 23.99).abs() < 0.05, "survivor still totals: {gb}");
805 });
806 assert!(
807 logs.contains("dropped=1"),
808 "the breadcrumb must report the dropped count, got: {logs}"
809 );
810 }
811
812 #[test]
813 fn vram_gb_from_smi_stdout_converts_mib_to_gb() {
814 let gb = vram_gb_from_smi_stdout("24564\n").unwrap();
816 assert!((gb - 23.99).abs() < 0.05, "got {gb}");
817 }
818
819 #[test]
820 fn vram_gb_from_smi_stdout_is_none_when_unparseable() {
821 assert_eq!(vram_gb_from_smi_stdout("\n[N/A]\n"), None);
822 }
823
824 #[test]
825 fn vram_gb_from_smi_stdout_emits_info_breadcrumb_on_success() {
826 let logs = crate::test_support::capture(|| {
827 let _ = vram_gb_from_smi_stdout("24564\n");
828 });
829 assert!(logs.contains("INFO"), "expected INFO level, got: {logs}");
830 assert!(
831 logs.contains("op=\"probe_vram\""),
832 "expected probe_vram op, got: {logs}"
833 );
834 assert!(
835 logs.contains("source=\"nvidia_smi\""),
836 "expected source=nvidia_smi, got: {logs}"
837 );
838 }
839}