From d5a9294eebcf303b63e585b4b668687f3146b471 Mon Sep 17 00:00:00 2001
From: Tempest <110401501+1-tempest@users.noreply.github.com>
Date: Fri, 17 Jan 2025 16:22:05 -0600
Subject: [PATCH] feat: Add additional env variables to ML container (#15398)

* Add additional variables to preload part ML models

* Add additional variables to preload part ML models

* Add additional variables to preload part ML models

* Add additional variables to preload part ML models

* Add additional variables to preload part ML models

* Add additional variables to preload part ML models

* Add additional variables to preload part ML models

* Add additional variables to preload part ML models

* Add additional variables to preload part ML models

* Update config.py

* Add additional variables to preload part ML models

* Add additional variables to preload part ML models

* Apply formatting

* minor update

* formatting

* root validator

* minor update

* minor update

* minor update

* change to support explicit models

* minor update

* minor change

* minor change

* minor change

* minor update

* add logs, resolve errors

* minor change

* add new enviornment variables

* minor revisons

* remove comments

* add additional variables to ML (fixed)

* add additional variables to ML (fixed)

* add additional variables to ML

* formatting

* remove comment

* remove mypy error

* remove unused module

* merge f strings
---
 machine-learning/app/config.py | 32 ++++++++++----------------------
 machine-learning/app/main.py   | 16 +++++++++++++++-
 2 files changed, 25 insertions(+), 23 deletions(-)

diff --git a/machine-learning/app/config.py b/machine-learning/app/config.py
index fa7dc61d47..e77531cbd6 100644
--- a/machine-learning/app/config.py
+++ b/machine-learning/app/config.py
@@ -25,31 +25,19 @@ class FacialRecognitionSettings(BaseModel):
 
 
 class PreloadModelData(BaseModel):
+    clip_fallback: str | None = os.getenv("MACHINE_LEARNING_PRELOAD__CLIP", None)
+    facial_recognition_fallback: str | None = os.getenv("MACHINE_LEARNING_PRELOAD__FACIAL_RECOGNITION", None)
+    if clip_fallback is not None:
+        os.environ["MACHINE_LEARNING_PRELOAD__CLIP__TEXTUAL"] = clip_fallback
+        os.environ["MACHINE_LEARNING_PRELOAD__CLIP__VISUAL"] = clip_fallback
+        del os.environ["MACHINE_LEARNING_PRELOAD__CLIP"]
+    if facial_recognition_fallback is not None:
+        os.environ["MACHINE_LEARNING_PRELOAD__FACIAL_RECOGNITION__RECOGNITION"] = facial_recognition_fallback
+        os.environ["MACHINE_LEARNING_PRELOAD__FACIAL_RECOGNITION__DETECTION"] = facial_recognition_fallback
+        del os.environ["MACHINE_LEARNING_PRELOAD__FACIAL_RECOGNITION"]
     clip: ClipSettings = ClipSettings()
     facial_recognition: FacialRecognitionSettings = FacialRecognitionSettings()
 
-    clip_model_fallback: str | None = os.getenv("MACHINE_LEARNING_PRELOAD__CLIP", None)
-    facial_recognition_model_fallback: str | None = os.getenv("MACHINE_LEARNING_PRELOAD__FACIAL_RECOGNITION", None)
-
-    def update_from_fallbacks(self) -> None:
-        if self.clip_model_fallback:
-            self.clip.textual = self.clip_model_fallback
-            self.clip.visual = self.clip_model_fallback
-            log.warning(
-                "Deprecated env variable: MACHINE_LEARNING_PRELOAD__CLIP. "
-                "Use MACHINE_LEARNING_PRELOAD__CLIP__TEXTUAL and "
-                "MACHINE_LEARNING_PRELOAD__CLIP__VISUAL instead."
-            )
-
-        if self.facial_recognition_model_fallback:
-            self.facial_recognition.recognition = self.facial_recognition_model_fallback
-            self.facial_recognition.detection = self.facial_recognition_model_fallback
-            log.warning(
-                "Deprecated environment variable: MACHINE_LEARNING_PRELOAD__FACIAL_RECOGNITION. "
-                "Use MACHINE_LEARNING_PRELOAD__FACIAL_RECOGNITION__RECOGNITION and "
-                "MACHINE_LEARNING_PRELOAD__FACIAL_RECOGNITION__DETECTION instead."
-            )
-
 
 class MaxBatchSize(BaseModel):
     facial_recognition: int | None = None
diff --git a/machine-learning/app/main.py b/machine-learning/app/main.py
index ee8881b6c6..fb6f84499a 100644
--- a/machine-learning/app/main.py
+++ b/machine-learning/app/main.py
@@ -75,7 +75,7 @@ async def lifespan(_: FastAPI) -> AsyncGenerator[None, None]:
 
 
 async def preload_models(preload: PreloadModelData) -> None:
-    log.info(f"Preloading models: {preload}")
+    log.info(f"Preloading models: clip:{preload.clip} facial_recognition:{preload.facial_recognition}")
 
     if preload.clip.textual is not None:
         model = await model_cache.get(preload.clip.textual, ModelType.TEXTUAL, ModelTask.SEARCH)
@@ -101,6 +101,20 @@ async def preload_models(preload: PreloadModelData) -> None:
         )
         await load(model)
 
+    if preload.clip_fallback is not None:
+        log.warning(
+            "Deprecated env variable: 'MACHINE_LEARNING_PRELOAD__CLIP'. "
+            "Use 'MACHINE_LEARNING_PRELOAD__CLIP__TEXTUAL' and "
+            "'MACHINE_LEARNING_PRELOAD__CLIP__VISUAL' instead."
+        )
+
+    if preload.facial_recognition_fallback is not None:
+        log.warning(
+            "Deprecated env variable: 'MACHINE_LEARNING_PRELOAD__FACIAL_RECOGNITION'. "
+            "Use 'MACHINE_LEARNING_PRELOAD__FACIAL_RECOGNITION__DETECTION' and "
+            "'MACHINE_LEARNING_PRELOAD__FACIAL_RECOGNITION__RECOGNITION' instead."
+        )
+
 
 def update_state() -> Iterator[None]:
     global active_requests, last_called