Repository: mlflow/mlflow
Pull request #15684: Remove active model id env var set
Author: login: serena-ruan
Labels: name: rn/none
description: List under Small Changes in Changelogs.
Created: 2025-05-12T09:55:37Z | Merged: 2025-05-13T08:39:23Z

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

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

#### Install mlflow from this PR

```
# mlflow
pip install git+https://github.com/mlflow/mlflow.git@refs/pull/15684/merge
# mlflow-skinny
pip install git+https://github.com/mlflow/mlflow.git@refs/pull/15684/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/15684/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?
Follow-up for this discussion: https://github.com/mlflow/mlflow/pull/15623#discussion_r2083707397

To avoid any issues caused by reading the env var from different threads, in this PR I use MLFLOW_ACTIVE_MODEL_ID as the default value for active_model. And users should `NEVER` modify this env var directly to set to a different model, they should only use `set_active_model` API.
Update `mlflow.models.predict` to pass this env var directly.
<!-- Please fill in changes proposed in this PR. -->

### How is this PR tested?

- [ ] Existing unit/integration tests
- [x] New unit/integration tests
- [ ] 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/environment_variables.py b/mlflow/environment_variables.py
index 1da052208c7dd..f83332d8fe85e 100644
--- a/mlflow/environment_variables.py
+++ b/mlflow/environment_variables.py
@@ -762,8 +762,9 @@ def get(self):
 )
 
 
-#: Specified the ID of the LoggedModel to link traces to.
-#: This should only by used by MLflow internally or in standalone environments such
-#: as Databricks serving.
+#: Default active LoggedModel ID.
+#: This should only by used by MLflow internally, users should always use
+#: `set_active_model` to set the active LoggedModel, and should not set
+#: this environment variable directly.
 #: (default: ``None``)
-MLFLOW_ACTIVE_MODEL_ID = _EnvironmentVariable("MLFLOW_ACTIVE_MODEL_ID", str, None)
+_MLFLOW_ACTIVE_MODEL_ID = _EnvironmentVariable("_MLFLOW_ACTIVE_MODEL_ID", str, None)
diff --git a/mlflow/tracking/fluent.py b/mlflow/tracking/fluent.py
index 600bd7ddafc42..ff8ae430c97d8 100644
--- a/mlflow/tracking/fluent.py
+++ b/mlflow/tracking/fluent.py
@@ -32,7 +32,7 @@
 )
 from mlflow.entities.lifecycle_stage import LifecycleStage
 from mlflow.environment_variables import (
-    MLFLOW_ACTIVE_MODEL_ID,
+    _MLFLOW_ACTIVE_MODEL_ID,
     MLFLOW_ENABLE_ASYNC_LOGGING,
     MLFLOW_ENABLE_SYSTEM_METRICS_LOGGING,
     MLFLOW_EXPERIMENT_ID,
@@ -3231,7 +3231,10 @@ class ActiveModelContext:
     """
 
     def __init__(self, model_id: Optional[str] = None, set_by_user: bool = False):
-        self._model_id = model_id
+        # use _MLFLOW_ACTIVE_MODEL_ID as the default value for model_id
+        # so that for subprocesses the default _ACTIVE_MODEL_CONTEXT.model_id
+        # is still valid, and we don't need to read from env var.
+        self._model_id = model_id or _MLFLOW_ACTIVE_MODEL_ID.get()
         self._set_by_user = set_by_user
 
     def __repr__(self):
@@ -3239,7 +3242,7 @@ def __repr__(self):
 
     @property
     def model_id(self) -> Optional[str]:
-        return self._model_id or MLFLOW_ACTIVE_MODEL_ID.get()
+        return self._model_id
 
     @property
     def set_by_user(self) -> bool:
@@ -3260,7 +3263,6 @@ class ActiveModel(LoggedModel):
     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()
         _set_active_model_id(self.model_id, set_by_user)
 
     def __enter__(self):
@@ -3268,7 +3270,6 @@ def __enter__(self):
 
     def __exit__(self, exc_type, exc_val, exc_tb):
         _ACTIVE_MODEL_CONTEXT.set(self.last_active_model_context)
-        _update_active_model_id_env_var(self.last_active_model_id_env_var)
 
 
 # NB: This function is only intended to be used publicly by users to set the
@@ -3280,8 +3281,7 @@ def set_active_model(*, name: Optional[str] = None, model_id: Optional[str] = No
     Set the active model with the specified name or model ID, and it will be used for linking
     traces that are generated during the lifecycle of the model. The return value can be used as
     a context manager within a ``with`` block; otherwise, you must call ``set_active_model()``
-    to update active model. Note that this function also sets the environment variable
-    ``MLFLOW_ACTIVE_MODEL_ID`` to the model ID of the active model.
+    to update active model.
 
     Args:
         name: The name of the :py:class:`mlflow.entities.LoggedModel` to set as active.
@@ -3370,7 +3370,6 @@ def _set_active_model_id(model_id: str, set_by_user: bool = False) -> None:
     """
     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:
@@ -3392,9 +3391,9 @@ def _get_active_model_context() -> ActiveModelContext:
 
 def get_active_model_id() -> Optional[str]:
     """
-    Get the active model ID. If no active model is set with ``set_active_model()``, this will
-    try to get the model ID from the environment variable ``MLFLOW_ACTIVE_MODEL_ID``.
-    If neither is set, return None.
+    Get the active model ID. If no active model is set with ``set_active_model()``, the
+    default active model is set using model ID from the environment variable
+    ``_MLFLOW_ACTIVE_MODEL_ID``. If neither is set, return None.
 
     Returns:
         The active model ID if set, otherwise None.
@@ -3406,12 +3405,5 @@ def _reset_active_model_context() -> None:
     """
     Should be called only for testing purposes.
     """
+    _MLFLOW_ACTIVE_MODEL_ID.unset()
     _ACTIVE_MODEL_CONTEXT.set(ActiveModelContext())
-    _update_active_model_id_env_var(None)
-
-
-def _update_active_model_id_env_var(value):
-    if value is None:
-        MLFLOW_ACTIVE_MODEL_ID.unset()
-    else:
-        MLFLOW_ACTIVE_MODEL_ID.set(value)
diff --git a/mlflow/utils/environment.py b/mlflow/utils/environment.py
index 7cdf6cb0f6aeb..3e51eec5c3c5f 100644
--- a/mlflow/utils/environment.py
+++ b/mlflow/utils/environment.py
@@ -15,6 +15,7 @@
 from packaging.version import Version
 
 from mlflow.environment_variables import (
+    _MLFLOW_ACTIVE_MODEL_ID,
     _MLFLOW_TESTING,
     MLFLOW_EXPERIMENT_ID,
     MLFLOW_INPUT_EXAMPLE_INFERENCE_TIMEOUT,
@@ -23,7 +24,7 @@
 from mlflow.exceptions import MlflowException
 from mlflow.protos.databricks_pb2 import INVALID_PARAMETER_VALUE
 from mlflow.tracking import get_tracking_uri
-from mlflow.tracking.fluent import _get_experiment_id
+from mlflow.tracking.fluent import _get_experiment_id, get_active_model_id
 from mlflow.utils import PYTHON_VERSION
 from mlflow.utils.databricks_utils import (
     _get_databricks_serverless_env_vars,
@@ -869,6 +870,8 @@ def execute(
             command_env.update(_get_databricks_serverless_env_vars())
         if exp_id := _get_experiment_id():
             command_env[MLFLOW_EXPERIMENT_ID.name] = exp_id
+        if active_model_id := get_active_model_id():
+            command_env[_MLFLOW_ACTIVE_MODEL_ID.name] = active_model_id
         command_env.update(self._extra_env)
         if not isinstance(command, list):
             command = [command]
diff --git a/tests/conftest.py b/tests/conftest.py
index d6b06c99fcae5..021b9096ffc1c 100644
--- a/tests/conftest.py
+++ b/tests/conftest.py
@@ -120,6 +120,13 @@ def reset_mlflow_uri():
     if "DISABLE_RESET_MLFLOW_URI_FIXTURE" not in os.environ:
         os.environ.pop("MLFLOW_TRACKING_URI", None)
         os.environ.pop("MLFLOW_REGISTRY_URI", None)
+        try:
+            from mlflow.tracking import set_registry_uri
+
+            set_registry_uri(None)
+        except ImportError:
+            # tracing sdk does not have the registry module
+            pass
 
 
 @pytest.fixture(autouse=True)
diff --git a/tests/models/test_python_api.py b/tests/models/test_python_api.py
index c540f36aac12f..55a9587594edf 100644
--- a/tests/models/test_python_api.py
+++ b/tests/models/test_python_api.py
@@ -16,8 +16,11 @@
     _CONTENT_TYPE_JSON,
     _serialize_input_data,
 )
+from mlflow.tracing.constant import SpanAttributeKey
 from mlflow.utils.env_manager import CONDA, LOCAL, UV, VIRTUALENV
 
+from tests.tracing.helper import get_traces
+
 
 @pytest.mark.parametrize(
     ("input_data", "expected_data", "content_type"),
@@ -368,3 +371,27 @@ def predict(self, context, model_input: list[str]):
     traces = client.search_traces(experiment_ids=[exp_id])
     assert len(traces) == 1
     assert json.loads(traces[0].data.request)["model_input"] == ["a", "b", "c"]
+
+
+def test_predict_traces_link_to_active_model():
+    model = mlflow.set_active_model(name="test_model")
+
+    class TestModel(mlflow.pyfunc.PythonModel):
+        @mlflow.trace
+        def predict(self, context, model_input: list[str]):
+            return model_input
+
+    with mlflow.start_run():
+        model_info = mlflow.pyfunc.log_model(
+            name="model",
+            python_model=TestModel(),
+        )
+
+    mlflow.models.predict(
+        model_uri=model_info.model_uri,
+        input_data=["a", "b", "c"],
+        env_manager=VIRTUALENV,
+    )
+    traces = get_traces()
+    assert len(traces) == 1
+    assert traces[0].info.request_metadata[SpanAttributeKey.MODEL_ID] == model.model_id
diff --git a/tests/tracking/fluent/test_fluent.py b/tests/tracking/fluent/test_fluent.py
index 0ecda0ed8c096..e6aa7db0a66b5 100644
--- a/tests/tracking/fluent/test_fluent.py
+++ b/tests/tracking/fluent/test_fluent.py
@@ -35,7 +35,7 @@
 )
 from mlflow.entities.logged_model_status import LoggedModelStatus
 from mlflow.environment_variables import (
-    MLFLOW_ACTIVE_MODEL_ID,
+    _MLFLOW_ACTIVE_MODEL_ID,
     MLFLOW_EXPERIMENT_ID,
     MLFLOW_EXPERIMENT_NAME,
     MLFLOW_REGISTRY_URI,
@@ -50,6 +50,8 @@
 from mlflow.store.tracking import SEARCH_MAX_RESULTS_DEFAULT
 from mlflow.tracing.constant import SpanAttributeKey
 from mlflow.tracking.fluent import (
+    _ACTIVE_MODEL_CONTEXT,
+    ActiveModelContext,
     _get_experiment_id,
     _get_experiment_id_from_env,
     _reset_last_logged_model_id,
@@ -1983,16 +1985,13 @@ def test_set_active_model():
 
     set_active_model(name=model.name)
     assert mlflow.get_active_model_id() == model.model_id
-    assert MLFLOW_ACTIVE_MODEL_ID.get() == model.model_id
 
     set_active_model(model_id=model.model_id)
     assert mlflow.get_active_model_id() == model.model_id
-    assert MLFLOW_ACTIVE_MODEL_ID.get() == model.model_id
 
     model2 = mlflow.create_external_model(name="test_model")
     set_active_model(name="test_model")
     assert mlflow.get_active_model_id() == model2.model_id
-    assert MLFLOW_ACTIVE_MODEL_ID.get() == model2.model_id
 
     set_active_model(name="new_model")
     logged_model = mlflow.search_logged_models(
@@ -2000,19 +1999,14 @@ def test_set_active_model():
     )[0]
     assert logged_model.name == "new_model"
     assert mlflow.get_active_model_id() == logged_model.model_id
-    assert MLFLOW_ACTIVE_MODEL_ID.get() == logged_model.model_id
 
     with set_active_model(model_id=model.model_id) as active_model:
         assert active_model.model_id == model.model_id
         assert mlflow.get_active_model_id() == model.model_id
-        assert MLFLOW_ACTIVE_MODEL_ID.get() == model.model_id
         with set_active_model(name="new_model"):
             assert mlflow.get_active_model_id() == logged_model.model_id
-            assert MLFLOW_ACTIVE_MODEL_ID.get() == logged_model.model_id
         assert mlflow.get_active_model_id() == model.model_id
-        assert MLFLOW_ACTIVE_MODEL_ID.get() == model.model_id
     assert mlflow.get_active_model_id() == logged_model.model_id
-    assert MLFLOW_ACTIVE_MODEL_ID.get() == logged_model.model_id
 
 
 def test_set_active_model_error():
@@ -2028,19 +2022,15 @@ def test_set_active_model_error():
 
 
 def test_set_active_model_env_var(monkeypatch):
-    monkeypatch.setenv(MLFLOW_ACTIVE_MODEL_ID.name, "1234")
+    monkeypatch.setenv(_MLFLOW_ACTIVE_MODEL_ID.name, "1234")
+    # mimic the behavior when mlflow is imported
+    _ACTIVE_MODEL_CONTEXT.set(ActiveModelContext())
     assert mlflow.get_active_model_id() == "1234"
 
-    with set_active_model(name="abc") as model:
-        model_id = mlflow.get_active_model_id()
-        assert model_id == model.model_id
-        assert MLFLOW_ACTIVE_MODEL_ID.get() == model_id
-    assert mlflow.get_active_model_id() == "1234"
-    assert MLFLOW_ACTIVE_MODEL_ID.get() == "1234"
-
-    monkeypatch.delenv(MLFLOW_ACTIVE_MODEL_ID.name)
+    monkeypatch.delenv(_MLFLOW_ACTIVE_MODEL_ID.name)
+    _ACTIVE_MODEL_CONTEXT.set(ActiveModelContext())
     assert mlflow.get_active_model_id() is None
-    assert MLFLOW_ACTIVE_MODEL_ID.get() is None
+    assert _MLFLOW_ACTIVE_MODEL_ID.get() is None
 
 
 def test_set_active_model_link_traces():


Reviews:
user: login: B-Step62
state: COMMENTED
author_association: COLLABORATOR
submitted_at: 2025-05-13T01:27:30Z
user: login: B-Step62
state: COMMENTED
author_association: COLLABORATOR
submitted_at: 2025-05-13T01:28:36Z
user: login: B-Step62
body: LGTM once https://github.com/mlflow/mlflow/pull/15684/files#r2085764428 is addressed!
state: APPROVED
author_association: COLLABORATOR
submitted_at: 2025-05-13T01:29:28Z
user: login: serena-ruan
state: COMMENTED
author_association: COLLABORATOR
submitted_at: 2025-05-13T01:40:04Z

Review comments:
diff_hunk: @@ -762,8 +762,9 @@ def get(self):
 )
 
 
-#: Specified the ID of the LoggedModel to link traces to.
+#: Default active LoggedModel ID.
 #: This should only by used by MLflow internally or in standalone environments such
-#: as Databricks serving.
+#: as Databricks serving, users should always use `set_active_model` to set the
+#: active LoggedModel, and should not set this environment variable directly.
 #: (default: ``None``)
 MLFLOW_ACTIVE_MODEL_ID = _EnvironmentVariable("MLFLOW_ACTIVE_MODEL_ID", str, None)
path: mlflow/environment_variables.py
user: login: B-Step62
body: ```suggestion
_MLFLOW_ACTIVE_MODEL_ID = _EnvironmentVariable("MLFLOW_ACTIVE_MODEL_ID", str, None)
```

nit: since this is an internal env variable now
created_at: 2025-05-13T01:27:29Z
author_association: COLLABORATOR
diff_hunk: @@ -3382,12 +3381,5 @@ def _reset_active_model_context() -> None:
     """
     Should be called only for testing purposes.
     """
+    MLFLOW_ACTIVE_MODEL_ID.unset()
path: mlflow/tracking/fluent.py
user: login: B-Step62
body: non-blocking: Can we make this an auto-use fixture?
created_at: 2025-05-13T01:28:35Z
author_association: COLLABORATOR
diff_hunk: @@ -3382,12 +3381,5 @@ def _reset_active_model_context() -> None:
     """
     Should be called only for testing purposes.
     """
+    MLFLOW_ACTIVE_MODEL_ID.unset()
path: mlflow/tracking/fluent.py
user: login: serena-ruan
body: Yes https://github.com/mlflow/mlflow/blob/df5dc56cf9083cd348d4b85ad4f43b812e884598/tests/conftest.py#L322 is the autouse fixture
created_at: 2025-05-13T01:40:04Z
author_association: COLLABORATOR

PR comments:
user: login: github-actions[bot]
created_at: 2025-05-12T09:56:13Z
body: <!-- documentation preview -->


Documentation preview for 0938ddd545b7dc57ee3d32ff080c06b3f9963fd9 will be available when [this CircleCI job](https://circleci.com/gh/mlflow/mlflow/158706)
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/6a9edee8-72a4-443c-bad4-4ce5de5adf09/artifacts/0/docs/build/latest/index.html)
- [Changed pages](https://output.circle-artifacts.com/output/job/6a9edee8-72a4-443c-bad4-4ce5de5adf09/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/14991242396.

</details>

author_association: CONTRIBUTOR
user: login: serena-ruan
created_at: 2025-05-13T06:51:35Z
body: https://github.com/mlflow/mlflow/actions/runs/14989921516/job/42110837631?pr=15684#step:11:38 shows that with only testing `test_wheeled_model.py` file, the test `test_serving_wheeled_model` succeeds. Then it proves the test failure in https://github.com/mlflow/mlflow/actions/runs/14989109882/job/42108505350#step:11:173 was caused by running the test together with other tests
author_association: COLLABORATOR
user: login: serena-ruan
created_at: 2025-05-13T07:30:57Z
body: The original issue can be repro with `pytest tests/models/test_model.py::test_model_log_with_unity_catalog_url tests/models/test_wheeled_model.py::test_serving_wheeled_model --serve-wheel` in local
author_association: COLLABORATOR