We are currently solving the following issue within our repository. Here is the issue text:
--- BEGIN ISSUE ---
Refactor for `subprocess.run`
### Summary

```diff
diff --git a/mlflow/models/wheeled_model.py b/mlflow/models/wheeled_model.py
index 0830208fe8..f64bf4194c 100644
--- a/mlflow/models/wheeled_model.py
+++ b/mlflow/models/wheeled_model.py
@@ -211,14 +211,23 @@ class WheeledModel:
             os.makedirs(dst_path)
 
         pip_wheel_options = MLFLOW_WHEELED_MODEL_PIP_DOWNLOAD_OPTIONS.get()
-
-        download_command = (
-            f"{sys.executable} -m pip wheel {pip_wheel_options} --wheel-dir={dst_path} -r"
-            f"{pip_requirements_path} --no-cache-dir"
-        )
-
         try:
-            subprocess.run([download_command], check=True, shell=True, capture_output=True)
+            subprocess.run(
+                [
+                    sys.executable,
+                    "-m",
+                    "pip",
+                    "wheel",
+                    pip_wheel_options,
+                    "--wheel-dir",
+                    dst_path,
+                    "-r",
+                    pip_requirements_path,
+                    "--no-cache-dir",
+                ],
+                check=True,
+                capture_output=True,
+            )
         except subprocess.CalledProcessError as e:
             raise MlflowException(
                 f"An error occurred while downloading the dependency wheels: {e.stderr}"
```

### Notes

- Make sure to open a PR from a **non-master** branch.
- Sign off the commit using the `-s` flag when making a commit:

  ```sh
  git commit -s -m "..."
           # ^^ make sure to use this
  ```

- Include `#{issue_number}` (e.g. `#123`) in the PR description when opening a PR.


--- END ISSUE ---

