Repository: mlflow/mlflow
Pull request #15623: Fix set active model during model loading
Author: login: serena-ruan
Labels: name: rn/none
description: List under Small Changes in Changelogs.
Created: 2025-05-08T05:26:47Z | Merged: 2025-05-12T07:06:09Z

<details><summary>&#x1F6E0 DevTools &#x1F6E0</summary>
<p>

[![Open in GitHub Codespaces](https://github.com/codespaces/badge.svg)](https://codespaces.new/serena-ruan/mlflow/pull/15623?quickstart=1)

#### Install mlflow from this PR

```
# mlflow
pip install git+https://github.com/mlflow/mlflow.git@refs/pull/15623/merge
# mlflow-skinny
pip install git+https://github.com/mlflow/mlflow.git@refs/pull/15623/merge#subdirectory=skinny
```

For Databricks, use the following command:

```
%sh curl -LsSf https://raw.githubusercontent.com/mlflow/mlflow/HEAD/dev/install-skinny.sh | sh -s pull/15623/merge
```

</p>
</details>

### Related Issues/PRs

<!-- Uncomment 'Resolve' if this PR can close the linked items. -->
<!-- Resolve --> #xxx

### What changes are proposed in this pull request?
When MLflow internally sets the active_model_id inside model loading, we don't need to check if the model exists or not. Besides set_active_model shouldn't block model loading even if it fails, so we wrap it with try catch in this PR.
Also move the set_active_model at the end of model loading, since users might set experiment inside the model definition.

<!-- Please fill in changes proposed in this PR. -->

### How is this PR tested?

- [ ] Existing unit/integration tests
- [ ] New unit/integration tests
- [x] Manual tests

<!-- Attach code, screenshot, video used for manual testing here. -->

### Does this PR require documentation update?

- [ ] No. You can skip the rest of this section.
- [ ] Yes. I've updated:
  - [ ] Examples
  - [ ] API references
  - [ ] Instructions

### Release Notes

#### Is this a user-facing change?

- [ ] No. You can skip the rest of this section.
- [ ] Yes. Give a description of this change to be included in the release notes for MLflow users.

<!-- Details in 1-2 sentences. You can just refer to another PR with a description if this PR is part of a larger change. -->

#### What component(s), interfaces, languages, and integrations does this PR affect?

Components

- [ ] `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/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

<!--
Insert an empty named anchor here to allow jumping to this section with a fragment URL
(e.g. https://github.com/mlflow/mlflow/pull/123#user-content-release-note-category).
Note that GitHub prefixes anchor names in markdown with "user-content-".
-->

<a name="release-note-category"></a>

#### How should the PR be classified in the release notes? Choose one:

- [x] `rn/none` - No description will be included. The PR will be mentioned only by the PR number in the "Small Bugfixes and Documentation Updates" section
- [ ] `rn/breaking-change` - The PR will be mentioned in the "Breaking Changes" section
- [ ] `rn/feature` - A new user-facing feature worth mentioning in the release notes
- [ ] `rn/bug-fix` - A user-facing bug fix worth mentioning in the release notes
- [ ] `rn/documentation` - A user-facing documentation change worth mentioning in the release notes

#### Should this PR be included in the next patch release?

`Yes` should be selected for bug fixes, documentation updates, and other small changes. `No` should be selected for new features and larger changes. If you're unsure about the release classification of this PR, leave this unchecked to let the maintainers decide.

<details>
<summary>What is a minor/patch release?</summary>

- Minor release: a release that increments the second part of the version number (e.g., 1.2.0 -> 1.3.0).
  Bug fixes, doc updates and new features usually go into minor releases.
- Patch release: a release that increments the third part of the version number (e.g., 1.2.0 -> 1.2.1).
  Bug fixes and doc updates usually go into patch releases.

</details>

<!-- patch -->

- [ ] Yes (this PR will be cherry-picked and included in the next patch release)
- [x] No (this PR will be included in the next minor release)



Diff:
diff --git a/mlflow/dspy/load.py b/mlflow/dspy/load.py
index 8495b09b093b4..3ed807518b73f 100644
--- a/mlflow/dspy/load.py
+++ b/mlflow/dspy/load.py
@@ -5,12 +5,9 @@
 
 from mlflow.models import Model
 from mlflow.models.dependencies_schemas import _get_dependencies_schema_from_model
+from mlflow.models.model import _update_active_model_id_based_on_mlflow_model
 from mlflow.tracing.provider import trace_disabled
 from mlflow.tracking.artifact_utils import _download_artifact_from_uri
-from mlflow.tracking.fluent import (
-    _get_active_model_context,
-    _set_active_model,
-)
 from mlflow.utils.annotations import experimental
 from mlflow.utils.autologging_utils import disable_autologging_globally
 from mlflow.utils.model_utils import (
@@ -40,17 +37,6 @@ def _set_dependency_schema_to_tracer(model_path, callbacks):
 def _load_model(model_uri, dst_path=None):
     local_model_path = _download_artifact_from_uri(artifact_uri=model_uri, output_path=dst_path)
     mlflow_model = Model.load(local_model_path)
-    if (
-        mlflow_model.model_id
-        and (amc := _get_active_model_context())
-        # only set the active model if the model is not set by the user
-        and not amc.set_by_user
-        and amc.model_id != mlflow_model.model_id
-    ):
-        _set_active_model(model_id=mlflow_model.model_id)
-        _logger.info(
-            "Use `mlflow.set_active_model` to set the active model to a different one if needed."
-        )
     flavor_conf = _get_flavor_configuration(model_path=local_model_path, flavor_name="dspy")
 
     _add_code_from_conf_to_system_path(local_model_path, flavor_conf)
@@ -59,6 +45,7 @@ def _load_model(model_uri, dst_path=None):
         loaded_wrapper = cloudpickle.load(f)
 
     _set_dependency_schema_to_tracer(local_model_path, loaded_wrapper.dspy_settings["callbacks"])
+    _update_active_model_id_based_on_mlflow_model(mlflow_model)
     return loaded_wrapper
 
 
diff --git a/mlflow/langchain/model.py b/mlflow/langchain/model.py
index 4ba3a98ea6857..8ac1586d6e598 100644
--- a/mlflow/langchain/model.py
+++ b/mlflow/langchain/model.py
@@ -47,7 +47,12 @@
     _get_dependencies_schema_from_model,
     _get_dependencies_schemas,
 )
-from mlflow.models.model import MLMODEL_FILE_NAME, MODEL_CODE_PATH, MODEL_CONFIG
+from mlflow.models.model import (
+    MLMODEL_FILE_NAME,
+    MODEL_CODE_PATH,
+    MODEL_CONFIG,
+    _update_active_model_id_based_on_mlflow_model,
+)
 from mlflow.models.resources import DatabricksFunction, Resource, _ResourceBuilder
 from mlflow.models.signature import _infer_signature_from_input_example
 from mlflow.models.utils import (
@@ -59,10 +64,6 @@
 from mlflow.tracing.provider import trace_disabled
 from mlflow.tracking._model_registry import DEFAULT_AWAIT_MAX_SLEEP_SECONDS
 from mlflow.tracking.artifact_utils import _download_artifact_from_uri
-from mlflow.tracking.fluent import (
-    _get_active_model_context,
-    _set_active_model,
-)
 from mlflow.types.schema import ColSpec, DataType, Schema
 from mlflow.utils.annotations import experimental
 from mlflow.utils.autologging_utils import (
@@ -847,17 +848,6 @@ def _load_pyfunc(path: str, model_config: Optional[dict[str, Any]] = None):  # n
 
 def _load_model_from_local_fs(local_model_path, model_config_overrides=None):
     mlflow_model = Model.load(local_model_path)
-    if (
-        mlflow_model.model_id
-        and (amc := _get_active_model_context())
-        # only set the active model if the model is not set by the user
-        and not amc.set_by_user
-        and amc.model_id != mlflow_model.model_id
-    ):
-        _set_active_model(model_id=mlflow_model.model_id)
-        logger.info(
-            "Use `mlflow.set_active_model` to set the active model to a different one if needed."
-        )
     flavor_conf = _get_flavor_configuration(model_path=local_model_path, flavor_name=FLAVOR_NAME)
     pyfunc_flavor_conf = _get_flavor_configuration(
         model_path=local_model_path, flavor_name=PYFUNC_FLAVOR_NAME
@@ -891,9 +881,12 @@ def _load_model_from_local_fs(local_model_path, model_config_overrides=None):
             # We would like to clean up the dependencies schema which is set to global
             # after loading the mode to avoid the schema being used in the next model loading
             _clear_dependencies_schemas()
-        return model
     else:
-        return _load_model(local_model_path, flavor_conf)
+        model = _load_model(local_model_path, flavor_conf)
+    # set active model after model loading since experiment ID might be set
+    # in the model loading process
+    _update_active_model_id_based_on_mlflow_model(mlflow_model)
+    return model
 
 
 @experimental
diff --git a/mlflow/llama_index/model.py b/mlflow/llama_index/model.py
index 99179a5220475..7a7954f5325b5 100644
--- a/mlflow/llama_index/model.py
+++ b/mlflow/llama_index/model.py
@@ -12,7 +12,12 @@
 from mlflow.llama_index.constant import FLAVOR_NAME
 from mlflow.llama_index.pyfunc_wrapper import create_pyfunc_wrapper
 from mlflow.models import Model, ModelInputExample, ModelSignature
-from mlflow.models.model import MLMODEL_FILE_NAME, MODEL_CODE_PATH, MODEL_CONFIG
+from mlflow.models.model import (
+    MLMODEL_FILE_NAME,
+    MODEL_CODE_PATH,
+    MODEL_CONFIG,
+    _update_active_model_id_based_on_mlflow_model,
+)
 from mlflow.models.signature import _infer_signature_from_input_example
 from mlflow.models.utils import (
     _load_model_code_path,
@@ -22,10 +27,6 @@
 from mlflow.tracing.provider import trace_disabled
 from mlflow.tracking._model_registry import DEFAULT_AWAIT_MAX_SLEEP_SECONDS
 from mlflow.tracking.artifact_utils import _download_artifact_from_uri
-from mlflow.tracking.fluent import (
-    _get_active_model_context,
-    _set_active_model,
-)
 from mlflow.utils.annotations import experimental
 from mlflow.utils.autologging_utils import (
     disable_autologging_globally,
@@ -561,23 +562,14 @@ def load_model(model_uri, dst_path=None):
 
     local_model_path = _download_artifact_from_uri(artifact_uri=model_uri, output_path=dst_path)
     mlflow_model = Model.load(local_model_path)
-    if (
-        mlflow_model.model_id
-        and (amc := _get_active_model_context())
-        # only set the active model if the model is not set by the user
-        and not amc.set_by_user
-        and amc.model_id != mlflow_model.model_id
-    ):
-        _set_active_model(model_id=mlflow_model.model_id)
-        _logger.info(
-            "Use `mlflow.set_active_model` to set the active model to a different one if needed."
-        )
     flavor_conf = _get_flavor_configuration(model_path=local_model_path, flavor_name=FLAVOR_NAME)
 
     settings_path = os.path.join(local_model_path, _SETTINGS_FILE)
     # NB: Settings is a singleton and can be loaded via llama_index.core.Settings
     deserialize_settings(settings_path)
-    return _load_llama_model(local_model_path, flavor_conf)
+    model = _load_llama_model(local_model_path, flavor_conf)
+    _update_active_model_id_based_on_mlflow_model(mlflow_model)
+    return model
 
 
 def _load_pyfunc(path, model_config: Optional[dict[str, Any]] = None):
diff --git a/mlflow/models/model.py b/mlflow/models/model.py
index 6a14bcded0b77..631c6fa99d13b 100644
--- a/mlflow/models/model.py
+++ b/mlflow/models/model.py
@@ -29,7 +29,11 @@
 from mlflow.tracking._model_registry import DEFAULT_AWAIT_MAX_SLEEP_SECONDS
 from mlflow.tracking._tracking_service.utils import _resolve_tracking_uri
 from mlflow.tracking.artifact_utils import _download_artifact_from_uri, _upload_artifact_to_uri
-from mlflow.tracking.fluent import _use_logged_model
+from mlflow.tracking.fluent import (
+    _get_active_model_context,
+    _set_active_model_id,
+    _use_logged_model,
+)
 from mlflow.utils.annotations import experimental
 from mlflow.utils.databricks_utils import (
     _construct_databricks_uc_registered_model_url,
@@ -1393,3 +1397,18 @@ def set_model(model) -> None:
             pass
 
     raise mlflow.MlflowException(SET_MODEL_ERROR)
+
+
+def _update_active_model_id_based_on_mlflow_model(mlflow_model: Model):
+    """
+    Update the current active model ID based on the provided MLflow model.
+    Only set the active model ID if it is not already set by the user.
+    This is useful for setting the active model ID when loading a model
+    to ensure traces generated are associated with the loaded model.
+    """
+    if mlflow_model.model_id is None:
+        return
+    amc = _get_active_model_context()
+    # only set the active model if the model is not set by the user
+    if amc.model_id != mlflow_model.model_id and not amc.set_by_user:
+        _set_active_model_id(model_id=mlflow_model.model_id)
diff --git a/mlflow/openai/model.py b/mlflow/openai/model.py
index a1e0eb0415318..b30e412e057fc 100644
--- a/mlflow/openai/model.py
+++ b/mlflow/openai/model.py
@@ -16,17 +16,13 @@
 from mlflow.environment_variables import MLFLOW_OPENAI_SECRET_SCOPE
 from mlflow.exceptions import MlflowException
 from mlflow.models import Model, ModelInputExample, ModelSignature
-from mlflow.models.model import MLMODEL_FILE_NAME
+from mlflow.models.model import MLMODEL_FILE_NAME, _update_active_model_id_based_on_mlflow_model
 from mlflow.models.signature import _infer_signature_from_input_example
 from mlflow.models.utils import _save_example
 from mlflow.openai.constant import FLAVOR_NAME
 from mlflow.protos.databricks_pb2 import INVALID_PARAMETER_VALUE
 from mlflow.tracking._model_registry import DEFAULT_AWAIT_MAX_SLEEP_SECONDS
 from mlflow.tracking.artifact_utils import _download_artifact_from_uri
-from mlflow.tracking.fluent import (
-    _get_active_model_context,
-    _set_active_model,
-)
 from mlflow.types import ColSpec, Schema, TensorSpec
 from mlflow.utils.annotations import experimental
 from mlflow.utils.autologging_utils import disable_autologging_globally
@@ -547,18 +543,7 @@ def _load_model(path):
     model_file_path = os.path.dirname(path)
     if os.path.exists(model_file_path):
         mlflow_model = Model.load(model_file_path)
-        if (
-            mlflow_model.model_id
-            and (amc := _get_active_model_context())
-            # only set the active model if the model is not set by the user
-            and not amc.set_by_user
-            and amc.model_id != mlflow_model.model_id
-        ):
-            _set_active_model(model_id=mlflow_model.model_id)
-            _logger.info(
-                "Use `mlflow.set_active_model` to set the active model "
-                "to a different one if needed."
-            )
+        _update_active_model_id_based_on_mlflow_model(mlflow_model)
     with open(path) as f:
         return yaml.safe_load(f)
 
diff --git a/mlflow/tracking/fluent.py b/mlflow/tracking/fluent.py
index add0da309aee4..71dca3154f97e 100644
--- a/mlflow/tracking/fluent.py
+++ b/mlflow/tracking/fluent.py
@@ -3213,9 +3213,7 @@ def __init__(self, logged_model: LoggedModel, set_by_user: bool):
         super().__init__(**logged_model.to_dictionary())
         self.last_active_model_context = _ACTIVE_MODEL_CONTEXT.get()
         self.last_active_model_id_env_var = MLFLOW_ACTIVE_MODEL_ID.get()
-        _ACTIVE_MODEL_CONTEXT.set(ActiveModelContext(self.model_id, set_by_user))
-        _update_active_model_id_env_var(self.model_id)
-        _logger.info(f"Active model set to model with ID: {self.model_id}")
+        _set_active_model_id(self.model_id, set_by_user)
 
     def __enter__(self):
         return self
@@ -3279,13 +3277,12 @@ def predict(model_input):
         traces = mlflow.search_traces(model_id=mlflow.get_active_model_id(), return_type="list")
         assert len(traces) == 1
     """
-    active_model = _set_active_model(name=name, model_id=model_id)
-    active_model_ctx = _ACTIVE_MODEL_CONTEXT.get()
-    _ACTIVE_MODEL_CONTEXT.set(ActiveModelContext(active_model_ctx.model_id, set_by_user=True))
-    return active_model
+    return _set_active_model(name=name, model_id=model_id, set_by_user=True)
 
 
-def _set_active_model(*, name: Optional[str] = None, model_id: Optional[str] = None) -> ActiveModel:
+def _set_active_model(
+    *, name: Optional[str] = None, model_id: Optional[str] = None, set_by_user: bool = False
+) -> ActiveModel:
     if name is None and model_id is None:
         raise MlflowException.invalid_parameter_value(
             message="Either name or model_id must be provided",
@@ -3312,7 +3309,29 @@ def _set_active_model(*, name: Optional[str] = None, model_id: Optional[str] = N
             logged_model = mlflow.create_external_model(name=name)
         else:
             logged_model = logged_models[0]
-    return ActiveModel(logged_model=logged_model, set_by_user=False)
+    return ActiveModel(logged_model=logged_model, set_by_user=set_by_user)
+
+
+def _set_active_model_id(model_id: str, set_by_user: bool = False) -> None:
+    """
+    Set the active model ID in the active model context and update the
+    corresponding environment variable. This should only be used when
+    we know the LoggedModel with the model_id exists.
+    This function should be used inside MLflow to set the active model
+    while not blocking other code execution.
+    """
+    try:
+        _ACTIVE_MODEL_CONTEXT.set(ActiveModelContext(model_id, set_by_user))
+        _update_active_model_id_env_var(model_id)
+    except Exception as e:
+        _logger.warning(f"Failed to set active model ID to {model_id}, error: {e}")
+    else:
+        _logger.info(f"Active model is set to the logged model with ID: {model_id}")
+        if not set_by_user:
+            _logger.info(
+                "Use `mlflow.set_active_model` to set the active model "
+                "to a different one if needed."
+            )
 
 
 def _get_active_model_context() -> ActiveModelContext:
diff --git a/tests/dspy/test_dspy_autolog.py b/tests/dspy/test_dspy_autolog.py
index 3d0c36cd1c41b..8ddb83c825f81 100644
--- a/tests/dspy/test_dspy_autolog.py
+++ b/tests/dspy/test_dspy_autolog.py
@@ -952,3 +952,23 @@ def test_autolog_link_traces_active_model():
         model_id = json.loads(trace.data.request)["args"][0]
         assert model_id != model.model_id
         assert trace.info.request_metadata[SpanAttributeKey.MODEL_ID] == model.model_id
+
+
+@skip_when_testing_trace_sdk
+def test_model_loading_set_active_model_id_without_fetching_logged_model():
+    mlflow.dspy.autolog()
+    dspy.settings.configure(
+        lm=DummyLM([{"answer": "test output", "reasoning": "No more responses"}])
+    )
+    dspy_model = CoT()
+
+    model_info = mlflow.dspy.log_model(dspy_model, name="model")
+
+    with mock.patch("mlflow.get_logged_model", side_effect=Exception("get_logged_model failed")):
+        loaded_model = mlflow.dspy.load_model(model_info.model_uri)
+    loaded_model(model_info.model_id)
+
+    traces = get_traces()
+    assert len(traces) == 1
+    model_id = json.loads(traces[0].data.request)["args"][0]
+    assert model_id == traces[0].info.request_metadata[SpanAttributeKey.MODEL_ID]
diff --git a/tests/langchain/test_langchain_autolog.py b/tests/langchain/test_langchain_autolog.py
index a79c311592d2c..dabaddb9aa893 100644
--- a/tests/langchain/test_langchain_autolog.py
+++ b/tests/langchain/test_langchain_autolog.py
@@ -1165,3 +1165,21 @@ def test_autolog_link_traces_to_active_model(model_infos):
         assert logged_model_id is not None
         assert trace.info.request_metadata[SpanAttributeKey.MODEL_ID] == model.model_id
         assert model.model_id != logged_model_id
+
+
+def test_model_loading_set_active_model_id_without_fetching_logged_model():
+    mlflow.langchain.autolog()
+
+    model = create_openai_runnable(temperature=0.9)
+    model_info = mlflow.langchain.log_model(
+        model, name="model", input_example={"product": "MLflow"}
+    )
+
+    with mock.patch("mlflow.get_logged_model", side_effect=Exception("get_logged_model failed")):
+        loaded_model = mlflow.langchain.load_model(model_info.model_uri)
+    loaded_model.invoke({"product": "MLflow"})
+
+    traces = get_traces()
+    assert len(traces) == 1
+    model_id = traces[0].info.request_metadata[SpanAttributeKey.MODEL_ID]
+    assert model_id == model_info.model_id
diff --git a/tests/llama_index/test_llama_index_autolog.py b/tests/llama_index/test_llama_index_autolog.py
index ad67d92fc9277..9b6d5373085a3 100644
--- a/tests/llama_index/test_llama_index_autolog.py
+++ b/tests/llama_index/test_llama_index_autolog.py
@@ -333,3 +333,27 @@ def test_autolog_link_traces_to_active_model():
     model_id = traces[0].info.request_metadata[SpanAttributeKey.MODEL_ID]
     assert model_id == model.model_id
     assert model_id != model_info.model_id
+
+
+@skip_when_testing_trace_sdk
+@pytest.mark.skipif(
+    llama_core_version < Version("0.11.0"),
+    reason="Workflow was introduced in 0.11.0",
+)
+def test_model_loading_set_active_model_id_without_fetching_logged_model():
+    mlflow.llama_index.autolog()
+    model_info = mlflow.llama_index.log_model(
+        "tests/llama_index/sample_code/simple_workflow.py",
+        name="model",
+        input_example={"topic": "Hello"},
+    )
+    with mock.patch("mlflow.get_logged_model", side_effect=Exception("get_logged_model failed")):
+        loaded_workflow = mlflow.pyfunc.load_model(model_info.model_uri)
+    loaded_workflow.predict({"topic": f"Hello {model_info.model_id}"})
+
+    traces = get_traces()
+    assert len(traces) == 1
+    span = traces[0].data.spans[0]
+    model_id = traces[0].info.request_metadata[SpanAttributeKey.MODEL_ID]
+    assert model_id is not None
+    assert span.inputs["kwargs"]["topic"] == f"Hello {model_id}"
diff --git a/tests/openai/test_openai_autolog.py b/tests/openai/test_openai_autolog.py
index 0cb7c9308ea2d..f262fd40e932e 100644
--- a/tests/openai/test_openai_autolog.py
+++ b/tests/openai/test_openai_autolog.py
@@ -754,3 +754,34 @@ def test_autolog_link_traces_to_active_model(monkeypatch, mock_openai):
         assert trace.info.request_metadata[SpanAttributeKey.MODEL_ID] == model.model_id
         logged_model_id = span.inputs["input"][0]
         assert logged_model_id != model.model_id
+
+
+@skip_when_testing_trace_sdk
+@pytest.mark.asyncio
+async def test_model_loading_set_active_model_id_without_fetching_logged_model(client):
+    mlflow.openai.autolog()
+
+    model_info = mlflow.openai.log_model(
+        "gpt-4o-mini",
+        "chat.completions",
+        name="model",
+        temperature=0.9,
+        messages=[{"role": "system", "content": "You are an MLflow expert."}],
+    )
+
+    with mock.patch("mlflow.get_logged_model", side_effect=Exception("get_logged_model failed")):
+        model_dict = mlflow.openai.load_model(model_info.model_uri)
+    resp = client.chat.completions.create(
+        messages=[{"role": "user", "content": f"test {model_info.model_id}"}],
+        model=model_dict["model"],
+        temperature=model_dict["temperature"],
+    )
+    if client._is_async:
+        await resp
+
+    traces = get_traces()
+    assert len(traces) == 1
+    span = traces[0].data.spans[0]
+    model_id = traces[0].info.request_metadata[SpanAttributeKey.MODEL_ID]
+    assert model_id is not None
+    assert span.inputs["messages"][0]["content"] == f"test {model_id}"


Reviews:
user: login: harupy
state: COMMENTED
author_association: MEMBER
submitted_at: 2025-05-12T01:45:40Z
user: login: B-Step62
state: COMMENTED
author_association: COLLABORATOR
submitted_at: 2025-05-12T01:58:35Z
user: login: serena-ruan
state: COMMENTED
author_association: COLLABORATOR
submitted_at: 2025-05-12T02:01:18Z
user: login: serena-ruan
state: COMMENTED
author_association: COLLABORATOR
submitted_at: 2025-05-12T02:07:58Z
user: login: B-Step62
state: COMMENTED
author_association: COLLABORATOR
submitted_at: 2025-05-12T02:37:22Z
user: login: serena-ruan
state: COMMENTED
author_association: COLLABORATOR
submitted_at: 2025-05-12T03:02:06Z
user: login: B-Step62
state: COMMENTED
author_association: COLLABORATOR
submitted_at: 2025-05-12T03:20:35Z
user: login: serena-ruan
state: COMMENTED
author_association: COLLABORATOR
submitted_at: 2025-05-12T05:14:05Z
user: login: B-Step62
state: COMMENTED
author_association: COLLABORATOR
submitted_at: 2025-05-12T05:28:03Z
user: login: B-Step62
body: LGTM once the comments are addressed!
state: APPROVED
author_association: COLLABORATOR
submitted_at: 2025-05-12T05:38:00Z

Review comments:
diff_hunk: @@ -333,3 +333,27 @@ def test_autolog_link_traces_to_active_model():
     model_id = traces[0].info.request_metadata[SpanAttributeKey.MODEL_ID]
     assert model_id == model.model_id
     assert model_id != model_info.model_id
+
+
+@skip_when_testing_trace_sdk
+@pytest.mark.skipif(
+    llama_core_version < Version("0.11.0"),
+    reason="Workflow was introduced in 0.11.0",
+)
+def test_model_loading_set_active_model_id_without_fetching_logged_model():
+    mlflow.llama_index.autolog()
+    model_info = mlflow.llama_index.log_model(
+        "tests/llama_index/sample_code/simple_workflow.py",
+        name="model",
+        input_example={"topic": "Hello"},
+    )
+    with mock.patch("mlflow.get_logged_model", side_effect=Exception):
+        loaded_workflow = mlflow.pyfunc.load_model(model_info.model_uri)
+        loaded_workflow.predict({"topic": f"Hello {model_info.model_id}"})
path: tests/llama_index/test_llama_index_autolog.py
user: login: harupy
body: ```suggestion
    with mock.patch("mlflow.get_logged_model", side_effect=Exception("get_logged_model failed")):
        loaded_workflow = mlflow.pyfunc.load_model(model_info.model_uri)
        loaded_workflow.predict({"topic": f"Hello {model_info.model_id}"})
```

both load_model and predict need mock?
created_at: 2025-05-12T01:45:40Z
author_association: MEMBER
diff_hunk: @@ -40,25 +40,25 @@ def _set_dependency_schema_to_tracer(model_path, callbacks):
 def _load_model(model_uri, dst_path=None):
     local_model_path = _download_artifact_from_uri(artifact_uri=model_uri, output_path=dst_path)
     mlflow_model = Model.load(local_model_path)
+    flavor_conf = _get_flavor_configuration(model_path=local_model_path, flavor_name="dspy")
+
+    _add_code_from_conf_to_system_path(local_model_path, flavor_conf)
+    model_path = flavor_conf.get("model_path", _DEFAULT_MODEL_PATH)
+    with open(os.path.join(local_model_path, model_path), "rb") as f:
+        loaded_wrapper = cloudpickle.load(f)
+
+    _set_dependency_schema_to_tracer(local_model_path, loaded_wrapper.dspy_settings["callbacks"])
     if (
         mlflow_model.model_id
path: mlflow/dspy/load.py
user: login: B-Step62
body: Can we extract this to a function and make it more readable? This code is duplicated across flavors so updating a piece of this logic is error prone.
created_at: 2025-05-12T01:53:31Z
author_association: COLLABORATOR
diff_hunk: @@ -3285,6 +3283,19 @@ def _set_active_model(*, name: Optional[str] = None, model_id: Optional[str] = N
     return ActiveModel(logged_model=logged_model, set_by_user=False)
 
 
+def _set_active_model_id(model_id: str, set_by_user: bool = False) -> None:
path: mlflow/tracking/fluent.py
user: login: B-Step62
body: Can we add try-catch for additional safety?
created_at: 2025-05-12T01:55:30Z
author_association: COLLABORATOR
diff_hunk: @@ -3285,6 +3283,19 @@ def _set_active_model(*, name: Optional[str] = None, model_id: Optional[str] = N
     return ActiveModel(logged_model=logged_model, set_by_user=False)
 
 
+def _set_active_model_id(model_id: str, set_by_user: bool = False) -> None:
+    """
+    Set the active model ID in the active model context and update the
+    corresponding environment variable. This should only be used when
+    we know the LoggedModel with the model_id exists.
+    This function should be used inside MLflow to set the active model
+    while not blocking other code execution.
+    """
+    _ACTIVE_MODEL_CONTEXT.set(ActiveModelContext(model_id, set_by_user))
+    _update_active_model_id_env_var(model_id)
path: mlflow/tracking/fluent.py
user: login: B-Step62
body: Q: Why do we need to update env var?
created_at: 2025-05-12T01:58:20Z
author_association: COLLABORATOR
diff_hunk: @@ -333,3 +333,27 @@ def test_autolog_link_traces_to_active_model():
     model_id = traces[0].info.request_metadata[SpanAttributeKey.MODEL_ID]
     assert model_id == model.model_id
     assert model_id != model_info.model_id
+
+
+@skip_when_testing_trace_sdk
+@pytest.mark.skipif(
+    llama_core_version < Version("0.11.0"),
+    reason="Workflow was introduced in 0.11.0",
+)
+def test_model_loading_set_active_model_id_without_fetching_logged_model():
+    mlflow.llama_index.autolog()
+    model_info = mlflow.llama_index.log_model(
+        "tests/llama_index/sample_code/simple_workflow.py",
+        name="model",
+        input_example={"topic": "Hello"},
+    )
+    with mock.patch("mlflow.get_logged_model", side_effect=Exception):
+        loaded_workflow = mlflow.pyfunc.load_model(model_info.model_uri)
+        loaded_workflow.predict({"topic": f"Hello {model_info.model_id}"})
path: tests/llama_index/test_llama_index_autolog.py
user: login: serena-ruan
body: No, only load_model needs the mock, let me update
created_at: 2025-05-12T02:01:18Z
author_association: COLLABORATOR
diff_hunk: @@ -3285,6 +3283,19 @@ def _set_active_model(*, name: Optional[str] = None, model_id: Optional[str] = N
     return ActiveModel(logged_model=logged_model, set_by_user=False)
 
 
+def _set_active_model_id(model_id: str, set_by_user: bool = False) -> None:
+    """
+    Set the active model ID in the active model context and update the
+    corresponding environment variable. This should only be used when
+    we know the LoggedModel with the model_id exists.
+    This function should be used inside MLflow to set the active model
+    while not blocking other code execution.
+    """
+    _ACTIVE_MODEL_CONTEXT.set(ActiveModelContext(model_id, set_by_user))
+    _update_active_model_id_env_var(model_id)
path: mlflow/tracking/fluent.py
user: login: serena-ruan
body: Because if the predict is executed in subprocess (e.g. mlflow.models.predict) then we need to attach the traces to the LoggedModel as well, so we use environment variable to sync this behavior
created_at: 2025-05-12T02:07:58Z
author_association: COLLABORATOR
diff_hunk: @@ -3285,6 +3283,19 @@ def _set_active_model(*, name: Optional[str] = None, model_id: Optional[str] = N
     return ActiveModel(logged_model=logged_model, set_by_user=False)
 
 
+def _set_active_model_id(model_id: str, set_by_user: bool = False) -> None:
+    """
+    Set the active model ID in the active model context and update the
+    corresponding environment variable. This should only be used when
+    we know the LoggedModel with the model_id exists.
+    This function should be used inside MLflow to set the active model
+    while not blocking other code execution.
+    """
+    _ACTIVE_MODEL_CONTEXT.set(ActiveModelContext(model_id, set_by_user))
+    _update_active_model_id_env_var(model_id)
path: mlflow/tracking/fluent.py
user: login: B-Step62
body: Is there any other use case than `mlflow.models.predict`? Fwiw we don't support multi-processing in MLflow tracing. So if the use case is `mlflow.models.predict` only, I'd suggest passing that env var directly inside `mlflow.models.predict` rather than always mutate an env var.
created_at: 2025-05-12T02:37:22Z
author_association: COLLABORATOR
diff_hunk: @@ -3285,6 +3283,19 @@ def _set_active_model(*, name: Optional[str] = None, model_id: Optional[str] = N
     return ActiveModel(logged_model=logged_model, set_by_user=False)
 
 
+def _set_active_model_id(model_id: str, set_by_user: bool = False) -> None:
+    """
+    Set the active model ID in the active model context and update the
+    corresponding environment variable. This should only be used when
+    we know the LoggedModel with the model_id exists.
+    This function should be used inside MLflow to set the active model
+    while not blocking other code execution.
+    """
+    _ACTIVE_MODEL_CONTEXT.set(ActiveModelContext(model_id, set_by_user))
+    _update_active_model_id_env_var(model_id)
path: mlflow/tracking/fluent.py
user: login: serena-ruan
body: So tracing doesn't work for spark_udf? I think `mlflow.models.predict` is the only use case I can think of now (inside MLflow), but is there any downside of mutating the env var?
created_at: 2025-05-12T03:02:06Z
author_association: COLLABORATOR
diff_hunk: @@ -3285,6 +3283,19 @@ def _set_active_model(*, name: Optional[str] = None, model_id: Optional[str] = N
     return ActiveModel(logged_model=logged_model, set_by_user=False)
 
 
+def _set_active_model_id(model_id: str, set_by_user: bool = False) -> None:
+    """
+    Set the active model ID in the active model context and update the
+    corresponding environment variable. This should only be used when
+    we know the LoggedModel with the model_id exists.
+    This function should be used inside MLflow to set the active model
+    while not blocking other code execution.
+    """
+    _ACTIVE_MODEL_CONTEXT.set(ActiveModelContext(model_id, set_by_user))
+    _update_active_model_id_env_var(model_id)
path: mlflow/tracking/fluent.py
user: login: B-Step62
body: It is supported as long as each distributor creates a single trace. We don't support sth like

```
@mlflow.trace
def predict():
   # Do something with subprocess
```

> is there any downside of mutating the env var?

Yes, it is not thread-safe and complicates testing. It is acceptable to set env var to a subprocess we spawn, but I don't think we should overwrite env var in main process in general.
created_at: 2025-05-12T03:20:35Z
author_association: COLLABORATOR
diff_hunk: @@ -3285,6 +3283,19 @@ def _set_active_model(*, name: Optional[str] = None, model_id: Optional[str] = N
     return ActiveModel(logged_model=logged_model, set_by_user=False)
 
 
+def _set_active_model_id(model_id: str, set_by_user: bool = False) -> None:
+    """
+    Set the active model ID in the active model context and update the
+    corresponding environment variable. This should only be used when
+    we know the LoggedModel with the model_id exists.
+    This function should be used inside MLflow to set the active model
+    while not blocking other code execution.
+    """
+    _ACTIVE_MODEL_CONTEXT.set(ActiveModelContext(model_id, set_by_user))
+    _update_active_model_id_env_var(model_id)
path: mlflow/tracking/fluent.py
user: login: serena-ruan
body: > it is not thread-safe

Yea seems like this happens if we try to fetch active_model_id using the env var... I think we should set the context var only and use the env var as the default value, and avoid letting users explicitly mutate the env var. Can I merge this PR in first and fix this in a follow-up PR?
created_at: 2025-05-12T05:14:04Z
author_association: COLLABORATOR
diff_hunk: @@ -3285,6 +3283,19 @@ def _set_active_model(*, name: Optional[str] = None, model_id: Optional[str] = N
     return ActiveModel(logged_model=logged_model, set_by_user=False)
 
 
+def _set_active_model_id(model_id: str, set_by_user: bool = False) -> None:
+    """
+    Set the active model ID in the active model context and update the
+    corresponding environment variable. This should only be used when
+    we know the LoggedModel with the model_id exists.
+    This function should be used inside MLflow to set the active model
+    while not blocking other code execution.
+    """
+    _ACTIVE_MODEL_CONTEXT.set(ActiveModelContext(model_id, set_by_user))
+    _update_active_model_id_env_var(model_id)
path: mlflow/tracking/fluent.py
user: login: B-Step62
body: Yea follow-up is fine:)
created_at: 2025-05-12T05:28:03Z
author_association: COLLABORATOR
diff_hunk: @@ -1393,3 +1397,23 @@ def set_model(model) -> None:
             pass
 
     raise mlflow.MlflowException(SET_MODEL_ERROR)
+
+
+def _update_active_model_id_based_on_mlflow_model(mlflow_model: Model):
+    """
+    Update the current active model ID based on the provided MLflow model.
+    Only set the active model ID if it is not already set by the user.
+    This is useful for setting the active model ID when loading a model
+    to ensure traces generated are associated with the loaded model.
+    """
+    if (
+        mlflow_model.model_id
+        and (amc := _get_active_model_context())
+        # only set the active model if the model is not set by the user
+        and not amc.set_by_user
+        and amc.model_id != mlflow_model.model_id
+    ):
+        _set_active_model_id(model_id=mlflow_model.model_id)
+        _logger.info(
+            "Use `mlflow.set_active_model` to set the active model to a different one if needed."
+        )
path: mlflow/models/model.py
user: login: B-Step62
body: ```suggestion
    if mlflow_model.model_id is None:
        return
   
    active_model = _get_active_model_context())
    if (
        active_model.model_id != mlflow_model.model_id
        # only set the active model if the model is not set by the user
        and not active_model.set_by_user
     ):
         _set_active_model_id(model_id=mlflow_model.model_id)
```

nit: The `if` condition is too complex to read imo, let's flatten it.
created_at: 2025-05-12T05:34:08Z
author_association: COLLABORATOR
diff_hunk: @@ -1393,3 +1397,23 @@ def set_model(model) -> None:
             pass
 
     raise mlflow.MlflowException(SET_MODEL_ERROR)
+
+
+def _update_active_model_id_based_on_mlflow_model(mlflow_model: Model):
+    """
+    Update the current active model ID based on the provided MLflow model.
+    Only set the active model ID if it is not already set by the user.
+    This is useful for setting the active model ID when loading a model
+    to ensure traces generated are associated with the loaded model.
+    """
+    if (
+        mlflow_model.model_id
+        and (amc := _get_active_model_context())
+        # only set the active model if the model is not set by the user
+        and not amc.set_by_user
+        and amc.model_id != mlflow_model.model_id
+    ):
+        _set_active_model_id(model_id=mlflow_model.model_id)
+        _logger.info(
+            "Use `mlflow.set_active_model` to set the active model to a different one if needed."
path: mlflow/models/model.py
user: login: B-Step62
body: Can we include this message in the logger call in `_set_active_model_id`? Showing this alone will be a bit strange.
created_at: 2025-05-12T05:35:29Z
author_association: COLLABORATOR
diff_hunk: @@ -3315,6 +3313,23 @@ def _set_active_model(*, name: Optional[str] = None, model_id: Optional[str] = N
     return ActiveModel(logged_model=logged_model, set_by_user=False)
 
 
+def _set_active_model_id(model_id: str, set_by_user: bool = False) -> None:
+    """
+    Set the active model ID in the active model context and update the
+    corresponding environment variable. This should only be used when
+    we know the LoggedModel with the model_id exists.
+    This function should be used inside MLflow to set the active model
+    while not blocking other code execution.
+    """
+    try:
+        _ACTIVE_MODEL_CONTEXT.set(ActiveModelContext(model_id, set_by_user))
+        _update_active_model_id_env_var(model_id)
+    except Exception as e:
+        _logger.warning(f"Failed to set active model ID to {model_id}, error: {e}")
+    else:
+        _logger.info(f"Active model set to model with ID: {model_id}")
path: mlflow/tracking/fluent.py
user: login: B-Step62
body: ```suggestion
        _logger.warning(f"Failed to set active model ID to {model_id}, error: {e}")
    else:
        _logger.info(f"Active model is set to the logged model with ID: {model_id}")
```
created_at: 2025-05-12T05:37:35Z
author_association: COLLABORATOR

PR comments:
user: login: github-actions[bot]
created_at: 2025-05-08T05:27:25Z
body: <!-- documentation preview -->


Documentation preview for 7b2d9ebf91bb2b8a9ca618af98de964f3d9ce2cd will be available when [this CircleCI job](https://circleci.com/gh/mlflow/mlflow/157999)
completes successfully. You may encounter a `{"message":"not found"}` error when reloading
a page. If so, add `/index.html` to the URL.

- [Top page](https://output.circle-artifacts.com/output/job/174f1adc-b10d-409b-9edf-d8c96067df96/artifacts/0/docs/build/latest/index.html)
- [Changed pages](https://output.circle-artifacts.com/output/job/174f1adc-b10d-409b-9edf-d8c96067df96/artifacts/0/docs/build/latest/diff.html) (⚠️ only MDX file changes are detected ⚠️)

<details>
<summary>More info</summary>

- Ignore this comment if this PR does not change the documentation.
- It takes a few minutes for the preview to be available.
- The preview is updated when a new commit is pushed to this PR.
- This comment was created by https://github.com/mlflow/mlflow/actions/runs/14965097859.

</details>

author_association: CONTRIBUTOR
user: login: serena-ruan
created_at: 2025-05-08T08:16:33Z
body: https://github.com/mlflow/mlflow/actions/runs/14899705134/job/41848931874?pr=15623#step:12:381 is irrelevant, I'll fix this on master branch
author_association: COLLABORATOR