MLFlow integration requirements duplicate fix (#3011)

* added copy * removed requirements * fixed integration test
zenml-io · Sep 19, 2024 · ea28592 · ea28592
1 parent aef636f
commit ea28592
Show file tree

Hide file tree

Showing 2 changed files with 17 additions and 18 deletions.
diff --git a/src/zenml/integrations/mlflow/__init__.py b/src/zenml/integrations/mlflow/__init__.py
@@ -37,20 +37,6 @@ class MlflowIntegration(Integration):
 
     NAME = MLFLOW
 
-    REQUIREMENTS = [
-        "mlflow>=2.1.1,<3",
-        # TODO: remove this requirement once rapidjson is fixed
-        "python-rapidjson<1.15",
-        # When you do:
-        # pip install zenml
-        # You get all our required dependencies. However, if you follow it with:
-        # zenml integration install mlflow
-        # This downgrades pydantic to v1 even though mlflow does not have
-        # any issues with v2. This is why we have to pin it here so a downgrade
-        # will not happen.
-        "pydantic>=2.8.0,<2.9.0",
-    ]
-
     REQUIREMENTS_IGNORED_ON_UNINSTALL = [
         "python-rapidjson",
         "pydantic",
@@ -71,7 +57,20 @@ def get_requirements(cls, target_os: Optional[str] = None) -> List[str]:
         from zenml.integrations.numpy import NumpyIntegration
         from zenml.integrations.pandas import PandasIntegration
 
-        reqs = cls.REQUIREMENTS
+        reqs = [
+            "mlflow>=2.1.1,<3",
+            # TODO: remove this requirement once rapidjson is fixed
+            "python-rapidjson<1.15",
+            # When you do:
+            #       pip install zenml
+            # You get all our required dependencies. However, if you follow it
+            # with:
+            #       zenml integration install mlflow
+            # This downgrades pydantic to v1 even though mlflow does not have
+            # any issues with v2. This is why we have to pin it here so a
+            # downgrade will not happen.
+            "pydantic>=2.8.0,<2.9.0",
+        ]
         if sys.version_info.minor >= 12:
             logger.debug(
                 "The MLflow integration on Python 3.12 and above is not yet "

diff --git a/tests/integration/functional/cli/test_integration.py b/tests/integration/functional/cli/test_integration.py
@@ -229,9 +229,9 @@ def test_integration_requirements_exporting(tmp_path) -> None:
     from zenml.integrations.mlflow import MlflowIntegration
 
     flow_integration_requirements = set(
-        AirflowIntegration.REQUIREMENTS
-        + KubeflowIntegration.REQUIREMENTS
-        + MlflowIntegration.REQUIREMENTS
+        AirflowIntegration.get_requirements()
+        + KubeflowIntegration.get_requirements()
+        + MlflowIntegration.get_requirements()
     )
 
     command = [