Below are some code segments, each from a relevant file. One or more of these files may contain bugs.
--- BEGIN FILES ---
Path: `mlflow/models/wheeled_model.py`
Content:
```
1 import os
2 import platform
3 import shutil
4 import subprocess
5 import sys
6 
7 import yaml
8 
9 import mlflow
10 from mlflow import MlflowClient
11 from mlflow.environment_variables import MLFLOW_WHEELED_MODEL_PIP_DOWNLOAD_OPTIONS
12 from mlflow.exceptions import MlflowException
13 from mlflow.protos.databricks_pb2 import BAD_REQUEST
14 from mlflow.pyfunc.model import MLMODEL_FILE_NAME, Model
15 from mlflow.store.artifact.utils.models import _parse_model_uri, get_model_name_and_version
16 from mlflow.tracking.artifact_utils import _download_artifact_from_uri
17 from mlflow.utils.annotations import experimental
18 from mlflow.utils.environment import (
19     _REQUIREMENTS_FILE_NAME,
20     _get_pip_deps,
21     _mlflow_additional_pip_env,
22     _overwrite_pip_deps,
23 )
24 from mlflow.utils.model_utils import _validate_and_prepare_target_save_path
25 from mlflow.utils.uri import get_databricks_profile_uri_from_artifact_uri
26 
27 _WHEELS_FOLDER_NAME = "wheels"
28 _ORIGINAL_REQ_FILE_NAME = "original_requirements.txt"
29 _PLATFORM = "platform"
30 
31 
32 @experimental
33 class WheeledModel:
34     """
35     Helper class to create a model with added dependency wheels from an existing registered model.
36     The `wheeled` model contains all the model dependencies as wheels stored as model artifacts.
37     .. note::
38         This utility only operates on a model that has been registered to the Model Registry.
39     """
40 
41     def __init__(self, model_uri):
42         self._model_uri = model_uri
43         databricks_profile_uri = (
44             get_databricks_profile_uri_from_artifact_uri(model_uri) or mlflow.get_registry_uri()
45         )
46         client = MlflowClient(registry_uri=databricks_profile_uri)
47         self._model_name, _ = get_model_name_and_version(client, model_uri)
48 
49     @classmethod
50     def log_model(cls, model_uri, registered_model_name=None):
51         """
52         Logs a registered model as an MLflow artifact for the current run. This only operates on
53         a model which has been registered to the Model Registry. Given a registered model_uri (
54         e.g. models:/<model_name>/<model_version>), this utility re-logs the model along with all
55         the required model libraries back to the Model Registry. The required model libraries are
56         stored along with the model as model artifacts. In addition, supporting files to the
57         model (e.g. conda.yaml, requirements.txt) are modified to use the added libraries.
58 
59         By default, this utility creates a new model version under the same registered model
60         specified by ``model_uri``. This behavior can be overridden by specifying the
61         ``registered_model_name`` argument.
62 
63         :param model_uri: A registered model uri in the Model Registry of the form
64                           models:/<model_name>/<model_version/stage/latest>
65         :param registered_model_name: The new model version (model with its libraries) is
66                                       registered under the inputted registered_model_name. If None,
67                                       a new version is logged to the existing model in the Model
68                                       Registry.
69 
70         .. code-block:: python
71             :caption: Example
72 
73             # Given a model uri, log the wheeled model
74             with mlflow.start_run():
75                 WheeledModel.log_model(model_uri)
76         """
77         parsed_uri = _parse_model_uri(model_uri)
78         return Model.log(
79             artifact_path=None,
80             flavor=WheeledModel(model_uri),
81             registered_model_name=registered_model_name or parsed_uri.name,
82         )
83 
84     def save_model(self, path, mlflow_model=None):
85         """
86         Given an existing registered model, saves the model along with it's dependencies stored as
87         wheels to a path on the local file system.
88 
89         This does not modify existing model behavior or existing model flavors. It simply downloads
90         the model dependencies as wheels and modifies the requirements.txt and conda.yaml file to
91         point to the downloaded wheels.
92 
93         The download_command defaults to downloading only binary packages using the
94         `--only-binary=:all:` option. This behavior can be overridden using an environment
95         variable `MLFLOW_WHEELED_MODEL_PIP_DOWNLOAD_OPTIONS`, which will allows setting
96         different options such as `--prefer-binary`, `--no-binary`, etc.
97         :param path: Local path where the model is to be saved.
98         :param mlflow_model: The new :py:mod:`mlflow.models.Model` metadata file to store the
99                              updated model metadata.
100         """
101         from mlflow.pyfunc import ENV, FLAVOR_NAME, _extract_conda_env
102 
103         path = os.path.abspath(path)
104         _validate_and_prepare_target_save_path(path)
105 
106         local_model_path = _download_artifact_from_uri(self._model_uri, output_path=path)
107 
108         wheels_dir = os.path.join(local_model_path, _WHEELS_FOLDER_NAME)
109         pip_requirements_path = os.path.join(local_model_path, _REQUIREMENTS_FILE_NAME)
110         model_metadata_path = os.path.join(local_model_path, MLMODEL_FILE_NAME)
111 
112         model_metadata = Model.load(model_metadata_path)
113 
114         # Check if the model file has `wheels` set to True
115         if model_metadata.__dict__.get(_WHEELS_FOLDER_NAME, None) is not None:
116             raise MlflowException("Model libraries are already added", BAD_REQUEST)
117 
118         conda_env = _extract_conda_env(model_metadata.flavors.get(FLAVOR_NAME, {}).get(ENV, None))
119         conda_env_path = os.path.join(local_model_path, conda_env)
120         if conda_env is None and not os.path.isfile(pip_requirements_path):
121             raise MlflowException(
122                 "Cannot add libraries for model with no logged dependencies.", BAD_REQUEST
123             )
124 
125         if not os.path.isfile(pip_requirements_path):
126             self._create_pip_requirement(conda_env_path, pip_requirements_path)
127 
128         WheeledModel._download_wheels(
129             pip_requirements_path=pip_requirements_path, dst_path=wheels_dir
130         )
131 
132         # Keep a copy of the original requirement.txt
133         shutil.copy2(pip_requirements_path, os.path.join(local_model_path, _ORIGINAL_REQ_FILE_NAME))
134 
135         # Update requirements.txt with wheels
136         pip_deps = self._overwrite_pip_requirements_with_wheels(
137             pip_requirements_path=pip_requirements_path, wheels_dir=wheels_dir
138         )
139 
140         # Update conda.yaml with wheels
141         self._update_conda_env(pip_deps, conda_env_path)
142 
143         # Update MLModel File
144         mlflow_model = self._update_mlflow_model(
145             original_model_metadata=model_metadata, mlflow_model=mlflow_model
146         )
147         mlflow_model.save(model_metadata_path)
148         return mlflow_model
149 
150     def _update_conda_env(self, new_pip_deps, conda_env_path):
151         """
152         Updates the list pip packages in the conda.yaml file to the list of wheels in the wheels
153         directory.
154         {
155             "name": "env",
156             "channels": [...],
157             "dependencies": [
158                 ...,
159                 "pip",
160                 {"pip": [...]},  <- Overwrite this with list of wheels
161             ],
162         }
163         :param new_pip_deps: List of pip dependencies as wheels
164         :param conda_env_path: Path to conda.yaml file in the model directory
165         """
166         with open(conda_env_path) as f:
167             conda_env = yaml.safe_load(f)
168 
169         new_conda_env = _overwrite_pip_deps(conda_env, new_pip_deps)
170 
171         with open(conda_env_path, "w") as out:
172             yaml.safe_dump(new_conda_env, stream=out, default_flow_style=False)
173 
174     def _update_mlflow_model(self, original_model_metadata, mlflow_model):
175         """
176         Modifies the MLModel file to reflect updated information such as the run_id,
177         utc_time_created. Additionally, this also adds `wheels` to the MLModel file to indicate that
178         this is a `wheeled` model.
179         :param original_model_file_path: The model metadata stored in the original MLmodel file.
180         :param mlflow_model: :py:mod:`mlflow.models.Model` configuration of the newly created
181                                 wheeled model
182         """
183 
184         run_id = mlflow.tracking.fluent._get_or_start_run().info.run_id
185         if mlflow_model is None:
186             mlflow_model = Model(run_id=run_id)
187 
188         original_model_metadata.__dict__.update(
189             {k: v for k, v in mlflow_model.__dict__.items() if v}
190         )
191         mlflow_model.__dict__.update(original_model_metadata.__dict__)
192         mlflow_model.artifact_path = WheeledModel.get_wheel_artifact_path(
193             mlflow_model.artifact_path
194         )
195 
196         mlflow_model.wheels = {_PLATFORM: platform.platform()}
197         return mlflow_model
198 
199     @classmethod
200     def _download_wheels(cls, pip_requirements_path, dst_path):
201         """
202         Downloads all the wheels of the dependencies specified in the requirements.txt file.
203         The pip wheel download_command defaults to downloading only binary packages using
204         the `--only-binary=:all:` option. This behavior can be overridden using an
205         environment variable `MLFLOW_WHEELED_MODEL_PIP_DOWNLOAD_OPTIONS`, which will allows
206         setting different options such as `--prefer-binary`, `--no-binary`, etc.
207         :param pip_requirements_path: Path to requirements.txt in the model directory
208         :param dst_path: Path to the directory where the wheels are to be downloaded
209         """
210         if not os.path.exists(dst_path):
211             os.makedirs(dst_path)
212 
213         pip_wheel_options = MLFLOW_WHEELED_MODEL_PIP_DOWNLOAD_OPTIONS.get()
214 
215         download_command = (
216             f"{sys.executable} -m pip wheel {pip_wheel_options} --wheel-dir={dst_path} -r"
217             f"{pip_requirements_path} --no-cache-dir"
218         )
219 
220         try:
221             subprocess.run([download_command], check=True, shell=True, capture_output=True)
222         except subprocess.CalledProcessError as e:
223             raise MlflowException(
224                 f"An error occurred while downloading the dependency wheels: {e.stderr}"
225             )
226 
227     def _overwrite_pip_requirements_with_wheels(self, pip_requirements_path, wheels_dir):
228         """
229         Overwrites the requirements.txt with the wheels of the required dependencies.
230         :param pip_requirements_path: Path to requirements.txt in the model directory
231         :param wheels_dir: Path to directory where wheels are stored
232         """
233         wheels = []
234         with open(pip_requirements_path, "w") as wheels_requirements:
235             for wheel_file in os.listdir(wheels_dir):
236                 if wheel_file.endswith(".whl"):
237                     complete_wheel_file = os.path.join(_WHEELS_FOLDER_NAME, wheel_file)
238                     wheels.append(complete_wheel_file)
239                     wheels_requirements.write(complete_wheel_file + "\n")
240         return wheels
241 
242     def _create_pip_requirement(self, conda_env_path, pip_requirements_path):
243         """
244         This method creates a requirements.txt file for the model dependencies if the file does not
245         already exist. If uses the pip dependencies found in the conda.yaml env file.
246         :param conda_env_path: Path to conda.yaml env file which contains the required pip
247                                 dependencies
248         :param pip_requirements_path: Path where the new requirements.txt will be created.
249         """
250         with open(conda_env_path) as f:
251             conda_env = yaml.safe_load(f)
252         pip_deps = _get_pip_deps(conda_env)
253         _mlflow_additional_pip_env(pip_deps, pip_requirements_path)
254 
255     @classmethod
256     def get_wheel_artifact_path(cls, original_artifact_path):
257         return original_artifact_path + "_" + _WHEELS_FOLDER_NAME
258
```
--- END FILES ---

