mlrun 1.7.2rc3__py3-none-any.whl → 1.8.0rc1__py3-none-any.whl

mlrun/frameworks/sklearn/__init__.py

@@ -11,10 +11,9 @@
 # WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
 # See the License for the specific language governing permissions and
 # limitations under the License.
-#
-# flake8: noqa  - this is until we take care of the F401 violations with respect to __all__ & sphinx
+
 import warnings
-from typing import Union
+from typing import Optional, Union
 
 import mlrun
 from mlrun.frameworks.sklearn.metric import Metric
@@ -36,26 +35,28 @@ def apply_mlrun(
     model: SKLearnTypes.ModelType = None,
     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,
-    custom_objects_directory: str = None,
+    model_path: Optional[str] = None,
+    modules_map: Optional[Union[dict[str, Union[None, str, list[str]]], str]] = None,
+    custom_objects_map: Optional[Union[dict[str, Union[str, list[str]]], str]] = None,
+    custom_objects_directory: Optional[str] = None,
     context: mlrun.MLClientCtx = None,
-    artifacts: Union[list[MLPlan], list[str], dict[str, dict]] = None,
-    metrics: Union[
-        list[Metric],
-        list[SKLearnTypes.MetricEntryType],
-        dict[str, SKLearnTypes.MetricEntryType],
+    artifacts: Optional[Union[list[MLPlan], list[str], dict[str, dict]]] = None,
+    metrics: Optional[
+        Union[
+            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,
-    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,
+    y_columns: Optional[Union[list[str], list[int]]] = None,
+    feature_vector: Optional[str] = None,
+    feature_weights: Optional[list[float]] = None,
+    labels: Optional[dict[str, Union[str, int, float]]] = None,
+    parameters: Optional[dict[str, Union[str, int, float]]] = None,
+    extra_data: Optional[dict[str, SKLearnTypes.ExtraDataType]] = None,
     auto_log: bool = True,
     **kwargs,
 ) -> SKLearnModelHandler:

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 Union
+from typing import Optional, 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: Optional[list[Metric]] = None,
     ):
         """
         Initialize an estimator with the given metrics. The estimator will log the calculated results using the given

mlrun/frameworks/sklearn/metric.py

@@ -15,7 +15,7 @@
 import importlib
 import json
 import sys
-from typing import Callable, Union
+from typing import Callable, Optional, Union
 
 import mlrun.errors
 
@@ -31,8 +31,8 @@ class Metric:
     def __init__(
         self,
         metric: Union[Callable, str],
-        name: str = None,
-        additional_arguments: dict = None,
+        name: Optional[str] = None,
+        additional_arguments: Optional[dict] = None,
         need_probabilities: bool = False,
     ):
         """

mlrun/frameworks/sklearn/metrics_library.py

@@ -13,7 +13,7 @@
 # limitations under the License.
 #
 from abc import ABC
-from typing import Union
+from typing import Optional, Union
 
 import sklearn
 from sklearn.preprocessing import LabelBinarizer
@@ -39,10 +39,12 @@ class MetricsLibrary(ABC):
     @classmethod
     def get_metrics(
         cls,
-        metrics: Union[
-            list[Metric],
-            list[SKLearnTypes.MetricEntryType],
-            dict[str, SKLearnTypes.MetricEntryType],
+        metrics: Optional[
+            Union[
+                list[Metric],
+                list[SKLearnTypes.MetricEntryType],
+                dict[str, SKLearnTypes.MetricEntryType],
+            ]
         ] = None,
         context: mlrun.MLClientCtx = None,
         include_default: bool = True,
@@ -262,7 +264,7 @@ class MetricsLibrary(ABC):
     def _to_metric_class(
         cls,
         metric_entry: SKLearnTypes.MetricEntryType,
-        metric_name: str = None,
+        metric_name: Optional[str] = None,
     ) -> Metric:
         """
         Create a Metric instance from a user given metric entry.

mlrun/frameworks/sklearn/mlrun_interface.py

@@ -13,6 +13,7 @@
 # limitations under the License.
 #
 from abc import ABC
+from typing import Optional
 
 import mlrun
 
@@ -161,8 +162,8 @@ class SKLearnMLRunInterface(MLRunInterface, ABC):
     def configure_logging(
         self,
         context: mlrun.MLClientCtx = None,
-        plans: list[MLPlan] = None,
-        metrics: list[Metric] = None,
+        plans: Optional[list[MLPlan]] = None,
+        metrics: Optional[list[Metric]] = None,
         x_test: SKLearnTypes.DatasetType = None,
         y_test: SKLearnTypes.DatasetType = None,
         model_handler: MLModelHandler = None,

mlrun/frameworks/sklearn/model_handler.py

@@ -14,6 +14,7 @@
 #
 import os
 import pickle
+from typing import Optional
 
 import cloudpickle
 
@@ -49,7 +50,7 @@ class SKLearnModelHandler(MLModelHandler):
             )
 
     @without_mlrun_interface(interface=SKLearnMLRunInterface)
-    def save(self, output_path: str = None, **kwargs):
+    def save(self, output_path: Optional[str] = None, **kwargs):
         """
         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,10 +82,10 @@ class SKLearnModelHandler(MLModelHandler):
 
     def to_onnx(
         self,
-        model_name: str = None,
+        model_name: Optional[str] = None,
         optimize: bool = True,
         input_sample: SKLearnTypes.DatasetType = None,
-        log: bool = None,
+        log: Optional[bool] = None,
     ):
         """
         Convert the model in this handler to an ONNX model. The inputs names are optional, they do not change the

mlrun/frameworks/tf_keras/__init__.py

@@ -11,9 +11,8 @@
 # WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
 # See the License for the specific language governing permissions and
 # 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, Union
+
+from typing import Any, Optional, Union
 
 from tensorflow import keras
 
@@ -29,20 +28,20 @@ from .utils import TFKerasTypes, TFKerasUtils
 
 def apply_mlrun(
     model: keras.Model = None,
-    model_name: str = None,
+    model_name: Optional[str] = None,
     tag: str = "",
-    model_path: str = None,
+    model_path: Optional[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,
-    custom_objects_directory: str = None,
+    modules_map: Optional[Union[dict[str, Union[None, str, list[str]]], str]] = None,
+    custom_objects_map: Optional[Union[dict[str, Union[str, list[str]]], str]] = None,
+    custom_objects_directory: Optional[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,
-    use_horovod: bool = None,
+    tensorboard_directory: Optional[str] = None,
+    mlrun_callback_kwargs: Optional[dict[str, Any]] = None,
+    tensorboard_callback_kwargs: Optional[dict[str, Any]] = None,
+    use_horovod: Optional[bool] = None,
     **kwargs,
 ) -> TFKerasModelHandler:
     """

mlrun/frameworks/tf_keras/callbacks/__init__.py

@@ -11,8 +11,7 @@
 # WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
 # See the License for the specific language governing permissions and
 # limitations under the License.
-#
-# flake8: noqa  - this is until we take care of the F401 violations with respect to __all__ & sphinx
+
 from .logging_callback import LoggingCallback
 from .mlrun_logging_callback import MLRunLoggingCallback
 from .tensorboard_logging_callback import TensorboardLoggingCallback

mlrun/frameworks/tf_keras/callbacks/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, Union
+from typing import Callable, Optional, Union
 
 import numpy as np
 import tensorflow as tf
@@ -36,11 +36,14 @@ class LoggingCallback(Callback):
     def __init__(
         self,
         context: mlrun.MLClientCtx = None,
-        dynamic_hyperparameters: dict[
-            str, Union[list[Union[str, int]], Callable[[], TFKerasTypes.TrackableType]]
+        dynamic_hyperparameters: Optional[
+            dict[
+                str,
+                Union[list[Union[str, int]], Callable[[], TFKerasTypes.TrackableType]],
+            ]
         ] = None,
-        static_hyperparameters: dict[
-            str, Union[TFKerasTypes.TrackableType, list[Union[str, int]]]
+        static_hyperparameters: Optional[
+            dict[str, Union[TFKerasTypes.TrackableType, list[Union[str, int]]]]
         ] = None,
         auto_log: bool = False,
     ):
@@ -175,7 +178,7 @@ class LoggingCallback(Callback):
         """
         return self._logger.validation_iterations
 
-    def on_train_begin(self, logs: dict = None):
+    def on_train_begin(self, logs: Optional[dict] = None):
         """
         Called once at the beginning of training process (one time call).
 
@@ -185,7 +188,7 @@ class LoggingCallback(Callback):
         self._is_training = True
         self._setup_run()
 
-    def on_test_begin(self, logs: dict = None):
+    def on_test_begin(self, logs: Optional[dict] = None):
         """
         Called at the beginning of evaluation or validation. Will be called on each epoch according to the validation
         per epoch configuration.
@@ -202,7 +205,7 @@ class LoggingCallback(Callback):
         if not self._is_training:
             self._setup_run()
 
-    def on_test_end(self, logs: dict = None):
+    def on_test_end(self, logs: Optional[dict] = None):
         """
         Called at the end of evaluation or validation. Will be called on each epoch according to the validation
         per epoch configuration. The recent evaluation / validation results will be summarized and logged.
@@ -220,7 +223,7 @@ class LoggingCallback(Callback):
                 result=float(sum(epoch_values[-1]) / len(epoch_values[-1])),
             )
 
-    def on_epoch_begin(self, epoch: int, logs: dict = None):
+    def on_epoch_begin(self, epoch: int, logs: Optional[dict] = None):
         """
         Called at the start of an epoch, logging it and appending a new epoch to the logger's dictionaries.
 
@@ -236,7 +239,7 @@ class LoggingCallback(Callback):
             for metric in sum_dictionary:
                 sum_dictionary[metric] = 0
 
-    def on_epoch_end(self, epoch: int, logs: dict = None):
+    def on_epoch_end(self, epoch: int, logs: Optional[dict] = None):
         """
         Called at the end of an epoch, logging the training summaries and the current dynamic hyperparameters values.
 
@@ -262,7 +265,7 @@ class LoggingCallback(Callback):
                     value=self._get_hyperparameter(key_chain=key_chain),
                 )
 
-    def on_train_batch_begin(self, batch: int, logs: dict = None):
+    def on_train_batch_begin(self, batch: int, logs: Optional[dict] = None):
         """
         Called at the beginning of a training batch in `fit` methods. The logger will check if this batch is needed to
         be logged according to the configuration. Note that if the `steps_per_execution` argument to `compile` in
@@ -274,7 +277,7 @@ class LoggingCallback(Callback):
         """
         self._logger.log_training_iteration()
 
-    def on_train_batch_end(self, batch: int, logs: dict = None):
+    def on_train_batch_end(self, batch: int, logs: Optional[dict] = None):
         """
         Called at the end of a training batch in `fit` methods. The batch metrics results will be logged. Note that if
         the `steps_per_execution` argument to `compile` in `tf.keras.Model` is set to `N`, this method will only be
@@ -289,7 +292,7 @@ class LoggingCallback(Callback):
             logs=logs,
         )
 
-    def on_test_batch_begin(self, batch: int, logs: dict = None):
+    def on_test_batch_begin(self, batch: int, logs: Optional[dict] = None):
         """
         Called at the beginning of a batch in `evaluate` methods. Also called at the beginning of a validation batch in
         the `fit` methods, if validation data is provided. The logger will check if this batch is needed to be logged
@@ -302,7 +305,7 @@ class LoggingCallback(Callback):
         """
         self._logger.log_validation_iteration()
 
-    def on_test_batch_end(self, batch: int, logs: dict = None):
+    def on_test_batch_end(self, batch: int, logs: Optional[dict] = None):
         """
         Called at the end of a batch in `evaluate` methods. Also called at the end of a validation batch in the `fit`
         methods, if validation data is provided. The batch metrics results will be logged. Note that if the

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, Union
+from typing import Callable, Optional, Union
 
 import mlrun
 from mlrun.artifacts import Artifact
@@ -50,16 +50,19 @@ 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[
-            str, Union[TFKerasTypes.TrackableType, Artifact]
+        log_model_labels: Optional[dict[str, TFKerasTypes.TrackableType]] = None,
+        log_model_parameters: Optional[dict[str, TFKerasTypes.TrackableType]] = None,
+        log_model_extra_data: Optional[
+            dict[str, Union[TFKerasTypes.TrackableType, Artifact]]
         ] = None,
-        dynamic_hyperparameters: dict[
-            str, Union[list[Union[str, int]], Callable[[], TFKerasTypes.TrackableType]]
+        dynamic_hyperparameters: Optional[
+            dict[
+                str,
+                Union[list[Union[str, int]], Callable[[], TFKerasTypes.TrackableType]],
+            ]
         ] = None,
-        static_hyperparameters: dict[
-            str, Union[TFKerasTypes, list[Union[str, int]]]
+        static_hyperparameters: Optional[
+            dict[str, Union[TFKerasTypes, list[Union[str, int]]]]
         ] = None,
         auto_log: bool = False,
     ):
@@ -116,7 +119,7 @@ class MLRunLoggingCallback(LoggingCallback):
         # Store the model handler:
         self._model_handler = model_handler
 
-    def on_train_end(self, logs: dict = None):
+    def on_train_end(self, logs: Optional[dict] = None):
         """
         Called at the end of training, logging the model and the summaries of this run.
 
@@ -125,7 +128,7 @@ class MLRunLoggingCallback(LoggingCallback):
         """
         self._end_run()
 
-    def on_test_end(self, logs: dict = None):
+    def on_test_end(self, logs: Optional[dict] = None):
         """
         Called at the end of evaluation or validation. Will be called on each epoch according to the validation
         per epoch configuration. The recent evaluation / validation results will be summarized and logged. If the logger
@@ -141,7 +144,7 @@ class MLRunLoggingCallback(LoggingCallback):
             self._logger.log_epoch_to_context(epoch=1)
             self._end_run()
 
-    def on_epoch_end(self, epoch: int, logs: dict = None):
+    def on_epoch_end(self, epoch: int, logs: Optional[dict] = None):
         """
         Called at the end of an epoch, logging the dynamic hyperparameters and results of this epoch via the stored
         context.

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, Union
+from typing import Callable, Optional, Union
 
 import tensorflow as tf
 from packaging import version
@@ -40,8 +40,8 @@ class _TFKerasTensorboardLogger(TensorboardLogger):
         self,
         statistics_functions: list[Callable[[Union[Variable]], Union[float, Variable]]],
         context: mlrun.MLClientCtx = None,
-        tensorboard_directory: str = None,
-        run_name: str = None,
+        tensorboard_directory: Optional[str] = None,
+        run_name: Optional[str] = None,
         update_frequency: Union[int, str] = "epoch",
     ):
         """
@@ -253,17 +253,20 @@ class TensorboardLoggingCallback(LoggingCallback):
     def __init__(
         self,
         context: mlrun.MLClientCtx = None,
-        tensorboard_directory: str = None,
-        run_name: str = None,
+        tensorboard_directory: Optional[str] = None,
+        run_name: Optional[str] = None,
         weights: Union[bool, list[str]] = False,
-        statistics_functions: list[
-            Callable[[Union[Variable, Tensor]], Union[float, Tensor]]
+        statistics_functions: Optional[
+            list[Callable[[Union[Variable, Tensor]], Union[float, Tensor]]]
         ] = None,
-        dynamic_hyperparameters: dict[
-            str, Union[list[Union[str, int]], Callable[[], TFKerasTypes.TrackableType]]
+        dynamic_hyperparameters: Optional[
+            dict[
+                str,
+                Union[list[Union[str, int]], Callable[[], TFKerasTypes.TrackableType]],
+            ]
         ] = None,
-        static_hyperparameters: dict[
-            str, Union[TFKerasTypes.TrackableType, list[Union[str, int]]]
+        static_hyperparameters: Optional[
+            dict[str, Union[TFKerasTypes.TrackableType, list[Union[str, int]]]]
         ] = None,
         update_frequency: Union[int, str] = "epoch",
         auto_log: bool = False,
@@ -370,7 +373,7 @@ class TensorboardLoggingCallback(LoggingCallback):
         """
         return self._logger.weight_statistics
 
-    def on_train_begin(self, logs: dict = None):
+    def on_train_begin(self, logs: Optional[dict] = None):
         """
         Called once at the beginning of training process (one time call). Will log the pre-training (epoch 0)
         hyperparameters and weights.
@@ -401,7 +404,7 @@ class TensorboardLoggingCallback(LoggingCallback):
         # Make sure all values were written to the directory logs:
         self._logger.flush()
 
-    def on_train_end(self, logs: dict = None):
+    def on_train_end(self, logs: Optional[dict] = None):
         """
         Called at the end of training, wrapping up the tensorboard logging session.
 
@@ -416,7 +419,7 @@ class TensorboardLoggingCallback(LoggingCallback):
         # Close the logger:
         self._logger.close()
 
-    def on_test_begin(self, logs: dict = None):
+    def on_test_begin(self, logs: Optional[dict] = None):
         """
         Called at the beginning of evaluation or validation. Will be called on each epoch according to the validation
         per epoch configuration. In case it is an evaluation, the epoch 0 will be logged.
@@ -445,7 +448,7 @@ class TensorboardLoggingCallback(LoggingCallback):
             # Make sure all values were written to the directory logs:
             self._logger.flush()
 
-    def on_test_end(self, logs: dict = None):
+    def on_test_end(self, logs: Optional[dict] = None):
         """
         Called at the end of evaluation or validation. Will be called on each epoch according to the validation
         per epoch configuration. The recent evaluation / validation results will be summarized and logged.
@@ -466,7 +469,7 @@ class TensorboardLoggingCallback(LoggingCallback):
             # Close the logger:
             self._logger.close()
 
-    def on_epoch_end(self, epoch: int, logs: dict = None):
+    def on_epoch_end(self, epoch: int, logs: Optional[dict] = None):
         """
         Called at the end of an epoch, logging the current dynamic hyperparameters values, summaries and weights to
         tensorboard.
@@ -504,7 +507,7 @@ class TensorboardLoggingCallback(LoggingCallback):
         # Make sure all values were written to the directory logs:
         self._logger.flush()
 
-    def on_train_batch_end(self, batch: int, logs: dict = None):
+    def on_train_batch_end(self, batch: int, logs: Optional[dict] = None):
         """
         Called at the end of a training batch in `fit` methods. The batch metrics results will be logged. If it is the
         first batch to end, the model architecture and hyperparameters will be logged as well. Note that if the
@@ -526,7 +529,7 @@ class TensorboardLoggingCallback(LoggingCallback):
             self._logged_hyperparameters = True
             self._logger.write_dynamic_hyperparameters()
 
-    def on_test_batch_end(self, batch: int, logs: dict = None):
+    def on_test_batch_end(self, batch: int, logs: Optional[dict] = None):
         """
         Called at the end of a batch in `evaluate` methods. Also called at the end of a validation batch in the `fit`
         methods, if validation data is provided. The batch metrics results will be logged. In case it is an evaluation

mlrun/frameworks/tf_keras/model_handler.py

@@ -15,7 +15,7 @@
 import os
 import shutil
 import zipfile
-from typing import Union
+from typing import Optional, Union
 
 import numpy as np
 import tensorflow as tf
@@ -63,13 +63,17 @@ class TFKerasModelHandler(DLModelHandler):
     def __init__(
         self,
         model: keras.Model = None,
-        model_path: str = None,
-        model_name: str = None,
+        model_path: Optional[str] = None,
+        model_name: Optional[str] = None,
         model_format: str = ModelFormats.SAVED_MODEL,
         context: mlrun.MLClientCtx = 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,
+        modules_map: Optional[
+            Union[dict[str, Union[None, str, list[str]]], str]
+        ] = None,
+        custom_objects_map: Optional[
+            Union[dict[str, Union[str, list[str]]], str]
+        ] = None,
+        custom_objects_directory: Optional[str] = None,
         save_traces: bool = False,
         **kwargs,
     ):
@@ -190,8 +194,8 @@ class TFKerasModelHandler(DLModelHandler):
 
     def set_labels(
         self,
-        to_add: dict[str, Union[str, int, float]] = None,
-        to_remove: list[str] = None,
+        to_add: Optional[dict[str, Union[str, int, float]]] = None,
+        to_remove: Optional[list[str]] = None,
     ):
         """
         Update the labels dictionary of this model artifact. There are required labels that cannot be edited or removed.
@@ -210,7 +214,7 @@ class TFKerasModelHandler(DLModelHandler):
     # TODO: output_path won't work well with logging artifacts. Need to look into changing the logic of 'log_artifact'.
     @without_mlrun_interface(interface=TFKerasMLRunInterface)
     def save(
-        self, output_path: str = None, **kwargs
+        self, output_path: Optional[str] = None, **kwargs
     ) -> 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
@@ -274,7 +278,7 @@ class TFKerasModelHandler(DLModelHandler):
 
         return artifacts if self._context is not None else None
 
-    def load(self, checkpoint: str = None, **kwargs):
+    def load(self, checkpoint: Optional[str] = None, **kwargs):
         """
         Load the specified model in this handler. If a checkpoint is required to be loaded, it can be given here
         according to the provided model path in the initialization of this handler. Additional parameters for the class
@@ -318,13 +322,13 @@ class TFKerasModelHandler(DLModelHandler):
 
     def to_onnx(
         self,
-        model_name: str = None,
+        model_name: Optional[str] = None,
         optimize: bool = True,
         input_signature: Union[
             list[tf.TensorSpec], list[np.ndarray], tf.TensorSpec, np.ndarray
         ] = None,
-        output_path: str = None,
-        log: bool = None,
+        output_path: Optional[str] = None,
+        log: Optional[bool] = None,
     ):
         """
         Convert the model in this handler to an ONNX model.

mlrun/frameworks/tf_keras/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, Union
+from typing import Any, Optional, Union
 
 import numpy as np
 from tensorflow import keras
@@ -32,16 +32,20 @@ class TFKerasModelServer(V2ModelServer):
     def __init__(
         self,
         context: mlrun.MLClientCtx = None,
-        name: str = None,
+        name: Optional[str] = None,
         model: keras.Model = None,
-        model_path: str = None,
-        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,
-        custom_objects_directory: str = None,
+        model_path: Optional[str] = None,
+        model_name: Optional[str] = None,
+        modules_map: Optional[
+            Union[dict[str, Union[None, str, list[str]]], str]
+        ] = None,
+        custom_objects_map: Optional[
+            Union[dict[str, Union[str, list[str]]], str]
+        ] = None,
+        custom_objects_directory: Optional[str] = None,
         model_format: str = TFKerasModelHandler.ModelFormats.SAVED_MODEL,
         to_list: bool = False,
-        protocol: str = None,
+        protocol: Optional[str] = None,
         **class_args,
     ):
         """