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

mlrun/frameworks/pytorch/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
 from torch import Tensor
@@ -59,15 +59,22 @@ class LoggingCallback(Callback):
     def __init__(
         self,
         context: mlrun.MLClientCtx = None,
-        dynamic_hyperparameters: dict[
-            str,
-            tuple[
+        dynamic_hyperparameters: Optional[
+            dict[
                 str,
-                Union[list[Union[str, int]], Callable[[], PyTorchTypes.TrackableType]],
-            ],
+                tuple[
+                    str,
+                    Union[
+                        list[Union[str, int]], Callable[[], PyTorchTypes.TrackableType]
+                    ],
+                ],
+            ]
         ] = None,
-        static_hyperparameters: dict[
-            str, Union[PyTorchTypes.TrackableType, tuple[str, list[Union[str, int]]]]
+        static_hyperparameters: Optional[
+            dict[
+                str,
+                Union[PyTorchTypes.TrackableType, tuple[str, list[Union[str, int]]]],
+            ]
         ] = None,
         auto_log: bool = False,
     ):

mlrun/frameworks/pytorch/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 torch
 from torch import Tensor
@@ -53,20 +53,27 @@ class MLRunLoggingCallback(LoggingCallback):
         context: mlrun.MLClientCtx,
         model_handler: PyTorchModelHandler,
         log_model_tag: str = "",
-        log_model_labels: dict[str, PyTorchTypes.TrackableType] = None,
-        log_model_parameters: dict[str, PyTorchTypes.TrackableType] = None,
-        log_model_extra_data: dict[
-            str, Union[PyTorchTypes.TrackableType, Artifact]
+        log_model_labels: Optional[dict[str, PyTorchTypes.TrackableType]] = None,
+        log_model_parameters: Optional[dict[str, PyTorchTypes.TrackableType]] = None,
+        log_model_extra_data: Optional[
+            dict[str, Union[PyTorchTypes.TrackableType, Artifact]]
         ] = None,
-        dynamic_hyperparameters: dict[
-            str,
-            tuple[
+        dynamic_hyperparameters: Optional[
+            dict[
                 str,
-                Union[list[Union[str, int]], Callable[[], PyTorchTypes.TrackableType]],
-            ],
+                tuple[
+                    str,
+                    Union[
+                        list[Union[str, int]], Callable[[], PyTorchTypes.TrackableType]
+                    ],
+                ],
+            ]
         ] = None,
-        static_hyperparameters: dict[
-            str, Union[PyTorchTypes.TrackableType, tuple[str, list[Union[str, int]]]]
+        static_hyperparameters: Optional[
+            dict[
+                str,
+                Union[PyTorchTypes.TrackableType, tuple[str, list[Union[str, int]]]],
+            ]
         ] = None,
         auto_log: bool = False,
     ):

mlrun/frameworks/pytorch/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 torch
 from torch import Tensor
@@ -67,8 +67,8 @@ class _PyTorchTensorboardLogger(TensorboardLogger):
             Callable[[Union[Parameter]], Union[float, Parameter]]
         ],
         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",
     ):
         """
@@ -247,21 +247,28 @@ 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[Parameter, Tensor]], Union[float, Tensor]]
+        statistics_functions: Optional[
+            list[Callable[[Union[Parameter, Tensor]], Union[float, Tensor]]]
         ] = None,
-        dynamic_hyperparameters: dict[
-            str,
-            tuple[
+        dynamic_hyperparameters: Optional[
+            dict[
                 str,
-                Union[list[Union[str, int]], Callable[[], PyTorchTypes.TrackableType]],
-            ],
+                tuple[
+                    str,
+                    Union[
+                        list[Union[str, int]], Callable[[], PyTorchTypes.TrackableType]
+                    ],
+                ],
+            ]
         ] = None,
-        static_hyperparameters: dict[
-            str, Union[PyTorchTypes.TrackableType, tuple[str, list[Union[str, int]]]]
+        static_hyperparameters: Optional[
+            dict[
+                str,
+                Union[PyTorchTypes.TrackableType, tuple[str, list[Union[str, int]]]],
+            ]
         ] = None,
         update_frequency: Union[int, str] = "epoch",
         auto_log: bool = False,
@@ -381,7 +388,7 @@ class TensorboardLoggingCallback(LoggingCallback):
         validation_set: DataLoader = None,
         loss_function: Module = None,
         optimizer: Optimizer = None,
-        metric_functions: list[PyTorchTypes.MetricFunctionType] = None,
+        metric_functions: Optional[list[PyTorchTypes.MetricFunctionType]] = None,
         scheduler=None,
     ):
         """

mlrun/frameworks/pytorch/callbacks_handler.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
 
 from torch import Tensor
 from torch.nn import Module
@@ -116,7 +116,7 @@ class CallbacksHandler:
         optimizer: Optimizer,
         metric_functions: list[PyTorchTypes.MetricFunctionType],
         scheduler,
-        callbacks: list[str] = None,
+        callbacks: Optional[list[str]] = None,
     ) -> bool:
         """
         Call the 'on_setup' method of every callback in the callbacks list. If the list is 'None' (not given), all
@@ -145,7 +145,7 @@ class CallbacksHandler:
             scheduler=scheduler,
         )
 
-    def on_run_begin(self, callbacks: list[str] = None) -> bool:
+    def on_run_begin(self, callbacks: Optional[list[str]] = None) -> bool:
         """
         Call the 'on_run_begin' method of every callback in the callbacks list. If the list is 'None' (not given), all
         callbacks will be called.
@@ -159,7 +159,7 @@ class CallbacksHandler:
             callbacks=self._parse_names(names=callbacks),
         )
 
