Repository: mlflow/mlflow

Issue:
[BUG] Transformers flavor expects both PyTorch and TensorFlow in environment
### Issues Policy acknowledgement

- [X] I have read and agree to submit bug reports in accordance with the [issues policy](https://www.github.com/mlflow/mlflow/blob/master/ISSUE_POLICY.md)

### Willingness to contribute

Yes. I would be willing to contribute a fix for this bug with guidance from the MLflow community.

### MLflow version

- Client: 2.4.0
- Tracking server: 2.4.0


### System information

- **OS Platform and Distribution (e.g., Linux Ubuntu 16.04)**: Ubuntu 22.04
- **Python version**: 3.9.11


### Describe the problem

The MLflow `transformers` flavor fails when logging a model should an environment not contain both PyTorch and TensorFlow. I would expect most user / CI environments favor one framework over the other, and would not install both. HuggingFace's `transformers` inherently has checks to see if either is installed, and does not require both.

### Tracking information

<!-- PLEASE KEEP BACKTICKS AND CHECK PREVIEW -->
N/A


### Code to reproduce issue

<!-- PLEASE KEEP BACKTICKS AND CHECK PREVIEW -->
With environment that has PyTorch and does not have TensorFlow:

```python
from transformers import AutoConfig, AutoModel, AutoTokenizer

import mlflow

model_name_or_path = "bert-large-cased-whole-word-masking-finetuned-squad"
transformers_config = AutoConfig.from_pretrained(model_name_or_path)
transformers_model = AutoModel.from_pretrained(model_name_or_path, config=transformers_config)
tokenizer =  AutoTokenizer.from_pretrained(model_name_or_path)

components = {"model": transformers_model , "tokenizer": tokenizer}

with mlflow.start_run():
    mlflow.transformers.log_model(components, registered_model_name=model_name_or_path, artifact_path="transformers_model")
```


### Stack trace

<!-- PLEASE KEEP BACKTICKS AND CHECK PREVIEW -->
Below is when I have only PyTorch and not TensorFlow installed.

When TF is installed, this error does not occur.
```
- This IS expected if you are initializing BertModel from the checkpoint of a model trained on another task or with another architecture (e.g. initializing a BertForSequenceClassification model from a BertForPreTraining model).
- This IS NOT expected if you are initializing BertModel from the checkpoint of a model that you expect to be exactly identical (initializing a BertForSequenceClassification model from a BertForSequenceClassification model).
The model 'BertModel' is not supported for question-answering. Supported models are ['AlbertForQuestionAnswering', 'BartForQuestionAnswering', 'BertForQuestionAnswering', 'BigBirdForQuestionAnswering', 'BigBirdPegasusForQuestionAnswering', 'BloomForQuestionAnswering', 'CamembertForQuestionAnswering', 'CanineForQuestionAnswering', 'ConvBertForQuestionAnswering', 'Data2VecTextForQuestionAnswering', 'DebertaForQuestionAnswering', 'DebertaV2ForQuestionAnswering', 'DistilBertForQuestionAnswering', 'ElectraForQuestionAnswering', 'ErnieForQuestionAnswering', 'ErnieMForQuestionAnswering', 'FlaubertForQuestionAnsweringSimple', 'FNetForQuestionAnswering', 'FunnelForQuestionAnswering', 'GPT2ForQuestionAnswering', 'GPTNeoForQuestionAnswering', 'GPTNeoXForQuestionAnswering', 'GPTJForQuestionAnswering', 'IBertForQuestionAnswering', 'LayoutLMv2ForQuestionAnswering', 'LayoutLMv3ForQuestionAnswering', 'LEDForQuestionAnswering', 'LiltForQuestionAnswering', 'LongformerForQuestionAnswering', 'LukeForQuestionAnswering', 'LxmertForQuestionAnswering', 'MarkupLMForQuestionAnswering', 'MBartForQuestionAnswering', 'MegaForQuestionAnswering', 'MegatronBertForQuestionAnswering', 'MobileBertForQuestionAnswering', 'MPNetForQuestionAnswering', 'MvpForQuestionAnswering', 'NezhaForQuestionAnswering', 'NystromformerForQuestionAnswering', 'OPTForQuestionAnswering', 'QDQBertForQuestionAnswering', 'ReformerForQuestionAnswering', 'RemBertForQuestionAnswering', 'RobertaForQuestionAnswering', 'RobertaPreLayerNormForQuestionAnswering', 'RoCBertForQuestionAnswering', 'RoFormerForQuestionAnswering', 'SplinterForQuestionAnswering', 'SqueezeBertForQuestionAnswering', 'XLMForQuestionAnsweringSimple', 'XLMRobertaForQuestionAnswering', 'XLMRobertaXLForQuestionAnswering', 'XLNetForQuestionAnsweringSimple', 'XmodForQuestionAnswering', 'YosoForQuestionAnswering'].
2023/06/09 10:48:39 WARNING mlflow.utils.requirements_utils: Found torch version (1.13.0+cu116) contains a local version label (+cu116). MLflow logged a pip requirement for this package as 'torch==1.13.0' without the local version label to make it installable from PyPI. To specify pip requirements containing local version labels, please use `conda_env` or `pip_requirements`.
2023/06/09 10:48:39 WARNING mlflow.utils.requirements_utils: Found torchvision version (0.14.0+cu116) contains a local version label (+cu116). MLflow logged a pip requirement for this package as 'torchvision==0.14.0' without the local version label to make it installable from PyPI. To specify pip requirements containing local version labels, please use `conda_env` or `pip_requirements`.
2023/06/09 10:48:39 WARNING mlflow.utils.requirements_utils: Found torch version (1.13.0+cu116) contains a local version label (+cu116). MLflow logged a pip requirement for this package as 'torch==1.13.0' without the local version label to make it installable from PyPI. To specify pip requirements containing local version labels, please use `conda_env` or `pip_requirements`.
2023/06/09 10:48:39 WARNING mlflow.utils.requirements_utils: Found torchvision version (0.14.0+cu116) contains a local version label (+cu116). MLflow logged a pip requirement for this package as 'torchvision==0.14.0' without the local version label to make it installable from PyPI. To specify pip requirements containing local version labels, please use `conda_env` or `pip_requirements`.
---------------------------------------------------------------------------
StopIteration                             Traceback (most recent call last)
File ~/.pyenv/versions/3.9.11/envs/mlflow/lib/python3.9/site-packages/importlib_metadata/__init__.py:409, in Distribution.from_name(cls, name)
    408 try:
--> 409     return next(iter(cls.discover(name=name)))
    410 except StopIteration:

StopIteration: 

During handling of the above exception, another exception occurred:

PackageNotFoundError                      Traceback (most recent call last)
File ~/.pyenv/versions/3.9.11/envs/mlflow/lib/python3.9/site-packages/mlflow/utils/requirements_utils.py:225, in _get_installed_version(package, module)
    224 try:
--> 225     version = importlib_metadata.version(package)
    226 except importlib_metadata.PackageNotFoundError:
    227     # Note `importlib_metadata.version(package)` is not necessarily equal to
    228     # `__import__(package).__version__`. See the example for pytorch below.
   (...)
    235     # $ python -c "import importlib_metadata; print(importlib_metadata.version('torch'))"
    236     # 1.9.0

File ~/.pyenv/versions/3.9.11/envs/mlflow/lib/python3.9/site-packages/importlib_metadata/__init__.py:910, in version(distribution_name)
    904 """Get the version string for the named package.
    905 
    906 :param distribution_name: The name of the distribution package to query.
    907 :return: The version string for the package as defined in the package's
    908     "Version" metadata key.
    909 """
--> 910 return distribution(distribution_name).version

File ~/.pyenv/versions/3.9.11/envs/mlflow/lib/python3.9/site-packages/importlib_metadata/__init__.py:883, in distribution(distribution_name)
    878 """Get the ``Distribution`` instance for the named package.
    879 
    880 :param distribution_name: The name of the distribution package as a string.
    881 :return: A ``Distribution`` instance (or subclass thereof).
    882 """
--> 883 return Distribution.from_name(distribution_name)

File ~/.pyenv/versions/3.9.11/envs/mlflow/lib/python3.9/site-packages/importlib_metadata/__init__.py:411, in Distribution.from_name(cls, name)
    410 except StopIteration:
--> 411     raise PackageNotFoundError(name)

PackageNotFoundError: No package metadata was found for accelerate

During handling of the above exception, another exception occurred:

ModuleNotFoundError                       Traceback (most recent call last)
File ~/.pyenv/versions/3.9.11/envs/mlflow/lib/python3.9/site-packages/mlflow/transformers.py:139, in get_default_pip_requirements(model)
    138     base_reqs = ["transformers", *_model_packages(model)]
--> 139     return [_get_pinned_requirement(module) for module in base_reqs]
    140 except Exception as e:

File ~/.pyenv/versions/3.9.11/envs/mlflow/lib/python3.9/site-packages/mlflow/transformers.py:139, in <listcomp>(.0)
    138     base_reqs = ["transformers", *_model_packages(model)]
--> 139     return [_get_pinned_requirement(module) for module in base_reqs]
    140 except Exception as e:

File ~/.pyenv/versions/3.9.11/envs/mlflow/lib/python3.9/site-packages/mlflow/utils/requirements_utils.py:476, in _get_pinned_requirement(package, version, module)
    475 if version is None:
--> 476     version_raw = _get_installed_version(package, module)
    477     local_version_label = _get_local_version_label(version_raw)

File ~/.pyenv/versions/3.9.11/envs/mlflow/lib/python3.9/site-packages/mlflow/utils/requirements_utils.py:237, in _get_installed_version(package, module)
    226 except importlib_metadata.PackageNotFoundError:
    227     # Note `importlib_metadata.version(package)` is not necessarily equal to
    228     # `__import__(package).__version__`. See the example for pytorch below.
   (...)
    235     # $ python -c "import importlib_metadata; print(importlib_metadata.version('torch'))"
    236     # 1.9.0
--> 237     version = __import__(module or package).__version__
    239 # Strip the suffix from `dev` versions of PySpark, which are not available for installation
    240 # from Anaconda or PyPI

ModuleNotFoundError: No module named 'accelerate'

During handling of the above exception, another exception occurred:

StopIteration                             Traceback (most recent call last)
File ~/.pyenv/versions/3.9.11/envs/mlflow/lib/python3.9/site-packages/importlib_metadata/__init__.py:409, in Distribution.from_name(cls, name)
    408 try:
--> 409     return next(iter(cls.discover(name=name)))
    410 except StopIteration:

StopIteration: 

During handling of the above exception, another exception occurred:

PackageNotFoundError                      Traceback (most recent call last)
File ~/.pyenv/versions/3.9.11/envs/mlflow/lib/python3.9/site-packages/mlflow/utils/requirements_utils.py:225, in _get_installed_version(package, module)
    224 try:
--> 225     version = importlib_metadata.version(package)
    226 except importlib_metadata.PackageNotFoundError:
    227     # Note `importlib_metadata.version(package)` is not necessarily equal to
    228     # `__import__(package).__version__`. See the example for pytorch below.
   (...)
    235     # $ python -c "import importlib_metadata; print(importlib_metadata.version('torch'))"
    236     # 1.9.0

File ~/.pyenv/versions/3.9.11/envs/mlflow/lib/python3.9/site-packages/importlib_metadata/__init__.py:910, in version(distribution_name)
    904 """Get the version string for the named package.
    905 
    906 :param distribution_name: The name of the distribution package to query.
    907 :return: The version string for the package as defined in the package's
    908     "Version" metadata key.
    909 """
--> 910 return distribution(distribution_name).version

File ~/.pyenv/versions/3.9.11/envs/mlflow/lib/python3.9/site-packages/importlib_metadata/__init__.py:883, in distribution(distribution_name)
    878 """Get the ``Distribution`` instance for the named package.
    879 
    880 :param distribution_name: The name of the distribution package as a string.
    881 :return: A ``Distribution`` instance (or subclass thereof).
    882 """
--> 883 return Distribution.from_name(distribution_name)

File ~/.pyenv/versions/3.9.11/envs/mlflow/lib/python3.9/site-packages/importlib_metadata/__init__.py:411, in Distribution.from_name(cls, name)
    410 except StopIteration:
--> 411     raise PackageNotFoundError(name)

PackageNotFoundError: No package metadata was found for tensorflow

During handling of the above exception, another exception occurred:

ModuleNotFoundError                       Traceback (most recent call last)
Cell In[29], line 13
     10 components = {"model": transformers_model , "tokenizer": tokenizer}
     12 with mlflow.start_run():
---> 13     mlflow.transformers.log_model(components, model_name=model_name_or_path, artifact_path="transformers_model")

File ~/.pyenv/versions/3.9.11/envs/mlflow/lib/python3.9/site-packages/mlflow/utils/docstring_utils.py:235, in docstring_version_compatibility_warning.<locals>.annotated_func.<locals>.version_func(*args, **kwargs)
    233 if installed_version < Version(min_ver) or installed_version > Version(max_ver):
    234     warnings.warn(notice, category=FutureWarning, stacklevel=2)
--> 235 return func(*args, **kwargs)

File ~/.pyenv/versions/3.9.11/envs/mlflow/lib/python3.9/site-packages/mlflow/transformers.py:723, in log_model(transformers_model, artifact_path, processor, task, model_card, inference_config, code_paths, registered_model_name, signature, input_example, await_registration_for, pip_requirements, extra_pip_requirements, conda_env, metadata, **kwargs)
    519 @experimental
    520 @docstring_version_compatibility_warning(integration_name=FLAVOR_NAME)
    521 @format_docstring(LOG_MODEL_PARAM_DOCS.format(package_name=FLAVOR_NAME))
   (...)
    538     **kwargs,
    539 ):
    540     """
    541     Log a ``transformers`` object as an MLflow artifact for the current run.
    542 
   (...)
    721     :param kwargs: Additional arguments for :py:class:`mlflow.models.model.Model`
    722     """
--> 723     return Model.log(
    724         artifact_path=artifact_path,
    725         flavor=mlflow.transformers,
    726         registered_model_name=registered_model_name,
    727         await_registration_for=await_registration_for,
    728         metadata=metadata,
    729         transformers_model=transformers_model,
    730         processor=processor,
    731         task=task,
    732         model_card=model_card,
    733         inference_config=inference_config,
    734         conda_env=conda_env,
    735         code_paths=code_paths,
    736         signature=signature,
    737         input_example=input_example,
    738         pip_requirements=pip_requirements,
    739         extra_pip_requirements=extra_pip_requirements,
    740         **kwargs,
    741     )

File ~/.pyenv/versions/3.9.11/envs/mlflow/lib/python3.9/site-packages/mlflow/models/model.py:562, in Model.log(cls, artifact_path, flavor, registered_model_name, await_registration_for, metadata, **kwargs)
    560 run_id = mlflow.tracking.fluent._get_or_start_run().info.run_id
    561 mlflow_model = cls(artifact_path=artifact_path, run_id=run_id, metadata=metadata)
--> 562 flavor.save_model(path=local_path, mlflow_model=mlflow_model, **kwargs)
    563 mlflow.tracking.fluent.log_artifacts(local_path, mlflow_model.artifact_path)
    564 tracking_uri = _resolve_tracking_uri()

File ~/.pyenv/versions/3.9.11/envs/mlflow/lib/python3.9/site-packages/mlflow/utils/docstring_utils.py:235, in docstring_version_compatibility_warning.<locals>.annotated_func.<locals>.version_func(*args, **kwargs)
    233 if installed_version < Version(min_ver) or installed_version > Version(max_ver):
    234     warnings.warn(notice, category=FutureWarning, stacklevel=2)
--> 235 return func(*args, **kwargs)

File ~/.pyenv/versions/3.9.11/envs/mlflow/lib/python3.9/site-packages/mlflow/transformers.py:497, in save_model(transformers_model, path, processor, task, model_card, inference_config, code_paths, mlflow_model, signature, input_example, pip_requirements, extra_pip_requirements, conda_env, metadata, **kwargs)
    495 if conda_env is None:
    496     if pip_requirements is None:
--> 497         default_reqs = get_default_pip_requirements(transformers_model.model)
    498         inferred_reqs = infer_pip_requirements(str(path), FLAVOR_NAME, fallback=default_reqs)
    499         default_reqs = sorted(set(inferred_reqs).union(default_reqs))

File ~/.pyenv/versions/3.9.11/envs/mlflow/lib/python3.9/site-packages/mlflow/transformers.py:141, in get_default_pip_requirements(model)
    139     return [_get_pinned_requirement(module) for module in base_reqs]
    140 except Exception as e:
--> 141     dependencies = [
    142         _get_pinned_requirement(module)
    143         for module in ["transformers", "torch", "torchvision", "tensorflow"]
    144     ]
    145     _logger.warning(
    146         "Could not infer model execution engine type due to huggingface_hub not "
    147         "being installed or unable to connect in online mode. Adding full "
    148         f"dependency chain: {dependencies}. \nFailure cause: {str(e)}"
    149     )
    150     return dependencies

File ~/.pyenv/versions/3.9.11/envs/mlflow/lib/python3.9/site-packages/mlflow/transformers.py:142, in <listcomp>(.0)
    139     return [_get_pinned_requirement(module) for module in base_reqs]
    140 except Exception as e:
    141     dependencies = [
--> 142         _get_pinned_requirement(module)
    143         for module in ["transformers", "torch", "torchvision", "tensorflow"]
    144     ]
    145     _logger.warning(
    146         "Could not infer model execution engine type due to huggingface_hub not "
    147         "being installed or unable to connect in online mode. Adding full "
    148         f"dependency chain: {dependencies}. \nFailure cause: {str(e)}"
    149     )
    150     return dependencies

File ~/.pyenv/versions/3.9.11/envs/mlflow/lib/python3.9/site-packages/mlflow/utils/requirements_utils.py:476, in _get_pinned_requirement(package, version, module)
    465 """
    466 Returns a string representing a pinned pip requirement to install the specified package and
    467 version (e.g. 'mlflow==1.2.3').
   (...)
    473                to `package`.
    474 """
    475 if version is None:
--> 476     version_raw = _get_installed_version(package, module)
    477     local_version_label = _get_local_version_label(version_raw)
    478     if local_version_label:

File ~/.pyenv/versions/3.9.11/envs/mlflow/lib/python3.9/site-packages/mlflow/utils/requirements_utils.py:237, in _get_installed_version(package, module)
    225     version = importlib_metadata.version(package)
    226 except importlib_metadata.PackageNotFoundError:
    227     # Note `importlib_metadata.version(package)` is not necessarily equal to
    228     # `__import__(package).__version__`. See the example for pytorch below.
   (...)
    235     # $ python -c "import importlib_metadata; print(importlib_metadata.version('torch'))"
    236     # 1.9.0
--> 237     version = __import__(module or package).__version__
    239 # Strip the suffix from `dev` versions of PySpark, which are not available for installation
    240 # from Anaconda or PyPI
    241 if package == "pyspark":

ModuleNotFoundError: No module named 'tensorflow'
```


### Other info / logs

<!-- PLEASE KEEP BACKTICKS AND CHECK PREVIEW -->
N/A


### What component(s) does this bug affect?

- [ ] `area/artifacts`: Artifact stores and artifact logging
- [ ] `area/build`: Build and test infrastructure for MLflow
- [ ] `area/docs`: MLflow documentation pages
- [ ] `area/examples`: Example code
- [ ] `area/model-registry`: Model Registry service, APIs, and the fluent client calls for Model Registry
- [X] `area/models`: MLmodel format, model serialization/deserialization, flavors
- [ ] `area/recipes`: Recipes, Recipe APIs, Recipe configs, Recipe Templates
- [ ] `area/projects`: MLproject format, project running backends
- [ ] `area/scoring`: MLflow Model server, model deployment tools, Spark UDFs
- [ ] `area/server-infra`: MLflow Tracking server backend
- [ ] `area/tracking`: Tracking Service, tracking client APIs, autologging

### What interface(s) does this bug affect?

- [ ] `area/uiux`: Front-end, user experience, plotting, JavaScript, JavaScript dev server
- [ ] `area/docker`: Docker use across MLflow's components, such as MLflow Projects and MLflow Models
- [ ] `area/sqlalchemy`: Use of SQLAlchemy in the Tracking Service or Model Registry
- [ ] `area/windows`: Windows support

### What language(s) does this bug affect?

- [ ] `language/r`: R APIs and clients
- [ ] `language/java`: Java APIs and clients
- [ ] `language/new`: Proposals for new client languages

### What integration(s) does this bug affect?

- [ ] `integrations/azure`: Azure and Azure ML integrations
- [ ] `integrations/sagemaker`: SageMaker integrations
- [ ] `integrations/databricks`: Databricks integrations


Discussion:
Looking at https://github.com/mlflow/mlflow/blob/master/mlflow/transformers.py#L139, I see that it was looking for `tensorflow` since `accelerate` was not in the environment (https://github.com/mlflow/mlflow/blob/master/mlflow/transformers.py#LL113C30-L113C30) . `accelerate` should be optional, right? Why is this required?
Thank you for reporting. We'll be sure to provide a fix for this.
<!-- assign-maintainer -->
@BenWilson2 @dbczumar @harupy @WeichenXu123 Please assign a maintainer and start triaging this issue.

Resolving patch:
diff --git a/mlflow/transformers.py b/mlflow/transformers.py
--- a/mlflow/transformers.py
+++ b/mlflow/transformers.py
@@ -115,7 +115,14 @@ def _model_packages(model) -> List[str]:
     """
     engine = _get_engine_type(model)
     if engine == "torch":
-        return ["torch", "torchvision", "accelerate"]
+        packages = ["torch", "torchvision"]
+        try:
+            import accelerate
+
+            packages.append("accelerate")
+        except ImportError:
+            pass
+        return packages
     else:
         return [engine]
 


Test patch:
diff --git a/tests/transformers/test_transformers_model_export.py b/tests/transformers/test_transformers_model_export.py
--- a/tests/transformers/test_transformers_model_export.py
+++ b/tests/transformers/test_transformers_model_export.py
@@ -1,6 +1,7 @@
 import base64
 from functools import wraps
 import gc
+import importlib.util
 import logging
 import json
 import time
@@ -1099,6 +1100,9 @@ def test_transformers_log_with_duplicate_extra_pip_requirements(small_multi_moda
     )
 
 
+@pytest.mark.skipif(
+    importlib.util.find_spec("accelerate") is not None, reason="fails when accelerate is installed"
+)
 def test_transformers_tf_model_save_without_conda_env_uses_default_env_with_expected_dependencies(
     small_seq2seq_pipeline, model_path
 ):
@@ -1109,6 +1113,7 @@ def test_transformers_tf_model_save_without_conda_env_uses_default_env_with_expe
     pip_requirements = _get_deps_from_requirement_file(model_path)
     assert "tensorflow" in pip_requirements
     assert "torch" not in pip_requirements
+    assert "accelerate" not in pip_requirements
 
 
 def test_transformers_pt_model_save_without_conda_env_uses_default_env_with_expected_dependencies(
@@ -1120,6 +1125,23 @@ def test_transformers_pt_model_save_without_conda_env_uses_default_env_with_expe
     )
     pip_requirements = _get_deps_from_requirement_file(model_path)
     assert "tensorflow" not in pip_requirements
+    assert "accelerate" in pip_requirements
+    assert "torch" in pip_requirements
+
+
+@pytest.mark.skipif(
+    importlib.util.find_spec("accelerate") is not None, reason="fails when accelerate is installed"
+)
+def test_transformers_pt_model_save_dependencies_without_accelerate(
+    translation_pipeline, model_path
+):
+    mlflow.transformers.save_model(translation_pipeline, model_path)
+    _assert_pip_requirements(
+        model_path, mlflow.transformers.get_default_pip_requirements(translation_pipeline.model)
+    )
+    pip_requirements = _get_deps_from_requirement_file(model_path)
+    assert "tensorflow" not in pip_requirements
+    assert "accelerate" not in pip_requirements
     assert "torch" in pip_requirements