Bug/gsk 254 auc test throws error for multiclass model with small dataset

giskard-ml-worker/logging_config.ini

@@ -1,5 +1,5 @@
 [loggers]
-keys = root
+keys = root,timer
 
 [handlers]
 keys = stream_handler
@@ -11,6 +11,13 @@ keys = formatter
 level = DEBUG
 handlers = stream_handler
 
+[logger_timer]
+level = DEBUG
+qualname=timer
+handlers = stream_handler
+propagate=0
+
+
 [handler_stream_handler]
 class = StreamHandler
 level = DEBUG

giskard-ml-worker/ml_worker/core/model.py

@@ -1,12 +1,15 @@
 import logging
+from timeit import timeit
 from typing import List, Any, Optional, Callable, Iterable, Union
+import time
 
 import numpy
 import pandas as pd
 from builtins import Exception
 from pydantic import BaseModel
 
 from ml_worker.core.giskard_dataset import GiskardDataset
+from ml_worker.utils.logging import timer, Timer
 
 
 class ModelPredictionResults(BaseModel):
@@ -31,6 +34,7 @@ def __init__(self,
         self.classification_labels = classification_labels
 
     def run_predict(self, dataset: GiskardDataset):
+        timer = Timer()
         df = self.prepare_dataframe(dataset)
         raw_prediction = self.prediction_function(df)
         if self.model_type == "regression":
@@ -62,6 +66,7 @@ def run_predict(self, dataset: GiskardDataset):
             raise ValueError(
                 f"Prediction task is not supported: {self.model_type}"
             )
+        timer.stop(f"Predicted dataset with shape {dataset.df.shape}")
         return result
 
     def prepare_dataframe(self, dataset):

giskard-ml-worker/ml_worker/server/ml_worker_service.py

@@ -1,5 +1,7 @@
 import logging
 import re
+import time
+
 
 import grpc
 import numpy as np
@@ -15,6 +17,7 @@
 from ml_worker.exceptions.IllegalArgumentError import IllegalArgumentError
 from ml_worker.exceptions.giskard_exception import GiskardException
 from ml_worker.utils.grpc_mapper import deserialize_model, deserialize_dataset
+from ml_worker.utils.logging import Timer
 
 logger = logging.getLogger()
 
@@ -25,7 +28,6 @@ def __init__(self) -> None:
 
     def runTest(self, request: RunTestRequest, context: grpc.ServicerContext) -> TestResultMessage:
         from ml_worker.testing.functions import GiskardTestFunctions
-
         model = deserialize_model(request.model)
 
         tests = GiskardTestFunctions()
@@ -38,7 +40,9 @@ def runTest(self, request: RunTestRequest, context: grpc.ServicerContext) -> Tes
         if request.actual_ds.serialized_df:
             _globals['actual_ds'] = deserialize_dataset(request.actual_ds)
         try:
+            timer = Timer()
             exec(request.code, _globals)
+            timer.stop(f"Test {tests.tests_results[0].name}")
         except NameError as e:
             missing_name = re.findall(r"name '(\w+)' is not defined", str(e))[0]
             if missing_name == 'reference_ds':

giskard-ml-worker/ml_worker/testing/metamorphic_tests.py

@@ -1,11 +1,13 @@
 import pandas as pd
-
+import time
+import logging
 from generated.ml_worker_pb2 import SingleTestResult, TestMessage, TestMessageType
 
 from ml_worker.core.giskard_dataset import GiskardDataset
 from ml_worker.core.model import GiskardModel
 from ml_worker.testing.abstract_test_collection import AbstractTestCollection
 from ml_worker.testing.utils import apply_perturbation_inplace
+from ml_worker.utils.logging import timer
 
 
 class MetamorphicTests(AbstractTestCollection):
@@ -40,6 +42,7 @@ def _perturb_and_predict(ds: GiskardDataset, model: GiskardModel, perturbation_d
             results_df["perturbed_prediction"] = results_df["prediction"]
         return results_df, len(modified_rows)
 
+    @timer("Compare and predict the data")
     def _compare_prediction(self, results_df, prediction_task, output_sensitivity=None, flag=None):
         if flag == 'Invariance':
             if prediction_task == 'classification':

giskard-ml-worker/ml_worker/testing/performance_tests.py

@@ -44,8 +44,12 @@ def test_auc(self, actual_slice: GiskardDataset, model: GiskardModel, threshold=
             metric = roc_auc_score(actual_slice.df[actual_slice.target],
                                    model.run_predict(actual_slice).raw_prediction)
         else:
-            metric = roc_auc_score(actual_slice.df[actual_slice.target],
-                                   model.run_predict(actual_slice).all_predictions, multi_class='ovr')
+            predictions = model.run_predict(actual_slice).all_predictions
+            non_declared_categories = set(predictions.columns) - set(actual_slice.df[actual_slice.target].unique())
+            assert not len(non_declared_categories), \
+                f"Predicted classes don't exist in the dataset \"{actual_slice.target}\" column: {non_declared_categories}"
+
+            metric = roc_auc_score(actual_slice.df[actual_slice.target], predictions, multi_class='ovr')
 
         return self.save_results(
             SingleTestResult(

giskard-ml-worker/ml_worker/testing/utils.py

@@ -1,6 +1,9 @@
 import pandas as pd
+import time
+import logging
 
 from generated.ml_worker_pb2 import SingleTestResult
+from ml_worker.utils.logging import timer
 
 
 def ge_result_to_test_result(result, passed=True) -> SingleTestResult:
@@ -25,6 +28,7 @@ def ge_result_to_test_result(result, passed=True) -> SingleTestResult:
     )
 
 
+@timer("Perturb data")
 def apply_perturbation_inplace(df: pd.DataFrame, perturbation_dict):
     modified_rows = []
     i = 0

giskard-ml-worker/ml_worker/utils/grpc_mapper.py

@@ -1,27 +1,34 @@
+import logging
 from io import BytesIO
 
 import cloudpickle
 import pandas as pd
+import time
 from zstandard import ZstdDecompressor
 from zstandard import decompress
 
+from generated.ml_worker_pb2 import SerializedGiskardModel, SerializedGiskardDataset
 from ml_worker.core.giskard_dataset import GiskardDataset
 from ml_worker.core.model import GiskardModel
-from generated.ml_worker_pb2 import SerializedGiskardModel, SerializedGiskardDataset
+from ml_worker.utils.logging import timer
 
 
+@timer()
 def deserialize_model(serialized_model: SerializedGiskardModel) -> GiskardModel:
-    return GiskardModel(
+    deserialized_model = GiskardModel(
         cloudpickle.load(ZstdDecompressor().stream_reader(serialized_model.serialized_prediction_function)),
         model_type=serialized_model.model_type,
         classification_threshold=serialized_model.threshold.value if serialized_model.HasField('threshold') else None,
         feature_names=list(serialized_model.feature_names),
         classification_labels=list(serialized_model.classification_labels)
     )
 
+    return deserialized_model
 
+
+@timer()
 def deserialize_dataset(serialized_dataset: SerializedGiskardDataset) -> GiskardDataset:
-    return GiskardDataset(
+    deserialized_dataset = GiskardDataset(
         df=pd.read_csv(
             BytesIO(decompress(serialized_dataset.serialized_df)),
             keep_default_na=False,
@@ -31,3 +38,5 @@ def deserialize_dataset(serialized_dataset: SerializedGiskardDataset) -> Giskard
         feature_types=dict(serialized_dataset.feature_types),
         column_types=dict(serialized_dataset.column_types)
     )
+
+    return deserialized_dataset

giskard-ml-worker/ml_worker/utils/logging.py

@@ -1,6 +1,11 @@
+import logging
 import os
-import sys
+from datetime import timedelta
+from functools import wraps
 from logging.config import fileConfig
+from timeit import default_timer
+
+import sys
 
 
 def resolve(filename):
@@ -21,3 +26,66 @@ def load_logging_config():
             print(f"Failed to load logging config from {config_path}")
     else:
         fileConfig(resolve('logging_config.ini'))
+
+
+timer_logger = logging.getLogger("timer")
+
+
+class Timer:
+    message_template: str
+
+    def __init__(self, message=None, start=True, level=logging.INFO) -> None:
+        if start:
+            self.start()
+        else:
+            self.__start_time = None
+        self.duration = None
+        self.message_template = message
+        self.message = None
+        self.level = level
+
+    def start(self):
+        self.__start_time = default_timer()
+
+    def stop(self, message=None, log=True):
+        if message:
+            self.message_template = message
+        if not self.__start_time:
+            timer_logger.error("Timer was not started")
+            return
+
+        self.duration = timedelta(seconds=default_timer() - self.__start_time)
+        if log:
+            self.message = self.create_message(self.duration)
+            timer_logger.log(self.level, self.message)
+        return self.duration
+
+    def __enter__(self):
+        if not self.__start_time:
+            self.start()
+
+    def __exit__(self, exc_type, exc_val, exc_tb):
+        self.stop()
+
+    def create_message(self, duration):
+        if self.message_template:
+            return self.message_template.strip() + f" executed in {duration}"
+        else:
+            return f"Executed in {duration}"
+
+    def prepare_message_template(self):
+        return self.message_template
+
+
+def timer(message=None):
+    @wraps(timer)
+    def timing_decorator(fn):
+        @wraps(timing_decorator)
+        def wrap(*args, **kw):
+            with Timer(message if message else f'{fn.__module__}.{fn.__qualname__}'):
+                result = fn(*args, **kw)
+            return result
+
+        return wrap
+
+    return timing_decorator

giskard-ml-worker/test/fixtures/enron_multilabel_classification.py

@@ -15,6 +15,7 @@
 
 from ml_worker.core.giskard_dataset import GiskardDataset
 from ml_worker.core.model import GiskardModel
+from ml_worker.utils.logging import timer, Timer
 from test import path
 
 input_types = {
@@ -51,7 +52,7 @@ def enron_test_data(enron_data):
 
 @pytest.fixture()
 def enron_model(enron_data) -> GiskardModel:
-    start = time.time()
+    timer = Timer()
 
     columns_to_scale = [key for key in input_types.keys() if input_types[key] == "numeric"]
 
@@ -87,9 +88,8 @@ def enron_model(enron_data) -> GiskardModel:
                                                                         stratify=Y)
     clf.fit(X_train, Y_train)
 
-    train_time = time.time() - start
     model_score = clf.score(X_test, Y_test)
-    logging.info(f"Trained model with score: {model_score} in {round(train_time * 1000)} ms")
+    timer.stop(f"Trained model with score: {model_score}")
 
     return GiskardModel(
         prediction_function=clf.predict_proba,