mlrun 1.6.4rc2__py3-none-any.whl → 1.7.0rc20__py3-none-any.whl

mlrun/frameworks/pytorch/model_handler.py

@@ -13,7 +13,7 @@
 # limitations under the License.
 #
 import os
-from typing import Dict, List, Tuple, Type, Union
+from typing import Union
 
 import numpy as np
 import torch
@@ -50,9 +50,9 @@ class PyTorchModelHandler(DLModelHandler):
         model: Module = None,
         model_path: str = None,
         model_name: str = None,
-        model_class: Union[Type[Module], str] = None,
-        modules_map: Union[Dict[str, Union[None, str, List[str]]], str] = None,
-        custom_objects_map: Union[Dict[str, Union[str, List[str]]], str] = None,
+        model_class: Union[type[Module], str] = None,
+        modules_map: Union[dict[str, Union[None, str, list[str]]], str] = None,
+        custom_objects_map: Union[dict[str, Union[str, list[str]]], str] = None,
         custom_objects_directory: str = None,
         context: mlrun.MLClientCtx = None,
         **kwargs,
@@ -136,7 +136,7 @@ class PyTorchModelHandler(DLModelHandler):
             )
 
         # Set up the base handler class:
-        super(PyTorchModelHandler, self).__init__(
+        super().__init__(
             model=model,
             model_path=model_path,
             model_name=model_name,
@@ -152,8 +152,8 @@ class PyTorchModelHandler(DLModelHandler):
 
     def set_labels(
         self,
-        to_add: Dict[str, Union[str, int, float]] = None,
-        to_remove: List[str] = None,
+        to_add: dict[str, Union[str, int, float]] = None,
+        to_remove: list[str] = None,
     ):
         """
         Update the labels dictionary of this model artifact. There are required labels that cannot be edited or removed.
@@ -162,14 +162,14 @@ class PyTorchModelHandler(DLModelHandler):
         :param to_remove: A list of labels keys to remove.
         """
         # Update the user's labels:
-        super(PyTorchModelHandler, self).set_labels(to_add=to_add, to_remove=to_remove)
+        super().set_labels(to_add=to_add, to_remove=to_remove)
 
         # Set the required labels:
         self._labels[self._LabelKeys.MODEL_CLASS_NAME] = self._model_class_name
 
     def save(
         self, output_path: str = None, **kwargs
-    ) -> Union[Dict[str, Artifact], None]:
+    ) -> Union[dict[str, Artifact], None]:
         """
         Save the handled model at the given output path.
 
@@ -182,7 +182,7 @@ class PyTorchModelHandler(DLModelHandler):
         :raise MLRunInvalidArgumentError: If an output path was not given, yet a context was not provided in
                                           initialization.
         """
-        super(PyTorchModelHandler, self).save(output_path=output_path)
+        super().save(output_path=output_path)
 
         # Set the output path:
         if output_path is None:
@@ -207,7 +207,7 @@ class PyTorchModelHandler(DLModelHandler):
 
         :raise MLRunInvalidArgumentError: If the model's class is not in the custom objects map.
         """
-        super(PyTorchModelHandler, self).load()
+        super().load()
 
         # Validate the model's class is in the custom objects map:
         if (
@@ -233,10 +233,10 @@ class PyTorchModelHandler(DLModelHandler):
     def to_onnx(
         self,
         model_name: str = None,
-        input_sample: Union[torch.Tensor, Tuple[torch.Tensor, ...]] = None,
-        input_layers_names: List[str] = None,
-        output_layers_names: List[str] = None,
-        dynamic_axes: Dict[str, Dict[int, str]] = None,
+        input_sample: Union[torch.Tensor, tuple[torch.Tensor, ...]] = None,
+        input_layers_names: list[str] = None,
+        output_layers_names: list[str] = None,
+        dynamic_axes: dict[str, dict[int, str]] = None,
         is_batched: bool = True,
         optimize: bool = True,
         output_path: str = None,
@@ -406,7 +406,7 @@ class PyTorchModelHandler(DLModelHandler):
         ]
 
         # Continue collecting from abstract class:
-        super(PyTorchModelHandler, self)._collect_files_from_store_object()
+        super()._collect_files_from_store_object()
 
     def _collect_files_from_local_path(self):
         """
@@ -443,7 +443,7 @@ class PyTorchModelHandler(DLModelHandler):
         """
         # Supported types:
         if isinstance(sample, np.ndarray):
-            return super(PyTorchModelHandler, self)._read_sample(sample=sample)
+            return super()._read_sample(sample=sample)
         elif isinstance(sample, torch.Tensor):
             return Feature(
                 value_type=PyTorchUtils.convert_torch_dtype_to_value_type(

mlrun/frameworks/pytorch/model_server.py

@@ -12,7 +12,7 @@
 # See the License for the specific language governing permissions and
 # limitations under the License.
 #
-from typing import Any, Dict, List, Type, Union
+from typing import Any, Union
 
 import numpy as np
 import torch
@@ -39,9 +39,9 @@ class PyTorchModelServer(V2ModelServer):
         model: Module = None,
         model_path: str = None,
         model_name: str = None,
-        model_class: Union[Type[Module], str] = None,
-        modules_map: Union[Dict[str, Union[None, str, List[str]]], str] = None,
-        custom_objects_map: Union[Dict[str, Union[str, List[str]]], str] = None,
+        model_class: Union[type[Module], str] = None,
+        modules_map: Union[dict[str, Union[None, str, list[str]]], str] = None,
+        custom_objects_map: Union[dict[str, Union[str, list[str]]], str] = None,
         custom_objects_directory: str = None,
         use_cuda: bool = True,
         to_list: bool = False,
@@ -106,7 +106,7 @@ class PyTorchModelServer(V2ModelServer):
         :param protocol:                 -
         :param class_args:               -
         """
-        super(PyTorchModelServer, self).__init__(
+        super().__init__(
             context=context,
             name=name,
             model_path=model_path,
@@ -158,7 +158,7 @@ class PyTorchModelServer(V2ModelServer):
             model=self._model_handler.model, context=self.context
         )
 
-    def predict(self, request: Dict[str, Any]) -> Union[Tensor, list]:
+    def predict(self, request: dict[str, Any]) -> Union[Tensor, list]:
         """
         Infer the inputs through the model using MLRun's PyTorch interface and return its output. The inferred data will
         be read from the "inputs" key of the request.
@@ -183,7 +183,7 @@ class PyTorchModelServer(V2ModelServer):
         # Return as list if required:
         return predictions if not self.to_list else predictions.tolist()
 
-    def explain(self, request: Dict[str, Any]) -> str:
+    def explain(self, request: dict[str, Any]) -> str:
         """
         Return a string explaining what model is being serve in this serving function and the function name.
 

mlrun/frameworks/sklearn/__init__.py

@@ -14,7 +14,7 @@
 #
 # flake8: noqa  - this is until we take care of the F401 violations with respect to __all__ & sphinx
 import warnings
-from typing import Dict, List, Union
+from typing import Union
 
 import mlrun
 from mlrun.frameworks.sklearn.metric import Metric
@@ -37,25 +37,25 @@ def apply_mlrun(
     model_name: str = "model",
     tag: str = "",
     model_path: str = None,
-    modules_map: Union[Dict[str, Union[None, str, List[str]]], str] = None,
-    custom_objects_map: Union[Dict[str, Union[str, List[str]]], str] = None,
+    modules_map: Union[dict[str, Union[None, str, list[str]]], str] = None,
+    custom_objects_map: Union[dict[str, Union[str, list[str]]], str] = None,
     custom_objects_directory: str = None,
     context: mlrun.MLClientCtx = None,
-    artifacts: Union[List[MLPlan], List[str], Dict[str, dict]] = None,
+    artifacts: Union[list[MLPlan], list[str], dict[str, dict]] = None,
     metrics: Union[
-        List[Metric],
-        List[SKLearnTypes.MetricEntryType],
-        Dict[str, SKLearnTypes.MetricEntryType],
+        list[Metric],
+        list[SKLearnTypes.MetricEntryType],
+        dict[str, SKLearnTypes.MetricEntryType],
     ] = None,
     x_test: SKLearnTypes.DatasetType = None,
     y_test: SKLearnTypes.DatasetType = None,
     sample_set: Union[SKLearnTypes.DatasetType, mlrun.DataItem, str] = None,
-    y_columns: Union[List[str], List[int]] = None,
+    y_columns: Union[list[str], list[int]] = None,
     feature_vector: str = None,
-    feature_weights: List[float] = None,
-    labels: Dict[str, Union[str, int, float]] = None,
-    parameters: Dict[str, Union[str, int, float]] = None,
-    extra_data: Dict[str, SKLearnTypes.ExtraDataType] = None,
+    feature_weights: list[float] = None,
+    labels: dict[str, Union[str, int, float]] = None,
+    parameters: dict[str, Union[str, int, float]] = None,
+    extra_data: dict[str, SKLearnTypes.ExtraDataType] = None,
     auto_log: bool = True,
     **kwargs,
 ) -> SKLearnModelHandler:
@@ -92,7 +92,7 @@ def apply_mlrun(
 
                                          {
                                              "/.../custom_model.py": "MyModel",
-                                             "/.../custom_objects.py": ["object1", "object2"]
+                                             "/.../custom_objects.py": ["object1", "object2"],
                                          }
 
                                      All the paths will be accessed from the given 'custom_objects_directory', meaning

mlrun/frameworks/sklearn/estimator.py

@@ -12,7 +12,7 @@
 # See the License for the specific language governing permissions and
 # limitations under the License.
 #
-from typing import Dict, List, Union
+from typing import Union
 
 import numpy as np
 import pandas as pd
@@ -32,7 +32,7 @@ class Estimator:
     def __init__(
         self,
         context: mlrun.MLClientCtx = None,
-        metrics: List[Metric] = None,
+        metrics: list[Metric] = None,
     ):
         """
         Initialize an estimator with the given metrics. The estimator will log the calculated results using the given
@@ -62,7 +62,7 @@ class Estimator:
         return self._context
 
     @property
-    def results(self) -> Dict[str, float]:
+    def results(self) -> dict[str, float]:
         """
         Get the logged results.
 
@@ -86,7 +86,7 @@ class Estimator:
         """
         self._context = context
 
-    def set_metrics(self, metrics: List[Metric]):
+    def set_metrics(self, metrics: list[Metric]):
         """
         Update the metrics of this logger to the given list of metrics here.
 

mlrun/frameworks/sklearn/metrics_library.py

@@ -13,7 +13,7 @@
 # limitations under the License.
 #
 from abc import ABC
-from typing import Dict, List, Union
+from typing import Union
 
 import sklearn
 from sklearn.preprocessing import LabelBinarizer
@@ -40,14 +40,14 @@ class MetricsLibrary(ABC):
     def get_metrics(
         cls,
         metrics: Union[
-            List[Metric],
-            List[SKLearnTypes.MetricEntryType],
-            Dict[str, SKLearnTypes.MetricEntryType],
+            list[Metric],
+            list[SKLearnTypes.MetricEntryType],
+            dict[str, SKLearnTypes.MetricEntryType],
         ] = None,
         context: mlrun.MLClientCtx = None,
         include_default: bool = True,
         **default_kwargs,
-    ) -> List[Metric]:
+    ) -> list[Metric]:
         """
         Get metrics for a run. The metrics will be taken from the provided metrics / configuration via code, from
         provided configuration via MLRun context and if the 'include_default' is True, from the metric library's
@@ -87,11 +87,11 @@ class MetricsLibrary(ABC):
     def _parse(
         cls,
         metrics: Union[
-            List[Metric],
-            List[SKLearnTypes.MetricEntryType],
-            Dict[str, SKLearnTypes.MetricEntryType],
+            list[Metric],
+            list[SKLearnTypes.MetricEntryType],
+            dict[str, SKLearnTypes.MetricEntryType],
         ],
-    ) -> List[Metric]:
+    ) -> list[Metric]:
         """
         Parse the given metrics by the possible rules of the framework implementing.
 
@@ -116,8 +116,8 @@ class MetricsLibrary(ABC):
 
     @classmethod
     def _from_list(
-        cls, metrics_list: List[Union[Metric, SKLearnTypes.MetricEntryType]]
-    ) -> List[Metric]:
+        cls, metrics_list: list[Union[Metric, SKLearnTypes.MetricEntryType]]
+    ) -> list[Metric]:
         """
         Collect the given metrics configurations from a list. The metrics names will be chosen by the following rules:
 
@@ -143,8 +143,8 @@ class MetricsLibrary(ABC):
 
     @classmethod
     def _from_dict(
-        cls, metrics_dictionary: Dict[str, SKLearnTypes.MetricEntryType]
-    ) -> List[Metric]:
+        cls, metrics_dictionary: dict[str, SKLearnTypes.MetricEntryType]
+    ) -> list[Metric]:
         """
         Collect the given metrics configurations from a dictionary.
 
@@ -165,7 +165,7 @@ class MetricsLibrary(ABC):
     @classmethod
     def _default(
         cls, model: SKLearnTypes.ModelType, y: SKLearnTypes.DatasetType = None
-    ) -> List[Metric]:
+    ) -> list[Metric]:
         """
         Get the default metrics list according to the algorithm functionality.
 

mlrun/frameworks/sklearn/mlrun_interface.py

@@ -13,7 +13,6 @@
 # limitations under the License.
 #
 from abc import ABC
-from typing import List
 
 import mlrun
 
@@ -75,9 +74,7 @@ class SKLearnMLRunInterface(MLRunInterface, ABC):
             cls._REPLACED_METHODS.remove("predict_proba")
 
         # Add the interface to the model:
-        super(SKLearnMLRunInterface, cls).add_interface(
-            obj=obj, restoration=restoration
-        )
+        super().add_interface(obj=obj, restoration=restoration)
 
         # Restore the '_REPLACED_METHODS' list for next models:
         if "predict_proba" not in cls._REPLACED_METHODS:
@@ -154,8 +151,8 @@ class SKLearnMLRunInterface(MLRunInterface, ABC):
     def configure_logging(
         self,
         context: mlrun.MLClientCtx = None,
-        plans: List[MLPlan] = None,
-        metrics: List[Metric] = None,
+        plans: list[MLPlan] = None,
+        metrics: list[Metric] = None,
         x_test: SKLearnTypes.DatasetType = None,
         y_test: SKLearnTypes.DatasetType = None,
         model_handler: MLModelHandler = None,

mlrun/frameworks/sklearn/model_handler.py

@@ -59,7 +59,7 @@ class SKLearnModelHandler(MLModelHandler):
 
         :return The saved model additional artifacts (if needed) dictionary if context is available and None otherwise.
         """
