diff --git a/mlflow/models/model.py b/mlflow/models/model.py index 76a9e3a6255b11..8f6f24e23c159a 100644 --- a/mlflow/models/model.py +++ b/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) diff --git a/mlflow/pyfunc/__init__.py b/mlflow/pyfunc/__init__.py index 7462cb13ace7fd..6af07e8c11898c 100644 --- a/mlflow/pyfunc/__init__.py +++ b/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: diff --git a/mlflow/pyfunc/model.py b/mlflow/pyfunc/model.py index 69902d064e8c50..917ab0943964b4 100644 --- a/mlflow/pyfunc/model.py +++ b/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: