Bugfix: Make vector database downloads thread safe

WolframResearch · Dec 3, 2024 · 0a5f6dd · 0a5f6dd
1 parent 6a8692a
commit 0a5f6dd
Showing 1 changed file with 32 additions and 17 deletions.
diff --git a/Source/Chatbook/PromptGenerators/VectorDatabases.wl b/Source/Chatbook/PromptGenerators/VectorDatabases.wl
@@ -128,21 +128,27 @@ downloadVectorDatabases[ ] :=
     downloadVectorDatabases[ $localVectorDBDirectory, $vectorDBDownloadURLs ];
 
 downloadVectorDatabases[ dir0_, urls_Association ] := Enclose[
-    Module[ { names, sizes, dir, tasks },
+    Module[ { dir, lock, names, sizes, tasks },
 
-        dir = ConfirmBy[ GeneralUtilities`EnsureDirectory @ dir0, DirectoryQ, "Directory" ];
+        dir   = ConfirmBy[ GeneralUtilities`EnsureDirectory @ dir0, DirectoryQ, "Directory" ];
+        lock  = FileNameJoin @ { dir, "download.lock" };
         names = ConfirmMatch[ Keys @ urls, { __String }, "Names" ];
         sizes = ConfirmMatch[ getDownloadSize /@ Values @ urls, { __? Positive }, "Sizes" ];
 
         $downloadProgress = AssociationMap[ 0 &, names ];
         $progressText = "Downloading semantic search indices\[Ellipsis]";
 
         evaluateWithProgress[
-
-            tasks = ConfirmMatch[ KeyValueMap[ downloadVectorDatabase @ dir, urls ], { __TaskObject }, "Download" ];
-            ConfirmMatch[ taskWait @ tasks, { __TaskObject }, "TaskWait" ];
-            $progressText = "Unpacking files\[Ellipsis]";
-            ConfirmBy[ unpackVectorDatabases @ dir, DirectoryQ, "Unpacked" ],
+            WithLock[
+                File @ lock
+                ,
+                tasks = ConfirmMatch[ KeyValueMap[ downloadVectorDatabase @ dir, urls ], { __TaskObject }, "Download" ];
+                ConfirmMatch[ taskWait @ tasks, { (_TaskObject|$Failed).. }, "TaskWait" ];
+                $progressText = "Unpacking files\[Ellipsis]";
+                ConfirmBy[ unpackVectorDatabases @ dir, DirectoryQ, "Unpacked" ]
+                ,
+                PersistenceTime -> 180
+            ],
 
             <|
                 "Text"             :> $progressText,
@@ -220,17 +226,26 @@ downloadVectorDatabase[ dir_ ] :=
     downloadVectorDatabase[ dir, ## ] &;
 
 downloadVectorDatabase[ dir_, name_String, url_String ] := Enclose[
-    Module[ { file },
+    Module[ { file, tmp },
+
         file = ConfirmBy[ FileNameJoin @ { dir, name<>".zip" }, StringQ, "File" ];
-        ConfirmMatch[
-            URLDownloadSubmit[
-                url,
-                file,
-                HandlerFunctions     -> <| "TaskProgress" -> setDownloadProgress @ name |>,
-                HandlerFunctionsKeys -> { "ByteCountDownloaded" }
-            ],
-            _TaskObject,
-            "Task"
+        tmp = file <> ".tmp";
+        Quiet[ DeleteFile /@ { file, tmp } ];
+
+        With[ { tmp = tmp, file = file },
+            ConfirmMatch[
+                URLDownloadSubmit[
+                    url,
+                    tmp,
+                    HandlerFunctions -> <|
+                        "TaskProgress" -> setDownloadProgress @ name,
+                        "TaskFinished" -> (RenameFile[ tmp, file ] &)
+                    |>,
+                    HandlerFunctionsKeys -> { "ByteCountDownloaded" }
+                ],
+                _TaskObject,
+                "Task"
+            ]
         ]
     ] // LogChatTiming[ { "DownloadVectorDatabase", name } ],
     throwInternalFailure