Context

I'm running an LLM pipeline on Databricks that distributes OpenAI API calls across Spark workers via mapInPandas. Getting mlflow.openai.autolog() to work on workers required solving three undocumented issues. Sharing here since I couldn't find this covered anywhere.

Issue 1: Workers need explicit MLflow context

The docs say to call mlflow.autolog() on workers. For mlflow.openai.autolog(), that's insufficient. Workers also need the tracking URI and experiment - they don't inherit either from the driver.

# Capture on driver
_tracking_uri = mlflow.get_tracking_uri()
_experiment_name = "/Shared/my-experiment"

# Inside mapInPandas partition function
mlflow.set_tracking_uri(_tracking_uri)
mlflow.set_experiment(_experiment_name)
mlflow.openai.autolog()
# Without all three, autolog silently produces zero traces. No error, no warning.

Issue 2: Span artifacts lost due to async export

Even with the correct setup, most traces appeared in the experiment list, but the "detailed trace view" was broken. Investigation showed that AsyncTraceExportQueue uses a daemon thread with `atexit` for flushing. mapInPandas worker processes are terminated (not exited) when the partition completes, so atexit never fires.

Result: trace metadata (inputs, outputs, tokens) is written synchronously and persists. Span artifacts are written asynchronously and are lost for most traces. In my test with 6 documents, 5 out of 6 had missing artifacts.

Fix:

import os
os.environ["MLFLOW_ENABLE_ASYNC_TRACE_LOGGING"] = "false"

Overhead is ~100-500ms per trace, negligible next to LLM latency.

Issue 3: No parent-child trace linking

Each chat.completions.create() call produces an independent trace. Autolog uses start_span_no_context() in mlflow/openai/autolog.py (line 287), which always creates root spans. There's no mechanism to attach autolog spans to a user-provided parent, even though start_span_no_context already accepts a parent_span parameter.

Processing 6 documents = 6 disconnected traces. Correlation is only possible by timestamp.

Complete pattern

_tracking_uri = mlflow.get_tracking_uri()
_experiment_name = "/Shared/my-experiment"

def process_partition(batch_iter):
import os, mlflow
os.environ["MLFLOW_ENABLE_ASYNC_TRACE_LOGGING"] = "false"
mlflow.set_tracking_uri(_tracking_uri)
mlflow.set_experiment(_experiment_name)
mlflow.openai.autolog()

client = DatabricksOpenAI(workspace_client=WorkspaceClient(host=_host, token=_token))
for batch_df in batch_iter:
for _, row in batch_df.iterrows():
client.chat.completions.create(model="endpoint", messages=[...])
yield batch_df

input_df.mapInPandas(process_partition, schema=schema).collect()

Side discovery: Spark re-evaluation

Autolog also exposed that certain mapInPandas materialization patterns cause Spark to re-evaluate the lazy plan multiple times. I saw 24 traces where 6 were expected - each document processed 4x. The createOrReplaceTempView + spark.table().cache() pattern doesn't guarantee single evaluation. Worth checking if you're seeing unexpected LLM costs.

Environment

• MLflow 3.10.1
• Databricks serverless compute
• databricks-openai / DatabricksOpenAI client
• Python 3.12

Curious if others have hit these. Any alternative approaches?