From 0fa0bf124aa39eccd8bafba49347eb9338992049 Mon Sep 17 00:00:00 2001
From: Manav Rathi <manav@mrmr.io>
Date: Wed, 3 Jul 2024 14:07:42 +0530
Subject: [PATCH] Only reset the idle duration on non-empty pulls

---
 .../new/photos/services/ml/embedding.ts       |  7 ++++
 web/packages/new/photos/services/ml/worker.ts | 32 +++++++++----------
 2 files changed, 23 insertions(+), 16 deletions(-)

diff --git a/web/packages/new/photos/services/ml/embedding.ts b/web/packages/new/photos/services/ml/embedding.ts
index 42047dd326..fdbc2ae24c 100644
--- a/web/packages/new/photos/services/ml/embedding.ts
+++ b/web/packages/new/photos/services/ml/embedding.ts
@@ -94,6 +94,8 @@ type RemoteEmbedding = z.infer<typeof RemoteEmbedding>;
  *
  * This function should be called only after we have synced files with remote.
  * See: [Note: Ignoring embeddings for unknown files].
+ *
+ * @returns true if at least one embedding was pulled, false otherwise.
  */
 const pullEmbeddings = async (
     model: EmbeddingModel,
@@ -127,6 +129,7 @@ const pullEmbeddings = async (
     //     scenario where this happens.
     const localFilesByID = new Map(localFiles.map((f) => [f.id, f]));
 
+    let didPull = false;
     let sinceTime = await embeddingSyncTime(model);
     // TODO: eslint has fixed this spurious warning, but we're not on the latest
     // version yet, so add a disable.
@@ -149,6 +152,7 @@ const pullEmbeddings = async (
                         file.key,
                     ),
                 );
+                didPull = true;
                 count++;
             } catch (e) {
                 log.warn(`Ignoring unparseable ${model} embedding`, e);
@@ -157,6 +161,7 @@ const pullEmbeddings = async (
         await saveEmbeddingSyncTime(sinceTime, model);
         log.info(`Fetched ${count} ${model} embeddings`);
     }
+    return didPull;
 };
 
 /**
@@ -291,6 +296,8 @@ export const putEmbedding = async (
  *
  * This function should be called only after we have synced files with remote.
  * See: [Note: Ignoring embeddings for unknown files].
+ *
+ * @returns true if at least one embedding was pulled, false otherwise.
  */
 export const pullFaceEmbeddings = () =>
     pullEmbeddings("file-ml-clip-face", (jsonString: string) =>
diff --git a/web/packages/new/photos/services/ml/worker.ts b/web/packages/new/photos/services/ml/worker.ts
index b8ec7558dd..1267672547 100644
--- a/web/packages/new/photos/services/ml/worker.ts
+++ b/web/packages/new/photos/services/ml/worker.ts
@@ -1,13 +1,13 @@
+import downloadManager from "@/new/photos/services/download";
 import { markIndexingFailed, saveFaceIndex } from "@/new/photos/services/ml/db";
 import type { FaceIndex } from "@/new/photos/services/ml/types";
 import type { EnteFile } from "@/new/photos/types/file";
-import log from "@/next/log";
-// import { expose } from "comlink";
-import downloadManager from "@/new/photos/services/download";
 import { getKVN } from "@/next/kv";
 import { ensureAuthToken } from "@/next/local-user";
+import log from "@/next/log";
 import { ensure } from "@/utils/ensure";
 import { wait } from "@/utils/promise";
+import { expose } from "comlink";
 import { syncWithLocalFilesAndGetFilesToIndex } from ".";
 import { fileLogID } from "../../utils/file";
 import { pullFaceEmbeddings, putFaceIndex } from "./embedding";
@@ -128,15 +128,15 @@ export class MLWorker {
             idleDuration: this.idleDuration,
         }));
 
-        const next = () => void setTimeout(() => this.tick(), 0);
-
         // If we've been asked to sync, do that irrespective of anything else.
         if (this.shouldSync) {
             this.shouldSync = false;
-            // Reset the idle duration as we start pulling.
-            this.idleDuration = idleDurationStart;
-            // Call tick again once the pull is done.
-            void this.pull().then(next);
+            void pull().then((didPull) => {
+                // Reset the idle duration if we did pull something.
+                if (didPull) this.idleDuration = idleDurationStart;
+                // Either ways, tick again.
+                void setTimeout(() => this.tick(), 0);
+            });
             // Return without waiting for the pull.
             return;
         }
@@ -148,7 +148,7 @@ export class MLWorker {
             // Everything is running smoothly. Reset the idle duration.
             this.idleDuration = idleDurationStart;
             // And tick again.
-            next();
+            void setTimeout(() => this.tick(), 0);
             return;
         }
 
@@ -165,14 +165,14 @@ export class MLWorker {
         this.idleDuration = Math.min(this.idleDuration * 2, idleDurationMax);
         this.idleTimeout = setTimeout(next, this.idleDuration * 1000);
     }
-
-    private async pull() {
-        await pullFaceEmbeddings();
-    }
 }
 
-// TODO-ML: Temorarily disable
-// expose(MLWorker);
+expose(MLWorker);
+
+/**
+ * Pull embeddings from remote.
+ */
+const pull = pullFaceEmbeddings;
 
 /**
  * Find out files which need to be indexed. Then index the next batch of them.