{"id":16697,"date":"2026-07-03T00:00:00","date_gmt":"2026-07-03T07:00:00","guid":{"rendered":"https:\/\/devblogs.microsoft.com\/ise\/?p=16697"},"modified":"2026-07-03T05:50:41","modified_gmt":"2026-07-03T12:50:41","slug":"mlflow-autolog-pyspark-workers","status":"publish","type":"post","link":"https:\/\/devblogs.microsoft.com\/ise\/mlflow-autolog-pyspark-workers\/","title":{"rendered":"Enabling MLflow OpenAI Autolog on PySpark Workers"},"content":{"rendered":"

Context<\/h2>\n

In a recent engagement, the team built an LLM-based contract intelligence pipeline on Azure Databricks. The goal was to extract entitlements from a large corpus of inconsistently formatted service-contract PDFs \u2014 what is covered, on which equipment, and under which terms \u2014 so downstream systems can tell what is in scope and what is billable.<\/p>\n

Rules or template-based extraction were not a realistic option given the variability in layout and wording across contracts, which made an LLM a good fit: it can absorb that variability, reason about context across a document, and emit structured output in a single pass.<\/p>\n

To parallelize the extraction, the pipeline fans those per-document LLM calls out across Spark workers. Per-call visibility into token spend, latency, and prompt\/response quality becomes essential to keep cost and output quality from drifting unnoticed.<\/p>\n

The natural tool for that is mlflow.openai.autolog()<\/code>. The catch:<\/strong> getting it to reliably emit traces in this setup takes more than the docs suggest. If you are running instrumented LLM calls across PySpark workers, the patterns below are what finally made tracing work end-to-end.<\/p>\n

Problem<\/h2>\n

The pipeline distributes OpenAI API calls across Spark workers using mapInPandas<\/code> to process documents in parallel, with mlflow.openai.autolog()<\/code> enabled for tracing. Traces from the driver looked exactly as expected; the workers produced none<\/strong>. The Databricks autologging docs<\/a> flag that autologging must be called explicitly on workers, but following that guidance alone was not enough. Three separate issues had to be solved before traces appeared reliably.<\/p>\n

Solution<\/h2>\n

Issue 1: Workers Need Full MLflow Context<\/h3>\n

Calling mlflow.openai.autolog()<\/code> on workers is necessary but not sufficient. Workers also need the MLflow tracking URI and experiment name – neither is inherited from the driver.<\/p>\n

# Capture on the driver before mapInPandas\r\n_tracking_uri = mlflow.get_tracking_uri()\r\n_experiment_name = \"\/Shared\/my-experiment\"\r\n\r\ndef process_partition(batch_iter):\r\n    import mlflow\r\n\r\n    # All three are required on workers\r\n    mlflow.set_tracking_uri(_tracking_uri)\r\n    mlflow.set_experiment(_experiment_name)\r\n    mlflow.openai.autolog()\r\n\r\n    # ... LLM calls here<\/code><\/pre>\n
Without set_tracking_uri<\/code> and set_experiment<\/code>, autolog silently discards traces. The tracking URI defaults to an empty local path on workers, and without an experiment, traces have nowhere to go. On Unity Catalog-enabled workspaces, workers also need access to the experiment path \u2014 permissions don’t carry over from the driver. The failure mode is the same: a silent drop.<\/p>\n
Since Spark can reuse worker processes across partitions, a module-level flag avoids redundant setup:<\/p>\n
_worker_initialized = False\r\n\r\ndef process_partition(batch_iter):\r\n    global _worker_initialized\r\n    if not _worker_initialized:\r\n        os.environ[\"MLFLOW_ENABLE_ASYNC_TRACE_LOGGING\"] = \"false\"\r\n        mlflow.set_tracking_uri(_tracking_uri)\r\n        mlflow.set_experiment(_experiment_name)\r\n        mlflow.openai.autolog()\r\n        _worker_initialized = True\r\n    # ... LLM calls<\/code><\/pre>\n
Issue 2: Span Artifacts Lost Due to Async Export<\/h3>\n
After fixing the context propagation, traces appeared in the experiment list with metadata (inputs, outputs, token counts), but the “detailed trace view” in the MLflow UI was broken for most traces. Investigation revealed that 5 out of 6 trace span artifacts were missing from storage.<\/p>\n
The root cause: MLflow’s AsyncTraceExportQueue<\/code> writes span artifacts via a background daemon thread, relying on atexit<\/code> to flush on shutdown. When running as a Databricks job task<\/strong>, the Python process exits shortly after the pipeline completes. The daemon thread races against process termination, and atexit<\/code> hooks may not complete in time.<\/p>\n
The fix is to disable async logging on workers:<\/p>\n
import os\r\nos.environ[\"MLFLOW_ENABLE_ASYNC_TRACE_LOGGING\"] = \"false\"\r\nmlflow.openai.autolog()<\/code><\/pre>\n
This forces synchronous trace export. The overhead is approximately 100-500ms per trace, negligible compared to 5-20 second LLM call latency. Alternatively, this can be set as a cluster or job environment variable<\/strong> so it applies to all workers without code changes.<\/p>\n
Issue 3: No Parent-Child Trace Linking<\/h3>\n
Each chat.completions.create()<\/code> call produces an independent trace. The MLflow autolog implementation uses start_span_no_context()<\/code> in mlflow\/openai\/autolog.py<\/code>, which creates root spans without checking for a parent context. There is currently no mechanism to group worker traces under a single parent span.<\/p>\n
Processing 6 documents produces 6 disconnected traces. Correlation is possible by timestamp and experiment, but no trace hierarchy exists in the UI. A feature request<\/a> has been filed with the MLflow team. Realistic workarounds include:<\/p>\n