Please first localize the bug based on the issue statement, and then generate a patch according to the `git diff` format fenced by three backticks.

Here is an example:

```diff
diff --git a/examples/server_async.py b/examples/server_async.py
--- a/examples/server_async.py
+++ b/examples/server_async.py
@@ -313,4 +313,4 @@
 
 
 if __name__ == "__main__":
-    asyncio.run(run_async_server("."), debug=True)
+    asyncio.run(run_async_server(), debug=True)
diff --git a/examples/server_sync.py b/examples/server_sync.py
--- a/examples/server_sync.py
+++ b/examples/server_sync.py
@@ -313,5 +313,5 @@
 
 
 if __name__ == "__main__":
-    server = run_sync_server(".")
+    server = run_sync_server()
     server.shutdown()

```

diff --git a/mlflow/models/wheeled_model.py b/mlflow/models/wheeled_model.py
--- a/mlflow/models/wheeled_model.py
+++ b/mlflow/models/wheeled_model.py
@@ -212,13 +212,23 @@
 
         pip_wheel_options = MLFLOW_WHEELED_MODEL_PIP_DOWNLOAD_OPTIONS.get()
 
-        download_command = (
-            f"{sys.executable} -m pip wheel {pip_wheel_options} --wheel-dir={dst_path} -r"
-            f"{pip_requirements_path} --no-cache-dir"
-        )
-
         try:
