fix: Windows and CUDA bindings (#254)

withcatai · Jun 30, 2024 · da00a17 · da00a17
1 parent 1fbbf72
commit da00a17
Show file tree

Hide file tree

Showing 2 changed files with 44 additions and 16 deletions.
diff --git a/src/bindings/getLlama.ts b/src/bindings/getLlama.ts
@@ -486,32 +486,28 @@ async function loadExistingLlamaBinary({
     }
 
     if (canUsePrebuiltBinaries) {
-        const prebuiltBinPath = await getPrebuiltBinaryPath(
+        const prebuiltBinDetails = await getPrebuiltBinaryPath(
             buildOptions,
             existingPrebuiltBinaryMustMatchBuildOptions
                 ? buildFolderName.withCustomCmakeOptions
                 : buildFolderName.withoutCustomCmakeOptions
         );
 
-        if (prebuiltBinPath != null) {
+        if (prebuiltBinDetails != null) {
             try {
-                const buildMetadata = await getPrebuiltBinaryBuildMetadata(
-                    existingPrebuiltBinaryMustMatchBuildOptions
-                        ? buildFolderName.withCustomCmakeOptions
-                        : buildFolderName.withoutCustomCmakeOptions
-                );
+                const buildMetadata = await getPrebuiltBinaryBuildMetadata(prebuiltBinDetails.folderPath, prebuiltBinDetails.folderName);
                 const shouldTestBinaryBeforeLoading = getShouldTestBinaryBeforeLoading({
                     isPrebuiltBinary: true,
                     platform,
                     platformInfo,
                     buildMetadata
                 });
                 const binaryCompatible = shouldTestBinaryBeforeLoading
-                    ? await testBindingBinary(prebuiltBinPath)
+                    ? await testBindingBinary(prebuiltBinDetails.binaryPath)
                     : true;
 
                 if (binaryCompatible) {
-                    const binding = loadBindingModule(prebuiltBinPath);
+                    const binding = loadBindingModule(prebuiltBinDetails.binaryPath);
 
                     return await Llama._create({
                         bindings: binding,

diff --git a/src/bindings/utils/compileLLamaCpp.ts b/src/bindings/utils/compileLLamaCpp.ts
@@ -141,18 +141,20 @@ export async function compileLlamaCpp(buildOptions: BuildOptions, compileOptions
 
                 for (const binFilesDirPath of binFilesDirPaths) {
                     if (await fs.pathExists(binFilesDirPath)) {
-                        const files = await fs.readdir(binFilesDirPath);
+                        const itemNames = await fs.readdir(binFilesDirPath);
 
                         await Promise.all(
-                            files.map((fileName) => (
-                                fs.copy(path.join(binFilesDirPath, fileName), path.join(compiledResultDirPath, fileName), {
+                            itemNames.map((itemName) => (
+                                fs.copy(path.join(binFilesDirPath, itemName), path.join(compiledResultDirPath, itemName), {
                                     overwrite: false
                                 })
                             ))
                         );
                     }
                 }
 
+                await applyResultDirFixes(compiledResultDirPath, path.join(outDirectory, "_temp"));
+
                 await fs.writeFile(path.join(compiledResultDirPath, buildMetadataFileName), JSON.stringify({
                     buildOptions: convertBuildOptionsToBuildOptionsJSON(buildOptions)
                 } satisfies BuildMetadataFile), "utf8");
@@ -282,7 +284,11 @@ export async function getPrebuiltBinaryPath(buildOptions: BuildOptions, folderNa
     const binaryPath = await resolvePrebuiltBinaryPath(localPrebuiltBinaryDirectoryPath);
 
     if (binaryPath != null)
-        return binaryPath;
+        return {
+            binaryPath,
+            folderName,
+            folderPath: localPrebuiltBinaryDirectoryPath
+        };
 
     const packagePrebuiltBinariesDirectoryPath = await getPrebuiltBinariesPackageDirectoryForBuildOptions(buildOptions);
     if (packagePrebuiltBinariesDirectoryPath == null)
@@ -292,13 +298,17 @@ export async function getPrebuiltBinaryPath(buildOptions: BuildOptions, folderNa
     const binaryPathFromPackage = await resolvePrebuiltBinaryPath(packagePrebuiltBinaryDirectoryPath);
 
     if (binaryPathFromPackage != null)
-        return binaryPathFromPackage;
+        return {
+            binaryPath: binaryPathFromPackage,
+            folderName,
+            folderPath: packagePrebuiltBinaryDirectoryPath
+        };
 
     return null;
 }
 
-export async function getPrebuiltBinaryBuildMetadata(folderName: string) {
-    const buildMetadataFilePath = path.join(llamaPrebuiltBinsDirectory, folderName, buildMetadataFileName);
+export async function getPrebuiltBinaryBuildMetadata(folderPath: string, folderName: string) {
+    const buildMetadataFilePath = path.join(folderPath, buildMetadataFileName);
 
     if (!(await fs.pathExists(buildMetadataFilePath)))
         throw new Error(`Could not find build metadata file for prebuilt build "${folderName}"`);
@@ -307,6 +317,28 @@ export async function getPrebuiltBinaryBuildMetadata(folderName: string) {
     return buildMetadata;
 }
 
+async function applyResultDirFixes(resultDirPath: string, tempDirPath: string) {
+    const releaseDirPath = path.join(resultDirPath, "Release");
+
+    if (await fs.pathExists(releaseDirPath)) {
+        await fs.remove(tempDirPath);
+        await fs.ensureDir(tempDirPath);
+        await fs.move(releaseDirPath, tempDirPath);
+
+        const itemNames = await fs.readdir(tempDirPath);
+
+        await Promise.all(
+            itemNames.map((itemName) => (
+                fs.move(path.join(tempDirPath, itemName), path.join(resultDirPath, itemName), {
+                    overwrite: true
+                })
+            ))
+        );
+
+        await fs.remove(tempDirPath);
+    }
+}
+
 async function resolvePrebuiltBinaryPath(prebuiltBinaryDirectoryPath: string) {
     const binaryPath = path.join(prebuiltBinaryDirectoryPath, "llama-addon.node");
     const buildMetadataFilePath = path.join(prebuiltBinaryDirectoryPath, buildMetadataFileName);