-    def on_run_end(self, callbacks: list[str] = None) -> bool:
+    def on_run_end(self, callbacks: Optional[list[str]] = None) -> bool:
         """
         Call the 'on_run_end' method of every callback in the callbacks list. If the list is 'None' (not given), all
         callbacks will be called.
@@ -173,7 +173,7 @@ class CallbacksHandler:
             callbacks=self._parse_names(names=callbacks),
         )
 
-    def on_epoch_begin(self, epoch: int, callbacks: list[str] = None) -> bool:
+    def on_epoch_begin(self, epoch: int, callbacks: Optional[list[str]] = None) -> bool:
         """
         Call the 'on_epoch_begin' method of every callback in the callbacks list. If the list is 'None' (not given), all
         callbacks will be called.
@@ -189,7 +189,7 @@ class CallbacksHandler:
             epoch=epoch,
         )
 
-    def on_epoch_end(self, epoch: int, callbacks: list[str] = None) -> bool:
+    def on_epoch_end(self, epoch: int, callbacks: Optional[list[str]] = None) -> bool:
         """
         Call the 'on_epoch_end' method of every callback in the callbacks list. If the list is 'None' (not given), all
         callbacks will be called.
@@ -205,7 +205,7 @@ class CallbacksHandler:
             epoch=epoch,
         )
 
-    def on_train_begin(self, callbacks: list[str] = None) -> bool:
+    def on_train_begin(self, callbacks: Optional[list[str]] = None) -> bool:
         """
         Call the 'on_train_begin' method of every callback in the callbacks list. If the list is 'None' (not given), all
         callbacks will be called.
@@ -219,7 +219,7 @@ class CallbacksHandler:
             callbacks=self._parse_names(names=callbacks),
         )
 
-    def on_train_end(self, callbacks: list[str] = None) -> bool:
+    def on_train_end(self, callbacks: Optional[list[str]] = None) -> bool:
         """
         Call the 'on_train_end' method of every callback in the callbacks list. If the list is 'None' (not given), all
         callbacks will be called.
@@ -233,7 +233,7 @@ class CallbacksHandler:
             callbacks=self._parse_names(names=callbacks),
         )
 
-    def on_validation_begin(self, callbacks: list[str] = None) -> bool:
+    def on_validation_begin(self, callbacks: Optional[list[str]] = None) -> bool:
         """
         Call the 'on_validation_begin' method of every callback in the callbacks list. If the list is 'None'
         (not given), all callbacks will be called.
@@ -251,7 +251,7 @@ class CallbacksHandler:
         self,
         loss_value: PyTorchTypes.MetricValueType,
         metric_values: list[float],
-        callbacks: list[str] = None,
+        callbacks: Optional[list[str]] = None,
     ) -> bool:
         """
         Call the 'on_validation_end' method of every callback in the callbacks list. If the list is 'None' (not given),
@@ -271,7 +271,7 @@ class CallbacksHandler:
         )
 
     def on_train_batch_begin(
-        self, batch: int, x, y_true: Tensor, callbacks: list[str] = None
+        self, batch: int, x, y_true: Tensor, callbacks: Optional[list[str]] = None
     ) -> bool:
         """
         Call the 'on_train_batch_begin' method of every callback in the callbacks list. If the list is 'None'
