From d5b616b2d835ded5abbecc23281be1b0eb2cf3e4 Mon Sep 17 00:00:00 2001
From: huangbin <huangbin58@huawei.com>
Date: Wed, 27 Nov 2024 17:32:45 +0800
Subject: [PATCH] add detect type for usad model

---
 anteater/model/detector/usad_detector.py | 19 ++++++++++++-------
 config/module/usad_model.job.json        |  3 ++-
 2 files changed, 14 insertions(+), 8 deletions(-)

diff --git a/anteater/model/detector/usad_detector.py b/anteater/model/detector/usad_detector.py
index 567a77c..142ccef 100644
--- a/anteater/model/detector/usad_detector.py
+++ b/anteater/model/detector/usad_detector.py
@@ -62,6 +62,7 @@ class UsadDetector(Detector):
         self.cause_list = None
         self.candidates = {}
         self.machine_training = {}
+        self.detect_type = config.params.get("detect_type", "machine")
 
     def detect_machine_kpis(self, kpis: List[KPI], features: List[Feature], machine_id, key):
         """Executes anomaly detection on kpis"""
@@ -301,13 +302,17 @@ class UsadDetector(Detector):
 
         entity_ids = []
         entity_keys = []
-        # machine_ids = self.data_loader.get_unique_machines(start, end, metrics)
-        # entity_keys.extend(['machine_id'] * len(machine_ids))
-        # entity_ids.extend(machine_ids)
-
-        pod_ids = self.data_loader.get_unique_pods(start, end, metrics)
-        entity_keys.extend(['pod_id'] * len(pod_ids))
-        entity_ids.extend(pod_ids)
+        
+        if self.detect_type == "machine":
+            # machine
+            machine_ids = self.data_loader.get_unique_machines(start, end, metrics)
+            entity_keys.extend(['machine_id'] * len(machine_ids))
+            entity_ids.extend(machine_ids)
+        else:
+            # pod
+            pod_ids = self.data_loader.get_unique_pods(start, end, metrics)
+            entity_keys.extend(['pod_id'] * len(pod_ids))
+            entity_ids.extend(pod_ids)
 
         if not entity_ids:
             logger.warning('Empty entity_ids, RETURN!')
diff --git a/config/module/usad_model.job.json b/config/module/usad_model.job.json
index a3ff4e6..54bb32c 100644
--- a/config/module/usad_model.job.json
+++ b/config/module/usad_model.job.json
@@ -73,6 +73,7 @@
   "model_config": {
     "name": "usad_model",
     "params": {
+      "detect_type": "machine",
       "th": 0.5,
       "max_error_rate": 0.7,
       "min_train_hours": 12,
@@ -419,4 +420,4 @@
     "metric": "gala_gopher_nic_tc_sent_drop"
   }
 ]
-}
\ No newline at end of file
+}
-- 
2.43.0

