Bundle native inference runtimes in releases.

crates/bds-release/src/packaging.rs

@@ -3,12 +3,13 @@ use apple_codesign::{SigningSettings, UnifiedSigner};
 use flate2::{Compression, write::GzEncoder};
 use icns::{IconFamily, Image};
 use plist::{Dictionary, Value};
+use sha2::{Digest, Sha256};
 use simple_file_manifest::FileEntry;
 use std::{
     error::Error,
     ffi::OsStr,
     fs::{self, File},
-    io::{self, BufReader},
+    io::{self, BufReader, Read, Seek},
     path::{Path, PathBuf},
 };
 use zip::{CompressionMethod, ZipWriter, write::SimpleFileOptions};
@@ -23,6 +24,8 @@ const LC_LOAD_WEAK_DYLIB: u32 = 0x8000_0018;
 const LC_REEXPORT_DYLIB: u32 = 0x8000_001f;
 const LC_LAZY_LOAD_DYLIB: u32 = 0x20;
 const LC_LOAD_UPWARD_DYLIB: u32 = 0x8000_0023;
+const DIRECTML_VERSION: &str = "1.15.4";
+const DIRECTML_SHA256: &str = "4e7cb7ddce8cf837a7a75dc029209b520ca0101470fcdf275c1f49736a3615b9";
 
 pub(crate) fn prepare_macos_linker_alias(directory: &Path) -> Result<(), Box<dyn Error>> {
     fs::create_dir_all(directory)?;
@@ -30,6 +33,84 @@ pub(crate) fn prepare_macos_linker_alias(directory: &Path) -> Result<(), Box<dyn
     Ok(())
 }
 
+pub(crate) fn prepare_windows_runtime(
+    directory: &Path,
+    sdk_libraries: &Path,
+) -> Result<(), Box<dyn Error>> {
+    fs::create_dir_all(directory)?;
+    let package_path = directory.join(format!("microsoft.ai.directml.{DIRECTML_VERSION}.nupkg"));
+    let result = (|| {
+        let url = format!(
+            "https://api.nuget.org/v3-flatcontainer/microsoft.ai.directml/{DIRECTML_VERSION}/microsoft.ai.directml.{DIRECTML_VERSION}.nupkg"
+        );
+        let mut response = reqwest::blocking::get(url)?.error_for_status()?;
+        let mut package = File::create(&package_path)?;
+        io::copy(&mut response, &mut package)?;
+        drop(package);
+
+        let mut package = BufReader::new(File::open(&package_path)?);
+        let mut digest = Sha256::new();
+        let mut buffer = [0_u8; 64 * 1024];
+        loop {
+            let read = package.read(&mut buffer)?;
+            if read == 0 {
+                break;
+            }
+            digest.update(&buffer[..read]);
+        }
+        let digest = digest.finalize();
+        if format!("{digest:x}") != DIRECTML_SHA256 {
+            return Err("DirectML package checksum mismatch".into());
+        }
+        prepare_windows_runtime_from_archive(File::open(&package_path)?, directory)?;
+        prepare_windows_sdk_linker_aliases(sdk_libraries, directory)
+    })();
+    if package_path.try_exists()? {
+        fs::remove_file(package_path)?;
+    }
+    result
+}
+
+fn prepare_windows_sdk_linker_aliases(
+    sdk_libraries: &Path,
+    directory: &Path,
+) -> Result<(), Box<dyn Error>> {
+    let source = fs::read_dir(sdk_libraries)?
+        .filter_map(Result::ok)
+        .find(|entry| {
+            entry
+                .file_name()
+                .to_str()
+                .is_some_and(|name| name.eq_ignore_ascii_case("pathcch.lib"))
+        })
+        .ok_or_else(|| {
+            format!(
+                "PathCch.lib is missing from Windows SDK directory: {}",
+                sdk_libraries.display()
+            )
+        })?
+        .path();
+    fs::copy(source, directory.join("PathCch.lib"))?;
+    Ok(())
+}
+
+fn prepare_windows_runtime_from_archive(
+    archive: impl Read + Seek,
+    directory: &Path,
+) -> Result<(), Box<dyn Error>> {
+    fs::create_dir_all(directory)?;
+    let mut archive = zip::ZipArchive::new(archive)?;
+    for (source, destination) in [
+        ("bin/x64-win/DirectML.lib", "DirectML.lib"),
+        ("bin/x64-win/DirectML.dll", "DirectML.dll"),
+        ("LICENSE.txt", "DirectML-LICENSE.txt"),
+    ] {
+        let mut source = archive.by_name(source)?;
+        io::copy(&mut source, &mut File::create(directory.join(destination))?)?;
+    }
+    Ok(())
+}
+
 #[derive(Clone, Copy, Debug, Eq, PartialEq)]
 enum Platform {
     MacOs,
@@ -177,6 +258,37 @@ fn runtime_files(
             .to_owned();
         files.push((path, name));
     }
+    let mut external_names = match platform {
+        Platform::MacOs => fs::read_dir(release_dir)?
+            .filter_map(Result::ok)
+            .map(|entry| entry.file_name())
+            .filter_map(|name| name.into_string().ok())
+            .filter(|name| name.starts_with("libonnxruntime") && name.contains(".dylib"))
+            .collect::<Vec<_>>(),
+        Platform::Linux => fs::read_dir(release_dir)?
+            .filter_map(Result::ok)
+            .map(|entry| entry.file_name())
+            .filter_map(|name| name.into_string().ok())
+            .filter(|name| name.starts_with("libonnxruntime.so"))
+            .collect::<Vec<_>>(),
+        Platform::Windows => ["onnxruntime.dll", "DirectML.dll", "DirectML-LICENSE.txt"]
+            .into_iter()
+            .map(str::to_owned)
+            .collect(),
+    };
+    external_names.sort();
+    if external_names.is_empty() {
+        return Err(format!(
+            "ONNX Runtime is missing from release directory: {}",
+            release_dir.display()
+        )
+        .into());
+    }
+    for name in external_names {
+        let path = release_dir.join(&name);
+        require_file(&path)?;
+        files.push((path, name));
+    }
     Ok(files)
 }
 