@@ -298,7 +298,7 @@ class CallbacksHandler:
         x,
         y_pred: Tensor,
         y_true: Tensor,
-        callbacks: list[str] = None,
+        callbacks: Optional[list[str]] = None,
     ) -> bool:
         """
         Call the 'on_train_batch_end' method of every callback in the callbacks list. If the list is 'None' (not given),
@@ -322,7 +322,7 @@ class CallbacksHandler:
         )
 
     def on_validation_batch_begin(
-        self, batch: int, x, y_true: Tensor, callbacks: list[str] = None
+        self, batch: int, x, y_true: Tensor, callbacks: Optional[list[str]] = None
     ) -> bool:
         """
         Call the 'on_validation_batch_begin' method of every callback in the callbacks list. If the list is 'None'
@@ -349,7 +349,7 @@ class CallbacksHandler:
         x,
         y_pred: Tensor,
         y_true: Tensor,
-        callbacks: list[str] = None,
+        callbacks: Optional[list[str]] = None,
     ) -> bool:
         """
         Call the 'on_validation_batch_end' method of every callback in the callbacks list. If the list is 'None'
@@ -375,7 +375,7 @@ class CallbacksHandler:
     def on_inference_begin(
         self,
         x,
-        callbacks: list[str] = None,
+        callbacks: Optional[list[str]] = None,
     ) -> bool:
         """
         Call the 'on_inference_begin' method of every callback in the callbacks list. If the list is 'None' (not given),
@@ -396,7 +396,7 @@ class CallbacksHandler:
         self,
         y_pred: Tensor,
         y_true: Tensor,
-        callbacks: list[str] = None,
+        callbacks: Optional[list[str]] = None,
     ) -> bool:
         """
         Call the 'on_inference_end' method of every callback in the callbacks list. If the list is 'None' (not given),
@@ -415,7 +415,7 @@ class CallbacksHandler:
             y_true=y_true,
         )
 
-    def on_train_loss_begin(self, callbacks: list[str] = None) -> bool:
+    def on_train_loss_begin(self, callbacks: Optional[list[str]] = None) -> bool:
         """
         Call the 'on_train_loss_begin' method of every callback in the callbacks list. If the list is 'None'
         (not given), all callbacks will be called.
@@ -430,7 +430,9 @@ class CallbacksHandler:
         )
 
     def on_train_loss_end(
-        self, loss_value: PyTorchTypes.MetricValueType, callbacks: list[str] = None
+        self,
+        loss_value: PyTorchTypes.MetricValueType,
+        callbacks: Optional[list[str]] = None,
     ) -> bool:
         """
         Call the 'on_train_loss_end' method of every callback in the callbacks list. If the list is 'None' (not given),
@@ -447,7 +449,7 @@ class CallbacksHandler:
             loss_value=loss_value,
         )
 
-    def on_validation_loss_begin(self, callbacks: list[str] = None) -> bool:
+    def on_validation_loss_begin(self, callbacks: Optional[list[str]] = None) -> bool:
         """
         Call the 'on_validation_loss_begin' method of every callback in the callbacks list. If the list is 'None'
         (not given), all callbacks will be called.
@@ -462,7 +464,9 @@ class CallbacksHandler:
         )
 
     def on_validation_loss_end(
-        self, loss_value: PyTorchTypes.MetricValueType, callbacks: list[str] = None
+        self,
+        loss_value: PyTorchTypes.MetricValueType,
+        callbacks: Optional[list[str]] = None,
     ) -> bool:
         """
         Call the 'on_validation_loss_end' method of every callback in the callbacks list. If the list is 'None'
@@ -479,7 +483,7 @@ class CallbacksHandler:
             loss_value=loss_value,
         )
 
-    def on_train_metrics_begin(self, callbacks: list[str] = None) -> bool:
+    def on_train_metrics_begin(self, callbacks: Optional[list[str]] = None) -> bool:
         """
         Call the 'on_train_metrics_begin' method of every callback in the callbacks list. If the list is 'None'
         (not given), all callbacks will be called.
