mlrun 1.6.4rc7__py3-none-any.whl → 1.7.0__py3-none-any.whl

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:
@@ -100,7 +97,7 @@ class SKLearnMLRunInterface(MLRunInterface, ABC):
 
         def wrapper(
             self: SKLearnTypes.ModelType,
-            X: SKLearnTypes.DatasetType,
+            X: SKLearnTypes.DatasetType,  # noqa: N803 - should be lowercase "x", kept for BC
             y: SKLearnTypes.DatasetType = None,
             *args,
             **kwargs,
@@ -127,7 +124,12 @@ class SKLearnMLRunInterface(MLRunInterface, ABC):
 
         return wrapper
 
-    def mlrun_predict(self, X: SKLearnTypes.DatasetType, *args, **kwargs):
+    def mlrun_predict(
+        self,
+        X: SKLearnTypes.DatasetType,  # noqa: N803 - should be lowercase "x", kept for BC
+        *args,
+        **kwargs,
+    ):
         """
         MLRun's wrapper for the common ML API predict method.
         """
@@ -139,7 +141,12 @@ class SKLearnMLRunInterface(MLRunInterface, ABC):
 
         return y_pred
 
-    def mlrun_predict_proba(self, X: SKLearnTypes.DatasetType, *args, **kwargs):
+    def mlrun_predict_proba(
+        self,
+        X: SKLearnTypes.DatasetType,  # noqa: N803 - should be lowercase "x", kept for BC
+        *args,
+        **kwargs,
+    ):
         """
         MLRun's wrapper for the common ML API predict_proba method.
         """
@@ -154,8 +161,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)

mlrun/frameworks/tf_keras/callbacks/tensorboard_logging_callback.py

@@ -13,7 +13,7 @@
 # limitations under the License.
 #
 from datetime import datetime
-from typing import Callable, Dict, List, Union
+from typing import Callable, Union
 
 import tensorflow as tf
 from packaging import version
