Fix critical bugs in MLE-STAR data leakage checker

[0xnavarro]

Summary

Fixed two critical bugs in the data leakage checker that caused crashes when use_data_leakage_checker=True.

Bugs Fixed

Bug 1: Missing functools.partial() wrapper

• Error: TypeError: replace_leakage_code() missing 1 required positional argument: 'prefix'
• Cause: after_model_callback was not wrapped with functools.partial() to pass the prefix parameter
• Fix: Added functools.partial() wrapper consistent with other callbacks in the same file

Bug 2: Undefined variable in exception handling

• Error: UnboundLocalError: cannot access local variable 'leakage_status' where it is not associated with a value
• Cause: leakage_status was not defined in the except block of update_extract_status()
• Fix: Added leakage_status = "Unknown" in the exception handler

Impact

• Enables data leakage checker to work properly when enabled
• Makes MLE-STAR more robust for production use
• Fixes crashes that prevented users from using this important feature

Testing

• Verified with use_data_leakage_checker=True on a task using this feature.
• No regressions when feature is disabled (default behavior unchanged)

[0xnavarro]

This change is required because, without explicitly setting leakage_status = "Unknown" in the exception block and passing prefix to replace_leakage_code via functools.partial, the system crashes during runtime in my environment. The previous implementation does not properly handle cases where exceptions occur or where the callback is invoked without the necessary context, leading to unstable behavior.

These improvements ensure the code is robust and can handle errors gracefully, making it reliable in production. Without this fix, the code does not work as intended and is unusable for our workflow. Please prioritize reviewing and merging this change so we can prevent further crashes and issues.

[ivanmkc]

@raj-sinha can you review since this is for MLE-STAR?

[0xnavarro]

@raj-sinha Could you please review it? Thank you.

[0xnavarro]

It’s a bit surprising that critical bug fixes stay open without review for days.
Is this repo meant to be just an example and not maintained for production use?