-        super(SKLearnModelHandler, self).save(output_path=output_path)
+        super().save(output_path=output_path)
 
         # Save the model pkl file:
         self._model_file = f"{self._model_name}.pkl"
@@ -73,7 +73,7 @@ class SKLearnModelHandler(MLModelHandler):
         Load the specified model in this handler. Additional parameters for the class initializer can be passed via the
         kwargs dictionary.
         """
-        super(SKLearnModelHandler, self).load()
+        super().load()
 
         # Load from a pkl file:
         with open(self._model_file, "rb") as pickle_file:

mlrun/frameworks/tf_keras/__init__.py

@@ -13,11 +13,12 @@
 # limitations under the License.
 #
 # flake8: noqa  - this is until we take care of the F401 violations with respect to __all__ & sphinx
-from typing import Any, Dict, List, Union
+from typing import Any, Union
 
 from tensorflow import keras
 
 import mlrun
+import mlrun.common.constants as mlrun_constants
 
 from .callbacks import MLRunLoggingCallback, TensorboardLoggingCallback
 from .mlrun_interface import TFKerasMLRunInterface
@@ -33,14 +34,14 @@ def apply_mlrun(
     model_path: str = None,
     model_format: str = TFKerasModelHandler.ModelFormats.SAVED_MODEL,
     save_traces: bool = False,
-    modules_map: Union[Dict[str, Union[None, str, List[str]]], str] = None,
-    custom_objects_map: Union[Dict[str, Union[str, List[str]]], str] = None,
+    modules_map: Union[dict[str, Union[None, str, list[str]]], str] = None,
+    custom_objects_map: Union[dict[str, Union[str, list[str]]], str] = None,
     custom_objects_directory: str = None,
     context: mlrun.MLClientCtx = None,
     auto_log: bool = True,
     tensorboard_directory: str = None,
-    mlrun_callback_kwargs: Dict[str, Any] = None,
-    tensorboard_callback_kwargs: Dict[str, Any] = None,
+    mlrun_callback_kwargs: dict[str, Any] = None,
+    tensorboard_callback_kwargs: dict[str, Any] = None,
     use_horovod: bool = None,
     **kwargs,
 ) -> TFKerasModelHandler:
@@ -85,7 +86,7 @@ def apply_mlrun(
 
                                             {
                                                 "/.../custom_optimizer.py": "optimizer",
-                                                "/.../custom_layers.py": ["layer1", "layer2"]
+                                                "/.../custom_layers.py": ["layer1", "layer2"],
                                             }
 
                                         All the paths will be accessed from the given 'custom_objects_directory',
@@ -126,7 +127,9 @@ def apply_mlrun(
     # # Use horovod:
     if use_horovod is None:
         use_horovod = (
-            context.labels.get("kind", "") == "mpijob" if context is not None else False
+            context.labels.get(mlrun_constants.MLRunInternalLabels.kind, "") == "mpijob"
+            if context is not None
+            else False
         )
 
     # Create a model handler:

mlrun/frameworks/tf_keras/callbacks/logging_callback.py

@@ -12,12 +12,12 @@
 # See the License for the specific language governing permissions and
 # limitations under the License.
 #
-from typing import Callable, Dict, List, Union
+from typing import Callable, Union
 
 import numpy as np
 import tensorflow as tf
 from tensorflow import Tensor, Variable
-from tensorflow.keras.callbacks import Callback
+from tensorflow.python.keras.callbacks import Callback
 
 import mlrun
 
@@ -36,11 +36,11 @@ class LoggingCallback(Callback):
     def __init__(
         self,
         context: mlrun.MLClientCtx = None,
-        dynamic_hyperparameters: Dict[
-            str, Union[List[Union[str, int]], Callable[[], TFKerasTypes.TrackableType]]
+        dynamic_hyperparameters: dict[
+            str, Union[list[Union[str, int]], Callable[[], TFKerasTypes.TrackableType]]
         ] = None,
-        static_hyperparameters: Dict[
-            str, Union[TFKerasTypes.TrackableType, List[Union[str, int]]]
+        static_hyperparameters: dict[
+            str, Union[TFKerasTypes.TrackableType, list[Union[str, int]]]
         ] = None,
         auto_log: bool = False,
     ):
@@ -70,7 +70,7 @@ class LoggingCallback(Callback):
         :param auto_log:                Whether or not to enable auto logging, trying to track common static and dynamic
                                         hyperparameters.
         """