@@ -38,7 +38,7 @@ class _TFKerasTensorboardLogger(TensorboardLogger):
 
     def __init__(
         self,
-        statistics_functions: List[Callable[[Union[Variable]], Union[float, Variable]]],
+        statistics_functions: list[Callable[[Union[Variable]], Union[float, Variable]]],
         context: mlrun.MLClientCtx = None,
         tensorboard_directory: str = None,
         run_name: str = None,
@@ -67,7 +67,7 @@ class _TFKerasTensorboardLogger(TensorboardLogger):
                                       update. Notice that writing to tensorboard too frequently may cause the training
                                       to be slower. Default: 'epoch'.
         """
-        super(_TFKerasTensorboardLogger, self).__init__(
+        super().__init__(
             statistics_functions=statistics_functions,
             context=context,
             tensorboard_directory=tensorboard_directory,
@@ -255,15 +255,15 @@ class TensorboardLoggingCallback(LoggingCallback):
         context: mlrun.MLClientCtx = None,
         tensorboard_directory: str = None,
         run_name: str = None,
-        weights: Union[bool, List[str]] = False,
-        statistics_functions: List[
+        weights: Union[bool, list[str]] = False,
+        statistics_functions: list[
             Callable[[Union[Variable, Tensor]], Union[float, Tensor]]
         ] = 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,
         update_frequency: Union[int, str] = "epoch",
         auto_log: bool = False,
@@ -325,7 +325,7 @@ class TensorboardLoggingCallback(LoggingCallback):
         :raise MLRunInvalidArgumentError: In case both 'context' and 'tensorboard_directory' parameters were not given
                                           or the 'update_frequency' was incorrect.
         """
-        super(TensorboardLoggingCallback, self).__init__(
+        super().__init__(
             dynamic_hyperparameters=dynamic_hyperparameters,
             static_hyperparameters=static_hyperparameters,
             auto_log=auto_log,
@@ -352,7 +352,7 @@ class TensorboardLoggingCallback(LoggingCallback):
         self._logged_model = False
         self._logged_hyperparameters = False
 
-    def get_weights(self) -> Dict[str, Variable]:
+    def get_weights(self) -> dict[str, Variable]:
         """
         Get the weights tensors tracked. The weights will be stored in a dictionary where each key is the weight's name
         and the value is the weight's parameter (tensor).
@@ -361,7 +361,7 @@ class TensorboardLoggingCallback(LoggingCallback):
         """
         return self._logger.weights
 
-    def get_weights_statistics(self) -> Dict[str, Dict[str, List[float]]]:
+    def get_weights_statistics(self) -> dict[str, dict[str, list[float]]]:
         """
         Get the weights mean results logged. The results will be stored in a dictionary where each key is the weight's
         name and the value is a list of mean values per epoch.
@@ -408,7 +408,7 @@ class TensorboardLoggingCallback(LoggingCallback):
         :param logs: Currently the output of the last call to `on_epoch_end()` is passed to this argument for this
                      method but that may change in the future.
         """
-        super(TensorboardLoggingCallback, self).on_train_end()
+        super().on_train_end()
 
         # Write the final run summary:
         self._logger.write_final_summary_text()
@@ -453,7 +453,7 @@ class TensorboardLoggingCallback(LoggingCallback):
         :param logs: Currently no data is passed to this argument for this method but that may change in the
                      future.
         """
-        super(TensorboardLoggingCallback, self).on_test_end(logs=logs)
+        super().on_test_end(logs=logs)
 
         # Check if needed to end the run (in case of evaluation and not training):
         if not self._is_training:
@@ -477,7 +477,7 @@ class TensorboardLoggingCallback(LoggingCallback):
                       `Model`'s metrics are returned. Example : `{'loss': 0.2, 'acc': 0.7}`.
         """
         # Update the dynamic hyperparameters
-        super(TensorboardLoggingCallback, self).on_epoch_end(epoch=epoch)
+        super().on_epoch_end(epoch=epoch)
 
         # Log the weights statistics:
         self._logger.log_weights_statistics()
@@ -515,9 +515,7 @@ class TensorboardLoggingCallback(LoggingCallback):
         :param logs:  Aggregated metric results up until this batch.
         """
         # Log the batch's results:
-        super(TensorboardLoggingCallback, self).on_train_batch_end(
-            batch=batch, logs=logs
-        )
+        super().on_train_batch_end(batch=batch, logs=logs)
 
         # Write the batch loss and metrics results to their graphs:
         self._logger.write_training_results()
@@ -540,9 +538,7 @@ class TensorboardLoggingCallback(LoggingCallback):
         :param logs:  Aggregated metric results up until this batch.
         """
         # Log the batch's results:
-        super(TensorboardLoggingCallback, self).on_test_batch_end(
-            batch=batch, logs=logs
-        )
+        super().on_test_batch_end(batch=batch, logs=logs)
 
         # Write the batch loss and metrics results to their graphs:
         self._logger.write_validation_results()
@@ -555,7 +551,7 @@ class TensorboardLoggingCallback(LoggingCallback):
 
     @staticmethod
     def get_default_weight_statistics_list() -> (
-        List[Callable[[Union[Variable, Tensor]], Union[float, Tensor]]]
+        list[Callable[[Union[Variable, Tensor]], Union[float, Tensor]]]
     ):
         """
         Get the default list of statistics functions being applied on the tracked weights each epoch.
@@ -569,7 +565,7 @@ class TensorboardLoggingCallback(LoggingCallback):
         After the trainer / evaluator run begins, this method will be called to setup the results, hyperparameters
         and weights dictionaries for logging.
         """
-        super(TensorboardLoggingCallback, self)._setup_run()
+        super()._setup_run()
 
         # Check if needed to track weights:
         if self._tracked_weights is False:

mlrun/frameworks/tf_keras/mlrun_interface.py

@@ -15,11 +15,12 @@
 import importlib
 import os
 from abc import ABC
-from typing import List, Tuple, Union
+from typing import Union
 
 import tensorflow as tf
 from tensorflow import keras
-from tensorflow.keras.callbacks import (
+from tensorflow.keras.optimizers import Optimizer
+from tensorflow.python.keras.callbacks import (
     BaseLogger,
     Callback,
     CSVLogger,
@@ -27,7 +28,6 @@ from tensorflow.keras.callbacks import (
     ProgbarLogger,
     TensorBoard,
 )
-from tensorflow.keras.optimizers import Optimizer
 
 import mlrun
 
@@ -88,9 +88,7 @@ class TFKerasMLRunInterface(MLRunInterface, ABC):
         :param restoration: Restoration information tuple as returned from 'remove_interface' in order to
                                         add the interface in a certain state.
         """
-        super(TFKerasMLRunInterface, cls).add_interface(
-            obj=obj, restoration=restoration
-        )
+        super().add_interface(obj=obj, restoration=restoration)
 
     def mlrun_compile(self, *args, **kwargs):
         """
@@ -237,7 +235,7 @@ class TFKerasMLRunInterface(MLRunInterface, ABC):
         """
         self._RANK_0_ONLY_CALLBACKS.add(callback_name)
 
-    def _pre_compile(self, optimizer: Optimizer) -> Tuple[Optimizer, Union[bool, None]]:
+    def _pre_compile(self, optimizer: Optimizer) -> tuple[Optimizer, Union[bool, None]]:
         """
         Method to call before calling 'compile' to setup the run and inputs for using horovod.
 
@@ -295,11 +293,11 @@ class TFKerasMLRunInterface(MLRunInterface, ABC):
 
     def _pre_fit(
         self,
-        callbacks: List[Callback],
+        callbacks: list[Callback],
         verbose: int,
         steps_per_epoch: Union[int, None],
         validation_steps: Union[int, None],
-    ) -> Tuple[List[Callback], int, Union[int, None], Union[int, None]]:
+    ) -> tuple[list[Callback], int, Union[int, None], Union[int, None]]:
         """
         Method to call before calling 'fit' to setup the run and inputs for using horovod.
 
@@ -366,9 +364,9 @@ class TFKerasMLRunInterface(MLRunInterface, ABC):
 
     def _pre_evaluate(
         self,
-        callbacks: List[Callback],
+        callbacks: list[Callback],
         steps: Union[int, None],
-    ) -> Tuple[List[Callback], Union[int, None]]:
+    ) -> tuple[list[Callback], Union[int, None]]:
         """
         Method to call before calling 'evaluate' to setup the run and inputs for using horovod.