Repository: mlflow/mlflow

Issue:
`mlflow.pyfunc.PythonModel`s referencing multiple artifacts with distinct uris but common sub-paths do not work as expected
Thank you for submitting an issue. Please refer to our [issue policy](https://www.github.com/mlflow/mlflow/blob/master/ISSUE_POLICY.md) for additional information about bug reports. For help with debugging your code, please refer to [Stack Overflow](https://stackoverflow.com/questions/tagged/mlflow).

### Willingness to contribute
The MLflow Community encourages bug fix contributions. Would you or another member of your organization be willing to contribute a fix for this bug to the MLflow code base?

- [ ] Yes. I can contribute a fix for this bug independently.
- [x] Yes. I would be willing to contribute a fix for this bug with guidance from the MLflow community.
- [ ] No. I cannot contribute a bug fix at this time.

### Describe the problem
Actual behaviour: 
A common pattern I have been working with is to use a pyfunc.PythonModel to wrap up a sequence of inference steps from separate models. The idea is illustrated below. The PythonModel references multiple artifacts that are distinguishable through their uris, but share a common artifact sub-path. When these are downloaded as part of a call to `mlflow.pyfunc.save_model` the artifacts are sequentially overwritten, and the result is that multiple distinct references are resolved to the same underlying artifacts. This all happens quietly, and could cause a real headache that might be hard to catch for practitioners.

```python
import mlflow

class TwoStepPyfuncModel(mlflow.pyfunc.PythonModel):
  def load_context(self, context):
    self._first_step_model = mlflow.keras.load_model(context.artifacts["first_step_model"])
    self._second_step_model = mlflow.keras.load_model(context.artifacts["second_step_model"])
    # Issue: context.artifacts["first_step_model"] = context.artifacts["second_step_model"]
    # Naive way to catch:
    if context.artifacts["first_step_model"] == context.artifacts["second_step_model"]:
      raise Exception("oops")

  def predict(self, context, model_input):    
    first_step_result = self._first_step_model(model_input, training=False)
    if condition(first_step_result):
      second_step_result = self.second_step_model(model_input, training=False)
      return {"first_step_result": first_step_result, "second_step_result": second_step_result}
    else:
      return {"first_step_result": first_step_result, "second_step_result": None}


two_step_model = TwoStepPyfuncModel()

mlflow.pyfunc.save_model(
  "./pyfunc_model", 
  python_model = two_step_model, 
  artifacts = {
    "first_step_model": 'runs:/abc/model', # Default artifact path for saving a model.
    "second_step_model": 'runs:/cde/model' # Default artifact path for saving a model.
  }
)

two_step_model_reloaded = mlflow.pyfunc.load_model("./pyfunc_model")

two_step_model_reloaded.predict(some_input)
# Exception: Oops
```
Desired behaviour:
It should not be possible for artifacts specified in `mlflow.pyfunc.save_model` to be quietly overwritten. I imagine there are reasons why it may be desirable for these artifacts to retain the same relative path as their originating sources, but at the very least we should tell the user that this has happened and raise an exception. An alternative would be to handle this situation by injecting a differentiator in the sub-paths of the artifacts downloaded for the new PythonModel, and to alert the user to this, though I don't have much visibility on whether this would cause side-effects elsewhere. I note that this is only an issue if the source artifacts in their separate runs happen to have the same relative path - but this is often the case, for the first model logged to an mlflow run, in my experience. 

### What component(s), interfaces, languages, and integrations does this bug affect?
Components 
- [x] `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/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

Interface 
- [ ] `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

Language 
- [ ] `language/r`: R APIs and clients
- [ ] `language/java`: Java APIs and clients
- [ ] `language/new`: Proposals for new client languages

Integrations
- [ ] `integrations/azure`: Azure and Azure ML integrations
- [ ] `integrations/sagemaker`: SageMaker integrations
- [ ] `integrations/databricks`: Databricks integrations



Discussion:
@kyle-jarvis is this just a typo here - you seem to be loading twice the same context (first_step_model)?
```
self._first_step_model = mlflow.keras.load_model(**context.artifacts["first_step_model"]**)
 self._second_step_model = mlflow.keras.load_model(**context.artifacts["first_step_model"]**)
```
@r3stl355 Yes, thanks for spotting. It is just a typo in the example I gave, it's corrected and the description of the issue is otherwise correct. 
In case helpful to illustrate the point further, I was able to work around by making the changes indicated below. This required back-tracking to the point at which the 'source' model artifacts were logged, and logging them to new uniquely-named paths in their respective runs, which then made the quiet overwriting issue in the derived `PythonModel` go away.
```python
mlflow.pyfunc.save_model(
  "./pyfunc_model", 
  python_model = two_step_model, 
  artifacts = {
    "first_step_model": 'runs:/abc/model_abc', # run_id appended to artifact path when saving this model.
    "second_step_model": 'runs:/cde/model_cde' # run_id appended to artifact path when saving this model.
  }
)

```
Each path of the downloaded artifact is explicitly captured and stored in `custom_model_config_kwargs[CONFIG_KEY_ARTIFACTS] = saved_artifacts_config` so adding another level in each path should not really matter (e.g. use indexer prefixes `0/`, '1/' and so on if there is more than one artifact. This will prevent the throwing and would still work with matching model names
    https://github.com/mlflow/mlflow/blob/6d597ff85c7c51ccdded5812ea95546f341ffda5/mlflow/pyfunc/model.py#L179-L190

For example, define `output_path` before L179 conditionally appending an index to it if there are multiple artifacts to download
@r3stl355 Thanks for the input, I have created a PR with a small patch that I'd be grateful for you to give a quick glance over and feedback and further suggestions. #5657 
@kyle-jarvis I think that should do but I don't have any power of formal review/approval

Resolving patch:
diff --git a/mlflow/pyfunc/model.py b/mlflow/pyfunc/model.py
--- a/mlflow/pyfunc/model.py
+++ b/mlflow/pyfunc/model.py
@@ -177,15 +177,9 @@ def _save_model_with_class_artifacts_params(
         with TempDir() as tmp_artifacts_dir:
             tmp_artifacts_config = {}
             saved_artifacts_dir_subpath = "artifacts"
-            for artifact_index, (artifact_name, artifact_uri) in enumerate(artifacts.items()):
-                artifact_download_output_path = (
-                    tmp_artifacts_dir.path(str(artifact_index))
-                    if len(artifacts) > 1
-                    else tmp_artifacts_dir.path()
-                )
-                os.makedirs(artifact_download_output_path, exist_ok=True)
+            for artifact_name, artifact_uri in artifacts.items():
                 tmp_artifact_path = _download_artifact_from_uri(
-                    artifact_uri=artifact_uri, output_path=artifact_download_output_path
+                    artifact_uri=artifact_uri, output_path=tmp_artifacts_dir.path()
                 )
                 tmp_artifacts_config[artifact_name] = tmp_artifact_path
                 saved_artifact_subpath = posixpath.join(


Test patch:
diff --git a/tests/pyfunc/test_model_export_with_class_and_artifacts.py b/tests/pyfunc/test_model_export_with_class_and_artifacts.py
--- a/tests/pyfunc/test_model_export_with_class_and_artifacts.py
+++ b/tests/pyfunc/test_model_export_with_class_and_artifacts.py
@@ -845,40 +845,6 @@ def predict(self, context, model_input):
     assert loaded_model.predict(iris_data[0])
 
 
-@pytest.mark.large
-def test_multiple_artifacts_with_common_sub_paths_resolve_correctly_when_saving_model(
-    sklearn_knn_model, model_path, iris_data
-):
-    default_model_path = "model"
-    with mlflow.start_run():
-        model_info = mlflow.sklearn.log_model(
-            sk_model=sklearn_knn_model, artifact_path=default_model_path
-        )
-        first_model_uri = model_info.model_uri
-
-    with mlflow.start_run():
-        model_info = mlflow.sklearn.log_model(
-            sk_model=sklearn_knn_model, artifact_path=default_model_path
-        )
-        second_model_uri = model_info.model_uri
-
-    class ArtifactValidationModel(mlflow.pyfunc.PythonModel):
-        def predict(self, context, model_input):
-            first_model_resolved_path = context.artifacts["first_model"]
-            second_model_resolved_path = context.artifacts["second_model"]
-            return first_model_resolved_path != second_model_resolved_path
-
-    mlflow.pyfunc.save_model(
-        path=model_path,
-        artifacts={"first_model": first_model_uri, "second_model": second_model_uri},
-        python_model=ArtifactValidationModel(),
-        conda_env=_conda_env(),
-    )
-
-    loaded_model = mlflow.pyfunc.load_model(model_uri=model_path)
-    assert loaded_model.predict(iris_data[0])
-
-
 @pytest.mark.large
 def test_save_model_with_no_artifacts_does_not_produce_artifacts_dir(model_path):
     mlflow.pyfunc.save_model(