-        super(LoggingCallback, self).__init__()
+        super().__init__()
         self._supports_tf_logs = True
 
         # Store the configurations:
@@ -93,7 +93,7 @@ class LoggingCallback(Callback):
         self._is_training = None  # type: bool
         self._auto_log = auto_log
 
-    def get_training_results(self) -> Dict[str, List[List[float]]]:
+    def get_training_results(self) -> dict[str, list[list[float]]]:
         """
         Get the training results logged. The results will be stored in a dictionary where each key is the metric name
         and the value is a list of lists of values. The first list is by epoch and the second list is by iteration
@@ -103,7 +103,7 @@ class LoggingCallback(Callback):
         """
         return self._logger.training_results
 
-    def get_validation_results(self) -> Dict[str, List[List[float]]]:
+    def get_validation_results(self) -> dict[str, list[list[float]]]:
         """
         Get the validation results logged. The results will be stored in a dictionary where each key is the metric name
         and the value is a list of lists of values. The first list is by epoch and the second list is by iteration
@@ -113,7 +113,7 @@ class LoggingCallback(Callback):
         """
         return self._logger.validation_results
 
-    def get_training_summaries(self) -> Dict[str, List[float]]:
+    def get_training_summaries(self) -> dict[str, list[float]]:
         """
         Get the training summaries of the metrics results. The summaries will be stored in a dictionary where each key
         is the metric names and the value is a list of all the summary values per epoch.
@@ -122,7 +122,7 @@ class LoggingCallback(Callback):
         """
         return self._logger.training_summaries
 
-    def get_validation_summaries(self) -> Dict[str, List[float]]:
+    def get_validation_summaries(self) -> dict[str, list[float]]:
         """
         Get the validation summaries of the metrics results. The summaries will be stored in a dictionary where each key
         is the metric names and the value is a list of all the summary values per epoch.
@@ -131,7 +131,7 @@ class LoggingCallback(Callback):
         """
         return self._logger.validation_summaries
 
-    def get_static_hyperparameters(self) -> Dict[str, TFKerasTypes.TrackableType]:
+    def get_static_hyperparameters(self) -> dict[str, TFKerasTypes.TrackableType]:
         """
         Get the static hyperparameters logged. The hyperparameters will be stored in a dictionary where each key is the
         hyperparameter name and the value is his logged value.
@@ -142,7 +142,7 @@ class LoggingCallback(Callback):
 
     def get_dynamic_hyperparameters(
         self,
-    ) -> Dict[str, List[TFKerasTypes.TrackableType]]:
+    ) -> dict[str, list[TFKerasTypes.TrackableType]]:
         """
         Get the dynamic hyperparameters logged. The hyperparameters will be stored in a dictionary where each key is the
         hyperparameter name and the value is a list of his logged values per epoch.
@@ -329,7 +329,7 @@ class LoggingCallback(Callback):
 
         # Static hyperparameters:
         for name, value in self._static_hyperparameters_keys.items():