@@ -496,7 +500,7 @@ class CallbacksHandler:
     def on_train_metrics_end(
         self,
         metric_values: list[PyTorchTypes.MetricValueType],
-        callbacks: list[str] = None,
+        callbacks: Optional[list[str]] = None,
     ) -> bool:
         """
         Call the 'on_train_metrics_end' method of every callback in the callbacks list. If the list is 'None'
@@ -513,7 +517,9 @@ class CallbacksHandler:
             metric_values=metric_values,
         )
 
-    def on_validation_metrics_begin(self, callbacks: list[str] = None) -> bool:
+    def on_validation_metrics_begin(
+        self, callbacks: Optional[list[str]] = None
+    ) -> bool:
         """
         Call the 'on_validation_metrics_begin' method of every callback in the callbacks list. If the list is 'None'
         (not given), all callbacks will be called.
@@ -530,7 +536,7 @@ class CallbacksHandler:
     def on_validation_metrics_end(
         self,
         metric_values: list[PyTorchTypes.MetricValueType],
-        callbacks: list[str] = None,
+        callbacks: Optional[list[str]] = None,
     ) -> bool:
         """
         Call the 'on_validation_metrics_end' method of every callback in the callbacks list. If the list is 'None'
@@ -547,7 +553,7 @@ class CallbacksHandler:
             metric_values=metric_values,
         )
 
-    def on_backward_begin(self, callbacks: list[str] = None) -> bool:
+    def on_backward_begin(self, callbacks: Optional[list[str]] = None) -> bool:
         """
         Call the 'on_backward_begin' method of every callback in the callbacks list. If the list is 'None' (not given),
         all callbacks will be called.
@@ -561,7 +567,7 @@ class CallbacksHandler:
             callbacks=self._parse_names(names=callbacks),
         )
 
-    def on_backward_end(self, callbacks: list[str] = None) -> bool:
+    def on_backward_end(self, callbacks: Optional[list[str]] = None) -> bool:
         """
         Call the 'on_backward_end' method of every callback in the callbacks list. If the list is 'None' (not given),
         all callbacks will be called.
@@ -575,7 +581,7 @@ class CallbacksHandler:
             callbacks=self._parse_names(names=callbacks),
         )
 
-    def on_optimizer_step_begin(self, callbacks: list[str] = None) -> bool:
+    def on_optimizer_step_begin(self, callbacks: Optional[list[str]] = None) -> bool:
         """
         Call the 'on_optimizer_step_begin' method of every callback in the callbacks list. If the list is 'None'
         (not given), all callbacks will be called.
@@ -589,7 +595,7 @@ class CallbacksHandler:
             callbacks=self._parse_names(names=callbacks),
         )
 
-    def on_optimizer_step_end(self, callbacks: list[str] = None) -> bool:
+    def on_optimizer_step_end(self, callbacks: Optional[list[str]] = None) -> bool:
         """
         Call the 'on_optimizer_step_end' method of every callback in the callbacks list. If the list is 'None'
         (not given), all callbacks will be called.
@@ -603,7 +609,7 @@ class CallbacksHandler:
             callbacks=self._parse_names(names=callbacks),
         )
 
-    def on_scheduler_step_begin(self, callbacks: list[str] = None) -> bool:
+    def on_scheduler_step_begin(self, callbacks: Optional[list[str]] = None) -> bool:
         """
         Call the 'on_scheduler_step_begin' method of every callback in the callbacks list. If the list is 'None'
         (not given), all callbacks will be called.
@@ -617,7 +623,7 @@ class CallbacksHandler:
             callbacks=self._parse_names(names=callbacks),
         )
 
-    def on_scheduler_step_end(self, callbacks: list[str] = None) -> bool:
+    def on_scheduler_step_end(self, callbacks: Optional[list[str]] = None) -> bool:
         """
         Call the 'on_scheduler_step_end' method of every callback in the callbacks list. If the list is 'None'
         (not given), all callbacks will be called.

mlrun/frameworks/pytorch/mlrun_interface.py

@@ -14,7 +14,7 @@
 #
 import importlib
 import sys
-from typing import Any, Union
+from typing import Any, Optional, Union
 
 import torch
 import torch.multiprocessing as mp
@@ -109,15 +109,15 @@ class PyTorchMLRunInterface:
         loss_function: Module,
         optimizer: Optimizer,
         validation_set: DataLoader = None,
-        metric_functions: list[PyTorchTypes.MetricFunctionType] = None,
+        metric_functions: Optional[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[Callback] = None,
+        training_iterations: Optional[int] = None,
+        validation_iterations: Optional[int] = None,
+        callbacks: Optional[list[Callback]] = None,
         use_cuda: bool = True,
-        use_horovod: bool = None,
+        use_horovod: Optional[bool] = None,
     ):
         """
         Initiate a training process on this interface configuration.
