- Mark as New
- Bookmark
- Subscribe
- Mute
- Subscribe to RSS Feed
- Permalink
- Report Inappropriate Content
10-27-2025 09:25 AM
Hi @shubham_lekhwar ,
This is a common context-passing issue when using Spark with MLflow.
The problem is that the nested=True flag in mlflow.start_run relies on an active run being present in the current process context. Your Parent_RUN is active on the driver node, but the build_tune_and_score_model function executes on worker nodes, which are separate processes and have no knowledge of the driver's active run. This causes the MLflow client on the worker to hang, waiting for a parent context that doesn't exist.
The solution is to manually pass the parent run's ID to the worker function and set the parent-child relationship using a tag.
You need to make two changes: one on the driver and one in your worker function.
On the Driver:
Get the parent_run_id before calling applyInPandas and use functools.partial to "bake" this ID into the function that Spark will distribute.
In the Worker Function (build_tune_and_score_model):
Modify the function signature to accept the new parent_run_id argument. Then, instead of nested=True, start a regular run and manually set the parent ID using mlflow.set_tag.