Skip to content

Commit 2d7e6a5

Browse files
committed
Add TF_RUST_DOWNLOAD_DIR environment variable
This controls the location for the downloaded prebuilt binary tarball.
1 parent ff7de73 commit 2d7e6a5

File tree

1 file changed

+8
-5
lines changed

1 file changed

+8
-5
lines changed

tensorflow-sys/build.rs

Lines changed: 8 additions & 5 deletions
Original file line numberDiff line numberDiff line change
@@ -82,11 +82,14 @@ fn install_prebuilt() {
8282
let mut base_name = short_file_name.to_string();
8383
remove_suffix(&mut base_name, ".tar.gz");
8484
log_var!(base_name);
85-
let target_dir = PathBuf::from(&get!("CARGO_MANIFEST_DIR")).join("target");
86-
if !target_dir.exists() {
87-
fs::create_dir(&target_dir).unwrap();
85+
let download_dir = match env::var("TF_RUST_DOWNLOAD_DIR") {
86+
Ok(s) => PathBuf::from(s),
87+
Err(_) => PathBuf::from(&get!("CARGO_MANIFEST_DIR")).join("target"),
88+
};
89+
if !download_dir.exists() {
90+
fs::create_dir(&download_dir).unwrap();
8891
}
89-
let file_name = target_dir.join(short_file_name);
92+
let file_name = download_dir.join(short_file_name);
9093
log_var!(file_name);
9194

9295
// Download the tarball.
@@ -107,7 +110,7 @@ fn install_prebuilt() {
107110
}
108111

109112
// Extract the tarball.
110-
let unpacked_dir = target_dir.join(base_name);
113+
let unpacked_dir = download_dir.join(base_name);
111114
let lib_dir = unpacked_dir.join("lib");
112115
if !lib_dir.join(format!("lib{}.so", LIBRARY)).exists() {
113116
extract(file_name, &unpacked_dir);

0 commit comments

Comments
 (0)