tensorflow_sys_runtime/
finder.rs

1use cfg_if::cfg_if;
2use std::env;
3
4/// Find the path to an Tensorflow library. This will try:
5/// - the `TENSORFLOW_DIR` environment variable with several subdirectories appended
6/// - the environment's library path (e.g. `LD_LIBRARY_PATH` in Linux)
7/// - Tensorflow's default installation paths for the OS
8pub fn find(library_name: &str) -> Option<PathBuf> {
9    let file = format!(
10        "{}{}{}",
11        env::consts::DLL_PREFIX,
12        library_name,
13        env::consts::DLL_SUFFIX
14    );
15    log::info!("Attempting to find library: {}", file);
16
17    // We search for the library in various different places and early-return if we find it.
18    macro_rules! check_and_return {
19        ($path: expr) => {
20            let path = $path;
21            log::debug!("Searching in: {}", path.display());
22            if path.is_file() {
23                log::info!("Found library at path: {}", path.display());
24                return Some(path);
25            }
26        };
27    }
28
29    // Search using the `TENSORFLOW_DIR` environment variable, this may be set by users.
30    if let Some(install_dir) = env::var_os(ENV_TENSORFLOW_DIR) {
31        check_and_return!(PathBuf::from(install_dir).join(&file));
32    }
33
34    // Search in the OS library path (i.e. `LD_LIBRARY_PATH` on Linux, `PATH` on Windows, and
35    // `DYLD_LIBRARY_PATH` on MacOS).
36    if let Some(path) = env::var_os(ENV_LIBRARY_PATH) {
37        for lib_dir in env::split_paths(&path) {
38            check_and_return!(lib_dir.join(&file));
39        }
40    }
41
42    // Search in Tensorflow's default installation directories (if they exist).
43    for default_dir in DEFAULT_INSTALLATION_DIRECTORIES
44        .iter()
45        .map(PathBuf::from)
46        .filter(|d| d.is_dir())
47    {
48        check_and_return!(default_dir.join(&file));
49    }
50
51    None
52}
53
54const ENV_TENSORFLOW_DIR: &'static str = "TENSORFLOW_DIR";
55
56cfg_if! {
57    if #[cfg(any(target_os = "linux"))] {
58        const ENV_LIBRARY_PATH: &'static str = "LD_LIBRARY_PATH";
59    } else if #[cfg(target_os = "macos")] {
60        const ENV_LIBRARY_PATH: &'static str = "DYLD_LIBRARY_PATH";
61    } else if #[cfg(target_os = "windows")] {
62        const ENV_LIBRARY_PATH: &'static str = "PATH";
63    } else {
64        // This may not work but seems like a sane default for target OS' not listed above.
65        const ENV_LIBRARY_PATH: &'static str = "LD_LIBRARY_PATH";
66    }
67}
68
69cfg_if! {
70    if #[cfg(any(target_os = "linux", target_os = "macos"))] {
71        const DEFAULT_INSTALLATION_DIRECTORIES: &'static [&'static str] =
72            &["/usr/local/lib", "/usr/local/lib/libtensorflow"];
73    } else if #[cfg(target_os = "windows")] {
74        const DEFAULT_INSTALLATION_DIRECTORIES: &'static [&'static str] = &[
75            "C:\\Program Files (x86)\\Tensorflow",
76            "C:\\Program Files (x86)\\tensorflow",
77        ];
78    } else {
79        const DEFAULT_INSTALLATION_DIRECTORIES: &'static [&'static str] = &[];
80    }
81}
82
tensorflow_sys_runtime/finder.rs

tensorflow_sys_runtime/
finder.rs
