Refine MVAD sample

sdk/anomalydetector/azure-ai-anomalydetector/samples/sample_multivariate_detect.py

@@ -21,15 +21,16 @@
 
 import os
 import time
-from datetime import datetime
+from datetime import datetime, timezone
 
 from azure.ai.anomalydetector import AnomalyDetectorClient
 from azure.ai.anomalydetector.models import DetectionRequest, ModelInfo
+from azure.ai.anomalydetector.models import ModelStatus, DetectionStatus
 from azure.core.credentials import AzureKeyCredential
 from azure.core.exceptions import HttpResponseError
 
 
-class MultivariateSample():
+class MultivariateSample:
 
     def __init__(self, subscription_key, anomaly_detector_endpoint, data_source=None):
         self.sub_key = subscription_key
@@ -43,40 +44,50 @@ def __init__(self, subscription_key, anomaly_detector_endpoint, data_source=None
 
         self.data_source = data_source
 
-    def train(self, start_time, end_time, max_tryout=500):
+    def train(self, start_time, end_time):
 
         # Number of models available now
         model_list = list(self.ad_client.list_multivariate_model(skip=0, top=10000))
         print("{:d} available models before training.".format(len(model_list)))
 
         # Use sample data to train the model
-        print("Training new model...")
+        print("Training new model...(it may take a few minutes)")
         data_feed = ModelInfo(start_time=start_time, end_time=end_time, source=self.data_source)
         response_header = \
             self.ad_client.train_multivariate_model(data_feed, cls=lambda *args: [args[i] for i in range(len(args))])[
                 -1]
         trained_model_id = response_header['Location'].split("/")[-1]
 
-        # Model list after training
-        new_model_list = list(self.ad_client.list_multivariate_model(skip=0, top=10000))
-
         # Wait until the model is ready. It usually takes several minutes
         model_status = None
-        tryout_count = 0
-        while (tryout_count < max_tryout and model_status != "READY"):
-            model_status = self.ad_client.get_multivariate_model(trained_model_id).model_info.status
-            tryout_count += 1
-            time.sleep(2)
 
-        assert model_status == "READY"
+        while model_status != ModelStatus.READY and model_status != ModelStatus.FAILED:
+            model_info = self.ad_client.get_multivariate_model(trained_model_id).model_info
+            model_status = model_info.status
+            time.sleep(1)
+
+        if model_status == ModelStatus.FAILED:
+            print("Creating model failed.")
+            print("Errors:")
+            if model_info.errors:
+                for error in model_info.errors:
+                    print("Error code: {}. Message: {}".format(error.code, error.message))
+            else:
+                print("None")
+            return None
+
+        if model_status == ModelStatus.READY:
+            # Model list after training
+            new_model_list = list(self.ad_client.list_multivariate_model(skip=0, top=10000))
+
+            print("Done.\n--------------------")
+            print("{:d} available models after training.".format(len(new_model_list)))
 
-        print("Done.", "\n--------------------")
-        print("{:d} available models after training.".format(len(new_model_list)))
+            # Return the latest model id
+            return trained_model_id
 
-        # Return the latest model id
-        return trained_model_id
 
-    def detect(self, model_id, start_time, end_time, max_tryout=500):
+    def detect(self, model_id, start_time, end_time):
 
         # Detect anomaly in the same data source (but a different interval)
         try:
@@ -87,14 +98,20 @@ def detect(self, model_id, start_time, end_time, max_tryout=500):
 
             # Get results (may need a few seconds)
             r = self.ad_client.get_detection_result(result_id)
-            tryout_count = 0
-            while r.summary.status != "READY" and tryout_count < max_tryout:
-                time.sleep(1)
+            print("Get detection result...(it may take a few seconds)")
+
+            while r.summary.status != DetectionStatus.READY and r.summary.status != DetectionStatus.FAILED:
                 r = self.ad_client.get_detection_result(result_id)
-                tryout_count += 1
+                time.sleep(1)
 
-            if r.summary.status != "READY":
-                print("Request timeout after %d tryouts.".format(max_tryout))
+            if r.summary.status == DetectionStatus.FAILED:
+                print("Detection failed.")
+                print("Errors:")
+                if r.summary.errors:
+                    for error in r.summary.errors:
+                        print("Error code: {}. Message: {}".format(error.code, error.message))
+                else:
+                    print("None")
                 return None
 
         except HttpResponseError as e:
@@ -138,10 +155,15 @@ def delete_model(self, model_id):
     sample = MultivariateSample(SUBSCRIPTION_KEY, ANOMALY_DETECTOR_ENDPOINT, data_source)
 
     # Train a new model
-    model_id = sample.train(datetime(2021, 1, 1, 0, 0, 0), datetime(2021, 1, 2, 12, 0, 0))
+    model_id = sample.train(datetime(2021, 1, 1, 0, 0, 0, tzinfo=timezone.utc),
+                            datetime(2021, 1, 2, 12, 0, 0, tzinfo=timezone.utc))
+    assert model_id is not None
 
     # Reference
-    result = sample.detect(model_id, datetime(2021, 1, 2, 12, 0, 0), datetime(2021, 1, 3, 0, 0, 0))
+    result = sample.detect(model_id, datetime(2021, 1, 2, 12, 0, 0, tzinfo=timezone.utc),
+                           datetime(2021, 1, 3, 0, 0, 0, tzinfo=timezone.utc))
+    assert result is not None
+
     print("Result ID:\t", result.result_id)
     print("Result summary:\t", result.summary)
     print("Result length:\t", len(result.results))