@@ -443,7 +555,7 @@ fn dmg_sectors(bundle_bytes: u64) -> Result<u32, Box<dyn Error>> {
 #[cfg(test)]
 mod tests {
     use super::*;
-    use std::io::Read;
+    use std::io::{Read, Write};
     use tempfile::tempdir;
 
     #[test]
@@ -477,6 +589,50 @@ mod tests {
         );
     }
 
+    #[test]
+    fn windows_runtime_extracts_only_redistributable_files() {
+        let temp = tempdir().unwrap();
+        let sdk = temp.path().join("sdk");
+        let runtime = temp.path().join("runtime");
+        fs::create_dir(&sdk).unwrap();
+        fs::create_dir(&runtime).unwrap();
+        fs::write(sdk.join("pathcch.lib"), b"sdk-import-library").unwrap();
+        let archive = std::io::Cursor::new(Vec::new());
+        let mut archive = ZipWriter::new(archive);
+        let options = SimpleFileOptions::default();
+        for (name, contents) in [
+            ("bin/x64-win/DirectML.lib", b"import-library".as_slice()),
+            ("bin/x64-win/DirectML.dll", b"runtime".as_slice()),
+            ("LICENSE.txt", b"license".as_slice()),
+            (
+                "bin/arm64-win/DirectML.dll",
+                b"wrong-architecture".as_slice(),
+            ),
+        ] {
+            archive.start_file(name, options).unwrap();
+            archive.write_all(contents).unwrap();
+        }
+        let archive = archive.finish().unwrap();
+
+        prepare_windows_runtime_from_archive(archive, &runtime).unwrap();
+        prepare_windows_sdk_linker_aliases(&sdk, &runtime).unwrap();
+
+        assert_eq!(
+            fs::read(runtime.join("DirectML.lib")).unwrap(),
+            b"import-library"
+        );
+        assert_eq!(fs::read(runtime.join("DirectML.dll")).unwrap(), b"runtime");
+        assert_eq!(
+            fs::read(runtime.join("DirectML-LICENSE.txt")).unwrap(),
+            b"license"
+        );
+        assert_eq!(
+            fs::read(runtime.join("PathCch.lib")).unwrap(),
+            b"sdk-import-library"
+        );
+        assert_eq!(fs::read_dir(runtime).unwrap().count(), 4);
+    }
+
     #[test]
     fn windows_package_contains_the_app_without_dynamic_std() {
         let temp = tempdir().unwrap();
@@ -491,6 +647,9 @@ mod tests {
             "bds-mcp.exe",
             "bds_core.dll",
             "bds_server.dll",
+            "onnxruntime.dll",
+            "DirectML.dll",
+            "DirectML-LICENSE.txt",
         ] {
             fs::write(release.join(file), format!("MZ-{file}")).unwrap();
         }
@@ -511,6 +670,14 @@ mod tests {
                 .by_name("bds-desktop-v1.2.3-windows-x64/bds_core.dll")
                 .is_ok()
         );
+        for runtime in ["onnxruntime.dll", "DirectML.dll", "DirectML-LICENSE.txt"] {
+            assert!(
+                archive
+                    .by_name(&format!("bds-desktop-v1.2.3-windows-x64/{runtime}"))
+                    .is_ok(),
+                "missing {runtime}"
+            );
+        }
     }
 
     #[test]
@@ -527,6 +694,7 @@ mod tests {
             "bds-mcp",
             "libbds_core.so",
             "libbds_server.so",
+            "libonnxruntime.so.1",
         ] {
             fs::write(release.join(file), file).unwrap();
         }
@@ -549,6 +717,11 @@ mod tests {
             .unwrap()
             .map(|entry| entry.unwrap().path().unwrap().into_owned())
             .collect::<Vec<_>>();
+        assert!(
+            names
+                .iter()
+                .any(|name| name.ends_with("libonnxruntime.so.1"))
+        );
         assert!(names.iter().any(|name| name.ends_with("bds-ui")));
         assert!(names.iter().any(|name| name.ends_with("libbds_core.so")));
     }