precommit

Signed-off-by: wenfeiy-db <[email protected]>

mlflow/langchain/__init__.py

@@ -48,7 +48,7 @@
     _PythonEnv,
     _validate_env_arguments,
 )
-from mlflow.utils.file_utils import write_to, get_total_file_size
+from mlflow.utils.file_utils import get_total_file_size, write_to
 from mlflow.utils.model_utils import (
     _add_code_from_conf_to_system_path,
     _get_flavor_configuration,
@@ -298,7 +298,7 @@ def load_retriever(persist_directory):
     try:
         mlflow_model.model_size_bytes = get_total_file_size(str(path))
     except Exception as e:
-        logger.info(f"Fail to get the total size of {str(path)} because of error :{e}")
+        logger.info(f"Fail to get the total size of {path!s} because of error :{e}")
     mlflow_model.save(os.path.join(path, MLMODEL_FILE_NAME))
 
     if conda_env is None:

mlflow/lightgbm/__init__.py

@@ -69,7 +69,7 @@
     _PythonEnv,
     _validate_env_arguments,
 )
-from mlflow.utils.file_utils import write_to, get_total_file_size
+from mlflow.utils.file_utils import get_total_file_size, write_to
 from mlflow.utils.mlflow_tags import (
     MLFLOW_DATASET_CONTEXT,
 )
