Fixes #1419: ensure correct run status on failure for MLflow, patch pre-execution NoneType crash

[dedyoc]

Fixes #1419

This PR resolves an issue where the MLflow plugin incorrectly reported failed Hamilton DAG as successful ("FINISHED"). Additionally, it patches a bug that caused the plugin to crash when no config or inputs were provided.

Changes

• Updated run_after_graph_execution in h_mlflow.py to explicitly pass the correct run status ("FINISHED" or "FAILED") to mlflow.end_run() as mlflow.end_run() defaulted to success.
• Added a fallback to the dictionary iteration in run_before_graph_execution to prevent an AttributeError when self.config or inputs evaluate to None.

How I tested this

1. Created a DAG using with MLFlowTracker adapter.
2. Intentionally raised a ValueError inside a node to simulate a crash.

import pandas as pd

def signups() -> pd.Series:
    return pd.Series([1, 10, 50, 100])

def spend() -> pd.Series:
    return pd.Series([10, 10, 20, 40])

def acquisition_cost(signups: pd.Series, spend: pd.Series) -> pd.Series:
    raise ValueError("intentional crash!!")
    return spend / signups

3. Verified via mlflow ui that the run correctly registered as FAILED.

5. Verified that running the DAG without external inputs no longer crashes the pre-execution hook.

Notes

• I detected the pre-execution crash while reproducing the issue with dummy DAGs that has no inputs.

Checklist

• PR has an informative and human-readable title (this will be pulled into the release notes)
• Changes are limited to a single goal (no scope creep)
• Code passed the pre-commit check & code is left cleaner/nicer than when first encountered.
• Any change in functionality is tested
• New functions are documented (with a description, list of inputs, and expected output)
• Placeholder code is flagged / future TODOs are captured in comments
• Project documentation has been updated if adding/changing functionality.

[skrawcz]

Great thanks. Is there a unit and/or integration test we could write for this?

[dedyoc]

I checked the tests/plugins directory and noticed there isn't an existing test folder for h_mlflow yet.
I'll go ahead and create one for it.

[skrawcz]

Test failed:
_______________ ERROR collecting tests/plugins/test_h_mlflow.py ________________
ImportError while importing test module '/home/runner/work/hamilton/hamilton/tests/plugins/test_h_mlflow.py'.
Hint: make sure your test modules/packages have valid Python names.
Traceback:
/opt/hostedtoolcache/Python/3.14.3/x64/lib/python3.14/importlib/init.py:88: in import_module
return _bootstrap._gcd_import(name[level:], package, level)
^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^
tests/plugins/test_h_mlflow.py:21: in
from hamilton.plugins.h_mlflow import MLFlowTracker
hamilton/plugins/h_mlflow.py:23: in
import mlflow
E ModuleNotFoundError: No module named 'mlflow'

[dedyoc]

I have added pytest.importorskip("mlflow") since mlflow isn't installed for Python 3.14

[dedyoc]

Hi @skrawcz ,

I've verified that the tests are now passing successfully on my fork's Action

Could you please take another look when you have a moment?

[skrawcz]

Cool. Mind fixing the ruff issues --- you should be able to run the pre-commit check for this:

- hook id: ruff-check
- exit code: 1
- files were modified by this hook

  B905 `zip()` without an explicit `strict=` parameter
     --> examples/LLM_Workflows/neo4j_graph_rag/embed_module.py:99:30
      |
   97 |             input=[item["text"] for item in batch],
   98 |         )
   99 |         for item, emb_obj in zip(batch, response.data):
      |                              ^^^^^^^^^^^^^^^^^^^^^^^^^
  100 |             results.append({"id": item["id"], "embedding": emb_obj.embedding})
      |
  help: Add explicit value for parameter `strict=`

  B007 Loop control variable `original` not used within loop body
     --> examples/LLM_Workflows/neo4j_graph_rag/retrieval_module.py:513:13
      |
  511 |     persons = resolved.get("persons", {})
  512 |     if persons:
  513 |         for original, canonical in persons.items():
      |             ^^^^^^^^
  514 |             lines.append(f'  Person: "{canonical}"')
      |
  help: Rename unused `original` to `_original`

  B007 Loop control variable `original` not used within loop body
     --> examples/LLM_Workflows/neo4j_graph_rag/retrieval_module.py:518:13
      |
  516 |     movies = resolved.get("movies", {})
  517 |     if movies:
  518 |         for original, canonical in movies.items():
      |             ^^^^^^^^
  519 |             lines.append(f'  Movie title: "{canonical}"')
      |
  help: Rename unused `original` to `_original`

  B007 Loop control variable `original` not used within loop body
     --> examples/LLM_Workflows/neo4j_graph_rag/retrieval_module.py:523:13
      |
  521 |     genres = resolved.get("genres", {})
  522 |     if genres:
  523 |         for original, canonical in genres.items():
      |             ^^^^^^^^
  524 |             lines.append(f'  Genre: "{canonical}"')
      |
  help: Rename unused `original` to `_original`

  B007 Loop control variable `original` not used within loop body
     --> examples/LLM_Workflows/neo4j_graph_rag/retrieval_module.py:528:13
      |
  526 |     companies = resolved.get("companies", {})
  527 |     if companies:
  528 |         for original, canonical in companies.items():
      |             ^^^^^^^^
  529 |             lines.append(f'  ProductionCompany: "{canonical}"')
      |
  help: Rename unused `original` to `_original`

  F841 Local variable `source` is assigned to but never used
     --> examples/LLM_Workflows/neo4j_graph_rag/retrieval_module.py:755:9
      |
  753 |     for row in merged_results:
  754 |         i += 1
  755 |         source = row.get("_source", "unknown")
      |         ^^^^^^
  756 |
  757 |         if "directors" in row:
      |
  help: Remove assignment to unused variable `source````

[dedyoc]

I thought they are a bit out of scope for this PR but sure, I will take care of em.

[skrawcz]

thanks for the fixes 🙏 appreciate it. looks good!

[skrawcz]

ugh @dedyoc sorry can you rebase and fix the conflicts?

[skrawcz]

maybe we didn't need to fix those issues because they're fixed on main? 🤔