-            subprocess.run([download_command], check=True, shell=True, capture_output=True)
+            subprocess.run(
+                [
+                    sys.executable,
+                    "-m",
+                    "pip",
+                    "wheel",
+                    pip_wheel_options,
+                    "--wheel-dir",
+                    dst_path,
+                    "-r",
+                    pip_requirements_path,
+                    "--no-cache-dir",
+                ],
+                check=True,
+                capture_output=True,
+            )
         except subprocess.CalledProcessError as e:
             raise MlflowException(
                 f"An error occurred while downloading the dependency wheels: {e.stderr}"

{"golden_diff": "diff --git a/mlflow/models/wheeled_model.py b/mlflow/models/wheeled_model.py\n--- a/mlflow/models/wheeled_model.py\n+++ b/mlflow/models/wheeled_model.py\n@@ -212,13 +212,23 @@\n \n         pip_wheel_options = MLFLOW_WHEELED_MODEL_PIP_DOWNLOAD_OPTIONS.get()\n \n-        download_command = (\n-            f\"{sys.executable} -m pip wheel {pip_wheel_options} --wheel-dir={dst_path} -r\"\n-            f\"{pip_requirements_path} --no-cache-dir\"\n-        )\n-\n         try:\n-            subprocess.run([download_command], check=True, shell=True, capture_output=True)\n+            subprocess.run(\n+                [\n+                    sys.executable,\n+                    \"-m\",\n+                    \"pip\",\n+                    \"wheel\",\n+                    pip_wheel_options,\n+                    \"--wheel-dir\",\n+                    dst_path,\n+                    \"-r\",\n+                    pip_requirements_path,\n+                    \"--no-cache-dir\",\n+                ],\n+                check=True,\n+                capture_output=True,\n+            )\n         except subprocess.CalledProcessError as e:\n             raise MlflowException(\n                 f\"An error occurred while downloading the dependency wheels: {e.stderr}\"\n", "issue": "Refactor for `subprocess.run`\n### Summary\n\n```diff\r\ndiff --git a/mlflow/models/wheeled_model.py b/mlflow/models/wheeled_model.py\r\nindex 0830208fe8..f64bf4194c 100644\r\n--- a/mlflow/models/wheeled_model.py\r\n+++ b/mlflow/models/wheeled_model.py\r\n@@ -211,14 +211,23 @@ class WheeledModel:\r\n             os.makedirs(dst_path)\r\n \r\n         pip_wheel_options = MLFLOW_WHEELED_MODEL_PIP_DOWNLOAD_OPTIONS.get()\r\n-\r\n-        download_command = (\r\n-            f\"{sys.executable} -m pip wheel {pip_wheel_options} --wheel-dir={dst_path} -r\"\r\n-            f\"{pip_requirements_path} --no-cache-dir\"\r\n-        )\r\n-\r\n         try:\r\n-            subprocess.run([download_command], check=True, shell=True, capture_output=True)\r\n+            subprocess.run(\r\n+                [\r\n+                    sys.executable,\r\n+                    \"-m\",\r\n+                    \"pip\",\r\n+                    \"wheel\",\r\n+                    pip_wheel_options,\r\n+                    \"--wheel-dir\",\r\n+                    dst_path,\r\n+                    \"-r\",\r\n+                    pip_requirements_path,\r\n+                    \"--no-cache-dir\",\r\n+                ],\r\n+                check=True,\r\n+                capture_output=True,\r\n+            )\r\n         except subprocess.CalledProcessError as e:\r\n             raise MlflowException(\r\n                 f\"An error occurred while downloading the dependency wheels: {e.stderr}\"\r\n```\n\n### Notes\n\n- Make sure to open a PR from a **non-master** branch.\r\n- Sign off the commit using the `-s` flag when making a commit:\r\n\r\n  ```sh\r\n  git commit -s -m \"...\"\r\n           # ^^ make sure to use this\r\n  ```\r\n\r\n- Include `#{issue_number}` (e.g. `#123`) in the PR description when opening a PR.\r\n\n", "before_files": [{"content": "import os\nimport platform\nimport shutil\nimport subprocess\nimport sys\n\nimport yaml\n\nimport mlflow\nfrom mlflow import MlflowClient\nfrom mlflow.environment_variables import MLFLOW_WHEELED_MODEL_PIP_DOWNLOAD_OPTIONS\nfrom mlflow.exceptions import MlflowException\nfrom mlflow.protos.databricks_pb2 import BAD_REQUEST\nfrom mlflow.pyfunc.model import MLMODEL_FILE_NAME, Model\nfrom mlflow.store.artifact.utils.models import _parse_model_uri, get_model_name_and_version\nfrom mlflow.tracking.artifact_utils import _download_artifact_from_uri\nfrom mlflow.utils.annotations import experimental\nfrom mlflow.utils.environment import (\n    _REQUIREMENTS_FILE_NAME,\n    _get_pip_deps,\n    _mlflow_additional_pip_env,\n    _overwrite_pip_deps,\n)\nfrom mlflow.utils.model_utils import _validate_and_prepare_target_save_path\nfrom mlflow.utils.uri import get_databricks_profile_uri_from_artifact_uri\n\n_WHEELS_FOLDER_NAME = \"wheels\"\n_ORIGINAL_REQ_FILE_NAME = \"original_requirements.txt\"\n_PLATFORM = \"platform\"\n\n\n@experimental\nclass WheeledModel:\n    \"\"\"\n    Helper class to create a model with added dependency wheels from an existing registered model.\n    The `wheeled` model contains all the model dependencies as wheels stored as model artifacts.\n    .. note::\n        This utility only operates on a model that has been registered to the Model Registry.\n    \"\"\"\n\n    def __init__(self, model_uri):\n        self._model_uri = model_uri\n        databricks_profile_uri = (\n            get_databricks_profile_uri_from_artifact_uri(model_uri) or mlflow.get_registry_uri()\n        )\n        client = MlflowClient(registry_uri=databricks_profile_uri)\n        self._model_name, _ = get_model_name_and_version(client, model_uri)\n\n    @classmethod\n    def log_model(cls, model_uri, registered_model_name=None):\n        \"\"\"\n        Logs a registered model as an MLflow artifact for the current run. This only operates on\n        a model which has been registered to the Model Registry. Given a registered model_uri (\n        e.g. models:/<model_name>/<model_version>), this utility re-logs the model along with all\n        the required model libraries back to the Model Registry. The required model libraries are\n        stored along with the model as model artifacts. In addition, supporting files to the\n        model (e.g. conda.yaml, requirements.txt) are modified to use the added libraries.\n\n        By default, this utility creates a new model version under the same registered model\n        specified by ``model_uri``. This behavior can be overridden by specifying the\n        ``registered_model_name`` argument.\n\n        :param model_uri: A registered model uri in the Model Registry of the form\n                          models:/<model_name>/<model_version/stage/latest>\n        :param registered_model_name: The new model version (model with its libraries) is\n                                      registered under the inputted registered_model_name. If None,\n                                      a new version is logged to the existing model in the Model\n                                      Registry.\n\n        .. code-block:: python\n            :caption: Example\n\n            # Given a model uri, log the wheeled model\n            with mlflow.start_run():\n                WheeledModel.log_model(model_uri)\n        \"\"\"\n        parsed_uri = _parse_model_uri(model_uri)\n        return Model.log(\n            artifact_path=None,\n            flavor=WheeledModel(model_uri),\n            registered_model_name=registered_model_name or parsed_uri.name,\n        )\n\n    def save_model(self, path, mlflow_model=None):\n        \"\"\"\n        Given an existing registered model, saves the model along with it's dependencies stored as\n        wheels to a path on the local file system.\n\n        This does not modify existing model behavior or existing model flavors. It simply downloads\n        the model dependencies as wheels and modifies the requirements.txt and conda.yaml file to\n        point to the downloaded wheels.\n\n        The download_command defaults to downloading only binary packages using the\n        `--only-binary=:all:` option. This behavior can be overridden using an environment\n        variable `MLFLOW_WHEELED_MODEL_PIP_DOWNLOAD_OPTIONS`, which will allows setting\n        different options such as `--prefer-binary`, `--no-binary`, etc.\n        :param path: Local path where the model is to be saved.\n        :param mlflow_model: The new :py:mod:`mlflow.models.Model` metadata file to store the\n                             updated model metadata.\n        \"\"\"\n        from mlflow.pyfunc import ENV, FLAVOR_NAME, _extract_conda_env\n\n        path = os.path.abspath(path)\n        _validate_and_prepare_target_save_path(path)\n\n        local_model_path = _download_artifact_from_uri(self._model_uri, output_path=path)\n\n        wheels_dir = os.path.join(local_model_path, _WHEELS_FOLDER_NAME)\n        pip_requirements_path = os.path.join(local_model_path, _REQUIREMENTS_FILE_NAME)\n        model_metadata_path = os.path.join(local_model_path, MLMODEL_FILE_NAME)\n\n        model_metadata = Model.load(model_metadata_path)\n\n        # Check if the model file has `wheels` set to True\n        if model_metadata.__dict__.get(_WHEELS_FOLDER_NAME, None) is not None:\n            raise MlflowException(\"Model libraries are already added\", BAD_REQUEST)\n\n        conda_env = _extract_conda_env(model_metadata.flavors.get(FLAVOR_NAME, {}).get(ENV, None))\n        conda_env_path = os.path.join(local_model_path, conda_env)\n        if conda_env is None and not os.path.isfile(pip_requirements_path):\n            raise MlflowException(\n                \"Cannot add libraries for model with no logged dependencies.\", BAD_REQUEST\n            )\n\n        if not os.path.isfile(pip_requirements_path):\n            self._create_pip_requirement(conda_env_path, pip_requirements_path)\n\n        WheeledModel._download_wheels(\n            pip_requirements_path=pip_requirements_path, dst_path=wheels_dir\n        )\n\n        # Keep a copy of the original requirement.txt\n        shutil.copy2(pip_requirements_path, os.path.join(local_model_path, _ORIGINAL_REQ_FILE_NAME))\n\n        # Update requirements.txt with wheels\n        pip_deps = self._overwrite_pip_requirements_with_wheels(\n            pip_requirements_path=pip_requirements_path, wheels_dir=wheels_dir\n        )\n\n        # Update conda.yaml with wheels\n        self._update_conda_env(pip_deps, conda_env_path)\n\n        # Update MLModel File\n        mlflow_model = self._update_mlflow_model(\n            original_model_metadata=model_metadata, mlflow_model=mlflow_model\n        )\n        mlflow_model.save(model_metadata_path)\n        return mlflow_model\n\n    def _update_conda_env(self, new_pip_deps, conda_env_path):\n        \"\"\"\n        Updates the list pip packages in the conda.yaml file to the list of wheels in the wheels\n        directory.\n        {\n            \"name\": \"env\",\n            \"channels\": [...],\n            \"dependencies\": [\n                ...,\n                \"pip\",\n                {\"pip\": [...]},  <- Overwrite this with list of wheels\n            ],\n        }\n        :param new_pip_deps: List of pip dependencies as wheels\n        :param conda_env_path: Path to conda.yaml file in the model directory\n        \"\"\"\n        with open(conda_env_path) as f:\n            conda_env = yaml.safe_load(f)\n\n        new_conda_env = _overwrite_pip_deps(conda_env, new_pip_deps)\n\n        with open(conda_env_path, \"w\") as out:\n            yaml.safe_dump(new_conda_env, stream=out, default_flow_style=False)\n\n    def _update_mlflow_model(self, original_model_metadata, mlflow_model):\n        \"\"\"\n        Modifies the MLModel file to reflect updated information such as the run_id,\n        utc_time_created. Additionally, this also adds `wheels` to the MLModel file to indicate that\n        this is a `wheeled` model.\n        :param original_model_file_path: The model metadata stored in the original MLmodel file.\n        :param mlflow_model: :py:mod:`mlflow.models.Model` configuration of the newly created\n                                wheeled model\n        \"\"\"\n\n        run_id = mlflow.tracking.fluent._get_or_start_run().info.run_id\n        if mlflow_model is None:\n            mlflow_model = Model(run_id=run_id)\n\n        original_model_metadata.__dict__.update(\n            {k: v for k, v in mlflow_model.__dict__.items() if v}\n        )\n        mlflow_model.__dict__.update(original_model_metadata.__dict__)\n        mlflow_model.artifact_path = WheeledModel.get_wheel_artifact_path(\n            mlflow_model.artifact_path\n        )\n\n        mlflow_model.wheels = {_PLATFORM: platform.platform()}\n        return mlflow_model\n\n    @classmethod\n    def _download_wheels(cls, pip_requirements_path, dst_path):\n        \"\"\"\n        Downloads all the wheels of the dependencies specified in the requirements.txt file.\n        The pip wheel download_command defaults to downloading only binary packages using\n        the `--only-binary=:all:` option. This behavior can be overridden using an\n        environment variable `MLFLOW_WHEELED_MODEL_PIP_DOWNLOAD_OPTIONS`, which will allows\n        setting different options such as `--prefer-binary`, `--no-binary`, etc.\n        :param pip_requirements_path: Path to requirements.txt in the model directory\n        :param dst_path: Path to the directory where the wheels are to be downloaded\n        \"\"\"\n        if not os.path.exists(dst_path):\n            os.makedirs(dst_path)\n\n        pip_wheel_options = MLFLOW_WHEELED_MODEL_PIP_DOWNLOAD_OPTIONS.get()\n\n        download_command = (\n            f\"{sys.executable} -m pip wheel {pip_wheel_options} --wheel-dir={dst_path} -r\"\n            f\"{pip_requirements_path} --no-cache-dir\"\n        )\n\n        try:\n            subprocess.run([download_command], check=True, shell=True, capture_output=True)\n        except subprocess.CalledProcessError as e:\n            raise MlflowException(\n                f\"An error occurred while downloading the dependency wheels: {e.stderr}\"\n            )\n\n    def _overwrite_pip_requirements_with_wheels(self, pip_requirements_path, wheels_dir):\n        \"\"\"\n        Overwrites the requirements.txt with the wheels of the required dependencies.\n        :param pip_requirements_path: Path to requirements.txt in the model directory\n        :param wheels_dir: Path to directory where wheels are stored\n        \"\"\"\n        wheels = []\n        with open(pip_requirements_path, \"w\") as wheels_requirements:\n            for wheel_file in os.listdir(wheels_dir):\n                if wheel_file.endswith(\".whl\"):\n                    complete_wheel_file = os.path.join(_WHEELS_FOLDER_NAME, wheel_file)\n                    wheels.append(complete_wheel_file)\n                    wheels_requirements.write(complete_wheel_file + \"\\n\")\n        return wheels\n\n    def _create_pip_requirement(self, conda_env_path, pip_requirements_path):\n        \"\"\"\n        This method creates a requirements.txt file for the model dependencies if the file does not\n        already exist. If uses the pip dependencies found in the conda.yaml env file.\n        :param conda_env_path: Path to conda.yaml env file which contains the required pip\n                                dependencies\n        :param pip_requirements_path: Path where the new requirements.txt will be created.\n        \"\"\"\n        with open(conda_env_path) as f:\n            conda_env = yaml.safe_load(f)\n        pip_deps = _get_pip_deps(conda_env)\n        _mlflow_additional_pip_env(pip_deps, pip_requirements_path)\n\n    @classmethod\n    def get_wheel_artifact_path(cls, original_artifact_path):\n        return original_artifact_path + \"_\" + _WHEELS_FOLDER_NAME\n", "path": "mlflow/models/wheeled_model.py"}], "after_files": [{"content": "import os\nimport platform\nimport shutil\nimport subprocess\nimport sys\n\nimport yaml\n\nimport mlflow\nfrom mlflow import MlflowClient\nfrom mlflow.environment_variables import MLFLOW_WHEELED_MODEL_PIP_DOWNLOAD_OPTIONS\nfrom mlflow.exceptions import MlflowException\nfrom mlflow.protos.databricks_pb2 import BAD_REQUEST\nfrom mlflow.pyfunc.model import MLMODEL_FILE_NAME, Model\nfrom mlflow.store.artifact.utils.models import _parse_model_uri, get_model_name_and_version\nfrom mlflow.tracking.artifact_utils import _download_artifact_from_uri\nfrom mlflow.utils.annotations import experimental\nfrom mlflow.utils.environment import (\n    _REQUIREMENTS_FILE_NAME,\n    _get_pip_deps,\n    _mlflow_additional_pip_env,\n    _overwrite_pip_deps,\n)\nfrom mlflow.utils.model_utils import _validate_and_prepare_target_save_path\nfrom mlflow.utils.uri import get_databricks_profile_uri_from_artifact_uri\n\n_WHEELS_FOLDER_NAME = \"wheels\"\n_ORIGINAL_REQ_FILE_NAME = \"original_requirements.txt\"\n_PLATFORM = \"platform\"\n\n\n@experimental\nclass WheeledModel:\n    \"\"\"\n    Helper class to create a model with added dependency wheels from an existing registered model.\n    The `wheeled` model contains all the model dependencies as wheels stored as model artifacts.\n    .. note::\n        This utility only operates on a model that has been registered to the Model Registry.\n    \"\"\"\n\n    def __init__(self, model_uri):\n        self._model_uri = model_uri\n        databricks_profile_uri = (\n            get_databricks_profile_uri_from_artifact_uri(model_uri) or mlflow.get_registry_uri()\n        )\n        client = MlflowClient(registry_uri=databricks_profile_uri)\n        self._model_name, _ = get_model_name_and_version(client, model_uri)\n\n    @classmethod\n    def log_model(cls, model_uri, registered_model_name=None):\n        \"\"\"\n        Logs a registered model as an MLflow artifact for the current run. This only operates on\n        a model which has been registered to the Model Registry. Given a registered model_uri (\n        e.g. models:/<model_name>/<model_version>), this utility re-logs the model along with all\n        the required model libraries back to the Model Registry. The required model libraries are\n        stored along with the model as model artifacts. In addition, supporting files to the\n        model (e.g. conda.yaml, requirements.txt) are modified to use the added libraries.\n\n        By default, this utility creates a new model version under the same registered model\n        specified by ``model_uri``. This behavior can be overridden by specifying the\n        ``registered_model_name`` argument.\n\n        :param model_uri: A registered model uri in the Model Registry of the form\n                          models:/<model_name>/<model_version/stage/latest>\n        :param registered_model_name: The new model version (model with its libraries) is\n                                      registered under the inputted registered_model_name. If None,\n                                      a new version is logged to the existing model in the Model\n                                      Registry.\n\n        .. code-block:: python\n            :caption: Example\n\n            # Given a model uri, log the wheeled model\n            with mlflow.start_run():\n                WheeledModel.log_model(model_uri)\n        \"\"\"\n        parsed_uri = _parse_model_uri(model_uri)\n        return Model.log(\n            artifact_path=None,\n            flavor=WheeledModel(model_uri),\n            registered_model_name=registered_model_name or parsed_uri.name,\n        )\n\n    def save_model(self, path, mlflow_model=None):\n        \"\"\"\n        Given an existing registered model, saves the model along with it's dependencies stored as\n        wheels to a path on the local file system.\n\n        This does not modify existing model behavior or existing model flavors. It simply downloads\n        the model dependencies as wheels and modifies the requirements.txt and conda.yaml file to\n        point to the downloaded wheels.\n\n        The download_command defaults to downloading only binary packages using the\n        `--only-binary=:all:` option. This behavior can be overridden using an environment\n        variable `MLFLOW_WHEELED_MODEL_PIP_DOWNLOAD_OPTIONS`, which will allows setting\n        different options such as `--prefer-binary`, `--no-binary`, etc.\n        :param path: Local path where the model is to be saved.\n        :param mlflow_model: The new :py:mod:`mlflow.models.Model` metadata file to store the\n                             updated model metadata.\n        \"\"\"\n        from mlflow.pyfunc import ENV, FLAVOR_NAME, _extract_conda_env\n\n        path = os.path.abspath(path)\n        _validate_and_prepare_target_save_path(path)\n\n        local_model_path = _download_artifact_from_uri(self._model_uri, output_path=path)\n\n        wheels_dir = os.path.join(local_model_path, _WHEELS_FOLDER_NAME)\n        pip_requirements_path = os.path.join(local_model_path, _REQUIREMENTS_FILE_NAME)\n        model_metadata_path = os.path.join(local_model_path, MLMODEL_FILE_NAME)\n\n        model_metadata = Model.load(model_metadata_path)\n\n        # Check if the model file has `wheels` set to True\n        if model_metadata.__dict__.get(_WHEELS_FOLDER_NAME, None) is not None:\n            raise MlflowException(\"Model libraries are already added\", BAD_REQUEST)\n\n        conda_env = _extract_conda_env(model_metadata.flavors.get(FLAVOR_NAME, {}).get(ENV, None))\n        conda_env_path = os.path.join(local_model_path, conda_env)\n        if conda_env is None and not os.path.isfile(pip_requirements_path):\n            raise MlflowException(\n                \"Cannot add libraries for model with no logged dependencies.\", BAD_REQUEST\n            )\n\n        if not os.path.isfile(pip_requirements_path):\n            self._create_pip_requirement(conda_env_path, pip_requirements_path)\n\n        WheeledModel._download_wheels(\n            pip_requirements_path=pip_requirements_path, dst_path=wheels_dir\n        )\n\n        # Keep a copy of the original requirement.txt\n        shutil.copy2(pip_requirements_path, os.path.join(local_model_path, _ORIGINAL_REQ_FILE_NAME))\n\n        # Update requirements.txt with wheels\n        pip_deps = self._overwrite_pip_requirements_with_wheels(\n            pip_requirements_path=pip_requirements_path, wheels_dir=wheels_dir\n        )\n\n        # Update conda.yaml with wheels\n        self._update_conda_env(pip_deps, conda_env_path)\n\n        # Update MLModel File\n        mlflow_model = self._update_mlflow_model(\n            original_model_metadata=model_metadata, mlflow_model=mlflow_model\n        )\n        mlflow_model.save(model_metadata_path)\n        return mlflow_model\n\n    def _update_conda_env(self, new_pip_deps, conda_env_path):\n        \"\"\"\n        Updates the list pip packages in the conda.yaml file to the list of wheels in the wheels\n        directory.\n        {\n            \"name\": \"env\",\n            \"channels\": [...],\n            \"dependencies\": [\n                ...,\n                \"pip\",\n                {\"pip\": [...]},  <- Overwrite this with list of wheels\n            ],\n        }\n        :param new_pip_deps: List of pip dependencies as wheels\n        :param conda_env_path: Path to conda.yaml file in the model directory\n        \"\"\"\n        with open(conda_env_path) as f:\n            conda_env = yaml.safe_load(f)\n\n        new_conda_env = _overwrite_pip_deps(conda_env, new_pip_deps)\n\n        with open(conda_env_path, \"w\") as out:\n            yaml.safe_dump(new_conda_env, stream=out, default_flow_style=False)\n\n    def _update_mlflow_model(self, original_model_metadata, mlflow_model):\n        \"\"\"\n        Modifies the MLModel file to reflect updated information such as the run_id,\n        utc_time_created. Additionally, this also adds `wheels` to the MLModel file to indicate that\n        this is a `wheeled` model.\n        :param original_model_file_path: The model metadata stored in the original MLmodel file.\n        :param mlflow_model: :py:mod:`mlflow.models.Model` configuration of the newly created\n                                wheeled model\n        \"\"\"\n\n        run_id = mlflow.tracking.fluent._get_or_start_run().info.run_id\n        if mlflow_model is None:\n            mlflow_model = Model(run_id=run_id)\n\n        original_model_metadata.__dict__.update(\n            {k: v for k, v in mlflow_model.__dict__.items() if v}\n        )\n        mlflow_model.__dict__.update(original_model_metadata.__dict__)\n        mlflow_model.artifact_path = WheeledModel.get_wheel_artifact_path(\n            mlflow_model.artifact_path\n        )\n\n        mlflow_model.wheels = {_PLATFORM: platform.platform()}\n        return mlflow_model\n\n    @classmethod\n    def _download_wheels(cls, pip_requirements_path, dst_path):\n        \"\"\"\n        Downloads all the wheels of the dependencies specified in the requirements.txt file.\n        The pip wheel download_command defaults to downloading only binary packages using\n        the `--only-binary=:all:` option. This behavior can be overridden using an\n        environment variable `MLFLOW_WHEELED_MODEL_PIP_DOWNLOAD_OPTIONS`, which will allows\n        setting different options such as `--prefer-binary`, `--no-binary`, etc.\n        :param pip_requirements_path: Path to requirements.txt in the model directory\n        :param dst_path: Path to the directory where the wheels are to be downloaded\n        \"\"\"\n        if not os.path.exists(dst_path):\n            os.makedirs(dst_path)\n\n        pip_wheel_options = MLFLOW_WHEELED_MODEL_PIP_DOWNLOAD_OPTIONS.get()\n\n        try:\n            subprocess.run(\n                [\n                    sys.executable,\n                    \"-m\",\n                    \"pip\",\n                    \"wheel\",\n                    pip_wheel_options,\n                    \"--wheel-dir\",\n                    dst_path,\n                    \"-r\",\n                    pip_requirements_path,\n                    \"--no-cache-dir\",\n                ],\n                check=True,\n                capture_output=True,\n            )\n        except subprocess.CalledProcessError as e:\n            raise MlflowException(\n                f\"An error occurred while downloading the dependency wheels: {e.stderr}\"\n            )\n\n    def _overwrite_pip_requirements_with_wheels(self, pip_requirements_path, wheels_dir):\n        \"\"\"\n        Overwrites the requirements.txt with the wheels of the required dependencies.\n        :param pip_requirements_path: Path to requirements.txt in the model directory\n        :param wheels_dir: Path to directory where wheels are stored\n        \"\"\"\n        wheels = []\n        with open(pip_requirements_path, \"w\") as wheels_requirements:\n            for wheel_file in os.listdir(wheels_dir):\n                if wheel_file.endswith(\".whl\"):\n                    complete_wheel_file = os.path.join(_WHEELS_FOLDER_NAME, wheel_file)\n                    wheels.append(complete_wheel_file)\n                    wheels_requirements.write(complete_wheel_file + \"\\n\")\n        return wheels\n\n    def _create_pip_requirement(self, conda_env_path, pip_requirements_path):\n        \"\"\"\n        This method creates a requirements.txt file for the model dependencies if the file does not\n        already exist. If uses the pip dependencies found in the conda.yaml env file.\n        :param conda_env_path: Path to conda.yaml env file which contains the required pip\n                                dependencies\n        :param pip_requirements_path: Path where the new requirements.txt will be created.\n        \"\"\"\n        with open(conda_env_path) as f:\n            conda_env = yaml.safe_load(f)\n        pip_deps = _get_pip_deps(conda_env)\n        _mlflow_additional_pip_env(pip_deps, pip_requirements_path)\n\n    @classmethod\n    def get_wheel_artifact_path(cls, original_artifact_path):\n        return original_artifact_path + \"_\" + _WHEELS_FOLDER_NAME\n", "path": "mlflow/models/wheeled_model.py"}]}