kiroshi 0.1.0

An opinionated image diffusion server for Rust
Documentation 
1
2
3
4
5
6
7
8
9
10
11
12
13
14
15
16
17
18
19
20
21
22
23
24
25
26
27
28
29
30
31
32
33
34
35
36
37
38
39
40
41
42
43
44
45
46
47
48
49
50
51
52
53
54
55
56
57
58
59
60
61
62
63
64
65
66
67
68
69
70
71
72
73
74
75
76
77
78
79
80
81
82
83
84
85
86
87
88
89
90
91
92
93
94
95
96
97
98
99
100
101
102
103
104
105
106
107
108
109
110
111
112
113
114

use crate::Error;

use tokio::process;

use std::fmt;

#[derive(Debug, Clone, Copy, PartialEq, Eq)]
pub struct Stats {
    pub vram_usage: Memory,
    pub gpu_temperature: Temperature,
}

impl Stats {
    pub async fn fetch() -> Result<Self, Error> {
        // TODO: Run on server
        let (vram_usage, gpu_temperature) = check_gpu().await?;

        Ok(Self {
            vram_usage,
            gpu_temperature,
        })
    }
}

#[derive(Debug, Clone, Copy, PartialEq, Eq)]
pub struct Memory {
    free: MBytes,
    total: MBytes,
}

impl Memory {
    pub fn ratio(self) -> f32 {
        self.used().0 as f32 / self.total.0 as f32
    }

    pub fn used(self) -> MBytes {
        MBytes(self.total.0 - self.free.0)
    }

    pub fn total(self) -> MBytes {
        self.total
    }
}

#[derive(Debug, Clone, Copy, PartialEq, Eq)]
pub struct Temperature {
    pub celsius: u64,
}

impl fmt::Display for Temperature {
    fn fmt(&self, f: &mut fmt::Formatter<'_>) -> fmt::Result {
        write!(f, "{} °C", self.celsius)
    }
}

#[derive(Debug, Clone, Copy, PartialEq, Eq)]
pub struct MBytes(usize);

impl MBytes {
    pub fn from_mebibytes(mebibytes: usize) -> Self {
        Self(mebibytes * 1000 / 1024)
    }
}

impl fmt::Display for MBytes {
    fn fmt(&self, f: &mut fmt::Formatter<'_>) -> fmt::Result {
        write!(f, "{} MB", self.0)
    }
}

async fn check_gpu() -> Result<(Memory, Temperature), Error> {
    let nvidia_smi = process::Command::new("nvidia-smi")
        .args(["-q", "-d", "MEMORY,TEMPERATURE"])
        .output()
        .await?;

    let unexpected_output = || Error::InvalidOutput("unexpected output by nvidia-smi".to_owned());

    let output = String::from_utf8(nvidia_smi.stdout).map_err(|_| unexpected_output())?;

    let total = output
        .lines()
        .find(|line| line.trim().starts_with("Total"))
        .and_then(|line| line.split_whitespace().rev().nth(1))
        .ok_or(unexpected_output())?
        .parse()
        .map(MBytes::from_mebibytes)
        .ok()
        .ok_or(unexpected_output())?;

    let free = output
        .lines()
        .find(|line| line.trim().starts_with("Free"))
        .and_then(|line| line.split_whitespace().rev().nth(1))
        .ok_or(unexpected_output())?
        .parse()
        .map(MBytes::from_mebibytes)
        .ok()
        .ok_or(unexpected_output())?;

    let memory = Memory { free, total };

    let temperature = output
        .lines()
        .find(|line| line.trim().starts_with("GPU Current Temp"))
        .and_then(|line| line.split_whitespace().rev().nth(1))
        .ok_or(unexpected_output())?
        .parse()
        .map(|celsius| Temperature { celsius })
        .ok()
        .ok_or(unexpected_output())?;

    Ok((memory, temperature))
}