From 9b580a4fd17fdbffe5efcad772b381b2eefe96d8 Mon Sep 17 00:00:00 2001
From: blaise <blaise@groundlight.ai>
Date: Mon, 22 Jan 2024 23:56:45 +0000
Subject: [PATCH] update edge inference model repository field

---
 app/core/edge_inference.py | 5 +++--
 deploy/bin/setup_db.sh     | 2 +-
 2 files changed, 4 insertions(+), 3 deletions(-)

diff --git a/app/core/edge_inference.py b/app/core/edge_inference.py
index 72e8afc9..6af5973c 100644
--- a/app/core/edge_inference.py
+++ b/app/core/edge_inference.py
@@ -23,6 +23,7 @@ class EdgeInferenceManager:
     INPUT_IMAGE_NAME = "image"
     MODEL_OUTPUTS = ["score", "confidence", "probability", "label"]
     INFERENCE_SERVER_URL = "inference-service:8000"
+    MODEL_REPOSITORY = MODEL_REPOSITORY_PATH
 
     def __init__(self, config: Dict[str, LocalInferenceConfig] | None, verbose: bool = False) -> None:
         """
@@ -163,7 +164,7 @@ def update_model(self, detector_id: str) -> bool:
         if cloud_binary_ksuid is None:
             logger.warning(f"No model binary ksuid returned for {detector_id}")
 
-        model_dir = os.path.join(MODEL_REPOSITORY_PATH, detector_id)
+        model_dir = os.path.join(self.MODEL_REPOSITORY, detector_id)
         edge_binary_ksuid = get_current_model_ksuid(model_dir)
         if edge_binary_ksuid and cloud_binary_ksuid is not None and cloud_binary_ksuid <= edge_binary_ksuid:
             logger.info(f"No new model available for {detector_id}")
@@ -179,7 +180,7 @@ def update_model(self, detector_id: str) -> bool:
             model_buffer,
             pipeline_config,
             binary_ksuid=cloud_binary_ksuid,
-            repository_root=MODEL_REPOSITORY_PATH,
+            repository_root=self.MODEL_REPOSITORY,
         )
         return True
 
diff --git a/deploy/bin/setup_db.sh b/deploy/bin/setup_db.sh
index 7d58a891..c8a7034f 100755
--- a/deploy/bin/setup_db.sh
+++ b/deploy/bin/setup_db.sh
@@ -1,4 +1,4 @@
-!/bin/bash
+#!/bin/bash
 
 # Sets up the SQLite database for the edge endpoint.
 # Expects the following environment variables: