Repository: mlflow/mlflow

Issue:
[BUG] Azure OpenAI model logged with deployment_id cannot be used as deployment not found
### 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?

Databricks

### Willingness to contribute

Yes. I can contribute a fix for this bug independently.

### MLflow version

- Client: >2.9.2


### System information

- **OS Platform and Distribution (e.g., Linux Ubuntu 16.04)**: Debian GNU/Linux 10 (buster)
- **Python version**: 3.10


### Describe the problem

I am using the notebook at https://www.mlflow.org/docs/2.8.0/llms/llm-evaluate/notebooks/question-answering-evaluation.html but I adapted it for Azure OpenAI as described at https://mlflow.org/docs/latest/llms/openai/guide/index.html#azure-openai-service-in-mlflow just with the deployment_id parameter  when logging the model:
```
basic_qa_model = mlflow.openai.log_model(
        model="gpt-35-turbo",
        deployment_id="gpt-35-turbo",

        task=openai.ChatCompletion,
        artifact_path="model",
        messages=[
            {"role": "system", "content": system_prompt},
            {"role": "user", "content": "{question}"},
        ],
    )
```

Evaluating the model:
```
results = mlflow.evaluate(
        basic_qa_model.model_uri,
        eval_df,
        targets="ground_truth",  # specify which column corresponds to the expected output
        model_type="question-answering",  # model type indicates which metrics are relevant for this task
        evaluators="default",
    )
```
or calling predict on the loaded model:
```
model = mlflow.pyfunc.load_model(basic_qa_model.model_uri)
model.predict(["question": "what is github?"}])
```
Fails with:
```
Request #1 failed with {'code': 'DeploymentNotFound', 'message': 'The API deployment for this resource does not exist. If you created the deployment within the last 5 minutes, please wait a moment and try again.'}
```
This seems to happen because when constructing the request url only self.envs is being used at https://github.com/mlflow/mlflow/blob/4d0af6f89f11e2bc27cebede01fe2d41e7e0f020/mlflow/openai/__init__.py#L714.
The self.envs explicitly excludes values that are part of the model, so recording a model with deployment_id as documented will prevent usage of the model.

### Tracking information

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


### Code to reproduce issue

<!-- PLEASE KEEP BACKTICKS AND CHECK PREVIEW -->
```
import os

os.environ.update(dict(
  OPENAI_API_KEY="xxx",
OPENAI_API_TYPE="azure",
OPENAI_API_VERSION="2023-05-15",
OPENAI_API_BASE="https://xyz.openai.azure.com/",
DEPLOYMENT_ID="gpt-35-turbo",
))

import openai
import pandas as pd

import mlflow


eval_df = pd.DataFrame(
    {
        "inputs": [
            "How does useEffect() work?",
            "What does the static keyword in a function mean?",
            "What does the 'finally' block in Python do?",
            "What is the difference between multiprocessing and multithreading?",
        ],
        "ground_truth": [
            "The useEffect() hook tells React that your component needs to do something after render. React will remember the function you passed (we’ll refer to it as our “effect”), and call it later after performing the DOM updates.",
            "Static members belongs to the class, rather than a specific instance. This means that only one instance of a static member exists, even if you create multiple objects of the class, or if you don't create any. It will be shared by all objects.",
            "'Finally' defines a block of code to run when the try... except...else block is final. The finally block will be executed no matter if the try block raises an error or not.",
            "Multithreading refers to the ability of a processor to execute multiple threads concurrently, where each thread runs a process. Whereas multiprocessing refers to the ability of a system to run multiple processors in parallel, where each processor can run one or more threads.",
        ],
    }
)


with mlflow.start_run() as run:
    system_prompt = "Answer the following question in two sentences"
    basic_qa_model = mlflow.openai.log_model(
        model="gpt-35-turbo",
        task=openai.ChatCompletion,
        artifact_path="model",
        messages=[
            {"role": "system", "content": system_prompt},
            {"role": "user", "content": "{question}"},
        ],
    )
    results = mlflow.evaluate(
        basic_qa_model.model_uri,
        eval_df,
        targets="ground_truth",  # specify which column corresponds to the expected output
        model_type="question-answering",  # model type indicates which metrics are relevant for this task
        evaluators="default",
    )
results.metrics
```


### Stack trace

<!-- PLEASE KEEP BACKTICKS AND CHECK PREVIEW -->
```
MlflowException: 4 tasks failed. See logs for details.
---------------------------------------------------------------------------
MlflowException                           Traceback (most recent call last)
File <command-2041551855651865>, line 13
      2     system_prompt = "Answer the following question in two sentences"
      3     basic_qa_model = mlflow.openai.log_model(
      4         model="gpt-35-turbo",
      5         deployment_id="gpt-35-turbo",
   (...)
     11         ],
     12     )
---> 13     results = mlflow.evaluate(
     14         basic_qa_model.model_uri,
     15         eval_df,
     16         targets="ground_truth",  # specify which column corresponds to the expected output
     17         model_type="question-answering",  # model type indicates which metrics are relevant for this task
     18         evaluators="default",
     19     )
     20 results.metrics

File /databricks/python/lib/python3.10/site-packages/mlflow/models/evaluation/base.py:1880, in evaluate(model, data, model_type, targets, predictions, dataset_path, feature_names, evaluators, evaluator_config, custom_metrics, extra_metrics, custom_artifacts, validation_thresholds, baseline_model, env_manager, model_config, baseline_config)
   1877 predictions_expected_in_model_output = predictions if model is not None else None
   1879 try:
-> 1880     evaluate_result = _evaluate(
   1881         model=model,
   1882         model_type=model_type,
   1883         dataset=dataset,
   1884         run_id=run_id,
   1885         evaluator_name_list=evaluator_name_list,
   1886         evaluator_name_to_conf_map=evaluator_name_to_conf_map,
   1887         custom_metrics=custom_metrics,
   1888         extra_metrics=extra_metrics,
   1889         custom_artifacts=custom_artifacts,
   1890         baseline_model=baseline_model,
   1891         predictions=predictions_expected_in_model_output,
   1892     )
   1893 finally:
   1894     if isinstance(model, _ServedPyFuncModel):

File /databricks/python_shell/dbruntime/MLWorkloadsInstrumentation/_evaluation.py:40, in apply_evaluate_patch.<locals>.patch_evaluate.<locals>.patched_evaluate(model, model_type, dataset, run_id, evaluator_name_list, evaluator_name_to_conf_map, **kwargs)
     38 try:
     39     original_succeeded = False
---> 40     original_result = original_evaluate_fn(
     41         model=model,
     42         model_type=model_type,
     43         run_id=run_id,
     44         dataset=dataset,
     45         evaluator_name_list=evaluator_name_list,
     46         evaluator_name_to_conf_map=evaluator_name_to_conf_map,
     47         **kwargs)
     48     original_succeeded = True
     49     return original_result

File /databricks/python/lib/python3.10/site-packages/mlflow/models/evaluation/base.py:1120, in _evaluate(model, model_type, dataset, run_id, evaluator_name_list, evaluator_name_to_conf_map, custom_metrics, extra_metrics, custom_artifacts, baseline_model, predictions)
   1118     if evaluator.can_evaluate(model_type=model_type, evaluator_config=config):
   1119         _logger.info(f"Evaluating the model with the {evaluator_name} evaluator.")
-> 1120         eval_result = evaluator.evaluate(
   1121             model=model,
   1122             model_type=model_type,
   1123             dataset=dataset,
   1124             run_id=run_id,
   1125             evaluator_config=config,
   1126             custom_metrics=custom_metrics,
   1127             extra_metrics=extra_metrics,
   1128             custom_artifacts=custom_artifacts,
   1129             baseline_model=baseline_model,
   1130             predictions=predictions,
   1131         )
   1132         eval_results.append(eval_result)
   1134 _last_failed_evaluator = None

File /databricks/python/lib/python3.10/site-packages/mlflow/models/evaluation/default_evaluator.py:1837, in DefaultEvaluator.evaluate(self, model_type, dataset, run_id, evaluator_config, model, custom_metrics, extra_metrics, custom_artifacts, baseline_model, predictions, **kwargs)
   1835     if baseline_model:
   1836         _logger.info("Evaluating candidate model:")
-> 1837     evaluation_result = self._evaluate(model, is_baseline_model=False)
   1839 if not baseline_model:
   1840     return evaluation_result

File /databricks/python/lib/python3.10/site-packages/mlflow/models/evaluation/default_evaluator.py:1706, in DefaultEvaluator._evaluate(self, model, is_baseline_model, **kwargs)
   1704         self.extra_metrics.remove(extra_metric)
   1705         break
-> 1706 self._generate_model_predictions(compute_latency=compute_latency)
   1707 if self.model_type in (_ModelType.CLASSIFIER, _ModelType.REGRESSOR):
   1708     self._compute_builtin_metrics()

File /databricks/python/lib/python3.10/site-packages/mlflow/models/evaluation/default_evaluator.py:1369, in DefaultEvaluator._generate_model_predictions(self, compute_latency)
   1367         model_predictions = predict_with_latency(X_copy)
   1368     else:
-> 1369         model_predictions = self.model.predict(X_copy)
   1370 else:
   1371     if self.dataset.predictions_data is None:

File /databricks/python/lib/python3.10/site-packages/mlflow/pyfunc/__init__.py:492, in PyFuncModel.predict(self, data, params)
    489         if _MLFLOW_TESTING.get():
    490             raise
--> 492 return _predict()

File /databricks/python/lib/python3.10/site-packages/mlflow/pyfunc/__init__.py:478, in PyFuncModel.predict.<locals>._predict()
    473 def _predict():
    474     # Models saved prior to MLflow 2.5.0 do not support `params` in the pyfunc `predict()`
    475     # function definition, nor do they support `**kwargs`. Accordingly, we only pass
    476     # `params` to the `predict()` method if it defines the `params` argument
    477     if inspect.signature(self._predict_fn).parameters.get("params"):
--> 478         return self._predict_fn(data, params=params)
    479     _log_warning_if_params_not_in_predict_signature(_logger, params)
    480     return self._predict_fn(data)

File /databricks/python/lib/python3.10/site-packages/mlflow/openai/__init__.py:801, in _OpenAIWrapper.predict(self, data, params)
    799 self.api_token.validate()
    800 if self.task == "chat.completions":
--> 801     return self._predict_chat(data, params or {})
    802 elif self.task == "completions":
    803     return self._predict_completions(data, params or {})

File /databricks/python/lib/python3.10/site-packages/mlflow/openai/__init__.py:724, in _OpenAIWrapper._predict_chat(self, data, params)
    721 requests = [{**self.model, **params, "messages": messages} for messages in messages_list]
    722 request_url = self._construct_request_url("chat/completions", REQUEST_URL_CHAT)
--> 724 results = process_api_requests(
    725     requests,
    726     request_url,
    727     api_token=self.api_token,
    728     max_requests_per_minute=self.api_config.max_requests_per_minute,
    729     max_tokens_per_minute=self.api_config.max_tokens_per_minute,
    730 )
    731 return [r["choices"][0]["message"]["content"] for r in results]

File /databricks/python/lib/python3.10/site-packages/mlflow/openai/api_request_parallel_processor.py:353, in process_api_requests(requests, request_url, api_token, max_requests_per_minute, max_tokens_per_minute, token_encoding_name, max_attempts, max_workers, throw_original_error)
    351     if throw_original_error and len(requests) == 1:
    352         raise status_tracker.error
--> 353     raise mlflow.MlflowException(
    354         f"{status_tracker.num_tasks_failed} tasks failed. See logs for details."
    355     )
    356 if status_tracker.num_rate_limit_errors > 0:
    357     _logger.debug(
    358         f"{status_tracker.num_rate_limit_errors} rate limit errors received. "
    359         "Consider running at a lower rate."
    360     )

MlflowException: 4 tasks failed. See logs for details.
```