@@ -221,11 +221,11 @@ class PyTorchMLRunInterface:
         self,
         dataset: DataLoader,
         loss_function: Module = None,
-        metric_functions: list[PyTorchTypes.MetricFunctionType] = None,
-        iterations: int = None,
-        callbacks: list[Callback] = None,
+        metric_functions: Optional[list[PyTorchTypes.MetricFunctionType]] = None,
+        iterations: Optional[int] = None,
+        callbacks: Optional[list[Callback]] = None,
         use_cuda: bool = True,
-        use_horovod: bool = None,
+        use_horovod: Optional[bool] = None,
     ) -> list[PyTorchTypes.MetricValueType]:
         """
         Initiate an evaluation process on this interface configuration.
@@ -303,9 +303,9 @@ class PyTorchMLRunInterface:
     def add_auto_logging_callbacks(
         self,
         add_mlrun_logger: bool = True,
-        mlrun_callback_kwargs: dict[str, Any] = None,
+        mlrun_callback_kwargs: Optional[dict[str, Any]] = None,
         add_tensorboard_logger: bool = True,
-        tensorboard_callback_kwargs: dict[str, Any] = None,
+        tensorboard_callback_kwargs: Optional[dict[str, Any]] = None,
     ):
         """
         Get automatic logging callbacks to both MLRun's context and Tensorboard. For further features of logging to both
@@ -402,15 +402,15 @@ class PyTorchMLRunInterface:
         loss_function: Module = None,
         optimizer: Optimizer = None,
         validation_set: DataLoader = None,
-        metric_functions: list[PyTorchTypes.MetricFunctionType] = None,
+        metric_functions: Optional[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[Callback] = None,
+        training_iterations: Optional[int] = None,
+        validation_iterations: Optional[int] = None,
+        callbacks: Optional[list[Callback]] = None,
         use_cuda: bool = True,
-        use_horovod: bool = None,
+        use_horovod: Optional[bool] = None,
     ):
         """
         Parse and store the given input so the interface can starting training / evaluating.

mlrun/frameworks/pytorch/model_handler.py

@@ -13,7 +13,7 @@
 # limitations under the License.
 #
 import os
-from typing import Union
+from typing import Optional, Union
 
 import numpy as np
 import torch
@@ -48,12 +48,16 @@ class PyTorchModelHandler(DLModelHandler):
     def __init__(
         self,
         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,
-        custom_objects_directory: str = None,
+        model_path: Optional[str] = None,
+        model_name: Optional[str] = None,
+        model_class: Optional[Union[type[Module], 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,
         **kwargs,
     ):
@@ -152,8 +156,8 @@ class PyTorchModelHandler(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.
@@ -168,7 +172,7 @@ class PyTorchModelHandler(DLModelHandler):
         self._labels[self._LabelKeys.MODEL_CLASS_NAME] = self._model_class_name
 
     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.
@@ -194,7 +198,7 @@ class PyTorchModelHandler(DLModelHandler):
 
         return 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
@@ -232,15 +236,15 @@ class PyTorchModelHandler(DLModelHandler):
 
     def to_onnx(
         self,
-        model_name: str = None,
+        model_name: Optional[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_layers_names: Optional[list[str]] = None,
+        output_layers_names: Optional[list[str]] = None,
+        dynamic_axes: Optional[dict[str, dict[int, str]]] = None,
         is_batched: bool = True,
         optimize: bool = True,
-        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. The layer names are optional, they do not change the

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, Union
+from typing import Any, Optional, Union
 
 import numpy as np
 import torch
@@ -35,17 +35,21 @@ class PyTorchModelServer(V2ModelServer):
     def __init__(
         self,
         context: mlrun.MLClientCtx = None,
-        name: str = None,
+        name: Optional[str] = None,
         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,
-        custom_objects_directory: str = None,
+        model_path: Optional[str] = None,
+        model_name: Optional[str] = None,
+        model_class: Optional[Union[type[Module], 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,
         use_cuda: bool = True,
         to_list: bool = False,
-        protocol: str = None,
+        protocol: Optional[str] = None,
         **class_args,
     ):
         """