Repository: mlflow/mlflow

Issue:
[BUG] Missing optional dependecies in requirements.txt
### 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)

### Where did you encounter this bug?

Local machine

### 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.1.1
- Tracking server: 2.1.1

### System information

- **OS Platform and Distribution**: Ubuntu 20.04 (WSL2)
- **Python version**: 3.10.13

### Describe the problem

I am running my script using similar command:
```
mlflow run git@gitlab.com:custom-repo.git#scripts/pv_model -e main -P parameter1=test --experiment-name test-model
```

Content of `MLproject`:
```
name: model_example

python_env: python_env.yaml

entry_points:
    ... # list of entry points
```

and content of `python_env.yaml`:
```
python: "3.10"
build_dependencies:
  - pip
dependencies:
  - mlflow==2.1.1
  - custom-lib[models]
```

All the training works fine, model is logged with all other stuff (including parameters, metrics, dependencies). The issue is that content of `requirements.txt` (same applies to `conda.yaml`) looks like this:
```
mlflow<3,>=2.1
cloudpickle==2.2.1
custom-lib==1.0.26
```
where optional dependencies/extras (**models**) are missing. This basically means that no one can install all models dependencies from such `requirements.txt` as those extras deps will be missing.

Is there some workaround or proper way how to deal with such extras dependencies or is it rather a bug that should be fixed?

### Tracking information

<!-- PLEASE KEEP BACKTICKS AND CHECK PREVIEW -->
```shell
REPLACE_ME
```


### Code to reproduce issue

Run `mlflow run` with `MLproject` pointing to `python_env.yaml` with a required optional dependencies like `custom-lib[extras]`.

### Stack trace

No stack trace (command runs correctly).

### Other info / logs

_No response_

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

- [X] `area/artifacts`: Artifact stores and artifact logging
- [ ] `area/build`: Build and test infrastructure for MLflow
- [ ] `area/deployments`: MLflow Deployments client APIs, server, and third-party Deployments integrations
- [ ] `area/docs`: MLflow documentation pages
- [ ] `area/examples`: Example code
- [ ] `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
- [X] `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:
Hi @sxooler! It's a little tricky to infer whether an extra is required for a model, but if you want to ensure that the extra is included in the `requirements.txt` file, you can use the `extra_pip_requirements` param when saving your model. Here's an example (you can also refer to [this page in the docs](https://mlflow.org/docs/latest/deployment/deploy-model-to-kubernetes/index.html#specifying-pip-requirements-using-pip-requirements-and-extra-pip-requirements) for more examples):

```python
import mlflow

class MyModel(mlflow.pyfunc.PythonModel):
    def predict(self, context, model_input):
        return model_input
    
mlflow.pyfunc.save_model(
    path="my-model", 
    python_model=MyModel(),
    extra_pip_requirements=["scikit-learn[openmp]"]
)
```

The `requirements.txt` file should have the package and extra listed:
```
mlflow==2.9.1
cloudpickle==2.2.1
scikit-learn[openmp]
```

Let me know if this solves your problem!
@daniellok-db thank you for your reply. I was not aware of this functionality (using `extra_pip_requirements`), so that's definitely useful, thank you! I have a few follow-up questions.

1. I was not able to make it work in my custom project, but it's a bit more complex one so I will try to find a way to make it work or submit another issue if I find a way how to reproduce it.

2. I was able to reproduce it with a simple project. But there are several issues:
- library version is not listed (as in your example, there is only `scikit-learn[openmp]` in the final `requirements.txt`, no version info)
- (I could provide version directly in `extra_pip_requirements`, I guess, but that would mean I need to update it every time when new lib is released, but my idea was that the latest lib version would be taken during training, logged and then re-used when deploying the trained model)
- should I keep it (e.g. `scikit-learn[openmp]`) in my `python_env.yaml` as well or only as `extra_pip_requirements`?

3. Which leads to another question - how to properly handle `python_env.yaml` and `extra_pip_requirements`? I understand it might be kind of tricky to correctly infer all the requirements including extras, but having a single place/file where I can put my dependencies (which was `python_env.yaml` in my case) is much easier to handle than adding another requirements directly in the code (and I usually have sth like multi-step workflow, therefore several steps share the same MLproject with its dependencies).
No worries glad it was helpful! For your follow up questions:

2. If you want to pin whatever version you have installed when saving your model without having to manually change the code all the time, you can do it like this:
```python
mlflow.pyfunc.save_model(
  ...
  extra_pip_requirements=[f"custom_lib[models]=={custom_lib.__version__}"]
)
```
You can keep `custom_lib[models]` in your `python_env.yaml` as well, there shouldn't be any issues with this.

3. I think this question depends on your workflow. Conceptually, the `requirements.txt` file saved with the model just contains all the dependencies required to load and perform inference with the model. The `python_env.yaml` file associated with your MLproject file is for dependencies required in your broader project (which may be more extensive than what your model requires for inference).

Hope this helps!

Great, this makes sense. To clarify one thing - do I understand it correctly that `python_env.yaml` is used as a source for project creation (including creating `venv` and installing listed packages) when running `mlflow run ...` and then when model is logged mlflow somehow decides which packages will go to `requirements.txt`?

I was confused several times about the content of the logged `requirements.txt` and the difference to `python_env.yaml` used in my workflow. Your explanation sheds some light on my confusion :) But I am still trying to understand what is going on under the hood. Do you have some reference on this topic (more details and so on)?

It seems to me that I should use `python_env.yaml` to provide dependencies to my workflow (all steps, not just inference) and if I cannot (or don't want to) rely on default handling of `requirements.txt` (that is somehow done by mlflow based on the provided dependencies and logged model afaik based on our conversation) I should rather provide my own requirements for inference (using `pip_requirements` and `extra_pip_requirements`). Is that correct?
Yup, you got it! If you're interested in the technical details of how we infer the requirements of a model, you can take a look at the [`infer_pip_requirements` function in the code](https://github.com/mlflow/mlflow/blob/51dd444ffb466ff5fce86b2024563d062f12c2e5/mlflow/utils/environment.py#L383). When you call `save_model()`, eventually this function gets called to infer the requirements of just the model.

I'll close the issue now since it seems more or less resolved, but feel free to let me know if you have anything else you want to know.
@daniellok-db I have one follow-up question - in some case package requirements are logged twice. E.g. in `python_env.yaml`, there is `mypkg[models]==1.0.1` (required for model trying and other parts of the pipeline to run) and I put `extra_pip_requirements=["mypkg[models]==1.0.1"]` to ensure it will get correctly logged as model requirements including extras. Then in `requirements.txt` of the logged model I can find following two lines:
```
mypkg==1.0.1
mypkg[models]==1.0.1
```
You said
> You can keep custom_lib[models] in your python_env.yaml as well, there shouldn't be any issues with this.

But is the described behavior expected? Because it gets logged twice (one without extras inferred by mlflow and one with extras because of used `extra_pip_requirements`).

Ah this may be a bug, I would expect that if a package is provided in extra_pip_requirements, and it's already in the inferred requirements, it should overwrite it. I'll check and get back to you!
<!-- assign-maintainer -->
@mlflow/mlflow-team Please assign a maintainer and start triaging this issue.
Confirmed that the extras should override the main package. I'll follow up with a fix, but probably only going to be able to get to it in a couple of weeks. Thank you for the report!

Resolving patch:
diff --git a/mlflow/utils/environment.py b/mlflow/utils/environment.py
--- a/mlflow/utils/environment.py
+++ b/mlflow/utils/environment.py
@@ -1,6 +1,8 @@
 import logging
 import os
 import re
+import subprocess
+import tempfile
 from collections import Counter
 
 import yaml
@@ -514,12 +516,14 @@ def _process_pip_requirements(
     if not _contains_mlflow_requirement(pip_reqs):
         pip_reqs.insert(0, _generate_mlflow_version_pinning())
 
+    sanitized_pip_reqs = _deduplicate_requirements(pip_reqs)
+
     if constraints:
-        pip_reqs.append(f"-c {_CONSTRAINTS_FILE_NAME}")
+        sanitized_pip_reqs.append(f"-c {_CONSTRAINTS_FILE_NAME}")
 
     # Set `install_mlflow` to False because `pip_reqs` already contains `mlflow`
-    conda_env = _mlflow_conda_env(additional_pip_deps=pip_reqs, install_mlflow=False)
-    return conda_env, pip_reqs, constraints
+    conda_env = _mlflow_conda_env(additional_pip_deps=sanitized_pip_reqs, install_mlflow=False)
+    return conda_env, sanitized_pip_reqs, constraints
 
 
 def _find_duplicate_requirements(requirements):
@@ -541,6 +545,136 @@ def _find_duplicate_requirements(requirements):
     return [package for package, count in package_counts.items() if count > 1]
 
 
+def _deduplicate_requirements(requirements):
+    """
+    De-duplicates a list of pip package requirements, handling complex scenarios such as merging
+    extras and combining version constraints.
+
+    This function processes a list of pip package requirements and de-duplicates them. It handles
+    standard PyPI packages and non-standard requirements (like URLs or local paths). The function
+    merges extras and combines version constraints for duplicate packages. The most restrictive
+    version specifications or the ones with extras are prioritized. If incompatible version
+    constraints are detected, it raises an MlflowException.
+
+    Args:
+        requirements (list of str): A list of pip package requirement strings.
+
+    Returns:
+        list of str: A deduplicated list of pip package requirements.
+
+    Raises:
+        MlflowException: If incompatible version constraints are detected among the provided
+                         requirements.
+
+    Examples:
+        - Input: ["packageA", "packageA==1.0"]
+          Output: ["packageA==1.0"]
+
+        - Input: ["packageX>1.0", "packageX[extras]", "packageX<2.0"]
+          Output: ["packageX[extras]>1.0,<2.0"]
+
+        - Input: ["markdown[extra1]>=3.5.1", "markdown[extra2]<4", "markdown"]
+          Output: ["markdown[extra1,extra2]>=3.5.1,<4"]
+
+        - Input: ["scikit-learn==1.1", "scikit-learn<1"]
+          Raises MlflowException indicating incompatible versions.
+
+    Note:
+        - Non-standard requirements (like URLs or file paths) are included as-is.
+        - If a requirement appears multiple times with different sets of extras, they are merged.
+        - The function uses `_validate_version_constraints` to check for incompatible version
+          constraints by doing a dry-run pip install of a requirements collection.
+    """
+    deduped_reqs = {}
+
+    for req in requirements:
+        try:
+            parsed_req = Requirement(req)
+            base_pkg = parsed_req.name
+
+            existing_req = deduped_reqs.get(base_pkg)
+
+            if not existing_req:
+                deduped_reqs[base_pkg] = parsed_req
+            else:
+                # Verify that there are not unresolvable constraints applied if set and combine
+                # if possible
+                if (
+                    existing_req.specifier
+                    and parsed_req.specifier
+                    and existing_req.specifier != parsed_req.specifier
+                ):
+                    _validate_version_constraints([str(existing_req), req])
+                    parsed_req.specifier = ",".join(
+                        [str(existing_req.specifier), str(parsed_req.specifier)]
+                    )
+
+                # Preserve existing specifiers
+                if existing_req.specifier and not parsed_req.specifier:
+                    parsed_req.specifier = existing_req.specifier
+
+                # Combine and apply extras if specified
+                if (
+                    existing_req.extras
+                    and parsed_req.extras
+                    and existing_req.extras != parsed_req.extras
+                ):
+                    parsed_req.extras = sorted(set(existing_req.extras).union(parsed_req.extras))
+                elif existing_req.extras and not parsed_req.extras:
+                    parsed_req.extras = existing_req.extras
+
+                deduped_reqs[base_pkg] = parsed_req
+
+        except InvalidRequirement:
+            # Include non-standard package strings as-is
+            if req not in deduped_reqs:
+                deduped_reqs[req] = req
+    return [str(req) for req in deduped_reqs.values()]
+
+
+def _validate_version_constraints(requirements):
+    """
+    Validates the version constraints of given Python package requirements using pip's resolver with
+    the `--dry-run` option enabled that performs validation only (will not install packages).
+
+    This function writes the requirements to a temporary file and then attempts to resolve
+    them using pip's `--dry-run` install option. If any version conflicts are detected, it
+    raises an MlflowException with details of the conflict.
+
+    Args:
+        requirements (list of str): A list of package requirements (e.g., `["pandas>=1.15",
+        "pandas<2"]`).
+
+    Raises:
+        MlflowException: If any version conflicts are detected among the provided requirements.
+
+    Returns:
+        None: This function does not return anything. It either completes successfully or raises
+        an MlflowException.
+
+    Example:
+        _validate_version_constraints(["tensorflow<2.0", "tensorflow>2.3"])
+        # This will raise an exception due to boundary validity.
+    """
+    with tempfile.NamedTemporaryFile(mode="w+", delete=False) as tmp_file:
+        tmp_file.write("\n".join(requirements))
+        tmp_file_name = tmp_file.name
+
+    try:
+        subprocess.run(
+            ["pip", "install", "--dry-run", "-r", tmp_file_name],
+            check=True,
+            capture_output=True,
+        )
+    except subprocess.CalledProcessError as e:
+        raise MlflowException.invalid_parameter_value(
+            "The specified requirements versions are incompatible. Detected "
+            f"conflicts: \n{e.stderr.decode()}"
+        )
+    finally:
+        os.remove(tmp_file_name)
+
+
 def _process_conda_env(conda_env):
     """
     Processes `conda_env` passed to `mlflow.*.save_model` or `mlflow.*.log_model`, and returns


Test patch:
diff --git a/tests/utils/test_environment.py b/tests/utils/test_environment.py
--- a/tests/utils/test_environment.py
+++ b/tests/utils/test_environment.py
@@ -3,8 +3,10 @@
 import pytest
 import yaml
 
+from mlflow.exceptions import MlflowException
 from mlflow.utils.environment import (
     _contains_mlflow_requirement,
+    _deduplicate_requirements,
     _find_duplicate_requirements,
     _get_pip_deps,
     _get_pip_requirement_specifier,
@@ -281,6 +283,19 @@ def test_process_pip_requirements(tmp_path):
     assert reqs == [expected_mlflow_ver, "b", "-c constraints.txt"]
     assert cons == ["c"]
 
+    conda_env, reqs, cons = _process_pip_requirements(["a"], extra_pip_requirements=["a[extras]"])
+    assert _get_pip_deps(conda_env) == [expected_mlflow_ver, "a[extras]"]
+    assert reqs == [expected_mlflow_ver, "a[extras]"]
+    assert cons == []
+
+    conda_env, reqs, cons = _process_pip_requirements(
+        ["mlflow==1.2.3", "b[extra1]", "a==1.2.3"],
+        extra_pip_requirements=["b[extra2]", "a[extras]"],
+    )
+    assert _get_pip_deps(conda_env) == ["mlflow==1.2.3", "b[extra1,extra2]", "a[extras]==1.2.3"]
+    assert reqs == ["mlflow==1.2.3", "b[extra1,extra2]", "a[extras]==1.2.3"]
+    assert cons == []
+
 
 def test_process_conda_env(tmp_path):
     def make_conda_env(pip_deps):
@@ -354,3 +369,64 @@ def test_duplicate_pip_requirements():
     ]
     evaluation = _find_duplicate_requirements(packages)
     assert sorted(evaluation) == ["fastapi", "numpy", "package"]
+
+
+@pytest.mark.parametrize(
+    ("input_requirements", "expected"),
+    [
+        # Simple cases
+        (["scikit-learn>1", "pandas"], ["scikit-learn>1", "pandas"]),
+        # Duplicates without extras, preserving version restrictions
+        (["packageA", "packageA==1.0"], ["packageA==1.0"]),
+        # Duplicates with extras
+        (["packageA", "packageA[extras]"], ["packageA[extras]"]),
+        # Mixed cases
+        (
+            ["packageA", "packageB", "packageA[extras]", "packageC<=2.0"],
+            ["packageA[extras]", "packageB", "packageC<=2.0"],
+        ),
+        # Mixed versions and extras
+        (["markdown>=3.5.1", "markdown[extras]", "markdown<4"], ["markdown[extras]>=3.5.1,<4"]),
+        # Overlapping extras
+        (
+            ["packageZ[extra1]", "packageZ[extra2]", "packageZ"],
+            ["packageZ[extra1,extra2]"],
+        ),
+        # No version on extras with final version on non-extras
+        (
+            ["markdown[extra1]", "markdown[extra2]", "markdown>3", "markdown<4"],
+            ["markdown[extra1,extra2]>3,<4"],
+        ),
+        # Version constraints with extras
+        (["markdown>1.0", "markdown[extras]<4"], ["markdown[extras]>1.0,<4"]),
+        # Verify duplicate specifiers are not preserved
+        (
+            ["markdown==3.5.1", "markdown[extras]==3.5.1", "markdown[extras]"],
+            ["markdown[extras]==3.5.1"],
+        ),
+        # Verify duplicate extras are not preserved
+        (["markdown[extras]", "markdown", "markdown[extras]"], ["markdown[extras]"]),
+    ],
+)
+def test_deduplicate_requirements_resolve_correctly(input_requirements, expected):
+    assert _deduplicate_requirements(input_requirements) == expected
+
+
+@pytest.mark.parametrize(
+    "input_requirements",
+    [
+        # Non-inclusive range with precise specifier
+        ["scikit-learn==1.1", "scikit-learn<1"],
+        # Incompatible ranges with extras
+        ["markdown[extras]==3.5.1", "markdown<3.4"],
+        # Invalid ranges
+        ["markdown<3", "markdown>3"],
+        # Conflicting versions
+        ["markdown==3.0", "markdown==3.5"],
+    ],
+)
+def test_invalid_requirements_raise(input_requirements):
+    with pytest.raises(
+        MlflowException, match="The specified requirements versions are incompatible"
+    ):
+        _deduplicate_requirements(input_requirements)