mlrun 1.6.0rc35__py3-none-any.whl → 1.7.0rc2__py3-none-any.whl

mlrun/frameworks/lgbm/callbacks/mlrun_logging_callback.py

@@ -12,7 +12,6 @@
 # See the License for the specific language governing permissions and
 # limitations under the License.
 #
-from typing import List
 
 import mlrun
 
@@ -34,8 +33,8 @@ class MLRunLoggingCallback(LoggingCallback):
     def __init__(
         self,
         context: mlrun.MLClientCtx,
-        dynamic_hyperparameters: List[str] = None,
-        static_hyperparameters: List[str] = None,
+        dynamic_hyperparameters: list[str] = None,
+        static_hyperparameters: list[str] = None,
         logging_frequency: int = 100,
     ):
         """
@@ -55,7 +54,7 @@ class MLRunLoggingCallback(LoggingCallback):
                                         them and the results to MLRun). Two low frequency may slow the training time.
                                         Default: 100.
         """
-        super(MLRunLoggingCallback, self).__init__(
+        super().__init__(
             dynamic_hyperparameters=dynamic_hyperparameters,
             static_hyperparameters=static_hyperparameters,
         )
@@ -75,7 +74,7 @@ class MLRunLoggingCallback(LoggingCallback):
                     information check the `Callback` doc string.
         """
         # Log the results and parameters:
-        super(MLRunLoggingCallback, self).__call__(env=env)
+        super().__call__(env=env)
 
         # Produce the artifacts (post iteration stage):
         if env.iteration % self._logging_frequency == 0:

mlrun/frameworks/lgbm/mlrun_interfaces/booster_mlrun_interface.py

@@ -43,6 +43,4 @@ class LGBMBoosterMLRunInterface(MLRunInterface, ABC):
         :param restoration: Restoration information tuple as returned from 'remove_interface' in order to add the
                             interface in a certain state.
         """
-        super(LGBMBoosterMLRunInterface, cls).add_interface(
-            obj=obj, restoration=restoration
-        )
+        super().add_interface(obj=obj, restoration=restoration)

mlrun/frameworks/lgbm/mlrun_interfaces/mlrun_interface.py

@@ -14,7 +14,7 @@
 #
 from abc import ABC
 from types import ModuleType
-from typing import Callable, List, Tuple, Union
+from typing import Callable, Union
 
 import lightgbm as lgb
 
@@ -88,7 +88,7 @@ class LGBMMLRunInterface(MLRunInterface, ABC):
                 globals().update({"lightgbm": lgb, "lgb": lgb})
 
         # Add the interface to the provided lightgbm module:
-        super(LGBMMLRunInterface, cls).add_interface(obj=obj, restoration=restoration)
+        super().add_interface(obj=obj, restoration=restoration)
 
     @staticmethod
     def mlrun_train(*args, **kwargs):
@@ -223,7 +223,7 @@ class LGBMMLRunInterface(MLRunInterface, ABC):
         pass
 
     @staticmethod
-    def _parse_callbacks(callbacks: List[Callable]):
+    def _parse_callbacks(callbacks: list[Callable]):
         """
         Parse the callbacks passed to the training API functions of LightGBM for adding logging and enabling the MLRun
         callbacks API.
@@ -259,9 +259,9 @@ class LGBMMLRunInterface(MLRunInterface, ABC):
     @staticmethod
     def _post_train(
         booster: lgb.Booster,
-        train_set: Tuple[MLTypes.DatasetType, Union[MLTypes.DatasetType, None]],
-        validation_sets: List[
-            Tuple[Tuple[MLTypes.DatasetType, Union[MLTypes.DatasetType, None]], str]
+        train_set: tuple[MLTypes.DatasetType, Union[MLTypes.DatasetType, None]],
+        validation_sets: list[
+            tuple[tuple[MLTypes.DatasetType, Union[MLTypes.DatasetType, None]], str]
         ],
     ):
         """

mlrun/frameworks/lgbm/model_handler.py

@@ -14,7 +14,7 @@
 #
 import os
 import pickle
-from typing import Dict, List, Union
+from typing import Union
 
 import cloudpickle
 import lightgbm as lgb
@@ -56,8 +56,8 @@ class LGBMModelHandler(MLModelHandler):
         model_name: str = None,
         model_path: str = None,
         model: LGBMTypes.ModelType = 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,
         model_format: str = ModelFormats.PKL,
@@ -139,7 +139,7 @@ class LGBMModelHandler(MLModelHandler):
         self._model_format = model_format
 
         # Set up the base handler class:
-        super(LGBMModelHandler, self).__init__(
+        super().__init__(
             model=model,
             model_path=model_path,
             model_name=model_name,
@@ -152,8 +152,8 @@ class LGBMModelHandler(MLModelHandler):
 
     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,7 +162,7 @@ class LGBMModelHandler(MLModelHandler):
         :param to_remove: A list of labels keys to remove.
         """
         # Update the user's labels:
-        super(LGBMModelHandler, 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_FORMAT] = self._model_format
@@ -193,7 +193,7 @@ class LGBMModelHandler(MLModelHandler):
 
         :return The saved model additional artifacts (if needed) dictionary if context is available and None otherwise.
         """
-        super(LGBMModelHandler, self).save(output_path=output_path)
+        super().save(output_path=output_path)
 
         if isinstance(self._model, lgb.LGBMModel):
             return self._save_lgbmmodel()
@@ -204,7 +204,7 @@ class LGBMModelHandler(MLModelHandler):
         Load the specified model in this handler. Additional parameters for the class initializer can be passed via the
         kwargs dictionary.
         """
-        super(LGBMModelHandler, self).load()
+        super().load()
 
         # ModelFormats.PKL - Load from a pkl file:
         if self._model_format == LGBMModelHandler.ModelFormats.PKL:

mlrun/frameworks/lgbm/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, Union
+from typing import Any, Union
 
 import numpy as np
 
@@ -37,8 +37,8 @@ class LGBMModelServer(V2ModelServer):
         model_path: LGBMTypes.PathType = None,
         model_name: str = None,
         model_format: str = LGBMModelHandler.ModelFormats.PKL,
-        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,
         to_list: bool = True,
         protocol: str = None,
@@ -100,7 +100,7 @@ class LGBMModelServer(V2ModelServer):
         :param protocol:                 -
         :param class_args:               -
         """
-        super(LGBMModelServer, self).__init__(
+        super().__init__(
             context=context,
             name=name,
             model_path=model_path,
@@ -139,7 +139,7 @@ class LGBMModelServer(V2ModelServer):
             self._model_handler.load()
         self.model = self._model_handler.model
 
-    def predict(self, request: Dict[str, Any]) -> Union[np.ndarray, list]:
+    def predict(self, request: dict[str, Any]) -> Union[np.ndarray, 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.
@@ -158,7 +158,7 @@ class LGBMModelServer(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 served in this serving function and the function name.
 

mlrun/frameworks/lgbm/utils.py

@@ -12,7 +12,7 @@
 # See the License for the specific language governing permissions and
 # limitations under the License.
 #
-from typing import List, Tuple, Union
+from typing import Union
 
 import lightgbm as lgb
 import numpy as np
@@ -36,13 +36,13 @@ class LGBMTypes(MLTypes):
 
     # An evaluation result as packaged by the training in LightGBM:
     EvaluationResultType = Union[
-        Tuple[str, str, float, bool],  # As packaged in `lightgbm.train`
-        Tuple[str, str, float, bool, float],  # As packaged in `lightgbm.cv`
+        tuple[str, str, float, bool],  # As packaged in `lightgbm.train`
+        tuple[str, str, float, bool, float],  # As packaged in `lightgbm.cv`
     ]
 
     # Detailed type for the named tuple `CallbackEnv` passed during LightGBM's training for the callbacks:
-    CallbackEnvType = Tuple[
-        lgb.Booster, dict, int, int, int, List[EvaluationResultType]
+    CallbackEnvType = tuple[
+        lgb.Booster, dict, int, int, int, list[EvaluationResultType]
     ]
 
 

mlrun/frameworks/onnx/dataset.py

@@ -13,7 +13,7 @@
 # limitations under the License.
 #
 import math
-from typing import Callable, List, Tuple, Union
+from typing import Callable, Union
 
 import numpy as np
 
@@ -25,11 +25,11 @@ class ONNXDataset:
 
     def __init__(
         self,
-        x: Union[np.ndarray, List[np.ndarray]],
-        y: Union[np.ndarray, List[np.ndarray]] = None,
+        x: Union[np.ndarray, list[np.ndarray]],
+        y: Union[np.ndarray, list[np.ndarray]] = None,
         batch_size: int = 1,
-        x_transforms: List[Callable[[np.ndarray], np.ndarray]] = None,
-        y_transforms: List[Callable[[np.ndarray], np.ndarray]] = None,
+        x_transforms: list[Callable[[np.ndarray], np.ndarray]] = None,
+        y_transforms: list[Callable[[np.ndarray], np.ndarray]] = None,
         is_batched_transforms: bool = False,
     ):
         """
@@ -71,7 +71,7 @@ class ONNXDataset:
         self._index = 0
         return self
 
-    def __next__(self) -> Union[np.ndarray, Tuple[np.ndarray, np.ndarray]]:
+    def __next__(self) -> Union[np.ndarray, tuple[np.ndarray, np.ndarray]]:
         """
         Get the next item in line (by the inner index) since calling '__iter__'. If ground truth was provided (y),
         a tuple of (x, y) will be returned. Otherwise x.
@@ -92,7 +92,7 @@ class ONNXDataset:
 
     def __getitem__(
         self, index: int
-    ) -> Union[np.ndarray, Tuple[np.ndarray, np.ndarray]]:
+    ) -> Union[np.ndarray, tuple[np.ndarray, np.ndarray]]:
         """
         Get the item at the given index. If ground truth was provided, a tuple of (x, y) will be returned. Otherwise x.
 
@@ -155,7 +155,7 @@ class ONNXDataset:
     def _call_transforms(
         self,
         items: np.ndarray,
-        transforms: List[Callable[[np.ndarray], np.ndarray]],
+        transforms: list[Callable[[np.ndarray], np.ndarray]],
         is_batched: bool,
     ):
         """

mlrun/frameworks/onnx/mlrun_interface.py

@@ -12,7 +12,7 @@
 # See the License for the specific language governing permissions and
 # limitations under the License.
 #
-from typing import Callable, List
+from typing import Callable
 
 import numpy as np
 import onnx
@@ -35,7 +35,7 @@ class ONNXMLRunInterface:
     def __init__(
         self,
         model: onnx.ModelProto,
-        execution_providers: List[str] = None,
+        execution_providers: list[str] = None,
         context: mlrun.MLClientCtx = None,
     ):
         # Set the context:
@@ -74,7 +74,7 @@ class ONNXMLRunInterface:
     def evaluate(
         self,
         dataset: ONNXDataset,
-        metrics: List[Callable[[np.ndarray, np.ndarray], float]],
+        metrics: list[Callable[[np.ndarray, np.ndarray], float]],
     ):
         pass
 

mlrun/frameworks/onnx/model_handler.py

@@ -13,7 +13,7 @@
 # limitations under the License.
 #
 import os
-from typing import Dict, List, Union
+from typing import Union
 
 import onnx
 import onnxoptimizer
@@ -60,7 +60,7 @@ class ONNXModelHandler(ModelHandler):
         :raise MLRunInvalidArgumentError: There was no model or model directory supplied.
         """
         # Setup the base handler class:
-        super(ONNXModelHandler, self).__init__(
+        super().__init__(
             model=model,
             model_path=model_path,
             model_name=model_name,
@@ -71,7 +71,7 @@ class ONNXModelHandler(ModelHandler):
     # TODO: output_path won't work well with logging artifacts. Need to look into changing the logic of 'log_artifact'.
     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. If a MLRun context is available, the saved model files will be
         logged and returned as artifacts.
@@ -81,7 +81,7 @@ class ONNXModelHandler(ModelHandler):
 
         :return The saved model additional artifacts (if needed) dictionary if context is available and None otherwise.
         """
-        super(ONNXModelHandler, self).save(output_path=output_path)
+        super().save(output_path=output_path)
 
         # Set the output path:
         if output_path is None:
@@ -97,7 +97,7 @@ class ONNXModelHandler(ModelHandler):
         """
         Load the specified model in this handler.
         """
-        super(ONNXModelHandler, self).load()
+        super().load()
 
         # Check that the model is well-formed:
         # TODO: Currently not working well with HuggingFace models so we skip it
@@ -106,7 +106,7 @@ class ONNXModelHandler(ModelHandler):
         # Load the ONNX model:
         self._model = onnx.load(self._model_file)
 
-    def optimize(self, optimizations: List[str] = None, fixed_point: bool = False):
+    def optimize(self, optimizations: list[str] = None, fixed_point: bool = False):
         """
         Use ONNX optimizer to optimize the ONNX model. The optimizations supported can be seen by calling
         'onnxoptimizer.get_available_passes()'

mlrun/frameworks/onnx/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, Tuple, Union
+from typing import Any, Union
 
 import numpy as np
 import onnx
@@ -37,7 +37,7 @@ class ONNXModelServer(V2ModelServer):
         model: onnx.ModelProto = None,
         model_path: str = None,
         model_name: str = None,
-        execution_providers: List[Union[str, Tuple[str, Dict[str, Any]]]] = None,
+        execution_providers: list[Union[str, tuple[str, dict[str, Any]]]] = None,
         protocol: str = None,
         **class_args,
     ):
@@ -76,7 +76,7 @@ class ONNXModelServer(V2ModelServer):
         :param protocol:            -
         :param class_args:          -
         """
-        super(ONNXModelServer, self).__init__(
+        super().__init__(
             context=context,
             name=name,
             model_path=model_path,
@@ -98,8 +98,8 @@ class ONNXModelServer(V2ModelServer):
         # Prepare inference parameters:
         self._model_handler: ONNXModelHandler = None
         self._inference_session: onnxruntime.InferenceSession = None
-        self._input_layers: List[str] = None
-        self._output_layers: List[str] = None
+        self._input_layers: list[str] = None
+        self._output_layers: list[str] = None
 
     def load(self):
         """
@@ -134,7 +134,7 @@ class ONNXModelServer(V2ModelServer):
             output_layer.name for output_layer in self._inference_session.get_outputs()
         ]
 
-    def predict(self, request: Dict[str, Any]) -> np.ndarray:
+    def predict(self, request: dict[str, Any]) -> np.ndarray:
         """
         Infer the inputs through the model using ONNXRunTime and return its output. The inferred data will be
         read from the "inputs" key of the request.
@@ -155,7 +155,7 @@ class ONNXModelServer(V2ModelServer):
             },
         )
 
-    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/parallel_coordinates.py

@@ -14,7 +14,7 @@
 #
 import datetime
 import os
-from typing import List, Union
+from typing import Union
 
 import numpy as np
 import pandas as pd
@@ -239,7 +239,7 @@ def _runs_list_to_df(runs_list, extend_iterations=False):
 
 @filter_warnings("ignore", FutureWarning)
 def compare_run_objects(
-    runs_list: Union[mlrun.model.RunObject, List[mlrun.model.RunObject]],
+    runs_list: Union[mlrun.model.RunObject, list[mlrun.model.RunObject]],
     hide_identical: bool = True,
     exclude: list = None,
     show: bool = None,

mlrun/frameworks/pytorch/__init__.py

@@ -13,7 +13,7 @@
 # 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, Tuple, Union
+from typing import Any, Union
 
 from torch.nn import Module
 from torch.optim import Optimizer
@@ -35,23 +35,23 @@ def train(
     loss_function: Module,
     optimizer: Optimizer,
     validation_set: DataLoader = None,
-    metric_functions: List[PyTorchTypes.MetricFunctionType] = None,
+    metric_functions: list[PyTorchTypes.MetricFunctionType] = None,
     scheduler=None,
     scheduler_step_frequency: Union[int, float, str] = "epoch",
     epochs: int = 1,
     training_iterations: int = None,
     validation_iterations: int = None,
-    callbacks_list: List[Callback] = None,
+    callbacks_list: list[Callback] = None,
     use_cuda: bool = True,
     use_horovod: bool = None,
     auto_log: bool = True,
     model_name: 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,
     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,
     context: mlrun.MLClientCtx = None,
 ) -> PyTorchModelHandler:
     """
@@ -205,19 +205,19 @@ def evaluate(
     dataset: DataLoader,
     model: Module = None,
     loss_function: Module = None,
-    metric_functions: List[PyTorchTypes.MetricFunctionType] = None,
+    metric_functions: list[PyTorchTypes.MetricFunctionType] = None,
     iterations: int = None,
-    callbacks_list: List[Callback] = None,
+    callbacks_list: list[Callback] = None,
     use_cuda: bool = True,
     use_horovod: bool = False,
     auto_log: bool = True,
     model_name: 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,
-    mlrun_callback_kwargs: Dict[str, Any] = None,
+    mlrun_callback_kwargs: dict[str, Any] = None,
     context: mlrun.MLClientCtx = None,
-) -> Tuple[PyTorchModelHandler, List[PyTorchTypes.MetricValueType]]:
+) -> tuple[PyTorchModelHandler, list[PyTorchTypes.MetricValueType]]:
     """
     Use MLRun's PyTorch interface to evaluate the model with the given parameters. For more information and further
     options regarding the auto logging, see 'PyTorchMLRunInterface' documentation. Notice for auto-logging: In order to
@@ -343,9 +343,9 @@ def evaluate(
 def _parse_callbacks_kwargs(
     handler: PyTorchModelHandler,
     tensorboard_directory: Union[str, None],
-    mlrun_callback_kwargs: Union[Dict[str, Any], None],
-    tensorboard_callback_kwargs: Union[Dict[str, Any], None],
-) -> Tuple[dict, dict]:
+    mlrun_callback_kwargs: Union[dict[str, Any], None],
+    tensorboard_callback_kwargs: Union[dict[str, Any], None],
+) -> tuple[dict, dict]:
     """
     Parse the given parameters into the MLRun and Tensorboard callbacks kwargs.
 

mlrun/frameworks/pytorch/callbacks/callback.py

@@ -13,7 +13,6 @@
 # limitations under the License.
 #
 from abc import ABC, abstractmethod
-from typing import List
 
 from torch import Tensor
 from torch.nn import Module
@@ -68,7 +67,7 @@ class Callback(ABC):
         validation_set: DataLoader = None,
         loss_function: Module = None,
         optimizer: Optimizer = None,
-        metric_functions: List[PyTorchTypes.MetricFunctionType] = None,
+        metric_functions: list[PyTorchTypes.MetricFunctionType] = None,
         scheduler=None,
     ):
         """
@@ -141,7 +140,7 @@ class Callback(ABC):
         pass
 
     def on_validation_end(
-        self, loss_value: PyTorchTypes.MetricValueType, metric_values: List[float]
+        self, loss_value: PyTorchTypes.MetricValueType, metric_values: list[float]
     ) -> bool:
         """
         Before the validation (in a training case it will be per epoch) ends, this method will be called.
@@ -258,7 +257,7 @@ class Callback(ABC):
         """
         pass
 
-    def on_train_metrics_end(self, metric_values: List[PyTorchTypes.MetricValueType]):
+    def on_train_metrics_end(self, metric_values: list[PyTorchTypes.MetricValueType]):
         """
         After the training calculation of the metrics, this method will be called.
 
@@ -273,7 +272,7 @@ class Callback(ABC):
         pass
 
     def on_validation_metrics_end(
-        self, metric_values: List[PyTorchTypes.MetricValueType]
+        self, metric_values: list[PyTorchTypes.MetricValueType]
     ):
         """
         After the validating calculation of the metrics, this method will be called.