feat: availability checks (#22185)

server/src/repositories/logging.repository.ts

@@ -142,6 +142,10 @@ export class LoggingRepository {
     this.handleMessage(LogLevel.Fatal, message, details);
   }
 
+  deprecate(message: string) {
+    this.warn(`[Deprecated] ${message}`);
+  }
+
   private handleFunction(level: LogLevel, message: LogFunction, details: LogDetails[]) {
     if (this.logger.isLevelEnabled(level)) {
       this.handleMessage(level, message(), details);

server/src/repositories/machine-learning.repository.ts

@@ -1,6 +1,7 @@
 import { Injectable } from '@nestjs/common';
+import { Duration } from 'luxon';
 import { readFile } from 'node:fs/promises';
-import { MACHINE_LEARNING_AVAILABILITY_BACKOFF_TIME, MACHINE_LEARNING_PING_TIMEOUT } from 'src/constants';
+import { MachineLearningConfig } from 'src/config';
 import { CLIPConfig } from 'src/dtos/model-config.dto';
 import { LoggingRepository } from 'src/repositories/logging.repository';
 
@@ -57,82 +58,100 @@ export type TextEncodingOptions = ModelOptions & { language?: string };
 
 @Injectable()
 export class MachineLearningRepository {
-  // Note that deleted URL's are not removed from this map (ie: they're leaked)
-  // Cleaning them up is low priority since there should be very few over a
-  // typical server uptime cycle
-  private urlAvailability: {
-    [url: string]:
-      | {
-          active: boolean;
-          lastChecked: number;
-        }
-      | undefined;
-  };
+  private healthyMap: Record<string, boolean> = {};
+  private interval?: ReturnType<typeof setInterval>;
+  private _config?: MachineLearningConfig;
+
+  private get config(): MachineLearningConfig {
+    if (!this._config) {
+      throw new Error('Machine learning repository not been setup');
+    }
+
+    return this._config;
+  }
 
   constructor(private logger: LoggingRepository) {
     this.logger.setContext(MachineLearningRepository.name);
-    this.urlAvailability = {};
   }
 
-  private setUrlAvailability(url: string, active: boolean) {
-    const current = this.urlAvailability[url];
-    if (current?.active !== active) {
-      this.logger.verbose(`Setting ${url} ML server to ${active ? 'active' : 'inactive'}.`);
+  setup(config: MachineLearningConfig) {
+    this._config = config;
+    this.teardown();
+
+    // delete old servers
+    for (const url of Object.keys(this.healthyMap)) {
+      if (!config.urls.includes(url)) {
+        delete this.healthyMap[url];
+      }
     }
-    this.urlAvailability[url] = {
-      active,
-      lastChecked: Date.now(),
-    };
+
+    if (!config.availabilityChecks.enabled) {
+      return;
+    }
+
+    this.tick();
+    this.interval = setInterval(
+      () => this.tick(),
+      Duration.fromObject({ milliseconds: config.availabilityChecks.interval }).as('milliseconds'),
+    );
   }
 
-  private async checkAvailability(url: string) {
-    let active = false;
+  teardown() {
+    if (this.interval) {
+      clearInterval(this.interval);
+    }
+  }
+
+  private tick() {
+    for (const url of this.config.urls) {
+      void this.check(url);
+    }
+  }
+
+  private async check(url: string) {
+    let healthy = false;
     try {
       const response = await fetch(new URL('/ping', url), {
-        signal: AbortSignal.timeout(MACHINE_LEARNING_PING_TIMEOUT),
+        signal: AbortSignal.timeout(this.config.availabilityChecks.timeout),
       });
-      active = response.ok;
+      if (response.ok) {
+        healthy = true;
+      }
     } catch {
       // nothing to do here
     }
-    this.setUrlAvailability(url, active);
-    return active;
+
+    this.setHealthy(url, healthy);
   }
 
-  private async shouldSkipUrl(url: string) {
-    const availability = this.urlAvailability[url];
-    if (availability === undefined) {
-      // If this is a new endpoint, then check inline and skip if it fails
-      if (!(await this.checkAvailability(url))) {
-        return true;
-      }
-      return false;
+  private setHealthy(url: string, healthy: boolean) {
+    if (this.healthyMap[url] !== healthy) {
+      this.logger.log(`Machine learning server became ${healthy ? 'healthy' : 'unhealthy'} (${url}).`);
     }
-    if (!availability.active && Date.now() - availability.lastChecked < MACHINE_LEARNING_AVAILABILITY_BACKOFF_TIME) {
-      // If this is an old inactive endpoint that hasn't been checked in a
-      // while then check but don't wait for the result, just skip it
-      // This avoids delays on every search whilst allowing higher priority
-      // ML servers to recover over time.
-      void this.checkAvailability(url);
+
+    this.healthyMap[url] = healthy;
+  }
+
+  private isHealthy(url: string) {
+    if (!this.config.availabilityChecks.enabled) {
       return true;
     }
-    return false;
+
+    return this.healthyMap[url];
   }
 
-  private async predict<T>(urls: string[], payload: ModelPayload, config: MachineLearningRequest): Promise<T> {
+  private async predict<T>(payload: ModelPayload, config: MachineLearningRequest): Promise<T> {
     const formData = await this.getFormData(payload, config);
-    let urlCounter = 0;
-    for (const url of urls) {
-      urlCounter++;
-      const isLast = urlCounter >= urls.length;
-      if (!isLast && (await this.shouldSkipUrl(url))) {
-        continue;
-      }
 
+    for (const url of [
+      // try healthy servers first
+      ...this.config.urls.filter((url) => this.isHealthy(url)),
+      ...this.config.urls.filter((url) => !this.isHealthy(url)),
+    ]) {
       try {
         const response = await fetch(new URL('/predict', url), { method: 'POST', body: formData });
         if (response.ok) {
-          this.setUrlAvailability(url, true);
+          this.setHealthy(url, true);
           return response.json();
         }
 
@@ -144,20 +163,21 @@ export class MachineLearningRepository {
           `Machine learning request to "${url}" failed: ${error instanceof Error ? error.message : error}`,
         );
       }
-      this.setUrlAvailability(url, false);
+
+      this.setHealthy(url, false);
     }
 
     throw new Error(`Machine learning request '${JSON.stringify(config)}' failed for all URLs`);
   }
 
-  async detectFaces(urls: string[], imagePath: string, { modelName, minScore }: FaceDetectionOptions) {
+  async detectFaces(imagePath: string, { modelName, minScore }: FaceDetectionOptions) {
     const request = {
       [ModelTask.FACIAL_RECOGNITION]: {
         [ModelType.DETECTION]: { modelName, options: { minScore } },
         [ModelType.RECOGNITION]: { modelName },
       },
     };
-    const response = await this.predict<FacialRecognitionResponse>(urls, { imagePath }, request);
+    const response = await this.predict<FacialRecognitionResponse>({ imagePath }, request);
     return {
       imageHeight: response.imageHeight,
       imageWidth: response.imageWidth,
@@ -165,15 +185,15 @@ export class MachineLearningRepository {
     };
   }
 
-  async encodeImage(urls: string[], imagePath: string, { modelName }: CLIPConfig) {
+  async encodeImage(imagePath: string, { modelName }: CLIPConfig) {
     const request = { [ModelTask.SEARCH]: { [ModelType.VISUAL]: { modelName } } };
-    const response = await this.predict<ClipVisualResponse>(urls, { imagePath }, request);
+    const response = await this.predict<ClipVisualResponse>({ imagePath }, request);
     return response[ModelTask.SEARCH];
   }
 
-  async encodeText(urls: string[], text: string, { language, modelName }: TextEncodingOptions) {
+  async encodeText(text: string, { language, modelName }: TextEncodingOptions) {
     const request = { [ModelTask.SEARCH]: { [ModelType.TEXTUAL]: { modelName, options: { language } } } };
-    const response = await this.predict<ClipTextualResponse>(urls, { text }, request);
+    const response = await this.predict<ClipTextualResponse>({ text }, request);
     return response[ModelTask.SEARCH];
   }