### Other info / logs

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


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

- [ ] `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
- [X] `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/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?

- [X] `integrations/azure`: Azure and Azure ML integrations
- [ ] `integrations/sagemaker`: SageMaker integrations
- [ ] `integrations/databricks`: Databricks integrations


Discussion:
Hi @sheigel2 The environment variable for deployment id should be `os.environ["OPENAI_DEPLOYMENT_NAME"] = "gpt-35-turbo"`, if setting this correctly you don't need to log the model with deployment_id. You're right the sample is not correct, I'll fix it.

Resolving patch:
diff --git a/examples/openai/azure_openai.py b/examples/openai/azure_openai.py
--- a/examples/openai/azure_openai.py
+++ b/examples/openai/azure_openai.py
@@ -1,5 +1,3 @@
-import os
-
 import openai
 import pandas as pd
 
@@ -14,7 +12,7 @@
 # OPENAI_API_VERSION e.g. 2023-05-15
 export OPENAI_API_VERSION="<AZURE OPENAI API VERSION>"
 export OPENAI_API_TYPE="azure"
-export DEPLOYMENT_ID="<AZURE OPENAI DEPLOYMENT ID OR NAME>"
+export OPENAI_DEPLOYMENT_NAME="<AZURE OPENAI DEPLOYMENT ID OR NAME>"
 """
 
 with mlflow.start_run():
@@ -24,7 +22,6 @@
         task=openai.ChatCompletion,
         artifact_path="model",
         messages=[{"role": "user", "content": "Tell me a joke about {animal}."}],
-        deployment_id=os.environ["DEPLOYMENT_ID"],
     )
 
 # Load native OpenAI model
diff --git a/mlflow/openai/__init__.py b/mlflow/openai/__init__.py
--- a/mlflow/openai/__init__.py
+++ b/mlflow/openai/__init__.py
@@ -653,15 +653,16 @@ def __init__(self, model):
         self.api_token = _OAITokenHolder(self.api_config.api_type)
         # If the same parameter exists in self.model & self.api_config,
         # we use the parameter from self.model
-        self.envs = {
-            x: getattr(self.api_config, x)
-            for x in ["api_base", "api_version", "api_type", "engine", "deployment_id"]
-            if getattr(self.api_config, x) is not None and x not in self.model
-        }
-        api_type = self.model.get("api_type") or self.envs.get("api_type")
-        if api_type in ("azure", "azure_ad", "azuread"):
-            deployment_id = self.model.get("deployment_id") or self.envs.get("deployment_id")
-            if self.model.get("engine") or self.envs.get("engine"):
+        self.request_configs = {}
+        for x in ["api_base", "api_version", "api_type", "engine", "deployment_id"]:
+            if x in self.model:
+                self.request_configs[x] = self.model.pop(x)
+            elif value := getattr(self.api_config, x):
+                self.request_configs[x] = value
+
+        if self.request_configs.get("api_type") in ("azure", "azure_ad", "azuread"):
+            deployment_id = self.request_configs.get("deployment_id")
+            if self.request_configs.get("engine"):
                 # Avoid using both parameters as they serve the same purpose
                 # Invalid inputs:
                 #   - Wrong engine + correct/wrong deployment_id
@@ -704,11 +705,11 @@ def get_params_list(self, data):
             return data[self.formater.variables].to_dict(orient="records")
 
     def _construct_request_url(self, task_url, default_url):
-        api_type = self.model.get("api_type") or self.envs.get("api_type")
+        api_type = self.request_configs.get("api_type")
         if api_type in ("azure", "azure_ad", "azuread"):
-            api_base = self.envs.get("api_base")
-            api_version = self.envs.get("api_version")
-            deployment_id = self.envs.get("deployment_id")
+            api_base = self.request_configs.get("api_base")
+            api_version = self.request_configs.get("api_version")
+            deployment_id = self.request_configs.get("deployment_id")
 
             return (
                 f"{api_base}/openai/deployments/{deployment_id}/"
diff --git a/mlflow/utils/openai_utils.py b/mlflow/utils/openai_utils.py
--- a/mlflow/utils/openai_utils.py
+++ b/mlflow/utils/openai_utils.py
@@ -211,11 +211,6 @@ def _validate_model_params(task, model, params):
         )
 
 
-def _exclude_params_from_envs(params, envs):
-    """Params passed at inference time should override envs."""
-    return {k: v for k, v in envs.items() if k not in params} if params else envs
-
-
 class _OAITokenHolder:
     def __init__(self, api_type):
         self._api_token = None


Test patch:
diff --git a/tests/openai/test_openai_model_export.py b/tests/openai/test_openai_model_export.py
--- a/tests/openai/test_openai_model_export.py
+++ b/tests/openai/test_openai_model_export.py
@@ -1,6 +1,5 @@
 import importlib
 import json
-from copy import deepcopy
 from unittest import mock
 
 import numpy as np
@@ -17,7 +16,6 @@
 from mlflow.models.signature import ModelSignature
 from mlflow.types.schema import ColSpec, ParamSchema, ParamSpec, Schema, TensorSpec
 from mlflow.utils.openai_utils import (
-    _exclude_params_from_envs,
     _mock_chat_completion_response,
     _mock_models_retrieve_response,
     _mock_request,
@@ -668,20 +666,3 @@ def test_engine_and_deployment_id_for_azure_openai(tmp_path, monkeypatch):
         MlflowException, match=r"Either engine or deployment_id must be set for Azure OpenAI API"
     ):
         mlflow.pyfunc.load_model(tmp_path)
-
-
-@pytest.mark.parametrize(
-    ("params", "envs"),
-    [
-        ({"a": None, "b": "b"}, {"a": "a", "c": "c"}),
-        ({"a": "a", "b": "b"}, {"a": "a", "d": "d"}),
-        ({}, {"a": "a", "b": "b"}),
-        ({"a": "a"}, {"b": "b"}),
-    ],
-)
-def test_exclude_params_from_envs(params, envs):
-    original_envs = deepcopy(envs)
-    result = _exclude_params_from_envs(params, envs)
-    assert envs == original_envs
-    assert not any(key in params for key in result)
-    assert all(key in envs for key in result)