adding container to update models on a schedule (#46)

* adding container to update models on a schedule

* Automatically reformatting code with black and isort

* dockerfile

* clean up

* update deployment

* udpate images

* Automatically reformatting code with black and isort

* update cicd

* clean up

* pr feedback

* reorganize stuff

* update image

* try/except block

* dict method error solved

* new image

* Automatically reformatting code with black and isort

* remove unnecessary change

* error log if refresh rates are different

* update image

---------

Co-authored-by: Auto-format Bot <autoformatbot@groundlight.ai>

Dockerfile

@@ -74,6 +74,10 @@ WORKDIR ${APP_ROOT}
 
 # Copy the remaining files
 COPY /app ${APP_ROOT}/app/
+
+# Copy model updating code 
+COPY model_updater ${APP_ROOT}/model_updater
+
 COPY --from=production-dependencies-build-stage ${APP_ROOT}/configs/nginx.conf /etc/nginx/nginx.conf
 
 # Remove default nginx config

app/core/configs.py

@@ -32,11 +32,11 @@ class LocalInferenceConfig(BaseModel):
     """
 
     enabled: bool = Field(False, description="Determines if local edge inference is enabled for a specific detector.")
-    refresh_every: float = Field(
-        3600.0,
+    refresh_rate: float = Field(
+        120.0,
         description=(
             "The refresh rate for the inference server (in seconds). This means how often to check for an updated model"
-            " binary (currently unused)."
+            " binary."
         ),
     )
 
@@ -80,7 +80,7 @@ def validate_templates(
                 'local_inference_templates': {
                     'default': LocalInferenceConfig(
                                     enabled=True,
-                                    refresh_every=3600.0
+                                    refresh_rate=120.0
                                 )
                 }
             }

app/main.py

@@ -18,19 +18,3 @@
 app.include_router(router=ping_router)
 
 app.state.app_state = AppState()
-
-
-@app.on_event("startup")
-async def on_startup():
-    """
-    On startup, update edge inference models.
-    """
-    if not os.environ.get("DEPLOY_DETECTOR_LEVEL_INFERENCE", None):
-        return
-
-    for detector_id, inference_config in app.state.app_state.edge_inference_manager.inference_config.items():
-        if inference_config.enabled:
-            try:
-                app.state.app_state.edge_inference_manager.update_model(detector_id)
-            except Exception:
-                logging.error(f"Failed to update model for {detector_id}", exc_info=True)

configs/edge-config.yaml

@@ -26,11 +26,7 @@ local_inference_templates:
   default:
     enabled: true
     # How often to fetch a new model binary (in seconds)
-    refresh_every: 3600
-
-  super-fast-refresh:
-    enabled: true
-    refresh_every: 60
+    refresh_rate: 120
 
   disabled:
     enabled: false

deploy/k3s/edge_deployment/edge_deployment.yaml

@@ -40,7 +40,7 @@ spec:
       serviceAccountName: edge-endpoint-service-account
       containers:
       - name: edge-endpoint
-        image: 723181461334.dkr.ecr.us-west-2.amazonaws.com/edge-endpoint:13361241c-state-management-via-k3s
+        image: 723181461334.dkr.ecr.us-west-2.amazonaws.com/edge-endpoint:6ef532de5-model-updater
         imagePullPolicy: IfNotPresent
         ports:
         - containerPort: 6717
@@ -56,12 +56,33 @@ spec:
               key: api-token
         volumeMounts:
         - name: edge-config-volume
-          # This is a path inside the container not the host
           mountPath: /etc/groundlight/edge-config
         - name: inference-deployment-template-volume
           mountPath: /etc/groundlight/inference-deployment 
         - name: model-repo
           mountPath: /mnt/models
+
+      - name: inference-model-updater
+        image: 723181461334.dkr.ecr.us-west-2.amazonaws.com/edge-endpoint:6ef532de5-model-updater
+        imagePullPolicy: IfNotPresent
+        command: ["/bin/bash", "-c"]
+        args: ["poetry run python -m model_updater.update_models"]
+        env:
+        - name: LOG_LEVEL
+          value: "INFO"
+        - name: DEPLOY_DETECTOR_LEVEL_INFERENCE
+          value: "True"
+        - name: GROUNDLIGHT_API_TOKEN
+          valueFrom:
+            secretKeyRef:
+              name: groundlight-secrets
+              key: api-token
+        volumeMounts:
+        - name: edge-config-volume
+          mountPath: /etc/groundlight/edge-config
+        - name: model-repo
+          mountPath: /mnt/models   
+
       imagePullSecrets:
       - name: registry-credentials
       volumes:

model_updater/update_models.py

@@ -0,0 +1,45 @@
+import os
+import logging
+import time
+from app.core.app_state import load_edge_config
+from app.core.configs import RootEdgeConfig
+from app.core.edge_inference import EdgeInferenceManager
+
+log_level = os.environ.get("LOG_LEVEL", "INFO").upper()
+logging.basicConfig(level=log_level)
+
+
+def update_models(edge_inference_manager: EdgeInferenceManager):
+    if not os.environ.get("DEPLOY_DETECTOR_LEVEL_INFERENCE", None) or not edge_inference_manager.inference_config:
+        return
+
+    inference_config = edge_inference_manager.inference_config
+
+    # All detectors should have the same refresh rate.
+    refresh_rates = [config.refresh_rate for config in inference_config.values()]
+    if set(refresh_rates) != 1:
+        logging.error(f"Detectors have different refresh rates.")
+
+    refresh_rate = refresh_rates[0]
+
+    while True:
+        for detector_id, config in inference_config.items():
+            if config.enabled:
+                try:
+                    edge_inference_manager.update_model(detector_id=detector_id)
+                except Exception as e:
+                    logging.error(f"Failed to update model for {detector_id}. {e}", exc_info=True)
+
+        time.sleep(refresh_rate)
+
+
+if __name__ == "__main__":
+    edge_config: RootEdgeConfig = load_edge_config()
+    edge_inference_templates = edge_config.local_inference_templates
+    inference_config = {
+        detector.detector_id: edge_inference_templates[detector.local_inference_template]
+        for detector in edge_config.detectors
+    }
+    edge_inference_manager = EdgeInferenceManager(config=inference_config, verbose=True)
+
+    update_models(edge_inference_manager=edge_inference_manager)

test/setup_inference_test_env.sh

@@ -20,6 +20,7 @@ motion_detection_templates:
 local_inference_templates:
   default:
     enabled: true 
+    refresh_rate: 120
   disabled:
     enabled: false