-            if isinstance(value, List):
+            if isinstance(value, list):
                 # Its a parameter that needed to be extracted via key chain.
                 self._logger.log_static_hyperparameter(
                     parameter_name=name,
@@ -398,7 +398,7 @@ class LoggingCallback(Callback):
     def _get_hyperparameter(
         self,
         key_chain: Union[
-            Callable[[], TFKerasTypes.TrackableType], List[Union[str, int]]
+            Callable[[], TFKerasTypes.TrackableType], list[Union[str, int]]
         ],
     ) -> TFKerasTypes.TrackableType:
         """

mlrun/frameworks/tf_keras/callbacks/mlrun_logging_callback.py

@@ -12,7 +12,7 @@
 # See the License for the specific language governing permissions and
 # limitations under the License.
 #
-from typing import Callable, Dict, List, Union
+from typing import Callable, Union
 
 import mlrun
 from mlrun.artifacts import Artifact
@@ -50,16 +50,16 @@ class MLRunLoggingCallback(LoggingCallback):
         context: mlrun.MLClientCtx,
         model_handler: TFKerasModelHandler,
         log_model_tag: str = "",
-        log_model_labels: Dict[str, TFKerasTypes.TrackableType] = None,
-        log_model_parameters: Dict[str, TFKerasTypes.TrackableType] = None,
-        log_model_extra_data: Dict[
+        log_model_labels: dict[str, TFKerasTypes.TrackableType] = None,
+        log_model_parameters: dict[str, TFKerasTypes.TrackableType] = None,
+        log_model_extra_data: dict[
             str, Union[TFKerasTypes.TrackableType, Artifact]
         ] = None,
-        dynamic_hyperparameters: Dict[
-            str, Union[List[Union[str, int]], Callable[[], TFKerasTypes.TrackableType]]
+        dynamic_hyperparameters: dict[
+            str, Union[list[Union[str, int]], Callable[[], TFKerasTypes.TrackableType]]
         ] = None,
-        static_hyperparameters: Dict[
-            str, Union[TFKerasTypes, List[Union[str, int]]]
+        static_hyperparameters: dict[
+            str, Union[TFKerasTypes, list[Union[str, int]]]
         ] = None,
         auto_log: bool = False,
     ):
@@ -97,7 +97,7 @@ class MLRunLoggingCallback(LoggingCallback):
                                          trying to track common static and dynamic hyperparameters such as learning
                                          rate.
         """
-        super(MLRunLoggingCallback, self).__init__(
+        super().__init__(
             dynamic_hyperparameters=dynamic_hyperparameters,
             static_hyperparameters=static_hyperparameters,
             auto_log=auto_log,
@@ -134,7 +134,7 @@ class MLRunLoggingCallback(LoggingCallback):
         :param logs: Currently no data is passed to this argument for this method but that may change in the
                      future.
         """
-        super(MLRunLoggingCallback, self).on_test_end(logs=logs)
+        super().on_test_end(logs=logs)
 
         # Check if its part of evaluation. If so, end the run:
         if self._logger.mode == LoggingMode.EVALUATION:
@@ -151,7 +151,7 @@ class MLRunLoggingCallback(LoggingCallback):
                       performed. Validation result keys are prefixed with `val_`. For training epoch, the values of the
                       `Model`'s metrics are returned. Example : `{'loss': 0.2, 'acc': 0.7}`.
         """
-        super(MLRunLoggingCallback, self).on_epoch_end(epoch=epoch)
+        super().on_epoch_end(epoch=epoch)
 
         # Log the current epoch's results:
         self._logger.log_epoch_to_context(epoch=epoch)