@@ -216,7 +216,7 @@ def save_model(
     try:
         mlflow_model.model_size_bytes = get_total_file_size(str(path))
     except Exception as e:
-        _logger.info(f"Fail to get the total size of {str(path)} because of error :{e}")
+        _logger.info(f"Fail to get the total size of {path!s} because of error :{e}")
     mlflow_model.save(os.path.join(path, MLMODEL_FILE_NAME))
 
     if conda_env is None:

mlflow/pyfunc/__init__.py

@@ -2183,7 +2183,7 @@ def _save_model_with_loader_module_and_data_path(
     try:
         mlflow_model.model_size_bytes = get_total_file_size(str(path))
     except Exception as e:
-        _logger.info(f"Fail to get the total size of {str(path)} because of error :{e}")
+        _logger.info(f"Fail to get the total size of {path!s} because of error :{e}")
     mlflow_model.save(os.path.join(path, MLMODEL_FILE_NAME))
 
     if conda_env is None:

mlflow/pytorch/__init__.py

@@ -50,7 +50,8 @@
 )
 from mlflow.utils.file_utils import (
     TempDir,
-    write_to, get_total_file_size,
+    get_total_file_size,
+    write_to,
 )
 from mlflow.utils.model_utils import (
     _add_code_from_conf_to_system_path,
@@ -551,7 +552,7 @@ class LinearNNModel(torch.nn.Module):
     try:
         mlflow_model.model_size_bytes = get_total_file_size(str(path))
     except Exception as e:
-        _logger.info(f"Fail to get the total size of {str(path)} because of error :{e}")
+        _logger.info(f"Fail to get the total size of {path!s} because of error :{e}")
     mlflow_model.save(os.path.join(path, MLMODEL_FILE_NAME))
 
     if conda_env is None:

mlflow/sentence_transformers/__init__.py

@@ -33,7 +33,7 @@
     _PythonEnv,
     _validate_env_arguments,
 )
-from mlflow.utils.file_utils import write_to, get_total_file_size
+from mlflow.utils.file_utils import get_total_file_size, write_to
 from mlflow.utils.model_utils import (
     _add_code_from_conf_to_system_path,
     _download_artifact_from_uri,
@@ -172,7 +172,7 @@ class that describes the model's inputs and outputs. If not specified but an
     try:
         mlflow_model.model_size_bytes = get_total_file_size(str(path))
     except Exception as e:
-        _logger.info(f"Fail to get the total size of {str(path)} because of error :{e}")
+        _logger.info(f"Fail to get the total size of {path!s} because of error :{e}")
     mlflow_model.save(str(path.joinpath(MLMODEL_FILE_NAME)))
 
     if inference_config:

mlflow/sklearn/__init__.py

@@ -65,7 +65,7 @@
     _PythonEnv,
     _validate_env_arguments,
 )
-from mlflow.utils.file_utils import write_to, get_total_file_size
+from mlflow.utils.file_utils import get_total_file_size, write_to
 from mlflow.utils.mlflow_tags import (
     MLFLOW_AUTOLOGGING,
     MLFLOW_DATASET_CONTEXT,
@@ -292,7 +292,7 @@ def save_model(
     try:
         mlflow_model.model_size_bytes = get_total_file_size(str(path))
     except Exception as e:
-        _logger.info(f"Fail to get the total size of {str(path)} because of error :{e}")
+        _logger.info(f"Fail to get the total size of {path!s} because of error :{e}")
     mlflow_model.save(os.path.join(path, MLMODEL_FILE_NAME))
 
     if conda_env is None:

mlflow/tensorflow/__init__.py

@@ -64,7 +64,7 @@
     _PythonEnv,
     _validate_env_arguments,
 )
-from mlflow.utils.file_utils import write_to, get_total_file_size
+from mlflow.utils.file_utils import get_total_file_size, write_to
 from mlflow.utils.model_utils import (
     _add_code_from_conf_to_system_path,
     _get_flavor_configuration,
@@ -459,7 +459,7 @@ def save_model(
     try:
         mlflow_model.model_size_bytes = get_total_file_size(str(path))
     except Exception as e:
-        _logger.info(f"Fail to get the total size of {str(path)} because of error :{e}")
+        _logger.info(f"Fail to get the total size of {path!s} because of error :{e}")
 
     # save mlflow_model to path/MLmodel
     mlflow_model.save(os.path.join(path, MLMODEL_FILE_NAME))

mlflow/transformers/__init__.py

@@ -532,7 +532,7 @@ def save_model(
     try:
         mlflow_model.model_size_bytes = get_total_file_size(str(path))
     except Exception as e:
-        _logger.info(f"Fail to get the total size of {str(path)} because of error :{e}")
+        _logger.info(f"Fail to get the total size of {path!s} because of error :{e}")
     mlflow_model.save(str(path.joinpath(MLMODEL_FILE_NAME)))
 
     if conda_env is None:

mlflow/xgboost/__init__.py

@@ -70,7 +70,7 @@
     _PythonEnv,
     _validate_env_arguments,
 )
-from mlflow.utils.file_utils import write_to, get_total_file_size
+from mlflow.utils.file_utils import get_total_file_size, write_to
 from mlflow.utils.mlflow_tags import (
     MLFLOW_DATASET_CONTEXT,
 )
@@ -187,7 +187,7 @@ def save_model(
     try:
         mlflow_model.model_size_bytes = get_total_file_size(str(path))
     except Exception as e:
-        _logger.info(f"Fail to get the total size of {str(path)} because of error :{e}")
+        _logger.info(f"Fail to get the total size of {path!s} because of error :{e}")
     mlflow_model.save(os.path.join(path, MLMODEL_FILE_NAME))
 
     if conda_env is None:

tests/helper_functions.py

@@ -12,8 +12,8 @@
 import time
 import uuid
 from contextlib import ExitStack, contextmanager
-from unittest import mock
 from functools import wraps
+from unittest import mock
 
 import pytest
 import requests
@@ -660,4 +660,4 @@ def decorated_func(*args, **kwargs):
 
         return decorated_func
 
-    return flaky_test_func
+    return flaky_test_func

tests/transformers/test_transformers_model_export.py

@@ -62,8 +62,8 @@
     _get_deps_from_requirement_file,
     _mlflow_major_version_string,
     assert_register_model_called_with_local_model_path,
+    flaky,
     pyfunc_serve_and_score_model,
-    flaky
 )
 
 pytestmark = pytest.mark.large

tests/utils/test_file_utils.py

@@ -4,18 +4,18 @@
 import json
 import os
 import pickle
-import cloudpickle
 import shutil
 import stat
 import tarfile
 
+import cloudpickle
 import jinja2.exceptions
 import pandas as pd
 import pytest
-import transformers
 import sklearn.linear_model as glm
-from sklearn import datasets
+import transformers
 from pyspark.sql import SparkSession
+from sklearn import datasets
 
 import mlflow
 from mlflow.exceptions import MissingConfigException, MlflowException
@@ -33,7 +33,7 @@
 )
 from mlflow.utils.os import is_windows
 
-from tests.helper_functions import random_file, random_int, safe_edit_yaml, flaky
+from tests.helper_functions import flaky, random_file, random_int, safe_edit_yaml
 from tests.projects.utils import TEST_PROJECT_DIR
 
 
@@ -391,11 +391,17 @@ def test_get_total_size_basic(tmp_path):
     assert get_total_file_size(subdir) == 22
 
     path_not_exists = os.path.join(root, "does_not_exist")
-    with pytest.raises(MlflowException, match=f"does not exist.",):
+    with pytest.raises(
+        MlflowException,
+        match="does not exist.",
+    ):
         get_total_file_size(path_not_exists)
 
     path_file = os.path.join(root, "file1.txt")
-    with pytest.raises(MlflowException, match=f"is not a directory.",):
+    with pytest.raises(
+        MlflowException,
+        match="is not a directory.",
+    ):
         get_total_file_size(path_file)