Repository: mlflow/mlflow
Pull request #15033: [Openai] create LoggedModel during autologging
Author: login: serena-ruan
Labels: name: rn/none
description: List under Small Changes in Changelogs.
Created: 2025-03-18T10:59:09Z | Merged: 2025-03-31T08:28:45Z

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

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

#### Install mlflow from this PR

```
pip install git+https://github.com/mlflow/mlflow.git@refs/pull/15033/merge
```

For Databricks, use the following command:

```
%sh
OPTIONS=$(if pip freeze | grep -q 'git+https://github.com/mlflow/mlflow.git'; then echo '--force-reinstall --no-deps'; fi)
pip install $OPTIONS git+https://github.com/mlflow/mlflow.git@refs/pull/15033/merge#subdirectory=skinny
```

</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?
1. remove log_models (and related) options in autolog
2. create LoggedModel if no existing models are linked to the model configs

<!-- Please fill in changes proposed in this PR. -->

### How is this PR tested?

- [ ] Existing unit/integration tests
- [ ] 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/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

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/openai/__init__.py b/mlflow/openai/__init__.py
index 566b1d67d0787..bf086f06d329a 100644
--- a/mlflow/openai/__init__.py
+++ b/mlflow/openai/__init__.py
@@ -68,6 +68,7 @@
 from mlflow.utils.autologging_utils import (
     autologging_integration,
     autologging_is_disabled,
+    disable_autologging_globally,
     safe_patch,
 )
 from mlflow.utils.databricks_utils import (
@@ -460,6 +461,7 @@ def save_model(
 
 @experimental
 @format_docstring(LOG_MODEL_PARAM_DOCS.format(package_name=FLAVOR_NAME))
+@disable_autologging_globally
 def log_model(
     model,
     task,
@@ -895,6 +897,7 @@ def autolog(
     disable_for_unsupported_versions=False,
     silent=False,
     log_traces=True,
+    create_logged_model=True,
 ):
     """
     Enables (or disables) and configures autologging from OpenAI to MLflow.
@@ -915,6 +918,13 @@ def autolog(
             autologging.
         log_traces: If ``True``, traces are logged for OpenAI models. If ``False``, no traces are
             collected during inference. Default to ``True``.
+        create_logged_model: If ``True``, automatically create a LoggedModel when the model
+            used for inference is not already logged. The created LoggedModel contains no model
+            artifacts, but it will be used to associate all traces generated by the model. If
+            ``False``, no LoggedModel is created and the traces will not be associated with any
+            model. Default to ``True``.
+            .. Note:: Experimental: This argument may change or be removed in a future release
+            without warning.
     """
     if Version(_get_openai_package_version()).major < 1:
         raise MlflowException("OpenAI autologging is only supported for openai >= 1.0.0")
@@ -927,6 +937,7 @@ def autolog(
         disable_for_unsupported_versions=disable_for_unsupported_versions,
         silent=silent,
         log_traces=log_traces,
+        create_logged_model=create_logged_model,
     )
 
     # Tracing OpenAI Agent SDK. This has to be done outside the function annotated with
@@ -960,6 +971,7 @@ def _autolog(
     disable_for_unsupported_versions=False,
     silent=False,
     log_traces=True,
+    create_logged_model=True,
 ):
     from openai.resources.chat.completions import AsyncCompletions as AsyncChatCompletions
     from openai.resources.chat.completions import Completions as ChatCompletions
diff --git a/mlflow/openai/_openai_autolog.py b/mlflow/openai/_openai_autolog.py
index 27ba666c2b3fb..6aa76844490bb 100644
--- a/mlflow/openai/_openai_autolog.py
+++ b/mlflow/openai/_openai_autolog.py
@@ -96,9 +96,12 @@ def patched_call(original, self, *args, **kwargs):
     task = mlflow.openai._get_task_name(self.__class__)
     model_identity = _generate_model_identity({"task": task, **kwargs})
     model_id = _MODEL_TRACKER.get(model_identity)
-    # TODO: create LoggedModel if model_id is None and set into _MODEL_TRACKER
-
     if config.log_traces:
+        if model_id is None and config.create_logged_model:
+            logged_model = mlflow.create_logged_model(name="openai")
+            _MODEL_TRACKER.set(model_identity, logged_model.model_id)
+            model_id = logged_model.model_id
+
         span = _start_span(mlflow_client, self, kwargs, run_id, model_id)
 
     # Execute the original function
@@ -125,6 +128,10 @@ async def async_patched_call(original, self, *args, **kwargs):
     model_identity = _generate_model_identity({"task": task, **kwargs})
     model_id = _MODEL_TRACKER.get(model_identity)
     if config.log_traces:
+        if model_id is None and config.create_logged_model:
+            logged_model = mlflow.create_logged_model(name="openai")
+            _MODEL_TRACKER.set(model_identity, logged_model.model_id)
+            model_id = logged_model.model_id
         span = _start_span(mlflow_client, self, kwargs, run_id, model_id)
 
     # Execute the original function
diff --git a/mlflow/utils/autologging_utils/config.py b/mlflow/utils/autologging_utils/config.py
index eee117dcb6302..abdf93e58c772 100644
--- a/mlflow/utils/autologging_utils/config.py
+++ b/mlflow/utils/autologging_utils/config.py
@@ -18,6 +18,7 @@ class AutoLoggingConfig:
     log_model_signatures: bool
     log_traces: bool
     extra_tags: Optional[dict] = None
+    create_logged_model: bool = True
 
     def should_log_optional_artifacts(self):
         """
@@ -36,4 +37,5 @@ def init(cls, flavor_name: str):
             log_model_signatures=config_dict.get("log_model_signatures", False),
             log_traces=config_dict.get("log_traces", True),
             extra_tags=config_dict.get("extra_tags", None),
+            create_logged_model=config_dict.get("create_logged_model", True),
         )
diff --git a/tests/openai/test_openai_autolog.py b/tests/openai/test_openai_autolog.py
index 4df5575c4046a..6409e6f3614a5 100644
--- a/tests/openai/test_openai_autolog.py
+++ b/tests/openai/test_openai_autolog.py
@@ -899,3 +899,57 @@ def test_generate_model_identity_same(model_dict1, model_dict2):
 )
 def test_generate_model_identity_different(model_dict1, model_dict2):
     assert _generate_model_identity(model_dict1) != _generate_model_identity(model_dict2)
+
+
+@pytest.mark.asyncio
+@pytest.mark.parametrize("create_logged_model", [True, False])
+async def test_autolog_create_logged_model_and_link_traces(client, create_logged_model):
+    mlflow.openai.autolog(create_logged_model=create_logged_model)
+
+    with mlflow.start_run() as run:
+        for _ in range(3):
+            response = client.chat.completions.create(
+                messages=[{"role": "user", "content": "test"}],
+                model="gpt-4o-mini",
+                temperature=0.1,
+            )
+            if client._is_async:
+                await response
+
+    traces = get_traces()
+    assert len(traces) == 3
+    logged_models = mlflow.search_logged_models(
+        filter_string=f"source_run_id='{run.info.run_id}'", output_format="list"
+    )
+    if create_logged_model:
+        assert len(logged_models) == 1
+        logged_model_id = logged_models[0].model_id
+        for i in range(3):
+            assert (
+                traces[i].data.spans[0].get_attribute(SpanAttributeKey.MODEL_ID) == logged_model_id
+            )
+    else:
+        assert len(logged_models) == 0
+        logged_model_id = None
+        for i in range(3):
+            assert SpanAttributeKey.MODEL_ID not in traces[i].data.spans[0].attributes
+
+    with mlflow.start_run():
+        model_info = mlflow.openai.log_model(
+            "gpt-4o-mini",
+            "chat.completions",
+            "model",
+            temperature=0.1,
+            messages=[{"role": "system", "content": "test"}],
+        )
+    assert model_info.model_id != logged_model_id
+    response = client.chat.completions.create(
+        messages=[{"role": "user", "content": "test"}],
+        model="gpt-4o-mini",
+        temperature=0.1,
+    )
+    if client._is_async:
+        await response
+    traces = get_traces()
+    assert len(traces) == 4
+    assert traces[0].data.spans[0].get_attribute(SpanAttributeKey.MODEL_ID) == model_info.model_id


Reviews:
user: login: harupy
state: COMMENTED
author_association: MEMBER
submitted_at: 2025-03-19T10:02:56Z
user: login: serena-ruan
state: COMMENTED
author_association: COLLABORATOR
submitted_at: 2025-03-19T11:06:42Z
user: login: serena-ruan
state: COMMENTED
author_association: COLLABORATOR
submitted_at: 2025-03-19T12:16:41Z
user: login: BenWilson2
state: COMMENTED
author_association: CONTRIBUTOR
submitted_at: 2025-03-19T23:29:54Z
user: login: harupy
body: LGTM!
state: APPROVED
author_association: MEMBER
submitted_at: 2025-03-31T07:46:29Z

Review comments:
diff_hunk: @@ -877,16 +881,10 @@ def load_model(model_uri, dst_path=None):
 
 @experimental
 def autolog(
-    log_input_examples=False,
-    log_model_signatures=False,
-    log_models=False,
-    log_datasets=False,
     disable=False,
     exclusive=False,
     disable_for_unsupported_versions=False,
     silent=False,
-    registered_model_name=None,
-    extra_tags=None,
path: mlflow/openai/__init__.py
user: login: harupy
body: This is a breaking change. Can we add a deprecation warning on master?
created_at: 2025-03-19T10:02:56Z
author_association: MEMBER
diff_hunk: @@ -877,16 +881,10 @@ def load_model(model_uri, dst_path=None):
 
 @experimental
 def autolog(
-    log_input_examples=False,
-    log_model_signatures=False,
-    log_models=False,
-    log_datasets=False,
     disable=False,
     exclusive=False,
     disable_for_unsupported_versions=False,
     silent=False,
-    registered_model_name=None,
-    extra_tags=None,
path: mlflow/openai/__init__.py
user: login: serena-ruan
body: Make sense! Let me file a PR
created_at: 2025-03-19T11:06:42Z
author_association: COLLABORATOR
diff_hunk: @@ -549,34 +550,37 @@ def log_model(
             model = mlflow.pyfunc.load_model(info.model_uri)
             print(model.predict(["hello", "world"]))
     """
-    mlflow_model = Model.log(
-        artifact_path=artifact_path,
-        name=name,
-        flavor=mlflow.openai,
-        registered_model_name=registered_model_name,
-        model=model,
-        task=task,
-        conda_env=conda_env,
-        code_paths=code_paths,
-        signature=signature,
-        input_example=input_example,
-        await_registration_for=await_registration_for,
-        pip_requirements=pip_requirements,
-        extra_pip_requirements=extra_pip_requirements,
-        metadata=metadata,
-        example_no_conversion=example_no_conversion,
-        prompts=prompts,
-        params=params,
-        tags=tags,
-        model_type=model_type,
-        step=step,
-        model_id=model_id,
-        **kwargs,
-    )
-    if mlflow_model.model_id:
+    with disable_autologging():
         model_identity = _generate_model_identity({"model": model, "task": task, **kwargs})
-        _MODEL_TRACKER.set(model_identity, mlflow_model.model_id)
-    return mlflow_model
+        if model_id is None and (existing_model_id := _MODEL_TRACKER.get(model_identity)):
+            model_id = existing_model_id
path: mlflow/openai/__init__.py
user: login: serena-ruan
body: This line (555) is added for case when a model is used for auto-tracing, we created a blank LoggedModel for linking traces, then later if users want to log the model we should use the existing model_id here. (test case -- test_autolog_create_logged_model_and_link_traces)
Below line (581) is added for case when a model is logged first, and then traces are generated, so we can link traces to the logged model. (test case -- test_autolog_link_traces_to_original_model_after_logging)

Possible issues with this setup:
1. If the model is loaded from a model_uri and we tracked it. If it's logged again, then we reuse the existing model_id
2. If the model is logged twice, then for the second logging time we reuse the first model_id

Probably it's fine for OpenAI because it's the same set of configs. But they'll become an issue for models like langchain.

Possible mitigation:
For case 1: Track if the model_id was set in load_model call, and don't reuse it in log_model
For case 2: Track if the model_id was set in log_model call, and don't reuse it in log_model
Basically we should track the source why the model_id is set, so we can determine when we should reuse it and when not to.

WDYT? @harupy 

created_at: 2025-03-19T12:16:40Z
author_association: COLLABORATOR
diff_hunk: @@ -549,34 +550,37 @@ def log_model(
             model = mlflow.pyfunc.load_model(info.model_uri)
             print(model.predict(["hello", "world"]))
     """
-    mlflow_model = Model.log(
-        artifact_path=artifact_path,
-        name=name,
-        flavor=mlflow.openai,
-        registered_model_name=registered_model_name,
-        model=model,
-        task=task,
-        conda_env=conda_env,
-        code_paths=code_paths,
-        signature=signature,
-        input_example=input_example,
-        await_registration_for=await_registration_for,
-        pip_requirements=pip_requirements,
-        extra_pip_requirements=extra_pip_requirements,
-        metadata=metadata,
-        example_no_conversion=example_no_conversion,
-        prompts=prompts,
-        params=params,
-        tags=tags,
-        model_type=model_type,
-        step=step,
-        model_id=model_id,
-        **kwargs,
-    )
-    if mlflow_model.model_id:
+    with disable_autologging():
         model_identity = _generate_model_identity({"model": model, "task": task, **kwargs})
-        _MODEL_TRACKER.set(model_identity, mlflow_model.model_id)
-    return mlflow_model
+        if model_id is None and (existing_model_id := _MODEL_TRACKER.get(model_identity)):
+            model_id = existing_model_id
path: mlflow/openai/__init__.py
user: login: BenWilson2
body: For case 1, we should, if the model is logged again, we should generate a new model_id and associate traces for inference of that object to the new logged model. 

For case 2, we should associate with last action (create a new model_id for the API call that was made and associate traces with new model_id) (last event state management)
created_at: 2025-03-19T23:29:54Z
author_association: CONTRIBUTOR

PR comments:
user: login: github-actions[bot]
created_at: 2025-03-18T10:59:42Z
body: <!-- documentation preview -->


Documentation preview for 62397f0f21cca86db7817367fd011e63916a49a5 will be available when [this CircleCI job](https://circleci.com/gh/mlflow/mlflow/151495)
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/cc22166a-a1e5-4422-803d-7b21e84ef376/artifacts/0/docs/build/latest/index.html)
- [Changed pages](https://output.circle-artifacts.com/output/job/cc22166a-a1e5-4422-803d-7b21e84ef376/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/14166479986.

</details>

author_association: CONTRIBUTOR
user: login: serena-ruan
created_at: 2025-03-20T07:30:26Z
body: Commit https://github.com/mlflow/mlflow/pull/15033/commits/ae05d3240c2d2cff271a6025e65faae70a115879 fixes issues pointed in [this comment](https://github.com/mlflow/mlflow/pull/15033#discussion_r2003192272) and added two test cases
author_association: COLLABORATOR
user: login: serena-ruan
created_at: 2025-03-25T06:07:45Z
body: Split this PR:
1. https://github.com/mlflow/mlflow/pull/15101 for removing log_models option
2. this PR for LoggedModel creation
author_association: COLLABORATOR