pyfunc

mlflow/models/model.py

@@ -598,6 +598,8 @@ def log(
             ):
                 _logger.warning(_LOG_MODEL_MISSING_SIGNATURE_WARNING)
             flavor.save_model(path=local_path, mlflow_model=mlflow_model, **kwargs)
+            # TODO: can we move saving MLModel file here? This way we can set model_size_bytes
+            #  here for all flavors here, instead of having to implement it in each save_model
             mlflow.tracking.fluent.log_artifacts(local_path, mlflow_model.artifact_path)
             try:
                 mlflow.tracking.fluent._record_logged_model(mlflow_model)

mlflow/pyfunc/__init__.py

@@ -288,7 +288,7 @@
     _copy_file_or_tree,
     get_or_create_nfs_tmp_dir,
     get_or_create_tmp_dir,
-    write_to,
+    write_to, get_total_size,
 )
 from mlflow.utils.model_utils import (
     _add_code_from_conf_to_system_path,
@@ -2179,6 +2179,7 @@ def _save_model_with_loader_module_and_data_path(
         python_env=_PYTHON_ENV_FILE_NAME,
         model_config=model_config,
     )
+    mlflow_model.model_size_bytes = get_total_size(str(path))
     mlflow_model.save(os.path.join(path, MLMODEL_FILE_NAME))
 
     if conda_env is None:

mlflow/pyfunc/model.py

@@ -32,7 +32,7 @@
     _process_pip_requirements,
     _PythonEnv,
 )
-from mlflow.utils.file_utils import TempDir, _copy_file_or_tree, write_to
+from mlflow.utils.file_utils import TempDir, _copy_file_or_tree, write_to, get_total_size
 from mlflow.utils.model_utils import _get_flavor_configuration
 from mlflow.utils.requirements_utils import _get_pinned_requirement
 
@@ -317,6 +317,7 @@ def _save_model_with_class_artifacts_params(
         model_config=model_config,
         **custom_model_config_kwargs,
     )
+    mlflow_model.model_size_bytes = get_total_size(str(path))
     mlflow_model.save(os.path.join(path, MLMODEL_FILE_NAME))
 
     if conda_env is None: