Separate explainers, predict and explorers in Units

[Felipedino]

This pull request introduces significant refactoring and improvements to the explainability and job execution pipeline in DashAI, focusing on modularizing job logic using units, improving SHAP model compatibility, and enhancing error handling and maintainability. The changes replace custom job logic with a unit-based approach, introduce a utility to wrap model predictors for SHAP, and register new units for explainability and prediction workflows.

Refactoring and Modularization of Job Logic:

• The job execution logic for exploration and prediction (explorer_job.py, predict_job.py) is refactored to use a unit-based workflow (ExecutionContext and units like LoadDatasetUnit, RunExplorationUnit, SaveExplorationUnit, etc.), replacing custom code for dataset loading, explorer/model instantiation, and result saving. This greatly improves maintainability and extensibility. [1] [2] [3] [4] [5] [6] [7]

• New units are registered in initial_components.py to support the new workflow, including units for building explainers, loading/saving models and datasets, running explorations, and generating explanations. [1] [2]

Explainability Improvements (SHAP Integration):

• Introduces as_shap_predictor in model_input.py, a utility to wrap model predictors as plain functions for SHAP, avoiding issues with read-only feature_names_in_ attributes in LightGBM/XGBoost wrappers. This ensures SHAP explainers work reliably across all supported models.

• Updates all SHAP-based explainers (contrastive_shap.py, kernel_shap.py, regression_kernel_shap.py) to use as_shap_predictor(self.model) instead of self.model.predict directly, ensuring compatibility and preventing attribute errors. [1] [2] [3] [4] [5] [6]

Error Handling and Cleanup:

• Improves error handling in jobs to ensure that failed runs correctly update the database row status to ERROR, preventing jobs from being stuck in an unfinished state. [1] [2] [3]

These changes collectively make the codebase more robust, modular, and maintainable, and ensure that explainability features work reliably with a wider range of models.

[Copilot]

🟡 Human review recommended

It performs a wide refactor across multiple core job execution paths (predict/explore/explain), so it warrants final human review despite the added regression and contract tests.

Pull request overview

This PR refactors DashAI’s prediction, exploration, and explainability job flows to be composed from reusable “units” wired through an ExecutionContext, and hardens SHAP integration by avoiding passing bound model methods to SHAP.

Changes:

• Refactors PredictJob, ExplorerJob, and ExplainerJob to orchestrate unit pipelines (load → run → save) via ExecutionContext.
• Adds new unit implementations for exploration, prediction, and explanation building/generation, and registers them in initial_components.py.
• Improves SHAP compatibility by introducing as_shap_predictor() and switching SHAP explainers to use it, plus adds targeted contract and end-to-end regression tests.

File summaries

File	Description
tests/back/units/test_unit_contracts.py	Adds a “reverse” contract test to ensure REQUIRES doesn’t include unread context keys.
tests/back/units/test_prediction_units.py	New unit-level contract tests for prediction units (model/dataset load, predict, save).
tests/back/units/test_exploration_units.py	New unit-level contract tests for exploration run/save units.
tests/back/units/test_explanation_units.py	New unit-level contract tests for explanation units (load/build/prepare/generate).
tests/back/explainers/test_shap_predictor_handover.py	New regression tests ensuring SHAP is not handed bound model.predict.
tests/back/api/test_units_api.py	Extends units API expectations to include newly registered units and their schemas.
tests/back/api/test_predict_job.py	New end-to-end regression net for PredictJob (status transitions, outputs, error cases).
tests/back/api/test_explorer_job.py	New end-to-end regression net for ExplorerJob (artifact paths, error cases).
tests/back/api/test_explainer_job.py	New end-to-end regression net for ExplainerJob (global/local artifacts, error cases).
DashAI/back/units/save_prediction_unit.py	New unit to persist prediction results to a unique results folder and publish results_path.
DashAI/back/units/save_exploration_unit.py	New unit to persist exploration artifacts under the notebook folder and publish exploration_path.
DashAI/back/units/run_exploration_unit.py	New unit to resolve/instantiate an explorer, prepare the dataset, run exploration, and publish outputs.
DashAI/back/units/prepare_explanation_data_unit.py	New unit to replay a run’s recorded split indexes and prepare data_x/data_y for explainers.
DashAI/back/units/predict_unit.py	New unit to run model prediction over selected input columns and decode via the task.
DashAI/back/units/load_training_dataset_unit.py	New unit to load the training dataset and publish both dataset and JSON-serializable types.
DashAI/back/units/load_trained_model_unit.py	New unit to restore a trained model from the Run row’s recorded artifact path.
DashAI/back/units/load_run_model_unit.py	New unit matching explainer-flow semantics: instantiate model with parameters then load artifact.
DashAI/back/units/generate_local_explanation_unit.py	New unit to select instances (split/rows/manual), fit explainer, generate plots, and publish paths.
DashAI/back/units/generate_global_explanation_unit.py	New unit to generate global explanation + plot artifacts and publish their paths.
DashAI/back/units/explanation_artifacts.py	New helper module for explainer instantiation and writing explanation artifacts (shared by units).
DashAI/back/units/build_manual_input_unit.py	New unit to build an in-memory dataset from user-typed manual input via the task.
DashAI/back/units/build_local_explainer_unit.py	New unit to instantiate a local explainer bound to the model in context.
DashAI/back/units/build_global_explainer_unit.py	New unit to instantiate a global explainer bound to the model in context.
DashAI/back/job/predict_job.py	Refactors prediction execution to a unit-based pipeline and improves error-path status marking.
DashAI/back/job/explorer_job.py	Refactors exploration execution to units and ensures failure paths mark the row as ERROR.
DashAI/back/job/explainer_job.py	Refactors explanation execution to units; centralizes artifact path publishing and row updates.
DashAI/back/initial_components.py	Registers new units so they are discoverable by the registry and exposed via units API.
DashAI/back/explainability/model_input.py	Adds as_shap_predictor() to provide SHAP a non-bound callable.
DashAI/back/explainability/explainers/regression_kernel_shap.py	Switches SHAP model handover to as_shap_predictor(self.model).
DashAI/back/explainability/explainers/kernel_shap.py	Switches SHAP model handover to as_shap_predictor(self.model).
DashAI/back/explainability/explainers/contrastive_shap.py	Switches SHAP model handover to as_shap_predictor(self.model).

Review details

• Files reviewed: 31/31 changed files
• Comments generated: 0
• Review effort level: Lite

We're testing this review assessment. Please use 👍 or 👎 to tell us if it's correct.