mlflow: [BUG] Lazy load issue with tensorflow

Issues Policy acknowledgement

• I have read and agree to submit bug reports in accordance with the issues policy

Where did you encounter this bug?

Local machine

Willingness to contribute

No. I cannot contribute a bug fix at this time.

MLflow version

• Client: 2.8.1
• Tracking server: 2.8.1

System information

Python 3.11 Linux

Describe the problem

When importing MLFlow, one can reference certain submodules referring to an ML framework without having them instakked

e.g. mlflow.sklearn, mlflow.pytorch mlflow.tensorflow

however, with mlflow.tensorflow the error ModuleNotFoundError: No module named 'tensorflow' occurs; it seems MLFlow has broken lazy loading for the tensorflow module.

Tracking information

No response

Code to reproduce issue

import mlflow
mlflow.tensorflow

Stack trace

Traceback (most recent call last):
  File "<stdin>", line 1, in <module>
  File "<frozen importlib._bootstrap>", line 310, in _module_repr
  File "/home/yaxattax/mlflow_repro/lib/python3.11/site-packages/mlflow/utils/lazy_load.py", line 41, in __getattr__
    module = self._load()
             ^^^^^^^^^^^^
  File "/home/yaxattax/mlflow_repro/lib/python3.11/site-packages/mlflow/utils/lazy_load.py", line 30, in _load
    module = importlib.import_module(self.__name__)
             ^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^
  File "/usr/lib64/python3.11/importlib/__init__.py", line 126, in import_module
    return _bootstrap._gcd_import(name[level:], package, level)
           ^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^
  File "<frozen importlib._bootstrap>", line 1204, in _gcd_import
  File "<frozen importlib._bootstrap>", line 1176, in _find_and_load
  File "<frozen importlib._bootstrap>", line 1147, in _find_and_load_unlocked
  File "<frozen importlib._bootstrap>", line 690, in _load_unlocked
  File "<frozen importlib._bootstrap_external>", line 940, in exec_module
  File "<frozen importlib._bootstrap>", line 241, in _call_with_frames_removed
  File "/home/yaxattax/mlflow_repro/lib/python3.11/site-packages/mlflow/tensorflow/__init__.py", line 35, in <module>
    from mlflow.tensorflow.callback import MLflowCallback  # noqa: F401
    ^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^
  File "/home/yaxattax/mlflow_repro/lib/python3.11/site-packages/mlflow/tensorflow/callback.py", line 1, in <module>
    from tensorflow import keras
ModuleNotFoundError: No module named 'tensorflow'

Other info / logs

No response

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/gateway: AI Gateway service, Gateway client APIs, third-party Gateway integrations
• area/model-registry: Model Registry service, APIs, and the fluent client calls for Model Registry
• 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