palace 0.1.0

A tool for mounting datasets into memory for fast loading in deep learning tasks.
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

use anyhow::{Context, Result};
use std::fs;
use std::path::Path;
use thiserror::Error;

#[derive(Error, Debug)]
pub enum PalaceError {
    #[error("Insufficient memory: required {required} bytes, available {available} bytes")]
    InsufficientMemory { required: u64, available: u64 },
    #[error("Invalid mount point: {0}")]
    InvalidMountPoint(String),
    #[error("I/O error: {0}")]
    IoError(#[from] std::io::Error),
}

pub fn get_available_memory() -> Result<u64> {
    let meminfo = fs::read_to_string("/proc/meminfo")?;
    let available = meminfo
        .lines()
        .find(|line| line.starts_with("MemAvailable:"))
        .and_then(|line| line.split_whitespace().nth(1))
        .and_then(|value| value.parse::<u64>().ok())
        .context("Failed to parse available memory")?;
    Ok(available * 1024) // Convert KB to bytes
}

pub fn get_directory_size(path: &Path) -> Result<u64> {
    let metadata = fs::metadata(path)?;
    if metadata.is_file() {
        return Ok(metadata.len());
    }
    let mut total_size = 0;
    for entry in fs::read_dir(path)? {
        let entry = entry?;
        let path = entry.path();
        if path.is_dir() {
            total_size += get_directory_size(&path)?;
        } else {
            total_size += entry.metadata()?.len();
        }
    }
    Ok(total_size)
}

pub fn check_mount_point(mount_point: &Path) -> Result<()> {
    if !mount_point.starts_with("/mnt") {
        return Err(PalaceError::InvalidMountPoint(mount_point.to_string_lossy().into_owned()).into());
    }
    Ok(())
}

pub fn ensure_directory(path: &Path) -> Result<()> {
    if !path.exists() {
        fs::create_dir_all(path).context("Failed to create directory")?;
    }
    Ok(())
}

pub fn check_memory_availability(required: u64, threshold_percentage: u8) -> Result<()> {
    let available = get_available_memory()?;
    let threshold = (available as f64 * (threshold_percentage as f64 / 100.0)) as u64;
    if required > threshold {
        Err(PalaceError::InsufficientMemory { required, available: threshold }.into())
    } else {
        Ok(())
    }
}

#[cfg(test)]
mod tests {
    use super::*;
    use tempfile::tempdir;

    #[test]
    fn test_get_directory_size() -> Result<()> {
        let dir = tempdir()?;
        let file_path = dir.path().join("test_file.txt");
        fs::write(&file_path, "Hello, world!")?;

        let size = get_directory_size(dir.path())?;
        assert_eq!(size, 13); // "Hello, world!" is 13 bytes
        Ok(())
    }

    #[test]
    fn test_check_mount_point() {
        assert!(check_mount_point(Path::new("/mnt/dataset/test")).is_ok());
        assert!(check_mount_point(Path::new("/tmp/test")).is_err());
    }

    #[test]
    fn test_ensure_directory() -> Result<()> {
        let dir = tempdir()?;
        let test_dir = dir.path().join("test_dir");
        ensure_directory(&test_dir)?;
        assert!(test_dir.exists());
        Ok(())
    }

    // Note: We can't reliably test get_available_memory and check_memory_availability
    // in a unit test environment, as they depend on system state.
}