From 0677b4b88d465532e7b32420ada761b6b1d804b4 Mon Sep 17 00:00:00 2001 From: Felipe Date: Mon, 3 Aug 2026 19:11:02 -0400 Subject: [PATCH 1/2] Add contract tests for exploration and prediction units - Introduced `test_exploration_units.py` to validate the functionality of exploration units, ensuring proper handling of datasets, explorers, and saving results. - Added `test_prediction_units.py` to test prediction units, focusing on model loading, dataset handling, and prediction saving. - Enhanced `test_unit_contracts.py` with a new test to ensure units do not require keys they do not read, preventing potential composability issues. --- .../explainers/contrastive_shap.py | 7 +- .../explainability/explainers/kernel_shap.py | 7 +- .../explainers/regression_kernel_shap.py | 7 +- DashAI/back/explainability/model_input.py | 40 +- DashAI/back/initial_components.py | 30 + DashAI/back/job/explainer_job.py | 487 +++--------- DashAI/back/job/explorer_job.py | 133 ++-- DashAI/back/job/predict_job.py | 158 ++-- .../back/units/build_global_explainer_unit.py | 63 ++ .../back/units/build_local_explainer_unit.py | 58 ++ DashAI/back/units/build_manual_input_unit.py | 143 ++++ DashAI/back/units/explanation_artifacts.py | 128 ++++ .../units/generate_global_explanation_unit.py | 88 +++ .../units/generate_local_explanation_unit.py | 445 +++++++++++ DashAI/back/units/load_run_model_unit.py | 124 ++++ DashAI/back/units/load_trained_model_unit.py | 114 +++ .../back/units/load_training_dataset_unit.py | 91 +++ DashAI/back/units/predict_unit.py | 148 ++++ .../units/prepare_explanation_data_unit.py | 188 +++++ DashAI/back/units/run_exploration_unit.py | 187 +++++ DashAI/back/units/save_exploration_unit.py | 122 +++ DashAI/back/units/save_prediction_unit.py | 140 ++++ tests/back/api/test_explainer_job.py | 701 ++++++++++++++++++ tests/back/api/test_explorer_job.py | 277 +++++++ tests/back/api/test_predict_job.py | 513 +++++++++++++ tests/back/api/test_units_api.py | 52 ++ .../test_shap_predictor_handover.py | 120 +++ tests/back/units/test_explanation_units.py | 663 +++++++++++++++++ tests/back/units/test_exploration_units.py | 326 ++++++++ tests/back/units/test_prediction_units.py | 423 +++++++++++ tests/back/units/test_unit_contracts.py | 26 + 31 files changed, 5433 insertions(+), 576 deletions(-) create mode 100644 DashAI/back/units/build_global_explainer_unit.py create mode 100644 DashAI/back/units/build_local_explainer_unit.py create mode 100644 DashAI/back/units/build_manual_input_unit.py create mode 100644 DashAI/back/units/explanation_artifacts.py create mode 100644 DashAI/back/units/generate_global_explanation_unit.py create mode 100644 DashAI/back/units/generate_local_explanation_unit.py create mode 100644 DashAI/back/units/load_run_model_unit.py create mode 100644 DashAI/back/units/load_trained_model_unit.py create mode 100644 DashAI/back/units/load_training_dataset_unit.py create mode 100644 DashAI/back/units/predict_unit.py create mode 100644 DashAI/back/units/prepare_explanation_data_unit.py create mode 100644 DashAI/back/units/run_exploration_unit.py create mode 100644 DashAI/back/units/save_exploration_unit.py create mode 100644 DashAI/back/units/save_prediction_unit.py create mode 100644 tests/back/api/test_explainer_job.py create mode 100644 tests/back/api/test_explorer_job.py create mode 100644 tests/back/api/test_predict_job.py create mode 100644 tests/back/explainers/test_shap_predictor_handover.py create mode 100644 tests/back/units/test_explanation_units.py create mode 100644 tests/back/units/test_exploration_units.py create mode 100644 tests/back/units/test_prediction_units.py diff --git a/DashAI/back/explainability/explainers/contrastive_shap.py b/DashAI/back/explainability/explainers/contrastive_shap.py index 3c6e8243d..fa6f1ad87 100644 --- a/DashAI/back/explainability/explainers/contrastive_shap.py +++ b/DashAI/back/explainability/explainers/contrastive_shap.py @@ -238,7 +238,10 @@ def fit( """ import shap - from DashAI.back.explainability.model_input import prepare_model_input + from DashAI.back.explainability.model_input import ( + as_shap_predictor, + prepare_model_input, + ) x, y = background_dataset # SHAP calls the model with perturbed frames, which skip the model @@ -254,7 +257,7 @@ def fit( background_data = shap.sample(background_data, n_samples) self.explainer = shap.KernelExplainer( - model=self.model.predict, + model=as_shap_predictor(self.model), data=background_data, feature_names=feature_names, ) diff --git a/DashAI/back/explainability/explainers/kernel_shap.py b/DashAI/back/explainability/explainers/kernel_shap.py index 69b0eb738..668d4ac27 100644 --- a/DashAI/back/explainability/explainers/kernel_shap.py +++ b/DashAI/back/explainability/explainers/kernel_shap.py @@ -333,7 +333,10 @@ def fit( """ sample_background_data = bool(sample_background_data) - from DashAI.back.explainability.model_input import prepare_model_input + from DashAI.back.explainability.model_input import ( + as_shap_predictor, + prepare_model_input, + ) x, y = background_dataset @@ -365,7 +368,7 @@ def fit( import shap self.explainer = shap.KernelExplainer( - model=self.model.predict, + model=as_shap_predictor(self.model), data=background_data, feature_names=feature_names, link=self.link, diff --git a/DashAI/back/explainability/explainers/regression_kernel_shap.py b/DashAI/back/explainability/explainers/regression_kernel_shap.py index 63661beaf..2fcedb437 100644 --- a/DashAI/back/explainability/explainers/regression_kernel_shap.py +++ b/DashAI/back/explainability/explainers/regression_kernel_shap.py @@ -179,7 +179,10 @@ def fit( """ import shap - from DashAI.back.explainability.model_input import prepare_model_input + from DashAI.back.explainability.model_input import ( + as_shap_predictor, + prepare_model_input, + ) x, y = background_dataset # SHAP calls the model with perturbed frames, which skip the model @@ -195,7 +198,7 @@ def fit( background_data = shap.sample(background_data, n_samples) self.explainer = shap.KernelExplainer( - model=self.model.predict, + model=as_shap_predictor(self.model), data=background_data, feature_names=feature_names, ) diff --git a/DashAI/back/explainability/model_input.py b/DashAI/back/explainability/model_input.py index 0387d3801..95c6913bc 100644 --- a/DashAI/back/explainability/model_input.py +++ b/DashAI/back/explainability/model_input.py @@ -14,12 +14,50 @@ that both live in the same space. """ -from typing import TYPE_CHECKING, Any +from typing import TYPE_CHECKING, Any, Callable if TYPE_CHECKING: from DashAI.back.dataloaders.classes.dashai_dataset import DashAIDataset +def as_shap_predictor(model: Any) -> Callable: + """Wrap ``model.predict`` so SHAP receives a plain function, not a method. + + SHAP suppresses scikit-learn's "X does not have valid feature names" + warning by blanking ``feature_names_in_`` on whatever object the callable + is bound to (``shap.utils._legacy.convert_to_model``). It reaches that + object through ``__self__``, so it only does this when handed a *bound + method*, and it assumes the attribute is writable. + + That assumption does not hold for every model DashAI ships: the LightGBM + and XGBoost wrappers inherit ``feature_names_in_`` from their upstream + estimator as a read-only ``property``, so the assignment raises + ``AttributeError: property 'feature_names_in_' ... has no setter`` and the + explanation fails before it starts. + + Handing over a plain closure instead leaves ``__self__`` absent, so SHAP + skips that step entirely — a function is SHAP's primary documented + interface for ``model``. The only thing lost is the suppression of a + cosmetic scikit-learn warning. + + Parameters + ---------- + model : Any + The trained model being explained. + + Returns + ------- + Callable + A one-argument function calling ``model.predict`` positionally, the + same way SHAP calls it today. + """ + + def predict(x): + return model.predict(x) + + return predict + + def prepare_model_input(model: Any, dataset: "DashAIDataset") -> "DashAIDataset": """Apply the model's own input preprocessing to a dataset. diff --git a/DashAI/back/initial_components.py b/DashAI/back/initial_components.py index 1cd493fe7..0b0d9f0ea 100644 --- a/DashAI/back/initial_components.py +++ b/DashAI/back/initial_components.py @@ -351,13 +351,30 @@ # Units from DashAI.back.units.apply_converter_unit import ApplyConverterUnit +from DashAI.back.units.build_global_explainer_unit import BuildGlobalExplainerUnit +from DashAI.back.units.build_local_explainer_unit import BuildLocalExplainerUnit +from DashAI.back.units.build_manual_input_unit import BuildManualInputUnit from DashAI.back.units.build_model_unit import BuildModelUnit from DashAI.back.units.evaluate_model_unit import EvaluateModelUnit from DashAI.back.units.fit_model_unit import FitModelUnit +from DashAI.back.units.generate_global_explanation_unit import ( + GenerateGlobalExplanationUnit, +) +from DashAI.back.units.generate_local_explanation_unit import ( + GenerateLocalExplanationUnit, +) from DashAI.back.units.load_dataset_unit import LoadDatasetUnit +from DashAI.back.units.load_run_model_unit import LoadRunModelUnit +from DashAI.back.units.load_trained_model_unit import LoadTrainedModelUnit +from DashAI.back.units.load_training_dataset_unit import LoadTrainingDatasetUnit +from DashAI.back.units.predict_unit import PredictUnit from DashAI.back.units.prepare_and_split_unit import PrepareAndSplitUnit +from DashAI.back.units.prepare_explanation_data_unit import PrepareExplanationDataUnit +from DashAI.back.units.run_exploration_unit import RunExplorationUnit from DashAI.back.units.save_dataset_unit import SaveDatasetUnit +from DashAI.back.units.save_exploration_unit import SaveExplorationUnit from DashAI.back.units.save_model_unit import SaveModelUnit +from DashAI.back.units.save_prediction_unit import SavePredictionUnit logging.basicConfig(level=logging.DEBUG) log = logging.getLogger(__name__) @@ -532,6 +549,19 @@ def get_initial_components(): SaveModelUnit, ApplyConverterUnit, SaveDatasetUnit, + RunExplorationUnit, + SaveExplorationUnit, + LoadTrainedModelUnit, + LoadTrainingDatasetUnit, + BuildManualInputUnit, + PredictUnit, + SavePredictionUnit, + LoadRunModelUnit, + BuildGlobalExplainerUnit, + BuildLocalExplainerUnit, + PrepareExplanationDataUnit, + GenerateGlobalExplanationUnit, + GenerateLocalExplanationUnit, # Explainers ContrastiveShap, DiceCounterfactual, diff --git a/DashAI/back/job/explainer_job.py b/DashAI/back/job/explainer_job.py index 390b6e149..e4328c03c 100644 --- a/DashAI/back/job/explainer_job.py +++ b/DashAI/back/job/explainer_job.py @@ -1,5 +1,5 @@ import logging -from typing import TYPE_CHECKING, Any, Dict, Tuple +from typing import TYPE_CHECKING from kink import inject from sqlalchemy import exc @@ -11,14 +11,21 @@ ModelSession, Run, ) -from DashAI.back.explainability.global_explainer import BaseGlobalExplainer -from DashAI.back.explainability.local_explainer import BaseLocalExplainer from DashAI.back.job.base_job import BaseJob, JobError -from DashAI.back.models.base_model import BaseModel -from DashAI.back.tasks.base_task import BaseTask +from DashAI.back.units.build_global_explainer_unit import BuildGlobalExplainerUnit +from DashAI.back.units.build_local_explainer_unit import BuildLocalExplainerUnit +from DashAI.back.units.context import ExecutionContext +from DashAI.back.units.generate_global_explanation_unit import ( + GenerateGlobalExplanationUnit, +) +from DashAI.back.units.generate_local_explanation_unit import ( + GenerateLocalExplanationUnit, +) +from DashAI.back.units.load_dataset_unit import LoadDatasetUnit +from DashAI.back.units.load_run_model_unit import LoadRunModelUnit +from DashAI.back.units.prepare_explanation_data_unit import PrepareExplanationDataUnit if TYPE_CHECKING: - from datasets import DatasetDict from sqlalchemy.orm import sessionmaker logging.basicConfig(level=logging.DEBUG) @@ -116,263 +123,6 @@ def get_job_name(self) -> str: return f"{explainer_scope.capitalize()} Explanation ({explainer_id})" - @inject - def _generate_global_explanation( - self, - explainer: BaseGlobalExplainer, - dataset=Tuple["DatasetDict", "DatasetDict"], - ) -> None: - import os - import pickle - - from kink import di - - from DashAI.back.core.artifacts import normalize_artifacts - - explainer_id: int = self.kwargs["explainer_id"] - session_factory = di["session_factory"] - config = di["config"] - with session_factory() as db: - try: - explanation = explainer.explain(dataset) - plot = normalize_artifacts(explainer.plot(explanation)) - except Exception as e: - log.exception(e) - raise JobError( - "Failed to generate the explanation", - ) from e - try: - explanation_filename = f"global_explanation_{explainer_id}.pickle" - explanation_path = os.path.join( - config["EXPLANATIONS_PATH"], explanation_filename - ) - with open(explanation_path, "wb") as file: - pickle.dump(explanation, file) - - plot_filename = f"global_explanation_plot_{explainer_id}.pickle" - plot_path = os.path.join(config["EXPLANATIONS_PATH"], plot_filename) - with open(plot_path, "wb") as file: - pickle.dump(plot, file) - - except Exception as e: - log.exception(e) - raise JobError( - "Explanation file saving failed", - ) from e - try: - self.explainer_db.explanation_path = explanation_path - self.explainer_db.plot_path = plot_path - self.explainer_db.plot_overrides = None - db.commit() - except Exception as e: - log.exception(e) - raise JobError( - "Explanation path saving failed", - ) from e - - @inject - def _generate_local_explanation( - self, - explainer: BaseLocalExplainer, - dataset: Tuple["DatasetDict", "DatasetDict"], - splits: Dict[str, Any], - task: BaseTask, - same_dataset: bool, - ) -> None: - import json - import os - import pickle - - from datasets import DatasetDict - from kink import di - - from DashAI.back.core.artifacts import normalize_artifacts - from DashAI.back.dataloaders.classes.dashai_dataset import ( - load_dataset, - prepare_for_model_session, - save_dataset, - select_columns, - split_dataset, - ) - - explainer_id: int = self.kwargs["explainer_id"] - session_factory = di["session_factory"] - config = di["config"] - - explainer.fit(dataset, **self.explainer_db.fit_parameters) - instance_id = self.explainer_db.dataset_id - with session_factory() as db: - instance: Dataset = db.get(Dataset, instance_id) - if not instance: - raise JobError( - f"Dataset {instance_id} to be explained does not exist in DB." - ) - try: - loaded_instance = load_dataset(f"{instance.file_path}/dataset") - except Exception as e: - log.exception(e) - raise JobError( - f"Can not load instance from path {instance.file_path}", - ) from e - try: - # The data source is selected via scope["mode"]. It defaults to - # "split" so explainers created before this field existed keep - # their original split + percentage behavior. - mode = self.explainer_db.scope.get("mode", "split") - - if mode == "manual": - # Build the instances from values the user typed in by hand, - # reusing the same conversion the manual prediction flow uses. - # The rows (and any image files rewritten by the job endpoint) - # travel in the job kwargs, not in scope. - manual_input_data = self.kwargs.get("manual_input_data") or [] - if not manual_input_data: - raise JobError( - "No manual input data provided for the explanation" - ) - prepared_instance = task.process_manual_input( - manual_input_data, - f"{instance.file_path}/dataset", - ) - # Manual input carries only the input columns (no target), so - # keep just those instead of the standard input/output split. - # select_columns returns a DashAIDataset (same shape the - # split path produces), which is what the explainers expect. - X = prepared_instance.select_columns(self.input_columns) - else: - prepared_instance = task.prepare_for_task( - loaded_instance, - input_columns=self.input_columns, - output_columns=self.output_columns, - ) - - if mode == "rows": - # Explain a set of rows the user marked in the table. - # Indexes are over the whole dataset (the split does not - # apply in this mode). - row_indexes = self.explainer_db.scope.get("row_indexes") or [] - valid_indexes = [ - i - for i in row_indexes - if isinstance(i, int) - and 0 <= i < prepared_instance.num_rows - ] - if row_indexes and not valid_indexes: - raise JobError( - "No valid row indexes provided for the explanation" - ) - if valid_indexes: - prepared_instance = prepared_instance.select(valid_indexes) - else: - split = self.explainer_db.scope.get("split") - if split not in ["train", "test", "val", "all"]: - raise JobError(f"{split} is not a valid split") - - if split != "all": - if not same_dataset: - if isinstance(splits, str): - splits = json.loads(splits) - ( - prepared_dataset_dict, - splits, - ) = prepare_for_model_session( - dataset=prepared_instance, - splits=splits, - output_columns=self.output_columns, - ) - split_key = "validation" if split == "val" else split - prepared_instance = prepared_dataset_dict[split_key] - else: - prepared_instance = split_dataset( - prepared_instance, - train_indexes=splits["train_indexes"], - test_indexes=splits["test_indexes"], - val_indexes=splits["val_indexes"], - ) - split_key = "validation" if split == "val" else split - prepared_instance = prepared_instance[split_key] - - n_rows = max( - 1, - int( - prepared_instance.num_rows - * self.explainer_db.scope.get("percentage") - / 100 - ), - ) - # When "shuffle" is set the percentage is taken as a random - # sample of the split; otherwise it is the leading rows. - if self.explainer_db.scope.get("shuffle"): - prepared_instance = prepared_instance.shuffle(seed=42) - prepared_instance = prepared_instance.select(range(n_rows)) - - prepared_instance = DatasetDict({"train": prepared_instance}) - X, _ = select_columns( - prepared_instance, - self.input_columns, - self.output_columns, - ) - # Persist the original selected rows (the model input for each - # explained instance) as a DashAIDataset before the model's own - # preprocessing runs, so the frontend can read them back with - # the existing dataset endpoints. - input_source = X["train"] if isinstance(X, DatasetDict) else X - input_dataset_path = os.path.join( - config["EXPLANATIONS_PATH"], - f"local_explanation_input_{explainer_id}", - ) - save_dataset(input_source, os.path.join(input_dataset_path, "dataset")) - # The instances are handed over unprepared, the same way the - # prediction job calls model.predict: the model applies its own - # preprocessing. Explainers that need the model feature space - # ask for it with prepare_model_input. - - except Exception as e: - log.exception(e) - raise JobError( - f"""Can not prepare Dataset with {instance_id} - to generate the local explanation.""", - ) from e - try: - explanation = explainer.explain_instance(X) - plots = normalize_artifacts( - explainer.plot(explanation), create_grouped=True - ) - except Exception as e: - log.exception(e) - raise JobError( - "Failed to generate the explanation", - ) from e - try: - explanation_filename = f"local_explanation_{explainer_id}.pickle" - explanation_path = os.path.join( - config["EXPLANATIONS_PATH"], explanation_filename - ) - with open(explanation_path, "wb") as file: - pickle.dump(explanation, file) - - plots_filename = f"local_explanation_plots_{explainer_id}.pickle" - plots_path = os.path.join(config["EXPLANATIONS_PATH"], plots_filename) - with open(plots_path, "wb") as file: - pickle.dump(plots, file) - - except Exception as e: - log.exception(e) - raise JobError( - "Explanation file saving failed", - ) from e - try: - self.explainer_db.explanation_path = explanation_path - self.explainer_db.plots_path = plots_path - self.explainer_db.input_dataset_path = input_dataset_path - self.explainer_db.plot_overrides = None - db.commit() - except Exception as e: - log.exception(e) - raise JobError( - "Explanation path saving failed", - ) from e - @inject def run( self, @@ -381,17 +131,13 @@ def run( from kink import di - from DashAI.back.dataloaders.classes.dashai_dataset import ( - load_dataset, - select_columns, - split_dataset, - ) - - component_registry = di["component_registry"] session_factory = di["session_factory"] explainer_id: int = self.kwargs["explainer_id"] explainer_scope: str = self.kwargs["explainer_scope"] + + ctx = ExecutionContext() + with session_factory() as db: if explainer_scope == "global": self.explainer_db: GlobalExplainer = db.get( @@ -402,6 +148,14 @@ def run( else: raise JobError(f"{explainer_scope} is an invalid explainer type") + if not self.explainer_db: + # Checked before the try below, whose handler would otherwise + # be the thing that crashes: it calls set_status_as_error on + # this very row. + raise JobError( + f"Explainer with id {explainer_id} does not exist in DB." + ) + try: run: Run = db.get(Run, self.explainer_db.run_id) if not run: @@ -415,118 +169,60 @@ def run( ) dataset: Dataset = db.get(Dataset, model_session.dataset_id) if not dataset: + # The id named here is the one that was looked up. It used + # to interpolate the explainer's own dataset_id, a column + # global explainers do not even have. raise JobError( - f"Dataset {self.explainer_db.dataset_id} does not exist in DB." + f"Dataset {model_session.dataset_id} does not exist in DB." ) - self.input_columns = model_session.input_columns - self.output_columns = model_session.output_columns - - try: - run_model_class = component_registry[run.model_name]["class"] - except Exception as e: - log.exception(e) - raise JobError( - f"Unable to find Model with name {run.model_name} in registry.", - ) from e - try: - model: BaseModel = run_model_class(**run.parameters) - except Exception as e: - log.exception(e) - raise JobError("Unable to instantiate model") from e - try: - trained_model = model.load(run.run_path) - except Exception as e: - log.exception(e) - raise JobError( - f"Can not load model from path {run.run_path}" - ) from e - try: - explainer_class = component_registry[ - self.explainer_db.explainer_name - ]["class"] - except Exception as e: - log.exception(e) - raise JobError( - f"""Unable to find the {explainer_scope} explainer with name - {self.explainer_db.explainer_name} in registry.""", - ) from e + self.explainer_db.huey_id = self.kwargs.get("huey_id", None) + db.commit() - try: - explainer = explainer_class( - model=trained_model, **self.explainer_db.parameters - ) - except Exception as e: - log.exception(e) - raise JobError( - f"Unable to instantiate {explainer_scope} explainer.", - ) from e - try: - loaded_dataset: "DatasetDict" = load_dataset( - f"{dataset.file_path}/dataset" - ) - except Exception as e: - log.exception(e) - raise JobError( - f"Can not load dataset from path {dataset.file_path}", - ) from e - try: - task: BaseTask = component_registry[model_session.task_name][ - "class" - ]() - except Exception as e: - log.exception(e) - raise JobError( - ( - f"Unable to find Task with name {model_session.task_name} " - "in registry" - ), - ) from e - try: - splits = json.loads(run.split_indexes) - loaded_dataset = split_dataset( - loaded_dataset, - train_indexes=splits["train_indexes"], - test_indexes=splits["test_indexes"], - val_indexes=splits["val_indexes"], - ) + input_columns = model_session.input_columns + output_columns = model_session.output_columns + + LoadRunModelUnit(run_id=run.id)(ctx) + + # How the explainer configuration is stored on the row — the + # component name and its parameters live in separate columns — + # rather than part of the explanation itself. + explainer_config = { + "component": self.explainer_db.explainer_name, + "params": self.explainer_db.parameters, + } + build_explainer = ( + BuildGlobalExplainerUnit + if explainer_scope == "global" + else BuildLocalExplainerUnit + ) + build_explainer(explainer=explainer_config)(ctx) - prepared_dataset = task.prepare_for_task( - dataset=loaded_dataset, - input_columns=self.input_columns, - output_columns=self.output_columns, - ) - data = select_columns( - prepared_dataset, - self.input_columns, - self.output_columns, - ) + LoadDatasetUnit(dataset_id=model_session.dataset_id)(ctx) - data_x = split_dataset( - data[0], - train_indexes=splits["train_indexes"], - test_indexes=splits["test_indexes"], - val_indexes=splits["val_indexes"], - ) - data_y = split_dataset( - data[1], - train_indexes=splits["train_indexes"], - test_indexes=splits["test_indexes"], - val_indexes=splits["val_indexes"], - ) - # Inputs stay unprepared (see the note in the local - # explanation path); targets are encoded because explainers - # compare them against the model's class indexes. - for split_name in data_y: - data_y[split_name] = trained_model.prepare_output( - data_y[split_name], is_fit=False - ) + prepare = PrepareExplanationDataUnit( + task_name=model_session.task_name, + input_columns=input_columns, + output_columns=output_columns, + ) + # Resolving the task outside the wrapper below keeps a missing + # task reported as a registry problem rather than a generic + # "cannot prepare" message. + prepare.validate(ctx) + try: + # Unpacking the JSON column is an artifact of how the row + # stores it, but it stays inside this block because a + # malformed value has always been reported as a + # preparation failure. + ctx.put_ref("split_indexes", json.loads(run.split_indexes)) + prepare(ctx) except Exception as e: log.exception(e) raise JobError( f"""Can not prepare dataset {dataset.id} for the explanation""", ) from e + try: self.explainer_db.set_status_as_started() db.commit() @@ -535,27 +231,44 @@ def run( raise JobError( "Connection with the database failed", ) from e - if explainer_scope == "global": - self._generate_global_explanation( - explainer=explainer, dataset=(data_x, data_y) - ) - elif explainer_scope == "local": + if explainer_scope == "global": + GenerateGlobalExplanationUnit(explainer_id=explainer_id)(ctx) + paths = { + "explanation_path": ctx.require("explanation_path"), + "plot_path": ctx.require("plot_path"), + } + else: same_dataset = ( model_session.dataset_id == self.explainer_db.dataset_id ) - if not same_dataset: - splits = model_session.splits - - self._generate_local_explanation( - explainer=explainer, - dataset=(data_x, data_y), - splits=splits, - task=task, + GenerateLocalExplanationUnit( + explainer_id=explainer_id, + instance_dataset_id=self.explainer_db.dataset_id, + scope=self.explainer_db.scope, + fit_parameters=self.explainer_db.fit_parameters, + input_columns=input_columns, + output_columns=output_columns, + manual_input_data=self.kwargs.get("manual_input_data"), same_dataset=same_dataset, - ) - else: - raise JobError(f"{explainer_scope} is an invalid explainer type") + session_splits=(None if same_dataset else model_session.splits), + )(ctx) + paths = { + "explanation_path": ctx.require("explanation_path"), + "plots_path": ctx.require("plots_path"), + "input_dataset_path": ctx.require("input_dataset_path"), + } + + try: + for column, value in paths.items(): + setattr(self.explainer_db, column, value) + self.explainer_db.plot_overrides = None + db.commit() + except Exception as e: + log.exception(e) + raise JobError( + "Explanation path saving failed", + ) from e self.explainer_db.set_status_as_finished() db.commit() @@ -564,3 +277,5 @@ def run( self.explainer_db.set_status_as_error() db.commit() raise e + finally: + ctx.clear_cache() diff --git a/DashAI/back/job/explorer_job.py b/DashAI/back/job/explorer_job.py index 97b98841e..9e336e781 100644 --- a/DashAI/back/job/explorer_job.py +++ b/DashAI/back/job/explorer_job.py @@ -1,12 +1,15 @@ import logging -from typing import TYPE_CHECKING, Type +from typing import TYPE_CHECKING from kink import inject from sqlalchemy import exc from DashAI.back.dependencies.database.models import Explorer, Notebook -from DashAI.back.exploration.base_explorer import BaseExplorer from DashAI.back.job.base_job import BaseJob, JobError +from DashAI.back.units.context import ExecutionContext +from DashAI.back.units.load_dataset_unit import LoadDatasetUnit +from DashAI.back.units.run_exploration_unit import RunExplorationUnit +from DashAI.back.units.save_exploration_unit import SaveExplorationUnit if TYPE_CHECKING: from sqlalchemy.orm import sessionmaker @@ -85,17 +88,13 @@ def get_job_name(self) -> str: def run( self, ) -> None: - import os - import pathlib - from kink import di - from DashAI.back.dataloaders.classes.dashai_dataset import load_dataset - - component_registry = di["component_registry"] session_factory = di["session_factory"] - config = di["config"] explorer_id: int = self.kwargs["explorer_id"] + + ctx = ExecutionContext() + with session_factory() as db: # Load the explorer information try: @@ -123,103 +122,59 @@ def run( explorer_info.set_status_as_error() db.commit() raise JobError("Error while loading the notebook info.") from e - - # Load the dataset from the notebook - try: - loaded_dataset = load_dataset(f"{notebook_info.file_path}/dataset") - except Exception as e: - log.exception(e) - explorer_info.set_status_as_error() - db.commit() - raise JobError( - f"Can not load dataset from path {notebook_info.file_path}", - ) from e - - # obtain the explorer component from the registry - try: - explorer_component_class: Type[BaseExplorer] = component_registry[ - explorer_info.exploration_type - ]["class"] - except KeyError as e: - log.exception(e) + except Exception: + # A notebook that is simply not there used to escape the + # SQLAlchemyError handler above and leave the row STARTED + # forever, because nothing else marks it: the Huey error signal + # writes only to its own task_copy table, and + # _execute_base_job calls run() with no handler at all. + # Re-raised as-is so the "not found" message survives. explorer_info.set_status_as_error() db.commit() - raise JobError( - ( - f"Explorer {explorer_info.exploration_type} " - "not found in the registry." - ) - ) from e + raise - # Instance the explorer (the explorer handles its validation) + # Load the dataset from the notebook: its own working copy, which + # is what the converters rewrite. try: - explorer_instance = explorer_component_class(**explorer_info.parameters) - assert isinstance(explorer_instance, BaseExplorer) + LoadDatasetUnit(notebook_id=notebook_info.id)(ctx) except Exception as e: + # Anything the load unit raises has to leave the row in ERROR. + # Nothing else marks it: the Huey error signal only writes to + # its own task_copy table, never to the Explorer row, so + # without this the exploration would stay STARTED forever. + # Re-raised as-is; the unit reports the same + # "Can not load dataset from path ..." message the job used to + # build here. log.exception(e) explorer_info.set_status_as_error() db.commit() - raise JobError( - f"Error instancing the explorer {explorer_info.exploration_type}." - ) from e - - # prepare the dataset + raise + + # How the exploration configuration is stored on the row — the + # component name and its parameters live in separate columns — + # rather than part of the exploration itself. + explorer = { + "component": explorer_info.exploration_type, + "params": explorer_info.parameters, + } + + # Run the exploration. The unit reports the registry, instancing, + # preparation and launch errors with the same texts the job used + # to build here; re-raised as-is so they reach the user intact. try: - prepared_dataset = explorer_instance.prepare_dataset( - loaded_dataset, explorer_info.columns - ) + RunExplorationUnit(explorer_id=explorer_id, explorer=explorer)(ctx) except Exception as e: log.exception(e) explorer_info.set_status_as_error() db.commit() - raise JobError( - ( - "Error preparing the dataset for the exploration " - f"{explorer_info.exploration_type}." - ) - ) from e - - # Launch the exploration - try: - result = explorer_instance.launch_exploration( - prepared_dataset, explorer_info - ) - except Exception as e: - log.exception(e) - explorer_info.set_status_as_error() - db.commit() - raise JobError( - f"Error launching the exploration {explorer_info.exploration_type}." - ) from e + raise # Save the result try: - # save in the notebook folder - save_path = pathlib.Path( - os.path.join( - config["NOTEBOOK_PATH"], - (f"{notebook_info.id}"), - ) - ) - if not save_path.exists(): - save_path.mkdir(parents=True) - - save_path = explorer_instance.save_notebook( - notebook_info, explorer_info, save_path, result - ) - if isinstance(save_path, str): - save_path = pathlib.Path(save_path) - if not isinstance(save_path, pathlib.Path): - raise JobError( - ( - f"Error while saving the exploration" - f" {explorer_info.exploration_type}" - f", save path is not a pathlib.Path." - ) - ) + SaveExplorationUnit(explorer_id=explorer_id)(ctx) # Update the explorer info - explorer_info.exploration_path = save_path.as_posix() + explorer_info.exploration_path = ctx.require("exploration_path") explorer_info.set_status_as_finished() db.commit() except Exception as e: @@ -232,3 +187,5 @@ def run( f"{explorer_info.exploration_type}." ) ) from e + finally: + ctx.clear_cache() diff --git a/DashAI/back/job/predict_job.py b/DashAI/back/job/predict_job.py index 3bb3c95f9..213eb1b0a 100644 --- a/DashAI/back/job/predict_job.py +++ b/DashAI/back/job/predict_job.py @@ -12,6 +12,13 @@ from DashAI.back.job.base_job import BaseJob, JobError from DashAI.back.models.base_model import BaseModel from DashAI.back.tasks.base_task import BaseTask +from DashAI.back.units.build_manual_input_unit import BuildManualInputUnit +from DashAI.back.units.context import ExecutionContext +from DashAI.back.units.load_dataset_unit import LoadDatasetUnit +from DashAI.back.units.load_trained_model_unit import LoadTrainedModelUnit +from DashAI.back.units.load_training_dataset_unit import LoadTrainingDatasetUnit +from DashAI.back.units.predict_unit import PredictUnit +from DashAI.back.units.save_prediction_unit import SavePredictionUnit if TYPE_CHECKING: from sqlalchemy.orm import sessionmaker @@ -274,18 +281,9 @@ def get_job_name(self) -> str: def run( self, ) -> List[Any]: - import uuid - from pathlib import Path - - from DashAI.back.dataloaders.classes.dashai_dataset import ( - load_dataset, - save_dataset, - to_dashai_dataset, - ) - - component_registry = di["component_registry"] session_factory = di["session_factory"] - config = di["config"] + + ctx = ExecutionContext() prediction_id: int = self.kwargs["prediction_id"] manual_input_data: List[dict] = self.kwargs.get("manual_input_data", []) @@ -329,18 +327,17 @@ def run( detail="Model session not found", ) - # Retrieve Dataset if dataset_id is provided - dataset: Dataset = None + # The dataset the model was trained on. The one to predict on, + # when there is one, is resolved by the unit that loads it. dataset_trained: Dataset = db.get(Dataset, model_session.dataset_id) if not dataset_trained: + prediction.set_status_as_error() + db.commit() raise HTTPException( status_code=status.HTTP_404_NOT_FOUND, detail="Training dataset not found", ) - if dataset_id: - dataset: Dataset = db.get(Dataset, dataset_id) - if not model_session.input_columns: prediction.set_status_as_error() db.commit() @@ -364,74 +361,66 @@ def run( detail="Internal database error", ) from e - # Retrieve Task + # The prediction step owns the task, and resolving it here — before + # the model is even looked up — is what keeps a missing task + # reported as a task problem instead of being overtaken by + # whatever fails next. Same shape as ModelJob validating the fit + # unit ahead of the status change. + predict = PredictUnit( + task_name=model_session.task_name, + input_columns=model_session.input_columns, + output_columns=model_session.output_columns, + ) try: - task: BaseTask = component_registry[model_session.task_name]["class"]() + predict.validate(ctx) except Exception as e: prediction.set_status_as_error() db.commit() log.exception(e) - raise JobError( - f"Task {model_session.task_name} not found in the registry", - ) from e + raise - # Load Model + # Load Model. The unit reports both the registry miss and the + # unreadable artifact with the same texts the job used to build + # here; re-raised as-is so they reach the user intact. try: - model = component_registry[prediction.run.model_name]["class"] - except KeyError as e: + LoadTrainedModelUnit(run_id=prediction.run_id)(ctx) + except Exception as e: prediction.set_status_as_error() db.commit() log.exception(e) - raise JobError( - f"Model {prediction.run.model_name} not found in the registry" - ) from e + raise + # Load training dataset for type info and label processing. Loaded + # before the dataset to predict on, which is the order the error + # messages depend on when both are unreadable. try: - trained_model: BaseModel = model.load(prediction.run.run_path) + LoadTrainingDatasetUnit( + train_dataset_file_path=dataset_trained.file_path + )(ctx) except Exception as e: + # This branch used to skip set_status_as_error, unlike every + # one around it, leaving the prediction STARTED forever. + # Re-raised as-is so the unit's specific message survives. prediction.set_status_as_error() db.commit() log.exception(e) - raise JobError( - f"Failed to load model {prediction.run.model_name} " - f"from path {prediction.run.run_path}" - ) from e - - # Load Dataset and make Predictions - try: - # Load training dataset for type info and label processing - train_dataset: "DashAIDataset" = load_dataset( - str(Path(f"{dataset_trained.file_path}/dataset/")) - ) - except Exception as e: - log.exception(e) - raise JobError( - f"Cannot load training dataset from " - f"{dataset_trained.file_path}/dataset/" - ) from e + raise try: - # Load or create prediction dataset + # Load or create prediction dataset. Both branches publish the + # same "dataset" key, so the prediction below cannot tell a + # dataset read from disk from one typed in by hand. if dataset_id: - loaded_dataset: "DashAIDataset" = load_dataset( - str(Path(f"{dataset.file_path}/dataset/")) - ) + LoadDatasetUnit(dataset_id=dataset_id)(ctx) else: - dataset_trained_path = str( - Path(f"{dataset_trained.file_path}/dataset/") - ) - loaded_dataset = task.process_manual_input( - manual_input_data, dataset_trained_path - ) + BuildManualInputUnit( + task_name=model_session.task_name, + train_dataset_file_path=dataset_trained.file_path, + manual_input_data=manual_input_data, + )(ctx) self.report_progress(0.4, "Running prediction") - _, y_pred = _run_prediction_pipeline( - task=task, - trained_model=trained_model, - train_dataset=train_dataset, - loaded_dataset=loaded_dataset, - model_session=model_session, - ) + predict(ctx) except ValueError as ve: prediction.set_status_as_error() @@ -442,6 +431,11 @@ def run( detail=f"Invalid input data: {str(ve)}", ) from ve except TypeError as te: + # Marked as failed like its ValueError neighbour: this branch + # used to return 400 without touching the row, which left the + # prediction STARTED forever. + prediction.set_status_as_error() + db.commit() log.error(f"Type Error: {te}") raise HTTPException( status_code=400, @@ -459,41 +453,13 @@ def run( # Save Predictions to Arrow file try: - # Create unique folder for predictions - path = str(Path(f"{config['DATASETS_PATH']}/predictions/")) - folder_name = str(uuid.uuid4()) - full_path = Path(path) / folder_name - full_path.mkdir(parents=True, exist_ok=True) - - output_col = model_session.output_columns[0] - base_columns = [ - col for col in loaded_dataset.column_names if col != output_col - ] - output_dataset = loaded_dataset.select_columns(base_columns) - dataset_with_prediction = to_dashai_dataset( - output_dataset.add_column(output_col, y_pred) - ) - - # Filter schema from trained dataset - trained_schema = train_dataset.types - filtered_schema = { - key: value.to_string() - for key, value in trained_schema.items() - if key in model_session.input_columns + model_session.output_columns - } - - # Store num of rows, columns, and column names - dataset_with_prediction.compute_base_metadata() - - # Save dataset with predictions - save_dataset( - dataset_with_prediction, - str(full_path / "dataset"), - filtered_schema, - ) + SavePredictionUnit( + input_columns=model_session.input_columns, + output_columns=model_session.output_columns, + )(ctx) # Update Prediction record - prediction.results_path = str(full_path) + prediction.results_path = ctx.require("results_path") prediction.set_status_as_finished() db.commit() except Exception as e: @@ -503,3 +469,5 @@ def run( raise JobError( "Can not save prediction to json file", ) from e + finally: + ctx.clear_cache() diff --git a/DashAI/back/units/build_global_explainer_unit.py b/DashAI/back/units/build_global_explainer_unit.py new file mode 100644 index 000000000..2f6e6c717 --- /dev/null +++ b/DashAI/back/units/build_global_explainer_unit.py @@ -0,0 +1,63 @@ +"""Unit that instantiates a global explainer bound to a trained model.""" + +import logging + +from DashAI.back.core.schema_fields import ( + BaseSchema, + component_field, + schema_field, +) +from DashAI.back.core.utils import MultilingualString +from DashAI.back.units.base_unit import BaseUnit +from DashAI.back.units.context import ExecutionContext +from DashAI.back.units.explanation_artifacts import build_explainer + +log = logging.getLogger(__name__) + + +class BuildGlobalExplainerSchema(BaseSchema): + explainer: schema_field( + component_field(parent="BaseGlobalExplainer"), + placeholder={"component": "PermutationFeatureImportance", "params": {}}, + description=MultilingualString( + en="Explainer for the model as a whole, together with its own " + "configuration.", + es="Explicador para el modelo completo, junto con su propia configuración.", + pt="Explicador para o modelo como um todo, junto com a sua própria " + "configuração.", + de="Erklärer für das gesamte Modell samt eigener Konfiguration.", + zh="针对整个模型的解释器及其自身配置。", + ), + alias=MultilingualString( + en="Global explainer", + es="Explicador global", + pt="Explicador global", + de="Globaler Erklärer", + zh="全局解释器", + ), + ) # type: ignore + + +class BuildGlobalExplainerUnit(BaseUnit): + """Instantiate a global explainer over an already trained model. + + Sibling of ``BuildLocalExplainerUnit`` rather than one unit with a scope + flag, even though the building step itself is identical — the two share it + through a helper. Global and local explainers are separate registries with + separate base classes, and a component field carries a single ``parent`` + hint that the front reads straight off the property to list the candidates. + One field covering both scopes would have to be optional, and an optional + component field is emitted as an ``anyOf``, which buries the hint where the + front does not look and leaves the user with no picker at all. + """ + + SCHEMA = BuildGlobalExplainerSchema + + REQUIRES = ("model",) + PROVIDES = ("explainer",) + + def execute(self, ctx: ExecutionContext) -> None: + explainer = build_explainer( + "global", self.config["explainer"], ctx.require("model") + ) + ctx.put("explainer", explainer) diff --git a/DashAI/back/units/build_local_explainer_unit.py b/DashAI/back/units/build_local_explainer_unit.py new file mode 100644 index 000000000..b303da357 --- /dev/null +++ b/DashAI/back/units/build_local_explainer_unit.py @@ -0,0 +1,58 @@ +"""Unit that instantiates a local explainer bound to a trained model.""" + +import logging + +from DashAI.back.core.schema_fields import ( + BaseSchema, + component_field, + schema_field, +) +from DashAI.back.core.utils import MultilingualString +from DashAI.back.units.base_unit import BaseUnit +from DashAI.back.units.context import ExecutionContext +from DashAI.back.units.explanation_artifacts import build_explainer + +log = logging.getLogger(__name__) + + +class BuildLocalExplainerSchema(BaseSchema): + explainer: schema_field( + component_field(parent="BaseLocalExplainer"), + placeholder={"component": "KernelShap", "params": {}}, + description=MultilingualString( + en="Explainer for individual instances, together with its own " + "configuration.", + es="Explicador para instancias individuales, junto con su propia " + "configuración.", + pt="Explicador para instâncias individuais, junto com a sua própria " + "configuração.", + de="Erklärer für einzelne Instanzen samt eigener Konfiguration.", + zh="针对单个实例的解释器及其自身配置。", + ), + alias=MultilingualString( + en="Local explainer", + es="Explicador local", + pt="Explicador local", + de="Lokaler Erklärer", + zh="局部解释器", + ), + ) # type: ignore + + +class BuildLocalExplainerUnit(BaseUnit): + """Instantiate a local explainer over an already trained model. + + See ``BuildGlobalExplainerUnit`` for why the two scopes are two units even + though they share their whole implementation. + """ + + SCHEMA = BuildLocalExplainerSchema + + REQUIRES = ("model",) + PROVIDES = ("explainer",) + + def execute(self, ctx: ExecutionContext) -> None: + explainer = build_explainer( + "local", self.config["explainer"], ctx.require("model") + ) + ctx.put("explainer", explainer) diff --git a/DashAI/back/units/build_manual_input_unit.py b/DashAI/back/units/build_manual_input_unit.py new file mode 100644 index 000000000..6cf31f19c --- /dev/null +++ b/DashAI/back/units/build_manual_input_unit.py @@ -0,0 +1,143 @@ +"""Unit that turns hand-typed rows into a dataset to predict on.""" + +import logging +from typing import TYPE_CHECKING + +from DashAI.back.core.schema_fields import ( + BaseSchema, + schema_field, + string_field, +) +from DashAI.back.core.utils import MultilingualString +from DashAI.back.job.base_job import JobError +from DashAI.back.units.base_unit import BaseUnit +from DashAI.back.units.context import ExecutionContext + +if TYPE_CHECKING: + from DashAI.back.tasks.base_task import BaseTask + +log = logging.getLogger(__name__) + + +class BuildManualInputSchema(BaseSchema): + task_name: schema_field( + string_field(), + placeholder="TabularClassificationTask", + description=MultilingualString( + en="Name of the task that validates and types the hand-typed rows.", + es="Nombre de la tarea que valida y tipa las filas ingresadas a mano.", + pt="Nome da tarefa que valida e tipa as linhas introduzidas à mão.", + de="Name der Aufgabe, die die manuell eingegebenen Zeilen prüft " + "und typisiert.", + zh="用于校验并确定手工输入行类型的任务名称。", + ), + alias=MultilingualString( + en="Task", es="Tarea", pt="Tarefa", de="Aufgabe", zh="任务" + ), + ) # type: ignore + train_dataset_file_path: schema_field( + string_field(), + placeholder="", + description=MultilingualString( + en="Folder of the dataset the model was trained on. Its column " + "specification is what the typed values are validated against.", + es="Carpeta del conjunto de datos con el que se entrenó el modelo. " + "Su especificación de columnas es contra lo que se validan los " + "valores ingresados.", + pt="Pasta do conjunto de dados com que o modelo foi treinado. A sua " + "especificação de colunas é aquilo contra o que os valores " + "introduzidos são validados.", + de="Ordner des Datensatzes, mit dem das Modell trainiert wurde. " + "Gegen dessen Spaltenspezifikation werden die eingegebenen Werte " + "geprüft.", + zh="模型训练所用数据集的文件夹。输入值将依据其列规格进行校验。", + ), + alias=MultilingualString( + en="Training dataset folder", + es="Carpeta del conjunto de entrenamiento", + pt="Pasta do conjunto de treino", + de="Ordner des Trainingsdatensatzes", + zh="训练数据集文件夹", + ), + ) # type: ignore + manual_input_data: schema_field( + list, + placeholder=[], + description=MultilingualString( + en="Rows to predict on, each a mapping from input column name to " + "value. Uploaded files arrive as a path reference instead of bytes.", + es="Filas a predecir, cada una un mapeo de nombre de columna de " + "entrada a valor. Los archivos subidos llegan como una referencia " + "a una ruta en vez de bytes.", + pt="Linhas a prever, cada uma um mapeamento de nome de coluna de " + "entrada para valor. Os ficheiros carregados chegam como uma " + "referência a um caminho em vez de bytes.", + de="Zu prognostizierende Zeilen, je eine Zuordnung von " + "Eingabespaltenname zu Wert. Hochgeladene Dateien kommen als " + "Pfadverweis statt als Bytes an.", + zh="要预测的行，每行是输入列名到值的映射。上传的文件以路径引用而非字节形式传入。", + ), + alias=MultilingualString( + en="Manual input", + es="Entrada manual", + pt="Entrada manual", + de="Manuelle Eingabe", + zh="手动输入", + ), + ) # type: ignore + + +class BuildManualInputUnit(BaseUnit): + """Build the dataset to predict on from values the user typed in. + + The counterpart of loading one from disk: it produces the same ``dataset`` + key, so whatever runs next cannot tell the two apart. That is what lets the + prediction step be written once for both sources. + + The task does the work — it is the task that knows the expected column + types and how to turn an uploaded file into a cell — so this unit only + resolves it and hands over the rows. + """ + + SCHEMA = BuildManualInputSchema + + PROVIDES = ("dataset",) + + def __init__(self, **config) -> None: + super().__init__(**config) + self._task = None + + def _resolve_task(self) -> "BaseTask": + """Instantiate the task from the registry, memoized on this unit.""" + if self._task is not None: + return self._task + + from kink import di + + component_registry = di["component_registry"] + task_name = self.config["task_name"] + + try: + task: "BaseTask" = component_registry[task_name]["class"]() + except Exception as e: + log.exception(e) + raise JobError(f"Task {task_name} not found in the registry") from e + + self._task = task + return task + + def validate(self, ctx: ExecutionContext) -> None: + """Resolve the task before anything observable happens.""" + self._resolve_task() + + def execute(self, ctx: ExecutionContext) -> None: + from pathlib import Path + + task = self._resolve_task() + + train_dataset_path = str( + Path(f"{self.config['train_dataset_file_path']}/dataset/") + ) + rows = self.config["manual_input_data"] + + ctx.put("dataset", task.process_manual_input(rows, train_dataset_path)) diff --git a/DashAI/back/units/explanation_artifacts.py b/DashAI/back/units/explanation_artifacts.py new file mode 100644 index 000000000..32608b742 --- /dev/null +++ b/DashAI/back/units/explanation_artifacts.py @@ -0,0 +1,128 @@ +"""Shared helpers for the two explanation-generating units. + +Not a unit: no configuration, no context, nothing to declare. It lives here +rather than in ``job/`` because importing from a job into a unit would invert +the dependency. +""" + +import logging +from typing import Any, Tuple + +from DashAI.back.job.base_job import JobError + +log = logging.getLogger(__name__) + + +def build_explainer(scope: str, selected: dict, trained_model: Any) -> Any: + """Resolve an explainer component and bind it to a trained model. + + Shared by the two scope-specific build units: building is identical either + way, only the registry the component comes from differs. + + Takes and returns plain values instead of touching the context. That keeps + every context write inside the unit itself, where the contract audit can + see it — a ``ctx.put`` hidden in a helper is invisible to the static check + and would let a broken ``PROVIDES`` through. + + Parameters + ---------- + scope : str + ``"global"`` or ``"local"``. Only decorates the error messages, which + are user-visible and worded per scope. + selected : dict + The ``{"component": ..., "params": ...}`` value of the unit's field. + trained_model : Any + The model the explainer explains. + + Returns + ------- + Any + The instantiated explainer. + + Raises + ------ + JobError + If the component is not registered or cannot be instantiated. + """ + from kink import di + + component_registry = di["component_registry"] + + explainer_name = selected["component"] + + try: + explainer_class = component_registry[explainer_name]["class"] + except Exception as e: + log.exception(e) + raise JobError( + f"""Unable to find the {scope} explainer with name + {explainer_name} in registry.""", + ) from e + + try: + return explainer_class(model=trained_model, **(selected.get("params") or {})) + except Exception as e: + log.exception(e) + raise JobError( + f"Unable to instantiate {scope} explainer.", + ) from e + + +def dump_explanation(explanation: Any, plots: Any, prefix: str, key: int) -> Tuple: + """Pickle an explanation and its plots under the explanations directory. + + Both files are named after the explanation they belong to, so re-running + one overwrites its own artifacts and never another's. + + Parameters + ---------- + explanation : Any + Whatever the explainer's ``explain``/``explain_instance`` returned. + plots : Any + The normalized artifacts produced from that explanation. + prefix : str + ``"global"`` or ``"local"``: the two scopes keep separate file names + because their ids come from separate tables and would otherwise clash. + key : int + Identifier of the explanation row. + + Returns + ------- + Tuple[str, str] + The explanation path and the plot path. + + Raises + ------ + JobError + If either file cannot be written. + """ + import os + import pickle + + from kink import di + + config = di["config"] + + plots_name = "plot" if prefix == "global" else "plots" + + try: + explanation_path = os.path.join( + config["EXPLANATIONS_PATH"], f"{prefix}_explanation_{key}.pickle" + ) + with open(explanation_path, "wb") as file: + pickle.dump(explanation, file) + + plot_path = os.path.join( + config["EXPLANATIONS_PATH"], + f"{prefix}_explanation_{plots_name}_{key}.pickle", + ) + with open(plot_path, "wb") as file: + pickle.dump(plots, file) + + except Exception as e: + log.exception(e) + raise JobError( + "Explanation file saving failed", + ) from e + + return explanation_path, plot_path diff --git a/DashAI/back/units/generate_global_explanation_unit.py b/DashAI/back/units/generate_global_explanation_unit.py new file mode 100644 index 000000000..677929b19 --- /dev/null +++ b/DashAI/back/units/generate_global_explanation_unit.py @@ -0,0 +1,88 @@ +"""Unit that explains a model as a whole and stores the result.""" + +import logging + +from DashAI.back.core.schema_fields import ( + BaseSchema, + int_field, + schema_field, +) +from DashAI.back.core.utils import MultilingualString +from DashAI.back.job.base_job import JobError +from DashAI.back.units.base_unit import BaseUnit +from DashAI.back.units.context import ExecutionContext +from DashAI.back.units.explanation_artifacts import dump_explanation + +log = logging.getLogger(__name__) + + +class GenerateGlobalExplanationSchema(BaseSchema): + explainer_id: schema_field( + int_field(gt=0), + placeholder=1, + description=MultilingualString( + en="Identifier of the global explanation being produced. It names " + "the files on disk, so a re-run overwrites its own artifacts and " + "never another explanation's.", + es="Identificador de la explicación global que se produce. Da " + "nombre a los archivos en disco, de modo que volver a ejecutarla " + "sobrescribe sus propios artefactos y nunca los de otra.", + pt="Identificador da explicação global a ser produzida. Dá nome aos " + "ficheiros em disco, pelo que uma nova execução substitui os seus " + "próprios artefactos e nunca os de outra.", + de="Kennung der erzeugten globalen Erklärung. Sie benennt die " + "Dateien auf der Festplatte, sodass ein erneuter Lauf nur die " + "eigenen Artefakte überschreibt.", + zh="所生成全局解释的标识符。它命名磁盘上的文件，因此重新运行只会覆盖自身产物。", + ), + alias=MultilingualString( + en="Explanation", + es="Explicación", + pt="Explicação", + de="Erklärung", + zh="解释", + ), + ) # type: ignore + + +class GenerateGlobalExplanationUnit(BaseUnit): + """Explain the model over the whole dataset and pickle the result. + + Sibling of ``GenerateLocalExplanationUnit`` rather than one unit with a + branch, for three reasons that all point the same way: the two produce + different outputs (a single plot here, a set of plots plus the explained + instances there), which a single ``PROVIDES`` could not describe since it + is checked unconditionally; the local path has steps this one does not + (fitting, selecting instances); and their configurations point at two + different component registries. + + The unit never touches the explanation row: it publishes where it wrote, + and the job owns the columns. + """ + + SCHEMA = GenerateGlobalExplanationSchema + + REQUIRES = ("explainer", "data_x", "data_y") + PROVIDES = ("explanation_path", "plot_path") + + def execute(self, ctx: ExecutionContext) -> None: + from DashAI.back.core.artifacts import normalize_artifacts + + explainer = ctx.require("explainer") + dataset = (ctx.require("data_x"), ctx.require("data_y")) + + try: + explanation = explainer.explain(dataset) + plot = normalize_artifacts(explainer.plot(explanation)) + except Exception as e: + log.exception(e) + raise JobError( + "Failed to generate the explanation", + ) from e + + explanation_path, plot_path = dump_explanation( + explanation, plot, "global", self.config["explainer_id"] + ) + + ctx.put_ref("explanation_path", explanation_path) + ctx.put_ref("plot_path", plot_path) diff --git a/DashAI/back/units/generate_local_explanation_unit.py b/DashAI/back/units/generate_local_explanation_unit.py new file mode 100644 index 000000000..7f12803ab --- /dev/null +++ b/DashAI/back/units/generate_local_explanation_unit.py @@ -0,0 +1,445 @@ +"""Unit that explains individual instances and stores the result.""" + +import logging + +from DashAI.back.core.schema_fields import ( + BaseSchema, + bool_field, + int_field, + list_field, + none_type, + schema_field, + string_field, +) +from DashAI.back.core.utils import MultilingualString +from DashAI.back.dependencies.database.models import Dataset +from DashAI.back.job.base_job import JobError +from DashAI.back.units.base_unit import BaseUnit +from DashAI.back.units.context import ExecutionContext +from DashAI.back.units.explanation_artifacts import dump_explanation + +log = logging.getLogger(__name__) + + +def _columns_field(alias: MultilingualString, description: MultilingualString): + return schema_field( + list_field(string_field(), min_items=1), + placeholder=[], + description=description, + alias=alias, + ) + + +class GenerateLocalExplanationSchema(BaseSchema): + explainer_id: schema_field( + int_field(gt=0), + placeholder=1, + description=MultilingualString( + en="Identifier of the local explanation being produced. It names " + "the files on disk, so a re-run overwrites its own artifacts.", + es="Identificador de la explicación local que se produce. Da nombre " + "a los archivos en disco, de modo que volver a ejecutarla " + "sobrescribe sus propios artefactos.", + pt="Identificador da explicação local a ser produzida. Dá nome aos " + "ficheiros em disco, pelo que uma nova execução substitui os seus " + "próprios artefactos.", + de="Kennung der erzeugten lokalen Erklärung. Sie benennt die " + "Dateien auf der Festplatte, sodass ein erneuter Lauf nur die " + "eigenen Artefakte überschreibt.", + zh="所生成局部解释的标识符。它命名磁盘上的文件，因此重新运行只会覆盖自身产物。", + ), + alias=MultilingualString( + en="Explanation", + es="Explicación", + pt="Explicação", + de="Erklärung", + zh="解释", + ), + ) # type: ignore + instance_dataset_id: schema_field( + int_field(gt=0), + placeholder=1, + description=MultilingualString( + en="Identifier of the dataset the explained instances come from. " + "It may differ from the one the model was trained on.", + es="Identificador del conjunto de datos del que provienen las " + "instancias explicadas. Puede diferir de aquel con el que se " + "entrenó el modelo.", + pt="Identificador do conjunto de dados de onde vêm as instâncias " + "explicadas. Pode diferir daquele com que o modelo foi treinado.", + de="Kennung des Datensatzes, aus dem die erklärten Instanzen " + "stammen. Er kann sich von dem des Trainings unterscheiden.", + zh="被解释实例所属数据集的标识符。它可能与模型训练所用的数据集不同。", + ), + alias=MultilingualString( + en="Instance dataset", + es="Conjunto de instancias", + pt="Conjunto de instâncias", + de="Instanzdatensatz", + zh="实例数据集", + ), + ) # type: ignore + scope: schema_field( + dict, + placeholder={"mode": "split", "split": "test", "percentage": 20}, + description=MultilingualString( + en="Which instances to explain. A 'mode' of 'split' takes a share " + "of one data split, 'rows' takes the row indexes the user marked, " + "and 'manual' takes hand-typed values. Defaults to 'split'.", + es="Qué instancias explicar. Un 'mode' de 'split' toma una " + "proporción de una partición, 'rows' toma los índices de fila que " + "marcó el usuario, y 'manual' toma valores ingresados a mano. Por " + "defecto es 'split'.", + pt="Que instâncias explicar. Um 'mode' de 'split' toma uma parte de " + "uma partição, 'rows' toma os índices de linha que o utilizador " + "marcou, e 'manual' toma valores introduzidos à mão. Por omissão é " + "'split'.", + de="Welche Instanzen erklärt werden. 'mode' 'split' nimmt einen " + "Anteil einer Teilmenge, 'rows' die vom Benutzer markierten " + "Zeilenindizes und 'manual' manuell eingegebene Werte. Standard " + "ist 'split'.", + zh="要解释哪些实例。'mode' 为 'split' 时取某个划分的一部分，" + "'rows' 取用户标记的行索引，'manual' 取手工输入的值。默认为 'split'。", + ), + alias=MultilingualString( + en="Scope", es="Alcance", pt="Âmbito", de="Umfang", zh="范围" + ), + ) # type: ignore + fit_parameters: schema_field( + dict, + placeholder={}, + description=MultilingualString( + en="Extra arguments handed to the explainer's fit step.", + es="Argumentos adicionales entregados al paso de ajuste del explicador.", + pt="Argumentos adicionais entregues ao passo de ajuste do explicador.", + de="Zusätzliche Argumente für den Fit-Schritt des Erklärers.", + zh="传递给解释器拟合步骤的额外参数。", + ), + alias=MultilingualString( + en="Fit parameters", + es="Parámetros de ajuste", + pt="Parâmetros de ajuste", + de="Fit-Parameter", + zh="拟合参数", + ), + ) # type: ignore + input_columns: _columns_field( + alias=MultilingualString( + en="Input columns", + es="Columnas de entrada", + pt="Colunas de entrada", + de="Eingabespalten", + zh="输入列", + ), + description=MultilingualString( + en="Names of the columns used as model input.", + es="Nombres de las columnas usadas como entrada del modelo.", + pt="Nomes das colunas usadas como entrada do modelo.", + de="Namen der als Modelleingabe verwendeten Spalten.", + zh="用作模型输入的列名。", + ), + ) # type: ignore + output_columns: _columns_field( + alias=MultilingualString( + en="Output columns", + es="Columnas de salida", + pt="Colunas de saída", + de="Ausgabespalten", + zh="输出列", + ), + description=MultilingualString( + en="Names of the columns the model predicts.", + es="Nombres de las columnas que el modelo predice.", + pt="Nomes das colunas que o modelo prevê.", + de="Namen der Spalten, die das Modell vorhersagt.", + zh="模型需要预测的列名。", + ), + ) # type: ignore + manual_input_data: schema_field( + none_type(list), + placeholder=None, + description=MultilingualString( + en="Rows to explain when the scope mode is 'manual'. Ignored otherwise.", + es="Filas a explicar cuando el modo del alcance es 'manual'. Se " + "ignora en otro caso.", + pt="Linhas a explicar quando o modo do âmbito é 'manual'. Ignorado " + "caso contrário.", + de="Zu erklärende Zeilen, wenn der Modus 'manual' ist. Sonst ignoriert.", + zh="范围模式为 'manual' 时要解释的行。其他情况下忽略。", + ), + alias=MultilingualString( + en="Manual input", + es="Entrada manual", + pt="Entrada manual", + de="Manuelle Eingabe", + zh="手动输入", + ), + ) # type: ignore + same_dataset: schema_field( + bool_field(), + placeholder=True, + description=MultilingualString( + en="Whether the instances come from the very dataset the model was " + "trained on. When they do not, the run's row indexes mean nothing " + "here and the split has to be recomputed.", + es="Si las instancias provienen del mismo conjunto de datos con el " + "que se entrenó el modelo. Si no, los índices de fila de la " + "ejecución no significan nada acá y la partición se recalcula.", + pt="Se as instâncias vêm do mesmo conjunto de dados com que o " + "modelo foi treinado. Se não, os índices de linha da execução não " + "significam nada aqui e a divisão tem de ser recalculada.", + de="Ob die Instanzen aus genau dem Datensatz stammen, mit dem das " + "Modell trainiert wurde. Andernfalls sind die Zeilenindizes des " + "Laufs hier bedeutungslos und der Split wird neu berechnet.", + zh="实例是否来自模型训练所用的同一数据集。若不是，则运行记录的行索引在此无意义，" + "需要重新计算划分。", + ), + alias=MultilingualString( + en="Same dataset", + es="Mismo conjunto", + pt="Mesmo conjunto", + de="Gleicher Datensatz", + zh="同一数据集", + ), + ) # type: ignore + session_splits: schema_field( + none_type(string_field()), + placeholder=None, + description=MultilingualString( + en="The model session's split configuration, used only when the " + "instances come from a different dataset and the split has to be " + "recomputed over it.", + es="La configuración de partición de la sesión del modelo, usada " + "solo cuando las instancias vienen de otro conjunto de datos y hay " + "que recalcular la partición sobre él.", + pt="A configuração de divisão da sessão do modelo, usada apenas " + "quando as instâncias vêm de outro conjunto de dados e a divisão " + "tem de ser recalculada sobre ele.", + de="Die Split-Konfiguration der Modellsitzung, nur verwendet, wenn " + "die Instanzen aus einem anderen Datensatz stammen und der Split " + "neu berechnet werden muss.", + zh="模型会话的划分配置，仅在实例来自其他数据集且需要在其上重新计算划分时使用。", + ), + alias=MultilingualString( + en="Session splits", + es="Particiones de la sesión", + pt="Partições da sessão", + de="Sitzungs-Splits", + zh="会话划分", + ), + ) # type: ignore + + +class GenerateLocalExplanationUnit(BaseUnit): + """Explain a selection of instances and pickle the result. + + Sibling of ``GenerateGlobalExplanationUnit``; see that class for why the + two scopes are two units and not one with a branch. + + The three ways of choosing instances — a share of a split, the rows the + user marked, or hand-typed values — stay one unit on purpose. They are + three branches of a single decision with a single output; as separate + nodes exactly one could ever run, which the contract cannot express. + + Row indexes are never taken on trust across datasets: when the instances + come from a dataset other than the one the run was trained on, the run's + indexes address rows that do not correspond, so the split is recomputed + from the session's ratios instead. That derived state is resolved here, + inside ``execute``, and never published. + """ + + SCHEMA = GenerateLocalExplanationSchema + + REQUIRES = ("explainer", "data_x", "data_y", "task", "split_indexes") + PROVIDES = ("explanation_path", "plots_path", "input_dataset_path") + + def _select_instances(self, prepared_instance, splits, instance, task): + """Narrow the loaded dataset down to the instances to explain.""" + import json + + from datasets import DatasetDict + + from DashAI.back.dataloaders.classes.dashai_dataset import ( + prepare_for_model_session, + select_columns, + split_dataset, + ) + + scope = self.config["scope"] or {} + input_columns = self.config["input_columns"] + output_columns = self.config["output_columns"] + + # The data source is selected via scope["mode"]. It defaults to + # "split" so explainers created before this field existed keep + # their original split + percentage behavior. + mode = scope.get("mode", "split") + + if mode == "manual": + # Build the instances from values the user typed in by hand, + # reusing the same conversion the manual prediction flow uses. + # The rows (and any image files rewritten by the job endpoint) + # travel in the job kwargs, not in scope. + manual_input_data = self.config.get("manual_input_data") or [] + if not manual_input_data: + raise JobError("No manual input data provided for the explanation") + prepared_instance = task.process_manual_input( + manual_input_data, + f"{instance.file_path}/dataset", + ) + # Manual input carries only the input columns (no target), so + # keep just those instead of the standard input/output split. + # select_columns returns a DashAIDataset (same shape the + # split path produces), which is what the explainers expect. + return prepared_instance.select_columns(input_columns) + + prepared_instance = task.prepare_for_task( + prepared_instance, + input_columns=input_columns, + output_columns=output_columns, + ) + + if mode == "rows": + # Explain a set of rows the user marked in the table. + # Indexes are over the whole dataset (the split does not + # apply in this mode). + row_indexes = scope.get("row_indexes") or [] + valid_indexes = [ + i + for i in row_indexes + if isinstance(i, int) and 0 <= i < prepared_instance.num_rows + ] + if row_indexes and not valid_indexes: + raise JobError("No valid row indexes provided for the explanation") + if valid_indexes: + prepared_instance = prepared_instance.select(valid_indexes) + else: + split = scope.get("split") + if split not in ["train", "test", "val", "all"]: + raise JobError(f"{split} is not a valid split") + + if split != "all": + if not self.config["same_dataset"]: + if isinstance(splits, str): + splits = json.loads(splits) + ( + prepared_dataset_dict, + splits, + ) = prepare_for_model_session( + dataset=prepared_instance, + splits=splits, + output_columns=output_columns, + ) + split_key = "validation" if split == "val" else split + prepared_instance = prepared_dataset_dict[split_key] + else: + prepared_instance = split_dataset( + prepared_instance, + train_indexes=splits["train_indexes"], + test_indexes=splits["test_indexes"], + val_indexes=splits["val_indexes"], + ) + split_key = "validation" if split == "val" else split + prepared_instance = prepared_instance[split_key] + + n_rows = max( + 1, + int(prepared_instance.num_rows * scope.get("percentage") / 100), + ) + # When "shuffle" is set the percentage is taken as a random + # sample of the split; otherwise it is the leading rows. + if scope.get("shuffle"): + prepared_instance = prepared_instance.shuffle(seed=42) + prepared_instance = prepared_instance.select(range(n_rows)) + + prepared_instance = DatasetDict({"train": prepared_instance}) + x, _ = select_columns(prepared_instance, input_columns, output_columns) + return x + + def execute(self, ctx: ExecutionContext) -> None: + import os + + from datasets import DatasetDict + from kink import di + + from DashAI.back.core.artifacts import normalize_artifacts + from DashAI.back.dataloaders.classes.dashai_dataset import ( + load_dataset, + save_dataset, + ) + + config = di["config"] + session_factory = di["session_factory"] + + explainer = ctx.require("explainer") + task = ctx.require("task") + dataset = (ctx.require("data_x"), ctx.require("data_y")) + splits = ctx.require("split_indexes") + + explainer_id = self.config["explainer_id"] + instance_id = self.config["instance_dataset_id"] + + # Fitting happens before the instances are even looked up, and is left + # unwrapped on purpose: the explainer's own error is what the user gets. + explainer.fit(dataset, **(self.config["fit_parameters"] or {})) + + if not self.config["same_dataset"]: + splits = self.config["session_splits"] + + with session_factory() as db: + instance: Dataset = db.get(Dataset, instance_id) + if not instance: + raise JobError( + f"Dataset {instance_id} to be explained does not exist in DB." + ) + + try: + loaded_instance = load_dataset(f"{instance.file_path}/dataset") + except Exception as e: + log.exception(e) + raise JobError( + f"Can not load instance from path {instance.file_path}", + ) from e + + try: + x = self._select_instances(loaded_instance, splits, instance, task) + + # Persist the original selected rows (the model input for each + # explained instance) as a DashAIDataset before the model's own + # preprocessing runs, so the frontend can read them back with + # the existing dataset endpoints. + input_source = x["train"] if isinstance(x, DatasetDict) else x + input_dataset_path = os.path.join( + config["EXPLANATIONS_PATH"], + f"local_explanation_input_{explainer_id}", + ) + save_dataset(input_source, os.path.join(input_dataset_path, "dataset")) + # The instances are handed over unprepared, the same way the + # prediction job calls model.predict: the model applies its own + # preprocessing. Explainers that need the model feature space + # ask for it with prepare_model_input. + + except Exception as e: + log.exception(e) + raise JobError( + f"""Can not prepare Dataset with {instance_id} + to generate the local explanation.""", + ) from e + + try: + explanation = explainer.explain_instance(x) + plots = normalize_artifacts( + explainer.plot(explanation), create_grouped=True + ) + except Exception as e: + log.exception(e) + raise JobError( + "Failed to generate the explanation", + ) from e + + explanation_path, plots_path = dump_explanation( + explanation, plots, "local", explainer_id + ) + + ctx.put_ref("explanation_path", explanation_path) + ctx.put_ref("plots_path", plots_path) + ctx.put_ref("input_dataset_path", input_dataset_path) diff --git a/DashAI/back/units/load_run_model_unit.py b/DashAI/back/units/load_run_model_unit.py new file mode 100644 index 000000000..3ed3970d7 --- /dev/null +++ b/DashAI/back/units/load_run_model_unit.py @@ -0,0 +1,124 @@ +"""Unit that restores a run's model the way the explanation flow expects.""" + +import logging + +from DashAI.back.core.schema_fields import ( + BaseSchema, + int_field, + schema_field, +) +from DashAI.back.core.utils import MultilingualString +from DashAI.back.dependencies.database.models import Run +from DashAI.back.job.base_job import JobError +from DashAI.back.units.base_unit import BaseUnit +from DashAI.back.units.context import ExecutionContext + +log = logging.getLogger(__name__) + + +class LoadRunModelSchema(BaseSchema): + run_id: schema_field( + int_field(gt=0), + placeholder=1, + description=MultilingualString( + en="Identifier of the run whose trained model is restored. The " + "model component, its parameters and the artifact path all come " + "from that row.", + es="Identificador de la ejecución cuyo modelo entrenado se " + "restaura. El componente del modelo, sus parámetros y la ruta del " + "artefacto salen de esa fila.", + pt="Identificador da execução cujo modelo treinado é restaurado. O " + "componente do modelo, os seus parâmetros e o caminho do artefacto " + "vêm todos dessa linha.", + de="Kennung des Laufs, dessen trainiertes Modell wiederhergestellt " + "wird. Modellkomponente, Parameter und Artefaktpfad stammen alle " + "aus dieser Zeile.", + zh="要恢复其已训练模型的运行标识符。模型组件、参数和产物路径都来自该行。", + ), + alias=MultilingualString( + en="Run", es="Ejecución", pt="Execução", de="Lauf", zh="运行" + ), + ) # type: ignore + + +class LoadRunModelUnit(BaseUnit): + """Restore a run's trained model for an explanation. + + Deliberately **not** ``LoadTrainedModelUnit``, and the difference is not + cosmetic to preserve even though it is very likely accidental: + + * this unit builds an instance with ``model_class(**run.parameters)`` and + only then calls ``load`` on it, the way the explanation flow always has; + * ``LoadTrainedModelUnit`` calls ``load`` straight on the class. + + Every concrete model in this codebase declares ``load`` as a + ``staticmethod`` or a ``classmethod`` that rebuilds the object from the + file, so the instance built here is thrown away and the extra step changes + nothing — and no explainer reads anything that ``__init__`` sets: the only + model attributes any of them touch (``one_hot_encoder``, + ``categorical_columns``, ``label_encoder``) are set during training and + restored from the artifact. + + The two units are kept apart because merging them would have to unify their + error messages, which are user-visible and differ word for word. Do not + collapse them into one with a flag without deciding that first. + """ + + SCHEMA = LoadRunModelSchema + + PROVIDES = ("model",) + + def __init__(self, **config) -> None: + super().__init__(**config) + self._model_class = None + + def _resolve_model_class(self, model_name: str) -> type: + """Resolve the model class from the registry, memoized on this unit.""" + if self._model_class is not None: + return self._model_class + + from kink import di + + component_registry = di["component_registry"] + + try: + model_class = component_registry[model_name]["class"] + except Exception as e: + log.exception(e) + raise JobError( + f"Unable to find Model with name {model_name} in registry.", + ) from e + + self._model_class = model_class + return model_class + + def execute(self, ctx: ExecutionContext) -> None: + from kink import di + + session_factory = di["session_factory"] + + run_id = self.config["run_id"] + + with session_factory() as db: + run: Run = db.get(Run, run_id) + if not run: + raise JobError(f"Run {run_id} does not exist in DB.") + model_name = run.model_name + run_path = run.run_path + parameters = dict(run.parameters or {}) + + run_model_class = self._resolve_model_class(model_name) + + try: + model = run_model_class(**parameters) + except Exception as e: + log.exception(e) + raise JobError("Unable to instantiate model") from e + + try: + trained_model = model.load(run_path) + except Exception as e: + log.exception(e) + raise JobError(f"Can not load model from path {run_path}") from e + + ctx.put("model", trained_model) diff --git a/DashAI/back/units/load_trained_model_unit.py b/DashAI/back/units/load_trained_model_unit.py new file mode 100644 index 000000000..c8e194028 --- /dev/null +++ b/DashAI/back/units/load_trained_model_unit.py @@ -0,0 +1,114 @@ +"""Unit that restores a trained model from the artifact a run left on disk.""" + +import logging + +from DashAI.back.core.schema_fields import ( + BaseSchema, + int_field, + schema_field, +) +from DashAI.back.core.utils import MultilingualString +from DashAI.back.dependencies.database.models import Run +from DashAI.back.job.base_job import JobError +from DashAI.back.units.base_unit import BaseUnit +from DashAI.back.units.context import ExecutionContext + +log = logging.getLogger(__name__) + + +class LoadTrainedModelSchema(BaseSchema): + run_id: schema_field( + int_field(gt=0), + placeholder=1, + description=MultilingualString( + en="Identifier of the run whose trained model is restored. Both " + "the model component and the artifact path come from that row.", + es="Identificador de la ejecución cuyo modelo entrenado se " + "restaura. Tanto el componente del modelo como la ruta del " + "artefacto salen de esa fila.", + pt="Identificador da execução cujo modelo treinado é restaurado. " + "Tanto o componente do modelo como o caminho do artefacto vêm " + "dessa linha.", + de="Kennung des Laufs, dessen trainiertes Modell wiederhergestellt " + "wird. Sowohl die Modellkomponente als auch der Artefaktpfad " + "stammen aus dieser Zeile.", + zh="要恢复其已训练模型的运行标识符。模型组件和产物路径都来自该行。", + ), + alias=MultilingualString( + en="Run", es="Ejecución", pt="Execução", de="Lauf", zh="运行" + ), + ) # type: ignore + + +class LoadTrainedModelUnit(BaseUnit): + """Rebuild a trained model from the run that produced it. + + Reads the model component name and the artifact path off the ``Run`` row + rather than taking them as configuration, so the model that is restored is + always the one that run actually saved. + + ``load`` is invoked on the model *class*, not on an instance, which is what + every concrete model in this codebase expects: they all declare it as a + ``staticmethod`` or a ``classmethod`` that rebuilds the object from the + file. ``ExplainerJob`` instead instantiates the class before calling + ``load``; that extra step has no effect for those models, and preserving + the difference is why ``LoadRunModelUnit`` exists separately instead of + this unit growing a flag. + """ + + SCHEMA = LoadTrainedModelSchema + + PROVIDES = ("model",) + + def __init__(self, **config) -> None: + super().__init__(**config) + self._model_class = None + + def _resolve_model_class(self, model_name: str) -> type: + """Resolve the model class from the registry, memoized on this unit. + + Memoized on the instance, not in the shared context: a context can hold + more than one model-loading node, and a context-global cache key would + make the second one silently reuse the first one's class. + """ + if self._model_class is not None: + return self._model_class + + from kink import di + + component_registry = di["component_registry"] + + try: + model_class = component_registry[model_name]["class"] + except KeyError as e: + log.exception(e) + raise JobError(f"Model {model_name} not found in the registry") from e + + self._model_class = model_class + return model_class + + def execute(self, ctx: ExecutionContext) -> None: + from kink import di + + session_factory = di["session_factory"] + + run_id = self.config["run_id"] + + with session_factory() as db: + run: Run = db.get(Run, run_id) + if not run: + raise JobError(f"Run {run_id} does not exist in DB.") + model_name = run.model_name + run_path = run.run_path + + model_class = self._resolve_model_class(model_name) + + try: + trained_model = model_class.load(run_path) + except Exception as e: + log.exception(e) + raise JobError( + f"Failed to load model {model_name} from path {run_path}" + ) from e + + ctx.put("model", trained_model) diff --git a/DashAI/back/units/load_training_dataset_unit.py b/DashAI/back/units/load_training_dataset_unit.py new file mode 100644 index 000000000..0cdd69f75 --- /dev/null +++ b/DashAI/back/units/load_training_dataset_unit.py @@ -0,0 +1,91 @@ +"""Unit that loads the dataset a model was trained on.""" + +import logging +from typing import TYPE_CHECKING + +from DashAI.back.core.schema_fields import ( + BaseSchema, + schema_field, + string_field, +) +from DashAI.back.core.utils import MultilingualString +from DashAI.back.job.base_job import JobError +from DashAI.back.units.base_unit import BaseUnit +from DashAI.back.units.context import ExecutionContext + +if TYPE_CHECKING: + from DashAI.back.dataloaders.classes.dashai_dataset import DashAIDataset + +log = logging.getLogger(__name__) + + +class LoadTrainingDatasetSchema(BaseSchema): + train_dataset_file_path: schema_field( + string_field(), + placeholder="", + description=MultilingualString( + en="Folder of the dataset the model was trained on — the stored " + "row's own path, not the inner dataset directory.", + es="Carpeta del conjunto de datos con el que se entrenó el " + "modelo: la ruta de la propia fila almacenada, no el directorio " + "interno del conjunto de datos.", + pt="Pasta do conjunto de dados com que o modelo foi treinado — o " + "caminho da própria linha armazenada, não o diretório interno do " + "conjunto de dados.", + de="Ordner des Datensatzes, mit dem das Modell trainiert wurde — " + "der Pfad der gespeicherten Zeile selbst, nicht das innere " + "Datensatzverzeichnis.", + zh="模型训练所用数据集的文件夹——已存储行自身的路径，而非内部数据集目录。", + ), + alias=MultilingualString( + en="Training dataset folder", + es="Carpeta del conjunto de entrenamiento", + pt="Pasta do conjunto de treino", + de="Ordner des Trainingsdatensatzes", + zh="训练数据集文件夹", + ), + ) # type: ignore + + +class LoadTrainingDatasetUnit(BaseUnit): + """Load the dataset a model was trained on, under a key of its own. + + Deliberately not ``LoadDatasetUnit``: this dataset is not the one being + transformed, it is a *reference* the prediction needs — the task decodes + predicted class indexes against its labels, and its declared types become + the schema of the saved result. Publishing it as ``dataset`` would collide + with the dataset actually being predicted on, since ``PROVIDES`` is fixed + per class and both would want the same key. + + Two outputs, with different rules on purpose: the live dataset is cached + for the prediction step, while the types travel as a plain JSON-serializable + mapping so the saving step never has to reopen the file. Nothing derived + from the dataset *being predicted on* crosses this boundary. + """ + + SCHEMA = LoadTrainingDatasetSchema + + PROVIDES = ("train_dataset", "train_dataset_types") + + def execute(self, ctx: ExecutionContext) -> None: + from pathlib import Path + + from DashAI.back.dataloaders.classes.dashai_dataset import load_dataset + + file_path = self.config["train_dataset_file_path"] + + try: + train_dataset: "DashAIDataset" = load_dataset( + str(Path(f"{file_path}/dataset/")) + ) + except Exception as e: + log.exception(e) + raise JobError( + f"Cannot load training dataset from {file_path}/dataset/" + ) from e + + ctx.put("train_dataset", train_dataset) + ctx.put_ref( + "train_dataset_types", + {name: kind.to_string() for name, kind in train_dataset.types.items()}, + ) diff --git a/DashAI/back/units/predict_unit.py b/DashAI/back/units/predict_unit.py new file mode 100644 index 000000000..4a6542fff --- /dev/null +++ b/DashAI/back/units/predict_unit.py @@ -0,0 +1,148 @@ +"""Unit that runs a trained model over a dataset and decodes its output.""" + +import logging +from typing import TYPE_CHECKING + +from DashAI.back.core.schema_fields import ( + BaseSchema, + list_field, + schema_field, + string_field, +) +from DashAI.back.core.utils import MultilingualString +from DashAI.back.job.base_job import JobError +from DashAI.back.units.base_unit import BaseUnit +from DashAI.back.units.context import ExecutionContext + +if TYPE_CHECKING: + from DashAI.back.tasks.base_task import BaseTask + +log = logging.getLogger(__name__) + + +class PredictSchema(BaseSchema): + task_name: schema_field( + string_field(), + placeholder="TabularClassificationTask", + description=MultilingualString( + en="Name of the task that turns raw model output into labels.", + es="Nombre de la tarea que convierte la salida cruda del modelo en " + "etiquetas.", + pt="Nome da tarefa que converte a saída bruta do modelo em rótulos.", + de="Name der Aufgabe, die die Rohausgabe des Modells in Labels umwandelt.", + zh="将模型原始输出转换为标签的任务名称。", + ), + alias=MultilingualString( + en="Task", es="Tarea", pt="Tarefa", de="Aufgabe", zh="任务" + ), + ) # type: ignore + input_columns: schema_field( + list_field(string_field(), min_items=1), + placeholder=[], + description=MultilingualString( + en="Names of the columns handed to the model as input.", + es="Nombres de las columnas entregadas al modelo como entrada.", + pt="Nomes das colunas entregues ao modelo como entrada.", + de="Namen der Spalten, die dem Modell als Eingabe übergeben werden.", + zh="作为输入交给模型的列名。", + ), + alias=MultilingualString( + en="Input columns", + es="Columnas de entrada", + pt="Colunas de entrada", + de="Eingabespalten", + zh="输入列", + ), + ) # type: ignore + output_columns: schema_field( + list_field(string_field(), min_items=1), + placeholder=[], + description=MultilingualString( + en="Names of the columns the model predicts. Only the first one is " + "used: it names the column the predictions are written to.", + es="Nombres de las columnas que el modelo predice. Solo se usa la " + "primera: da nombre a la columna donde se escriben las predicciones.", + pt="Nomes das colunas que o modelo prevê. Apenas a primeira é " + "usada: dá nome à coluna onde as previsões são escritas.", + de="Namen der Spalten, die das Modell vorhersagt. Nur die erste " + "wird verwendet: sie benennt die Spalte für die Vorhersagen.", + zh="模型预测的列名。仅使用第一个：它命名写入预测结果的列。", + ), + alias=MultilingualString( + en="Output columns", + es="Columnas de salida", + pt="Colunas de saída", + de="Ausgabespalten", + zh="输出列", + ), + ) # type: ignore + + +class PredictUnit(BaseUnit): + """Predict with a trained model and decode the result into labels. + + The input columns are selected against the dataset the context holds at + this moment, never against a column list captured earlier: whatever built + that dataset — a load from disk or hand-typed rows — is free to have + produced a different shape. + + The training dataset is required rather than reloaded because the task + decodes predicted class indexes against its labels. The model is handed the + selected columns unprepared: models apply their own preprocessing inside + ``predict``, and preparing beforehand would break the ones that replace + their input columns with derived features. + """ + + SCHEMA = PredictSchema + + REQUIRES = ("dataset", "model", "train_dataset") + PROVIDES = ("y_pred",) + + def __init__(self, **config) -> None: + super().__init__(**config) + self._task = None + + def _resolve_task(self) -> "BaseTask": + """Instantiate the task from the registry, memoized on this unit.""" + if self._task is not None: + return self._task + + from kink import di + + component_registry = di["component_registry"] + task_name = self.config["task_name"] + + try: + task: "BaseTask" = component_registry[task_name]["class"]() + except Exception as e: + log.exception(e) + raise JobError(f"Task {task_name} not found in the registry") from e + + self._task = task + return task + + def validate(self, ctx: ExecutionContext) -> None: + """Resolve the task before anything observable happens. + + The orchestrator calls this ahead of loading the model, which is what + keeps a missing task reported as a task problem rather than being + overtaken by whatever fails next. + """ + self._resolve_task() + + def execute(self, ctx: ExecutionContext) -> None: + import numpy as np + + task = self._resolve_task() + + dataset = ctx.require("dataset") + model = ctx.require("model") + train_dataset = ctx.require("train_dataset") + + prepared_dataset = dataset.select_columns(self.config["input_columns"]) + y_pred_proba = np.array(model.predict(prepared_dataset)) + y_pred = task.process_predictions( + train_dataset, y_pred_proba, self.config["output_columns"][0] + ) + + ctx.put("y_pred", y_pred) diff --git a/DashAI/back/units/prepare_explanation_data_unit.py b/DashAI/back/units/prepare_explanation_data_unit.py new file mode 100644 index 000000000..2ebc0ed20 --- /dev/null +++ b/DashAI/back/units/prepare_explanation_data_unit.py @@ -0,0 +1,188 @@ +"""Unit that rebuilds a run's train/test/val splits for an explanation.""" + +import logging +from typing import TYPE_CHECKING + +from DashAI.back.core.schema_fields import ( + BaseSchema, + list_field, + schema_field, + string_field, +) +from DashAI.back.core.utils import MultilingualString +from DashAI.back.job.base_job import JobError +from DashAI.back.units.base_unit import BaseUnit +from DashAI.back.units.context import ExecutionContext + +if TYPE_CHECKING: + from DashAI.back.tasks.base_task import BaseTask + +log = logging.getLogger(__name__) + + +def _columns_field(alias: MultilingualString, description: MultilingualString): + return schema_field( + list_field(string_field(), min_items=1), + placeholder=[], + description=description, + alias=alias, + ) + + +class PrepareExplanationDataSchema(BaseSchema): + task_name: schema_field( + string_field(), + placeholder="TabularClassificationTask", + description=MultilingualString( + en="Name of the task the dataset is prepared for.", + es="Nombre de la tarea para la que se prepara el conjunto de datos.", + pt="Nome da tarefa para a qual o conjunto de dados é preparado.", + de="Name der Aufgabe, für die der Datensatz vorbereitet wird.", + zh="数据集所准备的任务名称。", + ), + alias=MultilingualString( + en="Task", es="Tarea", pt="Tarefa", de="Aufgabe", zh="任务" + ), + ) # type: ignore + input_columns: _columns_field( + alias=MultilingualString( + en="Input columns", + es="Columnas de entrada", + pt="Colunas de entrada", + de="Eingabespalten", + zh="输入列", + ), + description=MultilingualString( + en="Names of the columns used as model input.", + es="Nombres de las columnas usadas como entrada del modelo.", + pt="Nomes das colunas usadas como entrada do modelo.", + de="Namen der als Modelleingabe verwendeten Spalten.", + zh="用作模型输入的列名。", + ), + ) # type: ignore + output_columns: _columns_field( + alias=MultilingualString( + en="Output columns", + es="Columnas de salida", + pt="Colunas de saída", + de="Ausgabespalten", + zh="输出列", + ), + description=MultilingualString( + en="Names of the columns the model predicts.", + es="Nombres de las columnas que el modelo predice.", + pt="Nomes das colunas que o modelo prevê.", + de="Namen der Spalten, die das Modell vorhersagt.", + zh="模型需要预测的列名。", + ), + ) # type: ignore + + +class PrepareExplanationDataUnit(BaseUnit): + """Rebuild the exact train/test/val split the run was trained on. + + Deliberately not ``PrepareAndSplitUnit``: that one *computes* a split from + a ratio configuration, which would hand the explainer different rows than + the model ever saw. This one replays the row indexes the run recorded, so + the explanation is about the model that exists. + + Features stay unprepared — models apply their own preprocessing inside + ``predict``, and preparing beforehand would break the ones that replace + their input columns with derived features — while targets are encoded, + because explainers compare them against the model's class indexes. + """ + + SCHEMA = PrepareExplanationDataSchema + + # Exactly the keys ``execute`` reads, and no more. ``dataset_id`` is + # deliberately absent: the unit does not name it anywhere, and every key + # listed here is demanded unconditionally by ``__call__``, so declaring an + # unused one would reject any upstream that publishes a dataset without an + # id — ``BuildManualInputUnit``, for one. + REQUIRES = ("dataset", "model", "split_indexes") + PROVIDES = ("data_x", "data_y", "task") + + def __init__(self, **config) -> None: + super().__init__(**config) + self._task = None + + def _resolve_task(self) -> "BaseTask": + """Instantiate the task from the registry, memoized on this unit.""" + if self._task is not None: + return self._task + + from kink import di + + component_registry = di["component_registry"] + task_name = self.config["task_name"] + + try: + task: "BaseTask" = component_registry[task_name]["class"]() + except Exception as e: + log.exception(e) + raise JobError( + (f"Unable to find Task with name {task_name} in registry"), + ) from e + + self._task = task + return task + + def validate(self, ctx: ExecutionContext) -> None: + """Resolve the task before anything observable happens. + + The orchestrator calls this outside the block that wraps preparation + failures, which is what keeps a missing task reported as a registry + problem instead of a generic "cannot prepare" message. + """ + self._resolve_task() + + def execute(self, ctx: ExecutionContext) -> None: + from DashAI.back.dataloaders.classes.dashai_dataset import ( + select_columns, + split_dataset, + ) + + task = self._resolve_task() + + loaded_dataset = ctx.require("dataset") + trained_model = ctx.require("model") + splits = ctx.require("split_indexes") + input_columns = self.config["input_columns"] + output_columns = self.config["output_columns"] + + loaded_dataset = split_dataset( + loaded_dataset, + train_indexes=splits["train_indexes"], + test_indexes=splits["test_indexes"], + val_indexes=splits["val_indexes"], + ) + + prepared_dataset = task.prepare_for_task( + dataset=loaded_dataset, + input_columns=input_columns, + output_columns=output_columns, + ) + data = select_columns(prepared_dataset, input_columns, output_columns) + + data_x = split_dataset( + data[0], + train_indexes=splits["train_indexes"], + test_indexes=splits["test_indexes"], + val_indexes=splits["val_indexes"], + ) + data_y = split_dataset( + data[1], + train_indexes=splits["train_indexes"], + test_indexes=splits["test_indexes"], + val_indexes=splits["val_indexes"], + ) + # Inputs stay unprepared (see the class docstring); targets are encoded + # because explainers compare them against the model's class indexes. + for split_name in data_y: + data_y[split_name] = trained_model.prepare_output( + data_y[split_name], is_fit=False + ) + + ctx.put("data_x", data_x) + ctx.put("data_y", data_y) + ctx.put("task", task) diff --git a/DashAI/back/units/run_exploration_unit.py b/DashAI/back/units/run_exploration_unit.py new file mode 100644 index 000000000..99f0ac5c2 --- /dev/null +++ b/DashAI/back/units/run_exploration_unit.py @@ -0,0 +1,187 @@ +"""Unit that runs one exploration over the dataset in the context.""" + +import logging +from typing import TYPE_CHECKING, Type + +from DashAI.back.core.schema_fields import ( + BaseSchema, + component_field, + int_field, + schema_field, +) +from DashAI.back.core.utils import MultilingualString +from DashAI.back.dependencies.database.models import Explorer +from DashAI.back.job.base_job import JobError +from DashAI.back.units.base_unit import BaseUnit +from DashAI.back.units.context import ExecutionContext + +if TYPE_CHECKING: + from DashAI.back.exploration.base_explorer import BaseExplorer + +log = logging.getLogger(__name__) + + +class RunExplorationSchema(BaseSchema): + explorer_id: schema_field( + int_field(gt=0), + placeholder=1, + description=MultilingualString( + en="Identifier of the exploration whose selected columns and " + "display name the explorer component reads.", + es="Identificador de la exploración cuyas columnas seleccionadas y " + "nombre para mostrar lee el componente de exploración.", + pt="Identificador da exploração cujas colunas selecionadas e nome " + "de exibição o componente de exploração lê.", + de="Kennung der Exploration, deren ausgewählte Spalten und " + "Anzeigename die Explorer-Komponente liest.", + zh="探索的标识符，探索组件从中读取所选列和显示名称。", + ), + alias=MultilingualString( + en="Exploration", + es="Exploración", + pt="Exploração", + de="Exploration", + zh="探索", + ), + ) # type: ignore + explorer: schema_field( + component_field(parent="BaseExplorer"), + placeholder={ + "component": "DescribeExplorer", + "params": {"percentiles": "25, 50, 75", "include": "all", "exclude": None}, + }, + description=MultilingualString( + en="Exploration to run, together with its own configuration.", + es="Exploración a ejecutar, junto con su propia configuración.", + pt="Exploração a executar, junto com a sua própria configuração.", + de="Auszuführende Exploration samt ihrer eigenen Konfiguration.", + zh="要运行的探索及其自身配置。", + ), + alias=MultilingualString( + en="Explorer", + es="Explorador", + pt="Explorador", + de="Explorer", + zh="探索器", + ), + ) # type: ignore + + +class RunExplorationUnit(BaseUnit): + """Instantiate an explorer component and run it over the dataset. + + Preparing the dataset and launching the exploration are one unit, not two: + ``prepare_dataset`` is a hook on the explorer component itself + (``BaseExplorer.prepare_dataset``), so it does not exist without an + instantiated explorer. Splitting them would force the live explorer + instance through the context, which is instance state wearing a context + key's clothes. + + The unit re-reads the ``Explorer`` row because the component API takes it: + ``prepare_dataset`` needs its ``columns`` and ``launch_exploration`` + receives the row itself. The read is strictly read-only — the row's status + belongs to the job. + + The explorer instance is published alongside the result because saving is + also a method on it (``save_notebook``). Handing over the same object, + rather than letting the save unit build its own from the same + configuration, is what keeps a stateful explorer working: ``CorrMatrix`` + and ``CovMatrix`` read ``self.plot`` while saving. Same shape as + ``BuildModelUnit`` publishing ``model`` for ``SaveModelUnit``. + """ + + SCHEMA = RunExplorationSchema + + PROVIDES = ("exploration_result", "explorer") + REQUIRES = ("dataset",) + + def __init__(self, **config) -> None: + super().__init__(**config) + self._explorer_class = None + + @property + def exploration_type(self) -> str: + return self.config["explorer"]["component"] + + @property + def parameters(self) -> dict: + return self.config["explorer"]["params"] + + def _resolve_explorer_class(self) -> Type["BaseExplorer"]: + """Resolve the explorer class from the registry, memoized on this unit. + + Memoized on the instance rather than in the shared context: a context + can hold more than one exploration node, and a context-global cache key + would make the second one silently reuse the first one's class. + """ + if self._explorer_class is not None: + return self._explorer_class + + from kink import di + + component_registry = di["component_registry"] + exploration_type = self.exploration_type + + try: + explorer_class = component_registry[exploration_type]["class"] + except KeyError as e: + log.exception(e) + raise JobError( + (f"Explorer {exploration_type} not found in the registry.") + ) from e + + self._explorer_class = explorer_class + return explorer_class + + def execute(self, ctx: ExecutionContext) -> None: + from kink import di + + from DashAI.back.exploration.base_explorer import BaseExplorer + + session_factory = di["session_factory"] + + explorer_id = self.config["explorer_id"] + exploration_type = self.exploration_type + loaded_dataset = ctx.require("dataset") + + explorer_component_class = self._resolve_explorer_class() + + with session_factory() as db: + explorer_info: Explorer = db.get(Explorer, explorer_id) + if explorer_info is None: + raise JobError(f"Explorer with id {explorer_id} not found.") + + try: + explorer_instance = explorer_component_class(**self.parameters) + assert isinstance(explorer_instance, BaseExplorer) + except Exception as e: + log.exception(e) + raise JobError( + f"Error instancing the explorer {exploration_type}." + ) from e + + try: + prepared_dataset = explorer_instance.prepare_dataset( + loaded_dataset, explorer_info.columns + ) + except Exception as e: + log.exception(e) + raise JobError( + ( + "Error preparing the dataset for the exploration " + f"{exploration_type}." + ) + ) from e + + try: + result = explorer_instance.launch_exploration( + prepared_dataset, explorer_info + ) + except Exception as e: + log.exception(e) + raise JobError( + f"Error launching the exploration {exploration_type}." + ) from e + + ctx.put("exploration_result", result) + ctx.put("explorer", explorer_instance) diff --git a/DashAI/back/units/save_exploration_unit.py b/DashAI/back/units/save_exploration_unit.py new file mode 100644 index 000000000..71146ac76 --- /dev/null +++ b/DashAI/back/units/save_exploration_unit.py @@ -0,0 +1,122 @@ +"""Unit that persists an exploration result under its notebook's folder.""" + +import logging + +from DashAI.back.core.schema_fields import ( + BaseSchema, + int_field, + schema_field, +) +from DashAI.back.core.utils import MultilingualString +from DashAI.back.dependencies.database.models import Explorer, Notebook +from DashAI.back.job.base_job import JobError +from DashAI.back.units.base_unit import BaseUnit +from DashAI.back.units.context import ExecutionContext + +log = logging.getLogger(__name__) + + +class SaveExplorationSchema(BaseSchema): + explorer_id: schema_field( + int_field(gt=0), + placeholder=1, + description=MultilingualString( + en="Identifier of the exploration being saved. It decides the " + "destination folder and the file name, so a re-run overwrites its " + "own artifact and never another exploration's.", + es="Identificador de la exploración que se guarda. Determina la " + "carpeta de destino y el nombre del archivo, de modo que volver a " + "ejecutarla sobrescribe su propio artefacto y nunca el de otra.", + pt="Identificador da exploração que está a ser guardada. Determina " + "a pasta de destino e o nome do ficheiro, pelo que uma nova " + "execução substitui o seu próprio artefacto e nunca o de outra.", + de="Kennung der zu speichernden Exploration. Sie bestimmt " + "Zielordner und Dateinamen, sodass ein erneuter Lauf nur das " + "eigene Artefakt überschreibt und nie das einer anderen.", + zh="要保存的探索的标识符。它决定目标文件夹和文件名，因此重新运行只会覆盖" + "自身的产物，而不会覆盖其他探索的产物。", + ), + alias=MultilingualString( + en="Exploration", + es="Exploración", + pt="Exploração", + de="Exploration", + zh="探索", + ), + ) # type: ignore + + +class SaveExplorationUnit(BaseUnit): + """Write an exploration result to disk and publish where it landed. + + How the result is serialised is the explorer component's own business — + a Plotly figure as JSON, a DataFrame as JSON, a word cloud as PNG — so the + unit delegates to ``save_notebook`` and only owns the destination: the + notebook's folder under ``NOTEBOOK_PATH``, keyed by the notebook id. + + The explorer arrives through the context rather than being rebuilt from a + configuration of its own, so the object that saves is the object that ran. + + Declares only ``exploration_path``: the result itself is on disk, and the + row that records the path belongs to the job. + """ + + SCHEMA = SaveExplorationSchema + + REQUIRES = ("exploration_result", "explorer") + PROVIDES = ("exploration_path",) + + def execute(self, ctx: ExecutionContext) -> None: + import os + import pathlib + + from kink import di + + config = di["config"] + session_factory = di["session_factory"] + + explorer_id = self.config["explorer_id"] + explorer_instance = ctx.require("explorer") + result = ctx.require("exploration_result") + + with session_factory() as db: + explorer_info: Explorer = db.get(Explorer, explorer_id) + if explorer_info is None: + raise JobError(f"Explorer with id {explorer_id} not found.") + + notebook_info: Notebook = db.get(Notebook, explorer_info.notebook_id) + if notebook_info is None: + raise JobError( + f"Notebook with id {explorer_info.notebook_id} not found." + ) + + # Read while the row is still attached: the error message below is + # built after the session is gone. + exploration_type = explorer_info.exploration_type + + # save in the notebook folder + save_path = pathlib.Path( + os.path.join( + config["NOTEBOOK_PATH"], + (f"{notebook_info.id}"), + ) + ) + if not save_path.exists(): + save_path.mkdir(parents=True) + + save_path = explorer_instance.save_notebook( + notebook_info, explorer_info, save_path, result + ) + + if isinstance(save_path, str): + save_path = pathlib.Path(save_path) + if not isinstance(save_path, pathlib.Path): + raise JobError( + ( + f"Error while saving the exploration" + f" {exploration_type}" + f", save path is not a pathlib.Path." + ) + ) + + ctx.put_ref("exploration_path", save_path.as_posix()) diff --git a/DashAI/back/units/save_prediction_unit.py b/DashAI/back/units/save_prediction_unit.py new file mode 100644 index 000000000..71e9625a7 --- /dev/null +++ b/DashAI/back/units/save_prediction_unit.py @@ -0,0 +1,140 @@ +"""Unit that stores a prediction alongside the data it was made on.""" + +import logging + +from DashAI.back.core.schema_fields import ( + BaseSchema, + list_field, + schema_field, + string_field, +) +from DashAI.back.core.utils import MultilingualString +from DashAI.back.units.base_unit import BaseUnit +from DashAI.back.units.context import ExecutionContext + +log = logging.getLogger(__name__) + + +def _columns_field(alias: MultilingualString, description: MultilingualString): + return schema_field( + list_field(string_field(), min_items=1), + placeholder=[], + description=description, + alias=alias, + ) + + +class SavePredictionSchema(BaseSchema): + input_columns: _columns_field( + alias=MultilingualString( + en="Input columns", + es="Columnas de entrada", + pt="Colunas de entrada", + de="Eingabespalten", + zh="输入列", + ), + description=MultilingualString( + en="Names of the model input columns. Together with the output " + "column they decide which declared types are kept in the schema " + "written next to the result.", + es="Nombres de las columnas de entrada del modelo. Junto con la " + "columna de salida deciden qué tipos declarados se conservan en el " + "esquema que se escribe junto al resultado.", + pt="Nomes das colunas de entrada do modelo. Juntamente com a coluna " + "de saída decidem que tipos declarados são mantidos no esquema " + "escrito junto ao resultado.", + de="Namen der Modelleingabespalten. Zusammen mit der Ausgabespalte " + "bestimmen sie, welche deklarierten Typen im Schema neben dem " + "Ergebnis erhalten bleiben.", + zh="模型输入列的列名。它们与输出列共同决定结果旁写入的模式中保留哪些声明类型。", + ), + ) # type: ignore + output_columns: _columns_field( + alias=MultilingualString( + en="Output columns", + es="Columnas de salida", + pt="Colunas de saída", + de="Ausgabespalten", + zh="输出列", + ), + description=MultilingualString( + en="Names of the predicted columns. Only the first one is used: it " + "names the column the predictions are written to.", + es="Nombres de las columnas predichas. Solo se usa la primera: da " + "nombre a la columna donde se escriben las predicciones.", + pt="Nomes das colunas previstas. Apenas a primeira é usada: dá nome " + "à coluna onde as previsões são escritas.", + de="Namen der vorhergesagten Spalten. Nur die erste wird verwendet: " + "sie benennt die Spalte für die Vorhersagen.", + zh="预测列的列名。仅使用第一个：它命名写入预测结果的列。", + ), + ) # type: ignore + + +class SavePredictionUnit(BaseUnit): + """Write the predicted column next to the data it was predicted from. + + The destination is a fresh folder under the datasets directory, named by a + generated identifier: a prediction has no natural key to overwrite, so + every run gets its own and no result can clobber another's. + + The columns to keep are resolved against the dataset the context holds + right now, at the top of this method. Publishing a column list earlier + would go stale the moment anything upstream renamed, added or dropped a + column. + """ + + SCHEMA = SavePredictionSchema + + REQUIRES = ("dataset", "y_pred", "train_dataset_types") + PROVIDES = ("results_path",) + + def execute(self, ctx: ExecutionContext) -> None: + import uuid + from pathlib import Path + + from kink import di + + from DashAI.back.dataloaders.classes.dashai_dataset import ( + save_dataset, + to_dashai_dataset, + ) + + config = di["config"] + + dataset = ctx.require("dataset") + y_pred = ctx.require("y_pred") + train_dataset_types = ctx.require("train_dataset_types") + + input_columns = self.config["input_columns"] + output_col = self.config["output_columns"][0] + + path = str(Path(f"{config['DATASETS_PATH']}/predictions/")) + folder_name = str(uuid.uuid4()) + full_path = Path(path) / folder_name + full_path.mkdir(parents=True, exist_ok=True) + + base_columns = [col for col in dataset.column_names if col != output_col] + output_dataset = dataset.select_columns(base_columns) + dataset_with_prediction = to_dashai_dataset( + output_dataset.add_column(output_col, y_pred) + ) + + # Only the columns the model session declares carry a type; anything + # the input dataset happened to bring along is left untyped. + filtered_schema = { + name: kind + for name, kind in train_dataset_types.items() + if name in input_columns + self.config["output_columns"] + } + + # Store num of rows, columns, and column names + dataset_with_prediction.compute_base_metadata() + + save_dataset( + dataset_with_prediction, + str(full_path / "dataset"), + filtered_schema, + ) + + ctx.put_ref("results_path", str(full_path)) diff --git a/tests/back/api/test_explainer_job.py b/tests/back/api/test_explainer_job.py new file mode 100644 index 000000000..2e55b422b --- /dev/null +++ b/tests/back/api/test_explainer_job.py @@ -0,0 +1,701 @@ +"""End-to-end regression net for ``ExplainerJob``. + +Written before the job is decomposed into atomic units, and asserted against the +monolithic implementation, so that the refactor has something to be measured +against. The assertions are deliberately explicit — exact status values, exact +columns written on the row, exact error message fragments, and which message is +the one the user actually sees versus which survives only as ``__cause__`` — +instead of the looser ``status in [1, 3]`` style used by +``test_explainer_jobs.py``, which cannot tell a unit that silently stopped doing +part of its work from one that did it. + +Tests named ``test_currently_*`` pin behaviour that is known to be wrong. They +exist so the refactor can be proven behaviour-preserving first; the fix lands +afterwards as its own change, which flips the assertion and renames the test. + +Lives under ``tests/back/api`` to reuse the ``client`` and ``dataset_1`` +fixtures from this package's ``conftest.py``. +""" + +import json +import shutil +from pathlib import Path + +import joblib +import pytest +from datasets import ClassLabel, Value +from fastapi.testclient import TestClient + +from DashAI.back.core.enums.status import ExplainerStatus +from DashAI.back.dependencies.database.models import ( + Dataset, + GlobalExplainer, + LocalExplainer, + ModelSession, + Run, +) +from DashAI.back.dependencies.registry import ComponentRegistry +from DashAI.back.explainability.global_explainer import BaseGlobalExplainer +from DashAI.back.explainability.local_explainer import BaseLocalExplainer +from DashAI.back.job.base_job import JobError +from DashAI.back.job.explainer_job import ExplainerJob +from DashAI.back.models.base_model import BaseModel +from DashAI.back.tasks.base_task import BaseTask + +INPUT_COLUMNS = ["SepalLengthCm", "SepalWidthCm", "PetalLengthCm", "PetalWidthCm"] +OUTPUT_COLUMNS = ["Species"] + +SPLIT_INDEXES = json.dumps( + { + "train_indexes": [0, 1, 2, 3, 4], + "test_indexes": [5, 6, 7, 8], + "val_indexes": [9, 10, 11, 12], + } +) + +SPLITS = json.dumps( + { + "train": 0.5, + "test": 0.2, + "validation": 0.3, + "is_random": True, + "has_changed": True, + "seed": 42, + "shuffle": True, + "stratify": False, + } +) + + +class DummyTask(BaseTask): + name: str = "DummyTask" + metadata: dict = { + "inputs_types": [ClassLabel, Value], + "outputs_types": [ClassLabel], + "inputs_cardinality": "n", + "outputs_cardinality": 1, + } + + def prepare_for_task(self, dataset, input_columns=None, output_columns=None): + return dataset + + +class DummyModel(BaseModel): + COMPATIBLE_COMPONENTS = ["DummyTask"] + + @classmethod + def get_schema(cls): + return {} + + def save(self, filename): + joblib.dump(self, filename) + + @staticmethod + def load(filename): + return DummyModel() + + def predict(self, x): + return {} + + def train(self, x_train, y_train, x_validation=None, y_validation=None): + return + + def prepare_dataset(self, dataset, is_fit=False): + return dataset + + def prepare_output(self, dataset, is_fit=False): + return dataset + + +class UninstantiableModel(DummyModel): + def __init__(self, *args, **kwargs): + raise RuntimeError("this model refuses to be built") + + +class UnloadableModel(DummyModel): + @staticmethod + def load(filename): + raise OSError("the artifact is not there") + + +class DummyGlobalExplainer(BaseGlobalExplainer): + COMPATIBLE_COMPONENTS = ["DummyTask"] + + def __init__(self, model: BaseModel) -> None: + self.model = model + self.explanation = None + + @classmethod + def get_schema(cls): + return {} + + def explain(self, dataset): + return {"importance": [1, 2, 3]} + + def plot(self, explanation): + return "a plot" + + +class ExplodingGlobalExplainer(DummyGlobalExplainer): + def explain(self, dataset): + raise RuntimeError("the explanation itself blew up") + + +class UninstantiableGlobalExplainer(DummyGlobalExplainer): + def __init__(self, model: BaseModel) -> None: + raise RuntimeError("this explainer refuses to be built") + + +class DummyLocalExplainer(BaseLocalExplainer): + COMPATIBLE_COMPONENTS = ["DummyTask"] + + def __init__(self, model: BaseModel) -> None: + self.model = model + self.explanation = None + + @classmethod + def get_schema(cls): + return {} + + def fit(self, dataset, **kwargs): + return self + + def explain_instance(self, instances): + return {"local": True} + + def plot(self, explanation): + return "a plot" + + +@pytest.fixture(autouse=True, name="test_registry") +def setup_test_registry(client, monkeypatch: pytest.MonkeyPatch): + container = client.app.container + + test_registry = ComponentRegistry( + initial_components=[ + DummyTask, + DummyModel, + UninstantiableModel, + UnloadableModel, + DummyGlobalExplainer, + ExplodingGlobalExplainer, + UninstantiableGlobalExplainer, + DummyLocalExplainer, + ExplainerJob, + ] + ) + + monkeypatch.setitem(container._services, "component_registry", test_registry) + return test_registry + + +@pytest.fixture(scope="module", name="model_session_id") +def create_model_session(client: TestClient, dataset_1: Dataset): + session_factory = client.app.container["session_factory"] + + with session_factory() as db: + model_session = ModelSession( + dataset_id=dataset_1.id, + name="ExplainerJobSession", + task_name="DummyTask", + input_columns=INPUT_COLUMNS, + output_columns=OUTPUT_COLUMNS, + splits=SPLITS, + ) + db.add(model_session) + db.commit() + db.refresh(model_session) + return model_session.id + + +@pytest.fixture(name="run_id") +def create_run(client: TestClient, model_session_id: int): + """Function scoped: the error-branch tests corrupt this row on purpose.""" + session_factory = client.app.container["session_factory"] + + with session_factory() as db: + run = Run( + model_session_id=model_session_id, + optimizer_name="OptunaOptimizer", + optimizer_parameters={}, + model_name="DummyModel", + parameters={}, + goal_metric="Accuracy", + name="ExplainerJobRun", + run_path="a/saved/model", + split_indexes=SPLIT_INDEXES, + ) + db.add(run) + db.commit() + db.refresh(run) + return run.id + + +def _create_global_explainer(client, run_id, explainer_name="DummyGlobalExplainer"): + session_factory = client.app.container["session_factory"] + with session_factory() as db: + explainer = GlobalExplainer( + run_id=run_id, + explainer_name=explainer_name, + parameters={}, + ) + db.add(explainer) + db.commit() + db.refresh(explainer) + return explainer.id + + +def _create_local_explainer(client, run_id, dataset_id, scope=None): + session_factory = client.app.container["session_factory"] + with session_factory() as db: + explainer = LocalExplainer( + run_id=run_id, + explainer_name="DummyLocalExplainer", + dataset_id=dataset_id, + scope=scope if scope is not None else {"split": "test", "percentage": 100}, + parameters={}, + fit_parameters={}, + ) + db.add(explainer) + db.commit() + db.refresh(explainer) + return explainer.id + + +def _stored(client, model, explainer_id): + session_factory = client.app.container["session_factory"] + with session_factory() as db: + row = db.get(model, explainer_id) + stored = { + "status": row.status, + "explanation_path": row.explanation_path, + "plot_overrides": row.plot_overrides, + "huey_id": row.huey_id, + } + if model is GlobalExplainer: + stored["plot_path"] = row.plot_path + else: + stored["plots_path"] = row.plots_path + stored["input_dataset_path"] = row.input_dataset_path + return stored + + +# --- happy paths -------------------------------------------------------- + + +def test_a_global_explanation_writes_both_pickles_and_finishes(client, run_id): + explainer_id = _create_global_explainer(client, run_id) + + ExplainerJob(explainer_id=explainer_id, explainer_scope="global").run() + + stored = _stored(client, GlobalExplainer, explainer_id) + assert stored["status"] == ExplainerStatus.FINISHED + assert Path(stored["explanation_path"]).name == ( + f"global_explanation_{explainer_id}.pickle" + ) + assert Path(stored["plot_path"]).name == ( + f"global_explanation_plot_{explainer_id}.pickle" + ) + assert Path(stored["explanation_path"]).exists() + assert Path(stored["plot_path"]).exists() + # Overrides belong to a previous result and must not survive a re-run. + assert stored["plot_overrides"] is None + + +def test_a_local_explanation_writes_its_three_paths_and_finishes( + client, run_id, dataset_1 +): + """The local row carries an extra artifact the global one does not: the + selected instances, saved so the frontend can read them back.""" + from DashAI.back.dataloaders.classes.dashai_dataset import load_dataset + + explainer_id = _create_local_explainer(client, run_id, dataset_1.id) + + ExplainerJob(explainer_id=explainer_id, explainer_scope="local").run() + + stored = _stored(client, LocalExplainer, explainer_id) + assert stored["status"] == ExplainerStatus.FINISHED + assert Path(stored["explanation_path"]).name == ( + f"local_explanation_{explainer_id}.pickle" + ) + assert Path(stored["plots_path"]).name == ( + f"local_explanation_plots_{explainer_id}.pickle" + ) + assert Path(stored["explanation_path"]).exists() + assert Path(stored["plots_path"]).exists() + assert stored["plot_overrides"] is None + + saved_input = load_dataset(str(Path(stored["input_dataset_path"]) / "dataset")) + assert saved_input.column_names == INPUT_COLUMNS + # scope percentage 100 over the four test indexes. + assert len(saved_input) == 4 + + +def test_a_rows_scope_explains_exactly_the_marked_rows(client, run_id, dataset_1): + """Row indexes address the whole dataset; the split does not apply.""" + from DashAI.back.dataloaders.classes.dashai_dataset import load_dataset + + explainer_id = _create_local_explainer( + client, run_id, dataset_1.id, scope={"mode": "rows", "row_indexes": [0, 7, 42]} + ) + + ExplainerJob(explainer_id=explainer_id, explainer_scope="local").run() + + stored = _stored(client, LocalExplainer, explainer_id) + assert stored["status"] == ExplainerStatus.FINISHED + + saved_input = load_dataset(str(Path(stored["input_dataset_path"]) / "dataset")) + assert len(saved_input) == 3 + + +# --- scope selection ---------------------------------------------------- + + +def test_an_invalid_scope_is_rejected_before_anything_is_touched(client, run_id): + explainer_id = _create_global_explainer(client, run_id) + + with pytest.raises(JobError, match="banana is an invalid explainer type"): + ExplainerJob(explainer_id=explainer_id, explainer_scope="banana").run() + + # Nothing ran, so the row is untouched. + assert _stored(client, GlobalExplainer, explainer_id)["status"] == ( + ExplainerStatus.NOT_STARTED + ) + + +def test_a_missing_explainer_row_is_reported_by_id(client): + """The missing row must be named, not crash the handler meant to mark it. + + The row is looked up before the guarded block, because the outer + ``except Exception`` calls ``set_status_as_error`` on that very row — so + without the check a bad id used to surface as an ``AttributeError`` raised + by the error handler itself. + """ + with pytest.raises( + JobError, match="Explainer with id 999999 does not exist in DB." + ): + ExplainerJob(explainer_id=999999, explainer_scope="global").run() + + +def test_the_huey_id_is_recorded_on_the_row(client, run_id): + """Like the other three jobs, the queue task id is stored on the row. + + Without it the explanation cannot be matched back to its queue entry. + """ + explainer_id = _create_global_explainer(client, run_id) + + ExplainerJob( + explainer_id=explainer_id, explainer_scope="global", huey_id="task-abc" + ).run() + + assert _stored(client, GlobalExplainer, explainer_id)["huey_id"] == "task-abc" + + +# --- loading errors ----------------------------------------------------- + + +def test_a_missing_run_is_reported_and_the_row_goes_to_error(client): + explainer_id = _create_global_explainer(client, run_id=999999) + + with pytest.raises(JobError, match="Run 999999 does not exist in DB."): + ExplainerJob(explainer_id=explainer_id, explainer_scope="global").run() + + # STARTED is set late, after everything is loaded, so a failure here takes + # the row straight from NOT_STARTED to ERROR. + assert _stored(client, GlobalExplainer, explainer_id)["status"] == ( + ExplainerStatus.ERROR + ) + + +def test_a_missing_model_session_is_reported_by_id(client, run_id): + explainer_id = _create_global_explainer(client, run_id) + + session_factory = client.app.container["session_factory"] + with session_factory() as db: + db.get(Run, run_id).model_session_id = 999999 + db.commit() + + with pytest.raises(JobError, match="Model session 999999 does not exist in DB."): + ExplainerJob(explainer_id=explainer_id, explainer_scope="global").run() + + assert _stored(client, GlobalExplainer, explainer_id)["status"] == ( + ExplainerStatus.ERROR + ) + + +def test_a_missing_training_dataset_names_the_id_that_was_looked_up( + client, run_id, model_session_id +): + """The message names the dataset the lookup actually used. + + It used to interpolate ``self.explainer_db.dataset_id`` while looking up + ``model_session.dataset_id`` — and that column exists on ``LocalExplainer`` + and *not* on ``GlobalExplainer``, so in the global scope the "does not + exist" error was never built at all: an ``AttributeError`` was raised while + formatting it. + """ + explainer_id = _create_global_explainer(client, run_id) + + session_factory = client.app.container["session_factory"] + with session_factory() as db: + db.get(ModelSession, model_session_id).dataset_id = 999999 + db.commit() + try: + with pytest.raises(JobError, match="Dataset 999999 does not exist in DB."): + ExplainerJob(explainer_id=explainer_id, explainer_scope="global").run() + + assert _stored(client, GlobalExplainer, explainer_id)["status"] == ( + ExplainerStatus.ERROR + ) + finally: + with session_factory() as db: + db.get(ModelSession, model_session_id).dataset_id = ( + db.query(Dataset).filter(Dataset.name == "test_csv_1").first().id + ) + db.commit() + + +def test_an_unknown_model_name_is_reported_by_name(client, run_id): + explainer_id = _create_global_explainer(client, run_id) + + session_factory = client.app.container["session_factory"] + with session_factory() as db: + db.get(Run, run_id).model_name = "NoSuchModel" + db.commit() + + with pytest.raises( + JobError, match="Unable to find Model with name NoSuchModel in registry." + ): + ExplainerJob(explainer_id=explainer_id, explainer_scope="global").run() + + assert _stored(client, GlobalExplainer, explainer_id)["status"] == ( + ExplainerStatus.ERROR + ) + + +def test_a_model_that_cannot_be_instantiated_is_reported(client, run_id): + """The job builds the model before loading it, unlike ``PredictJob``.""" + explainer_id = _create_global_explainer(client, run_id) + + session_factory = client.app.container["session_factory"] + with session_factory() as db: + db.get(Run, run_id).model_name = "UninstantiableModel" + db.commit() + + with pytest.raises(JobError, match="Unable to instantiate model"): + ExplainerJob(explainer_id=explainer_id, explainer_scope="global").run() + + +def test_a_model_that_cannot_be_loaded_names_the_path(client, run_id): + explainer_id = _create_global_explainer(client, run_id) + + session_factory = client.app.container["session_factory"] + with session_factory() as db: + run = db.get(Run, run_id) + run.model_name = "UnloadableModel" + run.run_path = "gone/from/disk" + db.commit() + + with pytest.raises(JobError, match="Can not load model from path gone/from/disk"): + ExplainerJob(explainer_id=explainer_id, explainer_scope="global").run() + + +def test_an_unknown_explainer_name_is_reported_with_the_multiline_message( + client, run_id +): + """The message is a triple-quoted f-string, so its newline and indentation + are literally part of the text the user sees. Pinned as-is.""" + explainer_id = _create_global_explainer( + client, run_id, explainer_name="NoSuchExplainer" + ) + + with pytest.raises(JobError) as excinfo: + ExplainerJob(explainer_id=explainer_id, explainer_scope="global").run() + + assert str(excinfo.value) == ( + "Unable to find the global explainer with name\n" + " NoSuchExplainer in registry." + ) + + +def test_an_explainer_that_cannot_be_instantiated_names_the_scope(client, run_id): + explainer_id = _create_global_explainer( + client, run_id, explainer_name="UninstantiableGlobalExplainer" + ) + + with pytest.raises(JobError, match="Unable to instantiate global explainer."): + ExplainerJob(explainer_id=explainer_id, explainer_scope="global").run() + + +def test_a_dataset_that_cannot_be_loaded_names_the_path( + client, run_id, dataset_1, tmp_path +): + explainer_id = _create_global_explainer(client, run_id) + + stored_folder = Path(dataset_1.file_path) / "dataset" + backup = tmp_path / "explainer-dataset-backup" + shutil.copytree(stored_folder, backup) + shutil.rmtree(stored_folder) + try: + with pytest.raises(JobError, match="Can not load dataset from path"): + ExplainerJob(explainer_id=explainer_id, explainer_scope="global").run() + + assert _stored(client, GlobalExplainer, explainer_id)["status"] == ( + ExplainerStatus.ERROR + ) + finally: + shutil.copytree(backup, stored_folder) + + +def test_an_unknown_task_name_is_reported_by_name(client, run_id, model_session_id): + explainer_id = _create_global_explainer(client, run_id) + + session_factory = client.app.container["session_factory"] + with session_factory() as db: + db.get(ModelSession, model_session_id).task_name = "NoSuchTask" + db.commit() + try: + with pytest.raises( + JobError, match="Unable to find Task with name NoSuchTask in registry" + ): + ExplainerJob(explainer_id=explainer_id, explainer_scope="global").run() + finally: + with session_factory() as db: + db.get(ModelSession, model_session_id).task_name = "DummyTask" + db.commit() + + +def test_incomplete_split_indexes_report_a_preparation_error(client, run_id, dataset_1): + """All three splits are read off the run; a missing one is a hard failure. + + The reads happen inside the block whose ``except Exception`` builds the + generic preparation message, so that wrapper is what the user sees. + """ + explainer_id = _create_global_explainer(client, run_id) + + session_factory = client.app.container["session_factory"] + with session_factory() as db: + db.get(Run, run_id).split_indexes = json.dumps({"train_indexes": [0, 1]}) + db.commit() + + with pytest.raises(JobError, match="Can not prepare dataset"): + ExplainerJob(explainer_id=explainer_id, explainer_scope="global").run() + + assert _stored(client, GlobalExplainer, explainer_id)["status"] == ( + ExplainerStatus.ERROR + ) + + +# --- generation errors -------------------------------------------------- + + +def test_a_failing_global_explanation_is_reported_and_errors(client, run_id): + explainer_id = _create_global_explainer( + client, run_id, explainer_name="ExplodingGlobalExplainer" + ) + + with pytest.raises(JobError, match="Failed to generate the explanation") as excinfo: + ExplainerJob(explainer_id=explainer_id, explainer_scope="global").run() + + assert "the explanation itself blew up" in str(excinfo.value.__cause__) + assert _stored(client, GlobalExplainer, explainer_id)["status"] == ( + ExplainerStatus.ERROR + ) + + +def test_an_invalid_split_is_swallowed_by_the_preparation_wrapper( + client, run_id, dataset_1 +): + """The specific complaint never reaches the user. + + ``"notasplit is not a valid split"`` is raised inside the block whose + ``except Exception`` replaces it with the generic wrapper, so it survives + only as ``__cause__``. Locking this in because it is an easy detail to + "fix" by accident while refactoring. + """ + explainer_id = _create_local_explainer( + client, run_id, dataset_1.id, scope={"split": "notasplit", "percentage": 100} + ) + + with pytest.raises(JobError, match="Can not prepare Dataset with") as excinfo: + ExplainerJob(explainer_id=explainer_id, explainer_scope="local").run() + + assert "notasplit is not a valid split" in str(excinfo.value.__cause__) + assert _stored(client, LocalExplainer, explainer_id)["status"] == ( + ExplainerStatus.ERROR + ) + + +def test_a_rows_scope_with_no_valid_index_is_swallowed_by_the_same_wrapper( + client, run_id, dataset_1 +): + explainer_id = _create_local_explainer( + client, + run_id, + dataset_1.id, + scope={"mode": "rows", "row_indexes": [10**9]}, + ) + + with pytest.raises(JobError, match="Can not prepare Dataset with") as excinfo: + ExplainerJob(explainer_id=explainer_id, explainer_scope="local").run() + + assert "No valid row indexes provided for the explanation" in str( + excinfo.value.__cause__ + ) + + +def test_a_manual_scope_with_no_rows_is_swallowed_by_the_same_wrapper( + client, run_id, dataset_1 +): + explainer_id = _create_local_explainer( + client, run_id, dataset_1.id, scope={"mode": "manual"} + ) + + with pytest.raises(JobError, match="Can not prepare Dataset with") as excinfo: + ExplainerJob(explainer_id=explainer_id, explainer_scope="local").run() + + assert "No manual input data provided for the explanation" in str( + excinfo.value.__cause__ + ) + + +def test_a_missing_instance_dataset_is_reported_by_id(client, run_id, dataset_1): + explainer_id = _create_local_explainer(client, run_id, dataset_1.id) + + session_factory = client.app.container["session_factory"] + with session_factory() as db: + db.get(LocalExplainer, explainer_id).dataset_id = 999999 + db.commit() + + with pytest.raises( + JobError, match="Dataset 999999 to be explained does not exist in DB." + ): + ExplainerJob(explainer_id=explainer_id, explainer_scope="local").run() + + assert _stored(client, LocalExplainer, explainer_id)["status"] == ( + ExplainerStatus.ERROR + ) + + +# --- delivery ----------------------------------------------------------- + + +def test_set_status_as_delivered_marks_the_right_row(client, run_id): + explainer_id = _create_global_explainer(client, run_id) + + ExplainerJob( + explainer_id=explainer_id, explainer_scope="global" + ).set_status_as_delivered() + + assert _stored(client, GlobalExplainer, explainer_id)["status"] == ( + ExplainerStatus.DELIVERED + ) + + +def test_set_status_as_delivered_rejects_an_invalid_scope(client, run_id): + explainer_id = _create_global_explainer(client, run_id) + + with pytest.raises(JobError, match="banana is an invalid explainer type"): + ExplainerJob( + explainer_id=explainer_id, explainer_scope="banana" + ).set_status_as_delivered() diff --git a/tests/back/api/test_explorer_job.py b/tests/back/api/test_explorer_job.py new file mode 100644 index 000000000..b591a06b3 --- /dev/null +++ b/tests/back/api/test_explorer_job.py @@ -0,0 +1,277 @@ +"""End-to-end regression net for ``ExplorerJob``. + +Written before the job is decomposed into atomic units, and asserted against the +monolithic implementation, so that the refactor has something to be measured +against. The assertions are deliberately explicit — exact status values, exact +files on disk, exact error message fragments — instead of the looser +``status in ["finished", "error"]`` style used elsewhere in this suite, which +cannot tell a unit that silently stopped doing part of its work from one that +did it. + +``ExplorerJob`` had no tests at all before this file. + +Tests named ``test_currently_*`` pin behaviour that is known to be wrong. They +exist so the refactor can be proven behaviour-preserving first; the fix lands +afterwards as its own change, which flips the assertion and renames the test. + +Lives under ``tests/back/api`` to reuse the ``client`` and ``dataset_1`` +fixtures from this package's ``conftest.py``. +""" + +import pathlib +import shutil + +import pytest +from fastapi.testclient import TestClient + +from DashAI.back.core.enums.status import ExplorerStatus +from DashAI.back.dependencies.database.models import Explorer +from DashAI.back.job.base_job import JobError +from DashAI.back.job.explorer_job import ExplorerJob + +#: ``DescribeExplorer.__init__`` reads these three keys unconditionally, and the +#: schema declares no defaults, so a valid configuration always carries all of +#: them. +DESCRIBE_PARAMETERS = {"percentiles": "25, 50, 75", "include": "all", "exclude": None} + +SEPAL_LENGTH = [{"columnName": "SepalLengthCm"}] + + +@pytest.fixture(name="notebook") +def create_notebook(client: TestClient, dataset_1): + """A notebook holding its own copy of the iris dataset. + + ``POST /notebook/`` copies the dataset folder, so an exploration always + reads the notebook's copy and never the source dataset. + """ + response = client.post( + "/api/v1/notebook/", + json={"dataset_id": dataset_1.id, "name": "explorer job test"}, + ) + assert response.status_code == 201, response.text + return response.json() + + +def _create_explorer(client, notebook_id, columns=None, parameters=None): + """Create an Explorer row through the API, which validates it.""" + response = client.post( + "/api/v1/explorer/", + json={ + "notebook_id": notebook_id, + "exploration_type": "DescribeExplorer", + "columns": columns if columns is not None else SEPAL_LENGTH, + "parameters": ( + parameters if parameters is not None else DESCRIBE_PARAMETERS + ), + }, + ) + assert response.status_code == 201, response.text + return response.json()["id"] + + +def _insert_explorer_row(client, **fields): + """Insert an Explorer row straight into the database. + + ``POST /explorer/`` validates the exploration type, the parameters and the + columns, so the branches of ``run()`` that react to an invalid row can only + be reached by writing the row directly. + """ + session_factory = client.app.container["session_factory"] + with session_factory() as db: + explorer = Explorer(**fields) + db.add(explorer) + db.commit() + db.refresh(explorer) + return explorer.id + + +def _stored_explorer(client, explorer_id): + """Read the Explorer row straight from the database.""" + session_factory = client.app.container["session_factory"] + with session_factory() as db: + explorer = db.get(Explorer, explorer_id) + return { + "status": explorer.status, + "exploration_path": explorer.exploration_path, + "start_time": explorer.start_time, + "end_time": explorer.end_time, + } + + +def _notebook_folder(client, notebook_id): + return pathlib.Path(client.app.container["config"]["NOTEBOOK_PATH"]) / str( + notebook_id + ) + + +def test_the_notebook_starts_as_a_readable_copy_of_the_dataset(client, notebook): + """Guards the fixture itself: the assertions below mean nothing if the + notebook copy is not a loadable iris dataset.""" + from DashAI.back.dataloaders.classes.dashai_dataset import load_dataset + + dataset = load_dataset(f"{notebook['file_path']}/dataset") + + assert "SepalLengthCm" in dataset.column_names + assert len(dataset) == 150 + + +def test_explorer_job_writes_the_result_and_finishes(client, notebook): + """The happy path, end to end: status transitions and the file on disk. + + ``DescribeExplorer`` writes ``{explorer_id}.json`` under the notebook's own + folder, and the row records that exact path. + """ + explorer_id = _create_explorer(client, notebook["id"]) + + ExplorerJob(explorer_id=explorer_id).run() + + stored = _stored_explorer(client, explorer_id) + assert stored["status"] == ExplorerStatus.FINISHED + assert stored["start_time"] is not None + assert stored["end_time"] is not None + + expected = _notebook_folder(client, notebook["id"]) / f"{explorer_id}.json" + assert expected.exists() + assert stored["exploration_path"] == expected.as_posix() + + +def test_two_explorations_on_one_notebook_keep_separate_files(client, notebook): + """The save path is keyed by explorer id, so runs never overwrite each + other's result. Any decomposition has to keep that key.""" + first = _create_explorer(client, notebook["id"]) + second = _create_explorer( + client, notebook["id"], columns=[{"columnName": "PetalWidthCm"}] + ) + + ExplorerJob(explorer_id=first).run() + ExplorerJob(explorer_id=second).run() + + first_path = _stored_explorer(client, first)["exploration_path"] + second_path = _stored_explorer(client, second)["exploration_path"] + + assert first_path != second_path + assert pathlib.Path(first_path).exists() + assert pathlib.Path(second_path).exists() + + +def test_a_missing_explorer_row_reports_it_by_id(client): + with pytest.raises(JobError, match="Explorer with id 999999 not found."): + ExplorerJob(explorer_id=999999).run() + + +def test_a_missing_notebook_leaves_the_row_in_error(client, notebook): + """A notebook that is gone must not leave the exploration stuck in STARTED. + + The "not found" error used to be raised inside a ``try`` whose only handler + was ``except exc.SQLAlchemyError``, so ``set_status_as_error`` never ran and + the UI showed the exploration as still running. Nothing else would have + fixed it: the Huey error signal writes only to its own ``task_copy`` table + and never touches the ``Explorer`` row, and ``_execute_base_job`` calls + ``job.run()`` with no handler at all. + + The message still has to be the specific one, not a generic wrapper. + """ + explorer_id = _create_explorer(client, notebook["id"]) + + session_factory = client.app.container["session_factory"] + with session_factory() as db: + explorer = db.get(Explorer, explorer_id) + explorer.notebook_id = 999999 + db.commit() + + with pytest.raises(JobError, match="Notebook with id 999999 not found."): + ExplorerJob(explorer_id=explorer_id).run() + + assert _stored_explorer(client, explorer_id)["status"] == ExplorerStatus.ERROR + + +def test_a_dataset_that_cannot_be_loaded_leaves_the_row_in_error(client, notebook): + """A load failure must not leave the explorer stuck in STARTED.""" + explorer_id = _create_explorer(client, notebook["id"]) + shutil.rmtree(f"{notebook['file_path']}/dataset") + + with pytest.raises(JobError, match="Can not load dataset from path"): + ExplorerJob(explorer_id=explorer_id).run() + + assert _stored_explorer(client, explorer_id)["status"] == ExplorerStatus.ERROR + + +def test_an_unknown_exploration_type_reports_it_and_errors(client, notebook): + """The registry lookup error names the culprit and reaches the user intact. + + Unlike ``ConverterJob``, nothing wraps this message on the way out. + """ + explorer_id = _insert_explorer_row( + client, + notebook_id=notebook["id"], + exploration_type="ThisExplorerDoesNotExist", + columns=SEPAL_LENGTH, + parameters={}, + ) + + with pytest.raises( + JobError, + match="Explorer ThisExplorerDoesNotExist not found in the registry.", + ): + ExplorerJob(explorer_id=explorer_id).run() + + assert _stored_explorer(client, explorer_id)["status"] == ExplorerStatus.ERROR + + +def test_parameters_the_explorer_cannot_accept_report_an_instancing_error( + client, notebook +): + """``DescribeExplorer.__init__`` reads its three keys unconditionally.""" + explorer_id = _insert_explorer_row( + client, + notebook_id=notebook["id"], + exploration_type="DescribeExplorer", + columns=SEPAL_LENGTH, + parameters={}, + ) + + with pytest.raises( + JobError, match="Error instancing the explorer DescribeExplorer." + ): + ExplorerJob(explorer_id=explorer_id).run() + + assert _stored_explorer(client, explorer_id)["status"] == ExplorerStatus.ERROR + + +def test_a_column_absent_from_the_dataset_reports_a_preparation_error(client, notebook): + """``prepare_dataset`` selects the requested columns and fails loudly. + + The column list is resolved against the dataset the job just loaded, which + is the behaviour any decomposition has to keep: nothing may resolve these + names ahead of time against a different dataset. + """ + explorer_id = _insert_explorer_row( + client, + notebook_id=notebook["id"], + exploration_type="DescribeExplorer", + columns=[{"columnName": "ThisColumnDoesNotExist"}], + parameters=DESCRIBE_PARAMETERS, + ) + + with pytest.raises( + JobError, + match="Error preparing the dataset for the exploration DescribeExplorer.", + ): + ExplorerJob(explorer_id=explorer_id).run() + + assert _stored_explorer(client, explorer_id)["status"] == ExplorerStatus.ERROR + + +def test_set_status_as_delivered_marks_the_row(client, notebook): + """The enqueue path marks the row before the worker ever picks it up.""" + explorer_id = _create_explorer(client, notebook["id"]) + + ExplorerJob(explorer_id=explorer_id).set_status_as_delivered() + + stored = _stored_explorer(client, explorer_id) + assert stored["status"] == ExplorerStatus.DELIVERED + + +def test_set_status_as_delivered_reports_a_missing_row(client): + with pytest.raises(JobError, match="Explorer with id 999999 not found."): + ExplorerJob(explorer_id=999999).set_status_as_delivered() diff --git a/tests/back/api/test_predict_job.py b/tests/back/api/test_predict_job.py new file mode 100644 index 000000000..dfafd836b --- /dev/null +++ b/tests/back/api/test_predict_job.py @@ -0,0 +1,513 @@ +"""End-to-end regression net for ``PredictJob``. + +Written before the job is decomposed into atomic units, and asserted against the +monolithic implementation, so that the refactor has something to be measured +against. The assertions are deliberately explicit — exact status values, exact +columns on disk, exact error message fragments — instead of the looser +``status in ["finished", "error"]`` style used elsewhere in this suite, which +cannot tell a unit that silently stopped doing part of its work from one that +did it. + +``test_predict_api.py`` does not cover this: it enqueues the job and then only +exercises the CRUD endpoints, never checking the job's outcome nor the +``Prediction`` row. + +Tests named ``test_currently_*`` pin behaviour that is known to be wrong. They +exist so the refactor can be proven behaviour-preserving first; the fix lands +afterwards as its own change, which flips the assertion and renames the test. + +Lives under ``tests/back/api`` to reuse the ``client`` and ``dataset_1`` +fixtures from this package's ``conftest.py``. +""" + +import json +import shutil +from pathlib import Path + +import pytest +from fastapi.exceptions import HTTPException +from fastapi.testclient import TestClient + +from DashAI.back.core.enums.status import PredictionStatus +from DashAI.back.dataloaders.classes.dashai_dataset import load_dataset +from DashAI.back.dependencies.database.models import ( + Dataset, + ModelSession, + Prediction, + Run, +) +from DashAI.back.job.base_job import JobError +from DashAI.back.job.model_job import ModelJob +from DashAI.back.job.predict_job import PredictJob + +INPUT_COLUMNS = [ + "SepalLengthCm", + "SepalWidthCm", + "PetalLengthCm", + "PetalWidthCm", +] +OUTPUT_COLUMN = "Species" +IRIS_ROWS = 150 + +SPLITS = json.dumps( + { + "train": 0.5, + "test": 0.2, + "validation": 0.3, + "is_random": True, + "has_changed": True, + "seed": 42, + "shuffle": True, + "stratify": False, + } +) + + +@pytest.fixture(scope="module", name="model_session_id") +def create_model_session(client: TestClient, dataset_1: Dataset): + session_factory = client.app.container["session_factory"] + + with session_factory() as db: + model_session = ModelSession( + dataset_id=dataset_1.id, + name="PredictJobSession", + task_name="TabularClassificationTask", + input_columns=INPUT_COLUMNS, + output_columns=[OUTPUT_COLUMN], + train_metrics=[], + validation_metrics=[], + test_metrics=[], + splits=SPLITS, + ) + db.add(model_session) + db.commit() + db.refresh(model_session) + return model_session.id + + +@pytest.fixture(scope="module", name="trained_run_id") +def create_trained_run(client: TestClient, model_session_id: int): + """A genuinely trained run: the prediction path loads the model from disk.""" + session_factory = client.app.container["session_factory"] + + with session_factory() as db: + run = Run( + model_session_id=model_session_id, + optimizer_name="OptunaOptimizer", + optimizer_parameters={ + "n_trials": 1, + "sampler": "TPESampler", + "pruner": "None", + }, + model_name="KNeighborsClassifier", + parameters={}, + name="PredictJobRun", + goal_metric="Accuracy", + ) + db.add(run) + db.commit() + db.refresh(run) + run_id = run.id + + ModelJob(run_id=run_id).run() + + with session_factory() as db: + run = db.get(Run, run_id) + assert run.run_path, "the run fixture did not produce a saved model" + return run_id + + +def _create_prediction(client, run_id, dataset_id=None): + response = client.post( + "/api/v1/predict/", + json={"run_id": run_id, "dataset_id": dataset_id}, + ) + assert response.status_code == 200, response.text + return response.json()["id"] + + +def _make_prediction_dataset(client, dataset_1: Dataset): + """A throwaway copy of the iris dataset, safe for a test to destroy.""" + import uuid + + config = client.app.container["config"] + session_factory = client.app.container["session_factory"] + + folder = Path(config["DATASETS_PATH"]) / f"predict-job-{uuid.uuid4()}" + shutil.copytree(Path(dataset_1.file_path), folder) + + with session_factory() as db: + row = Dataset(name=folder.name, file_path=str(folder)) + db.add(row) + db.commit() + db.refresh(row) + db.expunge(row) + return row + + +def _stored_prediction(client, prediction_id): + session_factory = client.app.container["session_factory"] + with session_factory() as db: + prediction = db.get(Prediction, prediction_id) + return { + "status": prediction.status, + "results_path": prediction.results_path, + "start_time": prediction.start_time, + "end_time": prediction.end_time, + } + + +@pytest.fixture(name="restore_run") +def fixture_restore_run(client: TestClient, trained_run_id: int): + """Let a test corrupt the module-scoped Run row and put it back after. + + The run and the model session are module scoped because training is slow; + without this the error-branch tests would poison every test after them. + """ + session_factory = client.app.container["session_factory"] + with session_factory() as db: + run = db.get(Run, trained_run_id) + original = {"model_name": run.model_name, "run_path": run.run_path} + + yield + + with session_factory() as db: + run = db.get(Run, trained_run_id) + run.model_name = original["model_name"] + run.run_path = original["run_path"] + db.commit() + + +@pytest.fixture(name="restore_model_session") +def fixture_restore_model_session(client: TestClient, model_session_id: int): + """Same idea for the module-scoped ModelSession row.""" + session_factory = client.app.container["session_factory"] + with session_factory() as db: + session_row = db.get(ModelSession, model_session_id) + original = { + "dataset_id": session_row.dataset_id, + "task_name": session_row.task_name, + "input_columns": list(session_row.input_columns), + "output_columns": list(session_row.output_columns), + } + + yield + + with session_factory() as db: + session_row = db.get(ModelSession, model_session_id) + for key, value in original.items(): + setattr(session_row, key, value) + db.commit() + + +def test_predict_job_writes_the_predictions_and_finishes( + client, trained_run_id, dataset_1 +): + """The happy path, end to end: status transitions and the dataset on disk. + + The saved dataset carries the input columns plus the predicted output + column, one row per row of the input. + """ + prediction_id = _create_prediction(client, trained_run_id, dataset_1.id) + + PredictJob(prediction_id=prediction_id).run() + + stored = _stored_prediction(client, prediction_id) + assert stored["status"] == PredictionStatus.FINISHED + assert stored["start_time"] is not None + assert stored["end_time"] is not None + assert stored["results_path"] is not None + + saved = load_dataset(str(Path(stored["results_path"]) / "dataset")) + assert saved.column_names == INPUT_COLUMNS + [OUTPUT_COLUMN] + assert len(saved) == IRIS_ROWS + + +def test_each_prediction_gets_its_own_results_folder(client, trained_run_id, dataset_1): + """The destination is a fresh uuid folder, so two runs never collide.""" + first = _create_prediction(client, trained_run_id, dataset_1.id) + second = _create_prediction(client, trained_run_id, dataset_1.id) + + PredictJob(prediction_id=first).run() + PredictJob(prediction_id=second).run() + + first_path = _stored_prediction(client, first)["results_path"] + second_path = _stored_prediction(client, second)["results_path"] + + assert first_path != second_path + assert Path(first_path).exists() + assert Path(second_path).exists() + + +def test_manual_input_predicts_without_a_dataset(client, trained_run_id): + """The manual branch builds the instances from typed values instead of disk.""" + prediction_id = _create_prediction(client, trained_run_id, dataset_id=None) + + PredictJob( + prediction_id=prediction_id, + manual_input_data=[ + { + "SepalLengthCm": 5.1, + "SepalWidthCm": 3.5, + "PetalLengthCm": 1.4, + "PetalWidthCm": 0.2, + } + ], + ).run() + + stored = _stored_prediction(client, prediction_id) + assert stored["status"] == PredictionStatus.FINISHED + + saved = load_dataset(str(Path(stored["results_path"]) / "dataset")) + assert saved.column_names == INPUT_COLUMNS + [OUTPUT_COLUMN] + assert len(saved) == 1 + + +def test_neither_a_dataset_nor_manual_input_is_rejected(client, trained_run_id): + prediction_id = _create_prediction(client, trained_run_id, dataset_id=None) + + with pytest.raises( + JobError, match="Either dataset_id or manual_input_data must be provided." + ): + PredictJob(prediction_id=prediction_id).run() + + assert _stored_prediction(client, prediction_id)["status"] == PredictionStatus.ERROR + + +def test_a_missing_prediction_row_is_a_404(client): + with pytest.raises(HTTPException) as excinfo: + PredictJob(prediction_id=999999).run() + + assert excinfo.value.status_code == 404 + assert excinfo.value.detail == "Prediction not found for id 999999" + + +def test_an_unknown_model_name_reports_it_and_errors( + client, trained_run_id, dataset_1, restore_run +): + prediction_id = _create_prediction(client, trained_run_id, dataset_1.id) + + session_factory = client.app.container["session_factory"] + with session_factory() as db: + db.get(Run, trained_run_id).model_name = "ThisModelDoesNotExist" + db.commit() + + with pytest.raises( + JobError, match="Model ThisModelDoesNotExist not found in the registry" + ): + PredictJob(prediction_id=prediction_id).run() + + assert _stored_prediction(client, prediction_id)["status"] == PredictionStatus.ERROR + + +def test_a_model_that_cannot_be_loaded_reports_the_path_and_errors( + client, trained_run_id, dataset_1, restore_run +): + prediction_id = _create_prediction(client, trained_run_id, dataset_1.id) + + session_factory = client.app.container["session_factory"] + with session_factory() as db: + db.get(Run, trained_run_id).run_path = "nowhere/at/all" + db.commit() + + with pytest.raises( + JobError, + match="Failed to load model KNeighborsClassifier from path nowhere/at/all", + ): + PredictJob(prediction_id=prediction_id).run() + + assert _stored_prediction(client, prediction_id)["status"] == PredictionStatus.ERROR + + +def test_an_unknown_task_name_reports_it_and_errors( + client, trained_run_id, dataset_1, restore_model_session, model_session_id +): + """The task is resolved before the model, so this is the first error seen.""" + prediction_id = _create_prediction(client, trained_run_id, dataset_1.id) + + session_factory = client.app.container["session_factory"] + with session_factory() as db: + db.get(ModelSession, model_session_id).task_name = "ThisTaskDoesNotExist" + db.commit() + + with pytest.raises( + JobError, match="Task ThisTaskDoesNotExist not found in the registry" + ): + PredictJob(prediction_id=prediction_id).run() + + assert _stored_prediction(client, prediction_id)["status"] == PredictionStatus.ERROR + + +def test_a_model_session_without_input_columns_is_a_422( + client, trained_run_id, dataset_1, restore_model_session, model_session_id +): + prediction_id = _create_prediction(client, trained_run_id, dataset_1.id) + + session_factory = client.app.container["session_factory"] + with session_factory() as db: + db.get(ModelSession, model_session_id).input_columns = [] + db.commit() + + with pytest.raises(HTTPException) as excinfo: + PredictJob(prediction_id=prediction_id).run() + + assert excinfo.value.status_code == 422 + assert excinfo.value.detail == "Model session has no input columns configured" + assert _stored_prediction(client, prediction_id)["status"] == PredictionStatus.ERROR + + +def test_a_missing_training_dataset_row_leaves_the_row_in_error( + client, trained_run_id, dataset_1, restore_model_session, model_session_id +): + """The 404 must also mark the prediction as failed. + + This branch used to skip ``set_status_as_error``, unlike every one around + it, so the row stayed STARTED forever — nothing else marks it, because the + Huey error signal only writes to its own ``task_copy`` table and + ``_execute_base_job`` calls ``run()`` with no handler. + """ + prediction_id = _create_prediction(client, trained_run_id, dataset_1.id) + + session_factory = client.app.container["session_factory"] + with session_factory() as db: + db.get(ModelSession, model_session_id).dataset_id = 999999 + db.commit() + + with pytest.raises(HTTPException) as excinfo: + PredictJob(prediction_id=prediction_id).run() + + assert excinfo.value.status_code == 404 + assert excinfo.value.detail == "Training dataset not found" + assert _stored_prediction(client, prediction_id)["status"] == PredictionStatus.ERROR + + +def test_an_unreadable_training_dataset_leaves_the_row_in_error( + client, trained_run_id, dataset_1, tmp_path +): + """Same omission as above, on the branch that reads the training dataset. + + The message still has to be the specific one, not the generic prediction + wrapper: this load happens before the dataset to predict on is even + touched, and that ordering is what the message depends on. + """ + prediction_id = _create_prediction(client, trained_run_id, dataset_1.id) + + stored_folder = Path(dataset_1.file_path) / "dataset" + backup = tmp_path / "training-dataset-backup" + shutil.copytree(stored_folder, backup) + shutil.rmtree(stored_folder) + try: + with pytest.raises(JobError, match="Cannot load training dataset from"): + PredictJob(prediction_id=prediction_id).run() + + assert ( + _stored_prediction(client, prediction_id)["status"] + == PredictionStatus.ERROR + ) + finally: + shutil.copytree(backup, stored_folder) + + +def test_an_unreadable_prediction_dataset_is_reported_as_a_prediction_failure( + client, trained_run_id, dataset_1, tmp_path +): + """Loading the dataset to predict on shares the "prediction failed" wrapper. + + Unlike the *training* dataset, which has a message of its own, the + inference dataset is loaded inside the same ``try`` as the prediction, so a + read failure surfaces as the generic message with the real cause attached. + Pinned because it is the exact spot a unit boundary lands on, and an + improved message here would be a silent behaviour change. + """ + prediction_dataset = _make_prediction_dataset(client, dataset_1) + prediction_id = _create_prediction(client, trained_run_id, prediction_dataset.id) + + shutil.rmtree(Path(prediction_dataset.file_path) / "dataset") + + with pytest.raises(JobError, match="Model prediction failed"): + PredictJob(prediction_id=prediction_id).run() + + assert _stored_prediction(client, prediction_id)["status"] == PredictionStatus.ERROR + + +def test_a_prediction_that_blows_up_leaves_the_row_in_error( + client, trained_run_id, dataset_1, monkeypatch +): + """Any unexpected failure while predicting is reported as one message.""" + from DashAI.back.models.scikit_learn.k_neighbors_classifier import ( + KNeighborsClassifier, + ) + + prediction_id = _create_prediction(client, trained_run_id, dataset_1.id) + + def _explode(self, x): + raise RuntimeError("the model itself blew up") + + monkeypatch.setattr(KNeighborsClassifier, "predict", _explode) + + with pytest.raises(JobError, match="Model prediction failed") as excinfo: + PredictJob(prediction_id=prediction_id).run() + + assert "the model itself blew up" in str(excinfo.value.__cause__) + assert _stored_prediction(client, prediction_id)["status"] == PredictionStatus.ERROR + + +def test_a_type_error_while_predicting_leaves_the_row_in_error( + client, trained_run_id, dataset_1, monkeypatch +): + """The ``TypeError`` branch is still a 400, but now it marks the row too. + + Its ``ValueError`` neighbour always did; this one did not, so a type + mismatch left the prediction STARTED forever. + """ + from DashAI.back.models.scikit_learn.k_neighbors_classifier import ( + KNeighborsClassifier, + ) + + prediction_id = _create_prediction(client, trained_run_id, dataset_1.id) + + def _wrong_type(self, x): + raise TypeError("bad type somewhere in the input") + + monkeypatch.setattr(KNeighborsClassifier, "predict", _wrong_type) + + with pytest.raises(HTTPException) as excinfo: + PredictJob(prediction_id=prediction_id).run() + + assert excinfo.value.status_code == 400 + assert "Type validation failed" in excinfo.value.detail + assert _stored_prediction(client, prediction_id)["status"] == PredictionStatus.ERROR + + +def test_a_value_error_while_predicting_leaves_the_row_in_error( + client, trained_run_id, dataset_1, monkeypatch +): + """The ``ValueError`` branch is reported as a 400 and does mark the row.""" + from DashAI.back.models.scikit_learn.k_neighbors_classifier import ( + KNeighborsClassifier, + ) + + prediction_id = _create_prediction(client, trained_run_id, dataset_1.id) + + def _bad_value(self, x): + raise ValueError("a value the model cannot use") + + monkeypatch.setattr(KNeighborsClassifier, "predict", _bad_value) + + with pytest.raises(HTTPException) as excinfo: + PredictJob(prediction_id=prediction_id).run() + + assert excinfo.value.status_code == 400 + assert "Invalid input data" in excinfo.value.detail + assert _stored_prediction(client, prediction_id)["status"] == PredictionStatus.ERROR + + +def test_set_status_as_delivered_marks_the_row(client, trained_run_id, dataset_1): + prediction_id = _create_prediction(client, trained_run_id, dataset_1.id) + + PredictJob(prediction_id=prediction_id).set_status_as_delivered() + + assert ( + _stored_prediction(client, prediction_id)["status"] + == PredictionStatus.DELIVERED + ) diff --git a/tests/back/api/test_units_api.py b/tests/back/api/test_units_api.py index 943662008..666379f18 100644 --- a/tests/back/api/test_units_api.py +++ b/tests/back/api/test_units_api.py @@ -12,6 +12,19 @@ "SaveModelUnit", "ApplyConverterUnit", "SaveDatasetUnit", + "RunExplorationUnit", + "SaveExplorationUnit", + "LoadTrainedModelUnit", + "LoadTrainingDatasetUnit", + "BuildManualInputUnit", + "PredictUnit", + "SavePredictionUnit", + "LoadRunModelUnit", + "BuildGlobalExplainerUnit", + "BuildLocalExplainerUnit", + "PrepareExplanationDataUnit", + "GenerateGlobalExplanationUnit", + "GenerateLocalExplanationUnit", } @@ -57,6 +70,32 @@ def test_unit_schemas_describe_their_configuration(units): } # SaveDatasetUnit is configuration-free: it saves where the load said. assert units["SaveDatasetUnit"]["schema"]["properties"] == {} + assert set(units["RunExplorationUnit"]["schema"]["properties"]) == { + "explorer_id", + "explorer", + } + # SaveExplorationUnit only picks the destination; how the result is + # serialised belongs to the explorer that produced it. + assert set(units["SaveExplorationUnit"]["schema"]["properties"]) == {"explorer_id"} + assert set(units["LoadTrainedModelUnit"]["schema"]["properties"]) == {"run_id"} + assert set(units["PredictUnit"]["schema"]["properties"]) == { + "task_name", + "input_columns", + "output_columns", + } + assert set(units["SavePredictionUnit"]["schema"]["properties"]) == { + "input_columns", + "output_columns", + } + assert set(units["BuildGlobalExplainerUnit"]["schema"]["properties"]) == { + "explainer" + } + assert set(units["BuildLocalExplainerUnit"]["schema"]["properties"]) == { + "explainer" + } + assert set(units["GenerateGlobalExplanationUnit"]["schema"]["properties"]) == { + "explainer_id" + } def test_component_fields_tell_the_front_which_components_to_offer(units): @@ -69,10 +108,23 @@ def test_component_fields_tell_the_front_which_components_to_offer(units): model = units["BuildModelUnit"]["schema"]["properties"]["model"] optimizer = units["FitModelUnit"]["schema"]["properties"]["optimizer"] converter = units["ApplyConverterUnit"]["schema"]["properties"]["converter"] + explorer = units["RunExplorationUnit"]["schema"]["properties"]["explorer"] assert model["parent"] == "BaseModel" assert optimizer["parent"] == "BaseOptimizer" assert converter["parent"] == "BaseConverter" + assert explorer["parent"] == "BaseExplorer" + + # Global and local explainers are separate registries with separate base + # classes, and a component field carries a single parent hint. Hence two + # sibling units with one required field each: making a single field cover + # both scopes would need it to be optional, and an optional component field + # is emitted as an anyOf, which hides the hint from the front — that is what + # the assertions below would catch. + global_explainer = units["BuildGlobalExplainerUnit"]["schema"]["properties"] + local_explainer = units["BuildLocalExplainerUnit"]["schema"]["properties"] + assert global_explainer["explainer"]["parent"] == "BaseGlobalExplainer" + assert local_explainer["explainer"]["parent"] == "BaseLocalExplainer" assert set(model["properties"]) == {"component", "params"} assert set(converter["properties"]) == {"component", "params"} diff --git a/tests/back/explainers/test_shap_predictor_handover.py b/tests/back/explainers/test_shap_predictor_handover.py new file mode 100644 index 000000000..b5bda46ed --- /dev/null +++ b/tests/back/explainers/test_shap_predictor_handover.py @@ -0,0 +1,120 @@ +"""SHAP must not be handed a bound method of the model. + +``shap.utils._legacy.convert_to_model`` suppresses scikit-learn's "X does not +have valid feature names" warning by blanking ``feature_names_in_`` on the +object the callable is bound to, reached through ``__self__``. It assumes that +attribute is writable. + +Two of the models DashAI ships inherit ``feature_names_in_`` from their upstream +estimator as a read-only ``property``, so that assignment raises and the +explanation dies before it starts: + + AttributeError: property 'feature_names_in_' of 'LGBMClassifier' object has + no setter + +``as_shap_predictor`` hands over a plain closure instead, which has no +``__self__``, so SHAP skips the step. These tests pin both halves: that the +wrappers really are read-only (otherwise the fix guards nothing), and that the +handover survives ``convert_to_model``. +""" + +import numpy as np +import pandas as pd +import pytest + +from DashAI.back.explainability.model_input import as_shap_predictor +from DashAI.back.models.scikit_learn.lightgbm_classifier import LGBMClassifier +from DashAI.back.models.scikit_learn.xgboost_classifier import XGBClassifier + +#: The models whose ``feature_names_in_`` cannot be assigned to. Every other +#: model stores it as a plain instance attribute, which is settable. +READ_ONLY_FEATURE_NAMES = [LGBMClassifier, XGBClassifier] + + +@pytest.fixture(name="frame") +def fixture_frame(): + rng = np.random.default_rng(0) + return pd.DataFrame({"a": rng.random(40), "b": rng.random(40)}), rng.integers( + 0, 2, 40 + ) + + +@pytest.mark.parametrize( + "model_class", READ_ONLY_FEATURE_NAMES, ids=lambda c: c.__name__ +) +def test_these_models_really_do_expose_feature_names_read_only(model_class, frame): + """Guards the premise: without this the tests below prove nothing. + + If an upstream release ever makes the attribute writable, this fails and the + workaround can be reconsidered. + """ + x, y = frame + model = model_class() + model.fit(x, y) + + assert hasattr(model, "feature_names_in_") + with pytest.raises(AttributeError, match="no setter"): + model.feature_names_in_ = None + + +@pytest.mark.parametrize( + "model_class", READ_ONLY_FEATURE_NAMES, ids=lambda c: c.__name__ +) +def test_a_bound_predict_breaks_shaps_model_conversion(model_class, frame): + """The failure this exists to prevent, reproduced directly. + + Pinned so the regression is recognisable if anyone reverts the handover to + ``model=self.model.predict``. + """ + from shap.utils._legacy import convert_to_model + + x, y = frame + model = model_class() + model.fit(x, y) + + with pytest.raises(AttributeError, match="feature_names_in_"): + convert_to_model(model.predict) + + +@pytest.mark.parametrize( + "model_class", READ_ONLY_FEATURE_NAMES, ids=lambda c: c.__name__ +) +def test_the_wrapped_predictor_survives_shaps_model_conversion(model_class, frame): + from shap.utils._legacy import convert_to_model + + x, y = frame + model = model_class() + model.fit(x, y) + + converted = convert_to_model(as_shap_predictor(model)) + + assert converted.f is not None + # The model itself must be left alone: SHAP deep-copies before blanking the + # attribute, but only on the branch we now skip. + assert list(model.feature_names_in_) == ["a", "b"] + + +def test_the_wrapped_predictor_forwards_to_predict_positionally(): + """SHAP calls the model with one positional argument; that must not change.""" + seen = {} + + class Model: + def predict(self, x): + seen["arg"] = x + return [0] + + predictor = as_shap_predictor(Model()) + assert predictor("the frame") == [0] + assert seen["arg"] == "the frame" + + +def test_the_wrapped_predictor_hides_the_model_from_shap(): + """The whole mechanism: no ``__self__`` means SHAP never reaches the model.""" + + class Model: + def predict(self, x): + return [0] + + model = Model() + assert getattr(model.predict, "__self__", None) is model + assert getattr(as_shap_predictor(model), "__self__", None) is None diff --git a/tests/back/units/test_explanation_units.py b/tests/back/units/test_explanation_units.py new file mode 100644 index 000000000..8d99124e1 --- /dev/null +++ b/tests/back/units/test_explanation_units.py @@ -0,0 +1,663 @@ +"""Contract tests for the explanation units, isolated from any orchestrating job. + +The context is built by hand rather than through a job, which is what exposes +composability mistakes: a job always wires the context "correctly", so an +end-to-end run cannot tell a real contract from a lucky one. +""" + +import json +import pickle +from pathlib import Path + +import pytest +from kink import di + +from DashAI.back.dataloaders.classes.dashai_dataset import ( + load_dataset, + save_dataset, + to_dashai_dataset, +) +from DashAI.back.job.base_job import JobError +from DashAI.back.units.build_global_explainer_unit import BuildGlobalExplainerUnit +from DashAI.back.units.build_local_explainer_unit import BuildLocalExplainerUnit +from DashAI.back.units.context import ExecutionContext, UnitContractError +from DashAI.back.units.generate_global_explanation_unit import ( + GenerateGlobalExplanationUnit, +) +from DashAI.back.units.generate_local_explanation_unit import ( + GenerateLocalExplanationUnit, +) +from DashAI.back.units.load_run_model_unit import LoadRunModelUnit +from DashAI.back.units.prepare_explanation_data_unit import PrepareExplanationDataUnit + +SPLITS = { + "train_indexes": [0, 1, 2], + "test_indexes": [3, 4], + "val_indexes": [5], +} + + +class _RunRow: + def __init__( + self, model_name="RecordingModel", run_path="somewhere", parameters=None + ): + self.id = 5 + self.model_name = model_name + self.run_path = run_path + self.parameters = parameters if parameters is not None else {"depth": 3} + + +class _DatasetRow: + def __init__(self, file_path): + self.file_path = file_path + + +class _FakeSession: + def __init__(self, rows): + self._rows = rows + + def get(self, model, row_id): + return self._rows.get(model.__name__, {}).get(row_id) + + def __enter__(self): + return self + + def __exit__(self, *exc_info): + return False + + +class _FakeSessionFactory: + """Stand-in for a ``sessionmaker``. + + A class rather than a lambda on purpose: kink invokes any registered lambda + with the container to resolve it. + """ + + def __init__(self, rows): + self._rows = rows + + def __call__(self): + return _FakeSession(self._rows) + + +class RecordingModel: + """Model that records how it was built and what it was asked to encode.""" + + built_with = None + + def __init__(self, **kwargs): + RecordingModel.built_with = kwargs + + @staticmethod + def load(filename): + # Bypasses __init__ on purpose: this is what every real model does — + # joblib or a checkpoint rebuilds the object, and the instance the unit + # constructed beforehand is thrown away. Going through __init__ here + # would overwrite the record of how that instance was built. + model = object.__new__(RecordingModel) + model.loaded_from = filename + return model + + def prepare_output(self, dataset, is_fit=False): + return dataset + + +class UninstantiableModel(RecordingModel): + def __init__(self, **kwargs): + raise RuntimeError("this model refuses to be built") + + +class UnloadableModel(RecordingModel): + @staticmethod + def load(filename): + raise OSError("the artifact is not there") + + +class RecordingTask: + def prepare_for_task(self, dataset, input_columns=None, output_columns=None): + return dataset + + def process_manual_input(self, rows, dataset_path): + import pandas as pd + import pyarrow as pa + + from DashAI.back.types.value_types import Integer + + frame = pd.DataFrame(rows) + types = {name: Integer(arrow_type=pa.int64()) for name in frame.columns} + return to_dashai_dataset(frame, types=types) + + +class RecordingGlobalExplainer: + def __init__(self, model, **kwargs): + self.model = model + self.kwargs = kwargs + self.seen = None + + def explain(self, dataset): + self.seen = dataset + return {"importance": [1, 2]} + + def plot(self, explanation): + return "a plot" + + +class ExplodingGlobalExplainer(RecordingGlobalExplainer): + def explain(self, dataset): + raise RuntimeError("the explanation itself blew up") + + +class RecordingLocalExplainer: + fitted_with = None + + def __init__(self, model, **kwargs): + self.model = model + self.explained_columns = None + + def fit(self, dataset, **kwargs): + RecordingLocalExplainer.fitted_with = kwargs + return self + + def explain_instance(self, instances): + columns = instances.column_names + if isinstance(columns, dict): + columns = [c for split in columns.values() for c in split] + self.explained_columns = columns + return {"local": True} + + def plot(self, explanation): + return "a plot" + + +@pytest.fixture(name="registry") +def fixture_registry(): + registry = { + "RecordingModel": {"class": RecordingModel}, + "UninstantiableModel": {"class": UninstantiableModel}, + "UnloadableModel": {"class": UnloadableModel}, + "RecordingTask": {"class": RecordingTask}, + "RecordingGlobalExplainer": {"class": RecordingGlobalExplainer}, + "ExplodingGlobalExplainer": {"class": ExplodingGlobalExplainer}, + "RecordingLocalExplainer": {"class": RecordingLocalExplainer}, + } + di["component_registry"] = registry + yield registry + del di["component_registry"] + + +def _dataset(rows=6): + import pandas as pd + import pyarrow as pa + + from DashAI.back.types.value_types import Integer + + frame = pd.DataFrame( + {"a": list(range(rows)), "b": list(range(rows)), "target": [0, 1] * (rows // 2)} + ) + types = {name: Integer(arrow_type=pa.int64()) for name in frame.columns} + return to_dashai_dataset(frame, types=types) + + +@pytest.fixture(name="stored_instances") +def fixture_stored_instances(tmp_path): + root = tmp_path / "instances" + save_dataset(_dataset(), str(root / "dataset")) + return root + + +@pytest.fixture(name="fake_db") +def fixture_fake_db(stored_instances): + rows = { + "Run": {5: _RunRow()}, + "Dataset": {9: _DatasetRow(str(stored_instances))}, + } + di["session_factory"] = _FakeSessionFactory(rows) + yield rows + del di["session_factory"] + + +@pytest.fixture(name="explanations_path") +def fixture_explanations_path(tmp_path): + target = tmp_path / "explanations" + target.mkdir() + di["config"] = {"EXPLANATIONS_PATH": target} + yield target + del di["config"] + + +# --- LoadRunModelUnit --------------------------------------------------- + + +def test_the_run_model_is_built_with_its_parameters_before_loading(registry, fake_db): + """The extra construction step is the difference from ``LoadTrainedModelUnit``. + + It has no effect for models whose ``load`` is a static or class method — all + of the real ones — but it is the inherited behaviour of this flow, and this + test is what would notice if the two units were quietly merged. + """ + ctx = ExecutionContext() + RecordingModel.built_with = None + + LoadRunModelUnit(run_id=5)(ctx) + + assert RecordingModel.built_with == {"depth": 3} + assert ctx.require("model").loaded_from == "somewhere" + + +def test_a_missing_run_is_reported_by_id(registry, fake_db): + with pytest.raises(JobError, match="Run 99 does not exist in DB."): + LoadRunModelUnit(run_id=99)(ExecutionContext()) + + +def test_an_unknown_model_name_uses_the_explanation_wording(registry, fake_db): + """Word for word different from ``LoadTrainedModelUnit``'s message, which is + why the two units are not merged.""" + fake_db["Run"][5].model_name = "NoSuchModel" + + with pytest.raises( + JobError, match="Unable to find Model with name NoSuchModel in registry." + ): + LoadRunModelUnit(run_id=5)(ExecutionContext()) + + +def test_a_model_that_cannot_be_built_is_reported_separately_from_loading( + registry, fake_db +): + fake_db["Run"][5].model_name = "UninstantiableModel" + + with pytest.raises(JobError, match="Unable to instantiate model") as excinfo: + LoadRunModelUnit(run_id=5)(ExecutionContext()) + + assert "refuses to be built" in str(excinfo.value.__cause__) + + +def test_a_model_that_cannot_be_loaded_names_the_path(registry, fake_db): + fake_db["Run"][5].model_name = "UnloadableModel" + fake_db["Run"][5].run_path = "gone" + + with pytest.raises(JobError, match="Can not load model from path gone"): + LoadRunModelUnit(run_id=5)(ExecutionContext()) + + +# --- the two build units ------------------------------------------------ + + +def test_the_global_build_unit_binds_the_model_from_the_context(registry): + ctx = ExecutionContext() + model = RecordingModel() + ctx.put("model", model) + + BuildGlobalExplainerUnit( + explainer={"component": "RecordingGlobalExplainer", "params": {"n": 5}} + )(ctx) + + explainer = ctx.require("explainer") + assert explainer.model is model + assert explainer.kwargs == {"n": 5} + + +def test_the_local_build_unit_produces_the_same_context_key(registry): + """Both scopes publish ``explainer``, so whatever generates the explanation + afterwards does not have to know which one ran.""" + ctx = ExecutionContext() + ctx.put("model", RecordingModel()) + + BuildLocalExplainerUnit( + explainer={"component": "RecordingLocalExplainer", "params": {}} + )(ctx) + + assert isinstance(ctx.require("explainer"), RecordingLocalExplainer) + + +def test_building_without_a_model_is_rejected_before_it_starts(registry): + with pytest.raises(UnitContractError, match="Context key 'model'"): + BuildGlobalExplainerUnit( + explainer={"component": "RecordingGlobalExplainer", "params": {}} + )(ExecutionContext()) + + +def test_each_build_unit_names_its_own_scope_in_its_errors(registry): + """The messages are worded per scope and are user-visible.""" + ctx = ExecutionContext() + ctx.put("model", RecordingModel()) + + with pytest.raises(JobError, match="Unable to find the global explainer with name"): + BuildGlobalExplainerUnit( + explainer={"component": "NoSuchExplainer", "params": {}} + )(ctx) + + with pytest.raises(JobError, match="Unable to find the local explainer with name"): + BuildLocalExplainerUnit( + explainer={"component": "NoSuchExplainer", "params": {}} + )(ctx) + + +# --- PrepareExplanationDataUnit ----------------------------------------- + + +def _prepared_context(registry): + ctx = ExecutionContext() + ctx.put("dataset", _dataset()) + ctx.put("model", RecordingModel()) + ctx.put_ref("dataset_id", 9) + ctx.put_ref("split_indexes", SPLITS) + return ctx + + +def _prepare_unit(**overrides): + config = { + "task_name": "RecordingTask", + "input_columns": ["a", "b"], + "output_columns": ["target"], + } + config.update(overrides) + return PrepareExplanationDataUnit(**config) + + +def test_prepare_replays_the_recorded_split(registry): + """The indexes come from the run, not from a ratio: the explanation has to + be about the rows the model actually saw.""" + ctx = _prepared_context(registry) + + _prepare_unit()(ctx) + + data_x = ctx.require("data_x") + assert sorted(data_x.keys()) == ["test", "train", "validation"] + assert len(data_x["train"]) == 3 + assert len(data_x["test"]) == 2 + assert len(data_x["validation"]) == 1 + assert data_x["train"].column_names == ["a", "b"] + assert ctx.require("data_y")["train"].column_names == ["target"] + + +def test_prepare_publishes_the_task_for_the_local_path(registry): + ctx = _prepared_context(registry) + + _prepare_unit()(ctx) + + assert isinstance(ctx.require("task"), RecordingTask) + + +def test_prepare_validates_the_task_before_it_runs(registry): + """``validate`` is called by the orchestrator outside the block that wraps + preparation failures, so a missing task stays a registry error.""" + with pytest.raises( + JobError, match="Unable to find Task with name NoSuchTask in registry" + ): + _prepare_unit(task_name="NoSuchTask").validate(ExecutionContext()) + + +def test_prepare_without_split_indexes_is_rejected_before_it_starts(registry): + """A missing key means "nothing published them", not "there is no split".""" + ctx = ExecutionContext() + ctx.put("dataset", _dataset()) + ctx.put("model", RecordingModel()) + ctx.put_ref("dataset_id", 9) + + with pytest.raises(UnitContractError, match="Context key 'split_indexes'"): + _prepare_unit()(ctx) + + +def test_prepare_without_a_model_is_rejected_before_it_starts(registry): + ctx = ExecutionContext() + ctx.put("dataset", _dataset()) + ctx.put_ref("dataset_id", 9) + ctx.put_ref("split_indexes", SPLITS) + + with pytest.raises(UnitContractError, match="Context key 'model'"): + _prepare_unit()(ctx) + + +def test_prepare_composes_after_a_loader_that_publishes_no_dataset_id(registry): + """A dataset with no id attached is enough to run. + + ``REQUIRES`` is demanded unconditionally, so listing a key the unit never + reads would silently restrict what it can follow. ``BuildManualInputUnit`` + publishes ``dataset`` alone, and this unit has to work after it. + """ + ctx = ExecutionContext() + ctx.put("dataset", _dataset()) + ctx.put("model", RecordingModel()) + ctx.put_ref("split_indexes", SPLITS) + + _prepare_unit()(ctx) + + assert not ctx.has("dataset_id") + assert len(ctx.require("data_x")["train"]) == 3 + + +# --- GenerateGlobalExplanationUnit -------------------------------------- + + +def test_the_global_explanation_pickles_both_artifacts(registry, explanations_path): + ctx = ExecutionContext() + explainer = RecordingGlobalExplainer(RecordingModel()) + ctx.put("explainer", explainer) + ctx.put("data_x", {"train": 1}) + ctx.put("data_y", {"train": 2}) + + GenerateGlobalExplanationUnit(explainer_id=7)(ctx) + + explanation_path = Path(ctx.require("explanation_path")) + plot_path = Path(ctx.require("plot_path")) + assert explanation_path.name == "global_explanation_7.pickle" + assert plot_path.name == "global_explanation_plot_7.pickle" + + with open(explanation_path, "rb") as handle: + assert pickle.load(handle) == {"importance": [1, 2]} + # The explainer receives the two halves as a pair, in order. + assert explainer.seen == ({"train": 1}, {"train": 2}) + + +def test_the_global_unit_never_writes_the_row(registry, explanations_path): + """It publishes where it wrote; the row belongs to the job. + + The unit takes only an id, so there is nothing for it to write a row with — + which is the point. + """ + assert set(GenerateGlobalExplanationUnit.SCHEMA.model_fields) == {"explainer_id"} + assert GenerateGlobalExplanationUnit.PROVIDES == ( + "explanation_path", + "plot_path", + ) + + +def test_a_failing_global_explanation_is_wrapped(registry, explanations_path): + ctx = ExecutionContext() + ctx.put("explainer", ExplodingGlobalExplainer(RecordingModel())) + ctx.put("data_x", {}) + ctx.put("data_y", {}) + + with pytest.raises(JobError, match="Failed to generate the explanation") as excinfo: + GenerateGlobalExplanationUnit(explainer_id=7)(ctx) + + assert "the explanation itself blew up" in str(excinfo.value.__cause__) + + +def test_generating_without_an_explainer_is_rejected_before_it_starts( + registry, explanations_path +): + ctx = ExecutionContext() + ctx.put("data_x", {}) + ctx.put("data_y", {}) + + with pytest.raises(UnitContractError, match="Context key 'explainer'"): + GenerateGlobalExplanationUnit(explainer_id=7)(ctx) + + +# --- GenerateLocalExplanationUnit --------------------------------------- + + +def _local_context(registry): + ctx = ExecutionContext() + ctx.put("explainer", RecordingLocalExplainer(RecordingModel())) + ctx.put("task", RecordingTask()) + ctx.put("data_x", {"train": 1}) + ctx.put("data_y", {"train": 2}) + ctx.put_ref("split_indexes", SPLITS) + return ctx + + +def _local_unit(**overrides): + config = { + "explainer_id": 7, + "instance_dataset_id": 9, + "scope": {"split": "test", "percentage": 100}, + "fit_parameters": {"nsamples": 10}, + "input_columns": ["a", "b"], + "output_columns": ["target"], + "manual_input_data": None, + "same_dataset": True, + "session_splits": None, + } + config.update(overrides) + return GenerateLocalExplanationUnit(**config) + + +def test_the_local_explanation_writes_three_artifacts( + registry, fake_db, explanations_path +): + ctx = _local_context(registry) + + _local_unit()(ctx) + + assert Path(ctx.require("explanation_path")).name == "local_explanation_7.pickle" + assert Path(ctx.require("plots_path")).name == "local_explanation_plots_7.pickle" + + saved = load_dataset(str(Path(ctx.require("input_dataset_path")) / "dataset")) + assert saved.column_names == ["a", "b"] + assert len(saved) == 2 + + +def test_the_local_explanation_forwards_its_fit_parameters( + registry, fake_db, explanations_path +): + ctx = _local_context(registry) + RecordingLocalExplainer.fitted_with = None + + _local_unit()(ctx) + + assert RecordingLocalExplainer.fitted_with == {"nsamples": 10} + + +def test_instances_from_another_dataset_recompute_the_split( + registry, fake_db, explanations_path +): + """The run's row indexes are meaningless over a different dataset. + + When the instances do not come from the dataset the model was trained on, + replaying ``split_indexes`` would address rows that do not correspond, so + the split is recomputed from the session's ratios over the dataset in hand. + That derived state is resolved inside ``execute`` and never published — this + is the branch that proves the recompute actually happens. + """ + ctx = _local_context(registry) + + _local_unit( + same_dataset=False, + session_splits=json.dumps( + { + "train": 0.5, + "test": 0.5, + "validation": 0.0, + "is_random": True, + "has_changed": True, + "seed": 42, + "shuffle": True, + "stratify": False, + } + ), + scope={"split": "test", "percentage": 100}, + )(ctx) + + saved = load_dataset(str(Path(ctx.require("input_dataset_path")) / "dataset")) + # Half of the six stored rows, not the two the run's test_indexes name. + assert len(saved) == 3 + # The recomputed split never leaks back into the context. + assert ctx.require("split_indexes") == SPLITS + + +def test_a_rows_scope_selects_exactly_the_valid_indexes( + registry, fake_db, explanations_path +): + ctx = _local_context(registry) + + _local_unit(scope={"mode": "rows", "row_indexes": [0, 3, 5]})(ctx) + + saved = load_dataset(str(Path(ctx.require("input_dataset_path")) / "dataset")) + assert len(saved) == 3 + + +def test_a_manual_scope_builds_the_instances_from_the_given_rows( + registry, fake_db, explanations_path +): + ctx = _local_context(registry) + + _local_unit( + scope={"mode": "manual"}, + manual_input_data=[{"a": 1, "b": 2}, {"a": 3, "b": 4}], + )(ctx) + + saved = load_dataset(str(Path(ctx.require("input_dataset_path")) / "dataset")) + assert saved.column_names == ["a", "b"] + assert len(saved) == 2 + + +def test_the_three_selection_complaints_are_swallowed_by_one_wrapper( + registry, fake_db, explanations_path +): + """All three modes report through the same message, keeping their own + complaint only as ``__cause__``. Pinned because it is an easy detail to + "fix" by accident.""" + cases = [ + ({"split": "notasplit", "percentage": 100}, None, "not a valid split"), + ( + {"mode": "rows", "row_indexes": [10**9]}, + None, + "No valid row indexes provided", + ), + ({"mode": "manual"}, None, "No manual input data provided"), + ] + + for scope, manual, cause in cases: + ctx = _local_context(registry) + with pytest.raises(JobError, match="Can not prepare Dataset with") as excinfo: + _local_unit(scope=scope, manual_input_data=manual)(ctx) + assert cause in str(excinfo.value.__cause__), scope + + +def test_a_missing_instance_dataset_is_reported_by_id( + registry, fake_db, explanations_path +): + ctx = _local_context(registry) + + with pytest.raises( + JobError, match="Dataset 99 to be explained does not exist in DB." + ): + _local_unit(instance_dataset_id=99)(ctx) + + +def test_the_local_unit_needs_the_task_and_says_so( + registry, fake_db, explanations_path +): + """The manual mode calls into the task, so it is a declared requirement even + though the other two modes barely touch it.""" + ctx = ExecutionContext() + ctx.put("explainer", RecordingLocalExplainer(RecordingModel())) + ctx.put("data_x", {"train": 1}) + ctx.put("data_y", {"train": 2}) + ctx.put_ref("split_indexes", SPLITS) + + with pytest.raises(UnitContractError, match="Context key 'task'"): + _local_unit()(ctx) + + +def test_the_published_paths_are_plain_strings(registry, fake_db, explanations_path): + """All three travel as refs, so they have to be JSON data.""" + ctx = _local_context(registry) + + _local_unit()(ctx) + + refs = ctx.to_dict() + for key in ("explanation_path", "plots_path", "input_dataset_path"): + assert isinstance(refs[key], str), key diff --git a/tests/back/units/test_exploration_units.py b/tests/back/units/test_exploration_units.py new file mode 100644 index 000000000..0a104fcb0 --- /dev/null +++ b/tests/back/units/test_exploration_units.py @@ -0,0 +1,326 @@ +"""Contract tests for the exploration units, isolated from any orchestrating job. + +The context is built by hand rather than through a job, which is what exposes +composability mistakes: a job always wires the context "correctly", so an +end-to-end run cannot tell a real contract from a lucky one. +""" + +import pathlib + +import pytest +from kink import di + +from DashAI.back.dataloaders.classes.dashai_dataset import to_dashai_dataset +from DashAI.back.exploration.base_explorer import BaseExplorer +from DashAI.back.job.base_job import JobError +from DashAI.back.units.context import ExecutionContext, UnitContractError +from DashAI.back.units.run_exploration_unit import RunExplorationUnit +from DashAI.back.units.save_exploration_unit import SaveExplorationUnit + + +class _ExplorerRow: + """Stand-in for an Explorer ORM row.""" + + def __init__( + self, notebook_id=3, columns=None, exploration_type="RecordingExplorer" + ): + self.id = 11 + self.notebook_id = notebook_id + self.columns = columns if columns is not None else [{"columnName": "a"}] + self.exploration_type = exploration_type + self.name = "an exploration" + + +class _NotebookRow: + """Stand-in for a Notebook ORM row.""" + + def __init__(self, notebook_id=3): + self.id = notebook_id + + +class _FakeSession: + def __init__(self, rows): + self._rows = rows + + def get(self, model, row_id): + return self._rows.get(model.__name__, {}).get(row_id) + + def __enter__(self): + return self + + def __exit__(self, *exc_info): + return False + + +class _FakeSessionFactory: + """Stand-in for a ``sessionmaker``. + + A class rather than a lambda on purpose: kink invokes any registered lambda + with the container to resolve it, so a lambda here would be called as a + service factory instead of being handed to the unit as one. + """ + + def __init__(self, rows): + self._rows = rows + + def __call__(self): + return _FakeSession(self._rows) + + +class RecordingExplorer(BaseExplorer): + """Explorer that records what the units hand it, and when.""" + + def __init__(self, **kwargs): + super().__init__(**kwargs) + self.marker = kwargs.get("marker") + self.seen_columns = None + self.seen_row_name = None + + def prepare_dataset(self, loaded_dataset, columns): + self.seen_columns = [column["columnName"] for column in columns] + return loaded_dataset.select_columns(self.seen_columns) + + def launch_exploration(self, dataset, explorer_info): + self.seen_row_name = explorer_info.name + return {"rows": len(dataset), "columns": dataset.column_names} + + def save_notebook(self, notebook_info, explorer_info, save_path, result): + # Reads instance state set at construction time, the way CorrMatrix + # reads self.plot, and returns a str the way DescribeExplorer does. + target = pathlib.Path(save_path) / f"{explorer_info.id}-{self.marker}.txt" + target.write_text(str(result), encoding="utf-8") + return target.as_posix() + + def get_results(self, exploration_path, options): + return [] + + +class BadPathExplorer(RecordingExplorer): + """Explorer whose save returns something that is not a path.""" + + def save_notebook(self, notebook_info, explorer_info, save_path, result): + return 42 + + +class ExplodingExplorer(RecordingExplorer): + def launch_exploration(self, dataset, explorer_info): + raise RuntimeError("the exploration itself blew up") + + +@pytest.fixture(name="registry") +def fixture_registry(): + registry = { + "RecordingExplorer": {"class": RecordingExplorer}, + "BadPathExplorer": {"class": BadPathExplorer}, + "ExplodingExplorer": {"class": ExplodingExplorer}, + } + di["component_registry"] = registry + yield registry + del di["component_registry"] + + +@pytest.fixture(name="fake_db") +def fixture_fake_db(): + rows = { + "Explorer": {11: _ExplorerRow()}, + "Notebook": {3: _NotebookRow()}, + } + di["session_factory"] = _FakeSessionFactory(rows) + yield rows + del di["session_factory"] + + +@pytest.fixture(name="notebook_path") +def fixture_notebook_path(tmp_path): + config = {"NOTEBOOK_PATH": tmp_path / "notebooks"} + di["config"] = config + yield config["NOTEBOOK_PATH"] + del di["config"] + + +@pytest.fixture(name="ctx") +def fixture_ctx(): + """A context holding a three-column dataset, as a loader would leave it.""" + import pandas as pd + import pyarrow as pa + + from DashAI.back.types.value_types import Integer + + frame = pd.DataFrame({"a": [1, 2, 3], "b": [4, 5, 6], "c": [7, 8, 9]}) + types = {name: Integer(arrow_type=pa.int64()) for name in ("a", "b", "c")} + + context = ExecutionContext() + context.put("dataset", to_dashai_dataset(frame, types=types)) + return context + + +def _explorer(component="RecordingExplorer", **params): + return {"component": component, "params": {"marker": "x", **params}} + + +# --- RunExplorationUnit ------------------------------------------------- + + +def test_run_exploration_publishes_the_result_and_the_explorer(ctx, registry, fake_db): + """Both outputs are declared, so both must be there. + + The explorer instance is an output and not a private detail: saving is a + method on it, and the save unit has to receive the object that ran. + """ + RunExplorationUnit(explorer_id=11, explorer=_explorer())(ctx) + + assert ctx.require("exploration_result") == {"rows": 3, "columns": ["a"]} + assert isinstance(ctx.require("explorer"), RecordingExplorer) + + +def test_run_exploration_narrows_the_dataset_to_the_rows_selected_columns( + ctx, registry, fake_db +): + """The column list comes from the row, resolved against the dataset the + context holds right now — never against a list captured earlier.""" + fake_db["Explorer"][11].columns = [{"columnName": "b"}, {"columnName": "c"}] + + RunExplorationUnit(explorer_id=11, explorer=_explorer())(ctx) + + assert ctx.require("explorer").seen_columns == ["b", "c"] + + +def test_run_exploration_hands_the_row_to_the_component(ctx, registry, fake_db): + """``launch_exploration`` takes the ORM row; the unit re-reads it itself.""" + RunExplorationUnit(explorer_id=11, explorer=_explorer())(ctx) + + assert ctx.require("explorer").seen_row_name == "an exploration" + + +def test_run_exploration_without_a_dataset_is_rejected_before_it_starts( + registry, fake_db +): + """REQUIRES is enforced by ``__call__``, so a wiring mistake is not + mistaken for an empty dataset.""" + with pytest.raises(UnitContractError, match="Context key 'dataset'"): + RunExplorationUnit(explorer_id=11, explorer=_explorer())(ExecutionContext()) + + +def test_an_unknown_explorer_is_reported_by_name(ctx, registry, fake_db): + with pytest.raises(JobError, match="Explorer NoSuchExplorer not found in the reg"): + RunExplorationUnit( + explorer_id=11, explorer=_explorer(component="NoSuchExplorer") + )(ctx) + + +def test_a_missing_explorer_row_is_reported_by_id(ctx, registry, fake_db): + with pytest.raises(JobError, match="Explorer with id 99 not found."): + RunExplorationUnit(explorer_id=99, explorer=_explorer())(ctx) + + +def test_a_column_absent_from_the_dataset_becomes_a_preparation_error( + ctx, registry, fake_db +): + fake_db["Explorer"][11].columns = [{"columnName": "nope"}] + + with pytest.raises( + JobError, match="Error preparing the dataset for the exploration" + ): + RunExplorationUnit(explorer_id=11, explorer=_explorer())(ctx) + + +def test_a_failing_exploration_is_wrapped_with_the_component_name( + ctx, registry, fake_db +): + with pytest.raises( + JobError, match="Error launching the exploration ExplodingExplorer." + ) as excinfo: + RunExplorationUnit( + explorer_id=11, explorer=_explorer(component="ExplodingExplorer") + )(ctx) + + assert "the exploration itself blew up" in str(excinfo.value.__cause__) + + +def test_two_exploration_units_do_not_share_a_resolved_class(ctx, registry, fake_db): + """The registry lookup is memoized on the instance, not in the context. + + Two exploration nodes in one context must each resolve their own component; + a context-global cache key would make the second silently reuse the first. + """ + first = RunExplorationUnit(explorer_id=11, explorer=_explorer()) + second = RunExplorationUnit( + explorer_id=11, explorer=_explorer(component="BadPathExplorer") + ) + + first(ctx) + second(ctx) + + assert first._explorer_class is RecordingExplorer + assert second._explorer_class is BadPathExplorer + + +# --- SaveExplorationUnit ------------------------------------------------ + + +def test_save_exploration_writes_under_the_notebook_folder( + ctx, registry, fake_db, notebook_path +): + RunExplorationUnit(explorer_id=11, explorer=_explorer())(ctx) + SaveExplorationUnit(explorer_id=11)(ctx) + + written = pathlib.Path(ctx.require("exploration_path")) + assert written.exists() + assert written.parent == notebook_path / "3" + assert written.name == "11-x.txt" + + +def test_save_exploration_uses_the_explorer_that_ran( + ctx, registry, fake_db, notebook_path +): + """Identity, not equality: the saved file name carries state the instance + was built with, so rebuilding a second explorer from the same config would + pass this by accident. Asserting ``is`` is what makes it real.""" + RunExplorationUnit(explorer_id=11, explorer=_explorer(marker="carried"))(ctx) + ran = ctx.require("explorer") + + SaveExplorationUnit(explorer_id=11)(ctx) + + assert ctx.require("explorer") is ran + assert pathlib.Path(ctx.require("exploration_path")).name == "11-carried.txt" + + +def test_save_exploration_creates_the_notebook_folder_when_absent( + ctx, registry, fake_db, notebook_path +): + assert not notebook_path.exists() + + RunExplorationUnit(explorer_id=11, explorer=_explorer())(ctx) + SaveExplorationUnit(explorer_id=11)(ctx) + + assert (notebook_path / "3").is_dir() + + +def test_save_exploration_without_a_result_is_rejected( + registry, fake_db, notebook_path +): + with pytest.raises(UnitContractError, match="Context key 'exploration_result'"): + SaveExplorationUnit(explorer_id=11)(ExecutionContext()) + + +def test_a_save_that_does_not_return_a_path_is_reported( + ctx, registry, fake_db, notebook_path +): + RunExplorationUnit(explorer_id=11, explorer=_explorer(component="BadPathExplorer"))( + ctx + ) + + with pytest.raises(JobError, match="save path is not a pathlib.Path"): + SaveExplorationUnit(explorer_id=11)(ctx) + + +def test_the_published_path_is_a_plain_string(ctx, registry, fake_db, notebook_path): + """``exploration_path`` travels as a ref, so it has to be JSON data. + + A ``pathlib.Path`` would raise on ``put_ref``; this pins that the unit + converts before publishing. + """ + RunExplorationUnit(explorer_id=11, explorer=_explorer())(ctx) + SaveExplorationUnit(explorer_id=11)(ctx) + + assert isinstance(ctx.to_dict()["exploration_path"], str) diff --git a/tests/back/units/test_prediction_units.py b/tests/back/units/test_prediction_units.py new file mode 100644 index 000000000..25a937b0f --- /dev/null +++ b/tests/back/units/test_prediction_units.py @@ -0,0 +1,423 @@ +"""Contract tests for the prediction units, isolated from any orchestrating job. + +The context is built by hand rather than through a job, which is what exposes +composability mistakes: a job always wires the context "correctly", so an +end-to-end run cannot tell a real contract from a lucky one. +""" + +from pathlib import Path + +import pytest +from kink import di + +from DashAI.back.dataloaders.classes.dashai_dataset import ( + load_dataset, + save_dataset, + to_dashai_dataset, +) +from DashAI.back.job.base_job import JobError +from DashAI.back.units.context import ExecutionContext, UnitContractError +from DashAI.back.units.load_trained_model_unit import LoadTrainedModelUnit +from DashAI.back.units.load_training_dataset_unit import LoadTrainingDatasetUnit +from DashAI.back.units.predict_unit import PredictUnit +from DashAI.back.units.save_prediction_unit import SavePredictionUnit + + +class _RunRow: + """Stand-in for a Run ORM row.""" + + def __init__(self, model_name="RecordingModel", run_path="somewhere"): + self.model_name = model_name + self.run_path = run_path + + +class _FakeSession: + def __init__(self, rows): + self._rows = rows + + def get(self, model, row_id): + return self._rows.get(model.__name__, {}).get(row_id) + + def __enter__(self): + return self + + def __exit__(self, *exc_info): + return False + + +class _FakeSessionFactory: + """Stand-in for a ``sessionmaker``. + + A class rather than a lambda on purpose: kink invokes any registered lambda + with the container to resolve it, so a lambda here would be called as a + service factory instead of being handed to the unit as one. + """ + + def __init__(self, rows): + self._rows = rows + + def __call__(self): + return _FakeSession(self._rows) + + +class RecordingModel: + """Model whose ``load`` is a staticmethod, the way every real one is.""" + + loaded_from = None + + def __init__(self): + self.seen_columns = None + + @staticmethod + def load(filename): + model = RecordingModel() + RecordingModel.loaded_from = filename + return model + + def predict(self, x): + self.seen_columns = x.column_names + return [0] * len(x) + + +class UnloadableModel(RecordingModel): + @staticmethod + def load(filename): + raise OSError("the artifact is not there") + + +class RecordingTask: + """Task that records the training dataset it was given for decoding.""" + + seen_train_columns = None + + def process_predictions(self, train_dataset, y_pred_proba, output_column): + RecordingTask.seen_train_columns = train_dataset.column_names + return [f"label-{int(value)}" for value in y_pred_proba] + + def process_manual_input(self, rows, dataset_path): + import pandas as pd + import pyarrow as pa + + from DashAI.back.types.value_types import Integer + + frame = pd.DataFrame(rows) + types = {name: Integer(arrow_type=pa.int64()) for name in frame.columns} + return to_dashai_dataset(frame, types=types) + + +@pytest.fixture(name="registry") +def fixture_registry(): + registry = { + "RecordingModel": {"class": RecordingModel}, + "UnloadableModel": {"class": UnloadableModel}, + "RecordingTask": {"class": RecordingTask}, + } + di["component_registry"] = registry + yield registry + del di["component_registry"] + + +@pytest.fixture(name="fake_db") +def fixture_fake_db(): + rows = {"Run": {5: _RunRow()}} + di["session_factory"] = _FakeSessionFactory(rows) + yield rows + del di["session_factory"] + + +def _dataset(**columns): + import pandas as pd + import pyarrow as pa + + from DashAI.back.types.value_types import Integer + + frame = pd.DataFrame(columns) + types = {name: Integer(arrow_type=pa.int64()) for name in frame.columns} + return to_dashai_dataset(frame, types=types) + + +@pytest.fixture(name="stored_training_dataset") +def fixture_stored_training_dataset(tmp_path): + """A training dataset on disk, laid out the way a Dataset row points at it.""" + root = tmp_path / "training" + save_dataset(_dataset(a=[1, 2, 3], b=[4, 5, 6]), str(root / "dataset")) + return root + + +@pytest.fixture(name="datasets_path") +def fixture_datasets_path(tmp_path): + config = {"DATASETS_PATH": tmp_path / "datasets"} + di["config"] = config + yield config["DATASETS_PATH"] + del di["config"] + + +# --- LoadTrainedModelUnit ----------------------------------------------- + + +def test_the_model_comes_from_the_path_the_run_recorded(registry, fake_db): + """Neither the component nor the path is configuration: both are read off + the run, so the model restored is always the one that run saved.""" + ctx = ExecutionContext() + fake_db["Run"][5].run_path = "the/recorded/path" + + LoadTrainedModelUnit(run_id=5)(ctx) + + assert isinstance(ctx.require("model"), RecordingModel) + assert RecordingModel.loaded_from == "the/recorded/path" + + +def test_a_missing_run_is_reported_by_id(registry, fake_db): + with pytest.raises(JobError, match="Run 99 does not exist in DB."): + LoadTrainedModelUnit(run_id=99)(ExecutionContext()) + + +def test_an_unknown_model_name_is_reported_by_name(registry, fake_db): + fake_db["Run"][5].model_name = "NoSuchModel" + + with pytest.raises(JobError, match="Model NoSuchModel not found in the registry"): + LoadTrainedModelUnit(run_id=5)(ExecutionContext()) + + +def test_an_artifact_that_cannot_be_read_names_the_model_and_the_path( + registry, fake_db +): + fake_db["Run"][5].model_name = "UnloadableModel" + fake_db["Run"][5].run_path = "gone" + + with pytest.raises( + JobError, match="Failed to load model UnloadableModel from path gone" + ) as excinfo: + LoadTrainedModelUnit(run_id=5)(ExecutionContext()) + + assert "the artifact is not there" in str(excinfo.value.__cause__) + + +def test_two_model_units_do_not_share_a_resolved_class(registry, fake_db): + """The registry lookup is memoized on the instance, not in the context.""" + fake_db["Run"][6] = _RunRow(model_name="UnloadableModel", run_path="gone") + + first = LoadTrainedModelUnit(run_id=5) + second = LoadTrainedModelUnit(run_id=6) + + first(ExecutionContext()) + with pytest.raises(JobError): + second(ExecutionContext()) + + assert first._model_class is RecordingModel + assert second._model_class is UnloadableModel + + +# --- LoadTrainingDatasetUnit -------------------------------------------- + + +def test_the_training_dataset_lands_under_its_own_key(stored_training_dataset): + """Not ``dataset``: this one is a reference for decoding and typing, and + would otherwise collide with the dataset actually being predicted on.""" + ctx = ExecutionContext() + + LoadTrainingDatasetUnit(train_dataset_file_path=str(stored_training_dataset))(ctx) + + assert ctx.require("train_dataset").column_names == ["a", "b"] + assert not ctx.has("dataset") + + +def test_the_training_dataset_can_coexist_with_the_one_being_predicted_on( + stored_training_dataset, +): + """The whole reason for the separate key: both datasets are live at once.""" + ctx = ExecutionContext() + ctx.put("dataset", _dataset(a=[9], b=[9])) + + LoadTrainingDatasetUnit(train_dataset_file_path=str(stored_training_dataset))(ctx) + + assert ctx.require("dataset")["a"] == [9] + assert ctx.require("train_dataset")["a"] == [1, 2, 3] + + +def test_the_declared_types_travel_as_plain_data(stored_training_dataset): + """``train_dataset_types`` is a ref so the saving step never reopens the + file. ``put_ref`` rejects anything that is not JSON data, which is what + stops a live type object from being smuggled across the boundary — note + ``to_string`` returns a dict despite its name. + """ + import json + + ctx = ExecutionContext() + + LoadTrainingDatasetUnit(train_dataset_file_path=str(stored_training_dataset))(ctx) + + types = ctx.to_dict()["train_dataset_types"] + assert types == { + "a": {"type": "Integer", "dtype": "int64"}, + "b": {"type": "Integer", "dtype": "int64"}, + } + json.dumps(types) + + +def test_an_unreadable_training_dataset_names_the_folder(tmp_path): + with pytest.raises(JobError, match="Cannot load training dataset from"): + LoadTrainingDatasetUnit(train_dataset_file_path=str(tmp_path / "nowhere"))( + ExecutionContext() + ) + + +# --- PredictUnit -------------------------------------------------------- + + +def _ready_context(stored_training_dataset, dataset=None): + ctx = ExecutionContext() + ctx.put("dataset", dataset if dataset is not None else _dataset(a=[1, 2], b=[3, 4])) + ctx.put("model", RecordingModel()) + LoadTrainingDatasetUnit(train_dataset_file_path=str(stored_training_dataset))(ctx) + return ctx + + +def _predict_unit(**overrides): + config = { + "task_name": "RecordingTask", + "input_columns": ["a"], + "output_columns": ["target"], + } + config.update(overrides) + return PredictUnit(**config) + + +def test_predict_publishes_decoded_labels(registry, stored_training_dataset): + ctx = _ready_context(stored_training_dataset) + + _predict_unit()(ctx) + + assert ctx.require("y_pred") == ["label-0", "label-0"] + + +def test_predict_hands_the_model_only_the_input_columns( + registry, stored_training_dataset +): + """Selected against the dataset in the context right now, so whatever + produced it — a load or hand-typed rows — is free to differ in shape.""" + ctx = _ready_context(stored_training_dataset) + + _predict_unit()(ctx) + + assert ctx.require("model").seen_columns == ["a"] + + +def test_predict_decodes_against_the_training_dataset( + registry, stored_training_dataset +): + ctx = _ready_context(stored_training_dataset) + + _predict_unit()(ctx) + + assert RecordingTask.seen_train_columns == ["a", "b"] + + +def test_predict_validates_the_task_before_it_runs(registry, stored_training_dataset): + """``validate`` is what the orchestrator calls early, so a missing task is + reported as a task problem rather than being overtaken by a later failure.""" + with pytest.raises(JobError, match="Task NoSuchTask not found in the registry"): + _predict_unit(task_name="NoSuchTask").validate(ExecutionContext()) + + +def test_predict_without_a_model_is_rejected_before_it_starts( + registry, stored_training_dataset +): + ctx = ExecutionContext() + ctx.put("dataset", _dataset(a=[1])) + LoadTrainingDatasetUnit(train_dataset_file_path=str(stored_training_dataset))(ctx) + + with pytest.raises(UnitContractError, match="Context key 'model'"): + _predict_unit()(ctx) + + +def test_predict_without_a_training_dataset_is_rejected_before_it_starts(registry): + """A missing key means "the loader did not run", not "there is nothing to + decode against" — so it has to fail loudly instead of predicting anyway.""" + ctx = ExecutionContext() + ctx.put("dataset", _dataset(a=[1])) + ctx.put("model", RecordingModel()) + + with pytest.raises(UnitContractError, match="Context key 'train_dataset'"): + _predict_unit()(ctx) + + +# --- SavePredictionUnit ------------------------------------------------- + + +def _save_unit(**overrides): + config = {"input_columns": ["a"], "output_columns": ["target"]} + config.update(overrides) + return SavePredictionUnit(**config) + + +def test_save_writes_the_inputs_plus_the_predicted_column( + registry, stored_training_dataset, datasets_path +): + ctx = _ready_context(stored_training_dataset) + _predict_unit()(ctx) + + _save_unit()(ctx) + + saved = load_dataset(str(Path(ctx.require("results_path")) / "dataset")) + assert saved.column_names == ["a", "b", "target"] + assert saved["target"] == ["label-0", "label-0"] + + +def test_save_resolves_the_columns_against_the_dataset_it_is_handed( + registry, stored_training_dataset, datasets_path +): + """The column list is read at the top of execute, never published earlier. + + Here the dataset already carries a column named like the output one; it has + to be replaced, not duplicated — which only works if the names are resolved + from the dataset in hand. + """ + ctx = _ready_context( + stored_training_dataset, dataset=_dataset(a=[1, 2], target=[7, 8]) + ) + _predict_unit()(ctx) + + _save_unit()(ctx) + + saved = load_dataset(str(Path(ctx.require("results_path")) / "dataset")) + assert saved.column_names == ["a", "target"] + assert saved["target"] == ["label-0", "label-0"] + + +def test_two_saves_never_collide(registry, stored_training_dataset, datasets_path): + """A prediction has no natural key to overwrite, so each run gets a folder.""" + ctx = _ready_context(stored_training_dataset) + _predict_unit()(ctx) + + _save_unit()(ctx) + first = ctx.require("results_path") + _save_unit()(ctx) + second = ctx.require("results_path") + + assert first != second + assert Path(first).exists() + assert Path(second).exists() + + +def test_save_without_a_prediction_is_rejected_before_it_starts( + registry, stored_training_dataset, datasets_path +): + ctx = ExecutionContext() + ctx.put("dataset", _dataset(a=[1])) + ctx.put_ref("train_dataset_types", {}) + + with pytest.raises(UnitContractError, match="Context key 'y_pred'"): + _save_unit()(ctx) + + +def test_the_published_results_path_is_a_plain_string( + registry, stored_training_dataset, datasets_path +): + """``results_path`` travels as a ref, so it has to be JSON data.""" + ctx = _ready_context(stored_training_dataset) + _predict_unit()(ctx) + + _save_unit()(ctx) + + assert isinstance(ctx.to_dict()["results_path"], str) diff --git a/tests/back/units/test_unit_contracts.py b/tests/back/units/test_unit_contracts.py index b472a885a..7649bdf6e 100644 --- a/tests/back/units/test_unit_contracts.py +++ b/tests/back/units/test_unit_contracts.py @@ -105,6 +105,32 @@ def test_every_context_key_a_unit_reads_is_declared_in_requires(name, cls): ) +@pytest.mark.parametrize(("name", "cls"), UNITS, ids=[name for name, _ in UNITS]) +def test_a_unit_does_not_require_a_key_it_never_reads(name, cls): + """The mirror of the test above, and just as load-bearing. + + ``__call__`` demands every key in ``REQUIRES`` unconditionally, so a key + listed but never read is not harmless documentation: it rejects any upstream + that does not happen to publish it. ``PrepareExplanationDataUnit`` used to + require ``dataset_id`` — left over from an error message that moved to the + job — which would have made it impossible to compose after + ``BuildManualInputUnit``, whose ``PROVIDES`` is just ``("dataset",)``. + + It passed every end-to-end test because the one job wiring it happened to + run a loader that publishes the id first. That is exactly the class of + mistake this file exists to catch. + """ + required = _declared(cls, "REQUIRES") + read = _context_calls(cls, {"require", "get", "has"}) + + unread = required - read + assert not unread, ( + f"{name} declares {sorted(unread)} in REQUIRES but never reads them. " + "Every declared key is demanded before the unit runs, so an unused one " + "only narrows what the unit can be composed after." + ) + + @pytest.mark.parametrize(("name", "cls"), UNITS, ids=[name for name, _ in UNITS]) def test_every_key_a_unit_promises_is_actually_written(name, cls): written = _context_calls(cls, {"put", "put_ref"}) From 5a880cf0cabd84b75f82e5bbca2553fe76684bbe Mon Sep 17 00:00:00 2001 From: Felipe Date: Wed, 5 Aug 2026 17:38:39 -0400 Subject: [PATCH 2/2] feat: Refactor manual prediction to use shared units and add comprehensive tests for preview endpoint --- DashAI/back/api/api_v1/endpoints/predict.py | 2 - DashAI/back/job/predict_job.py | 124 ++--- tests/back/api/test_predict_preview_api.py | 490 ++++++++++++++++++++ 3 files changed, 553 insertions(+), 63 deletions(-) create mode 100644 tests/back/api/test_predict_preview_api.py diff --git a/DashAI/back/api/api_v1/endpoints/predict.py b/DashAI/back/api/api_v1/endpoints/predict.py index 23da1a336..3fffa2f04 100644 --- a/DashAI/back/api/api_v1/endpoints/predict.py +++ b/DashAI/back/api/api_v1/endpoints/predict.py @@ -252,7 +252,6 @@ async def delete_prediction( @inject async def preview_manual_prediction( request: Request, - component_registry: "ComponentRegistry" = Depends(lambda: di["component_registry"]), session_factory: "sessionmaker" = Depends(lambda: di["session_factory"]), ): """Run a synchronous manual prediction and return results without persisting. @@ -335,7 +334,6 @@ async def preview_manual_prediction( run_manual_prediction, run_id=run_id_int, manual_input_data=rows_data, - component_registry=component_registry, session_factory=session_factory, ) return {"columns": columns, "rows": rows} diff --git a/DashAI/back/job/predict_job.py b/DashAI/back/job/predict_job.py index 213eb1b0a..a6d8ef953 100644 --- a/DashAI/back/job/predict_job.py +++ b/DashAI/back/job/predict_job.py @@ -1,5 +1,4 @@ import logging -from pathlib import Path from typing import TYPE_CHECKING, Any, Dict, List, Tuple from fastapi import status @@ -10,8 +9,6 @@ from DashAI.back.dependencies.database.models import Dataset, ModelSession, Prediction from DashAI.back.job.base_job import BaseJob, JobError -from DashAI.back.models.base_model import BaseModel -from DashAI.back.tasks.base_task import BaseTask from DashAI.back.units.build_manual_input_unit import BuildManualInputUnit from DashAI.back.units.context import ExecutionContext from DashAI.back.units.load_dataset_unit import LoadDatasetUnit @@ -29,24 +26,6 @@ log = logging.getLogger(__name__) -def _run_prediction_pipeline( - task: BaseTask, - trained_model: BaseModel, - train_dataset: "DashAIDataset", - loaded_dataset: "DashAIDataset", - model_session: ModelSession, -) -> Tuple["DashAIDataset", Any]: - """Run shared prediction steps from prepared input data to final predictions.""" - import numpy as np - - prepared_dataset = loaded_dataset.select_columns(model_session.input_columns) - y_pred_proba = np.array(trained_model.predict(prepared_dataset)) - y_pred = task.process_predictions( - train_dataset, y_pred_proba, model_session.output_columns[0] - ) - return prepared_dataset, y_pred - - def _build_preview_rows( prepared_dataset: "DashAIDataset", input_columns: List[str], @@ -85,21 +64,31 @@ def _to_native(v: Any) -> Any: def run_manual_prediction( run_id: int, manual_input_data: List[Dict], - component_registry: Any, session_factory: "sessionmaker", ) -> Tuple[List[str], List[List]]: """Execute a manual prediction synchronously without persisting results. + Composes the same units ``PredictJob`` does, so there is one definition of + what predicting means. The difference is entirely in the orchestration: no + state row to advance, nothing written to disk, and failures reported as + ``HTTPException`` because this runs inside a request instead of a worker. + + Each unit call sits in its own ``try`` so the HTTP response is decided by + *which step* failed, never by matching on an error message. That is what + keeps the endpoint's contract — eleven distinct responses across four status + codes — independent of how the units happen to word their errors. + Parameters ---------- run_id : int The ID of the trained run. manual_input_data : List[Dict] List of row dicts keyed by input column name. - component_registry : Any - The DashAI component registry. session_factory : sessionmaker - SQLAlchemy session factory. + SQLAlchemy session factory, used for this function's own row reads. It + must be the container's: the units resolve ``session_factory`` and + ``component_registry`` from the DI container themselves, so a different + one passed here would leave the two halves reading different databases. Returns ------- @@ -112,10 +101,11 @@ def run_manual_prediction( HTTPException On missing run, model session, or prediction failure. """ - with session_factory() as db: - from DashAI.back.dataloaders.classes.dashai_dataset import load_dataset - from DashAI.back.dependencies.database.models import Run + from DashAI.back.dependencies.database.models import Run + # Read everything this function needs off the rows, then let the session go: + # nothing here writes, and the units open their own sessions. + with session_factory() as db: run = db.get(Run, run_id) if not run: raise HTTPException( @@ -149,54 +139,65 @@ def run_manual_prediction( detail="Model session has no output columns configured", ) - try: - task: BaseTask = component_registry[model_session.task_name]["class"]() - except Exception as e: - raise HTTPException( - status_code=status.HTTP_500_INTERNAL_SERVER_ERROR, - detail=f"Task {model_session.task_name} not found in the registry", - ) from e + task_name = model_session.task_name + input_columns = list(model_session.input_columns) + output_columns = list(model_session.output_columns) + train_dataset_file_path = dataset_trained.file_path + + ctx = ExecutionContext() + + build_input = BuildManualInputUnit( + task_name=task_name, + train_dataset_file_path=train_dataset_file_path, + manual_input_data=manual_input_data, + ) + predict = PredictUnit( + task_name=task_name, + input_columns=input_columns, + output_columns=output_columns, + ) + try: + # Both units resolve the task; validating up front keeps a missing task + # reported as a task problem, ahead of the model, the way it always was. try: - model_cls = component_registry[run.model_name]["class"] - except KeyError as e: + build_input.validate(ctx) + predict.validate(ctx) + except JobError as e: raise HTTPException( status_code=status.HTTP_500_INTERNAL_SERVER_ERROR, - detail=f"Model {run.model_name} not found in the registry", + detail=f"Task {task_name} not found in the registry", ) from e try: - trained_model: BaseModel = model_cls.load(run.run_path) - except Exception as e: + LoadTrainedModelUnit(run_id=run_id)(ctx) + except JobError as e: + # The unit distinguishes "not in the registry" from "cannot be read + # from disk" with the same two texts this endpoint has always + # returned, so its message is forwarded rather than rebuilt. raise HTTPException( status_code=status.HTTP_500_INTERNAL_SERVER_ERROR, - detail=( - f"Failed to load model {run.model_name} from path {run.run_path}" - ), + detail=str(e), ) from e try: - train_dataset: "DashAIDataset" = load_dataset( - str(Path(f"{dataset_trained.file_path}/dataset/")) + LoadTrainingDatasetUnit(train_dataset_file_path=train_dataset_file_path)( + ctx ) - except Exception as e: + except JobError as e: + # Not forwarded: the unit names the path, this endpoint never has. raise HTTPException( status_code=status.HTTP_500_INTERNAL_SERVER_ERROR, detail="Cannot load training dataset", ) from e try: - dataset_trained_path = str(Path(f"{dataset_trained.file_path}/dataset/")) - loaded_dataset: "DashAIDataset" = task.process_manual_input( - manual_input_data, dataset_trained_path - ) - prepared_dataset, y_pred = _run_prediction_pipeline( - task=task, - trained_model=trained_model, - train_dataset=train_dataset, - loaded_dataset=loaded_dataset, - model_session=model_session, - ) + build_input(ctx) + predict(ctx) + # Re-derived here rather than published by the unit: a narrowed view + # of the dataset is exactly the kind of derived value that must not + # cross a unit boundary, since anything upstream may reshape it. + prepared_dataset = ctx.require("dataset").select_columns(input_columns) except (ValueError, TypeError) as e: logging.exception("Manual prediction input error: %s", e) raise HTTPException( @@ -209,13 +210,14 @@ def run_manual_prediction( detail="Model prediction failed", ) from e - output_col = model_session.output_columns[0] return _build_preview_rows( prepared_dataset=prepared_dataset, - input_columns=list(model_session.input_columns), - output_col=output_col, - y_pred=y_pred, + input_columns=input_columns, + output_col=output_columns[0], + y_pred=ctx.require("y_pred"), ) + finally: + ctx.clear_cache() class PredictJob(BaseJob): diff --git a/tests/back/api/test_predict_preview_api.py b/tests/back/api/test_predict_preview_api.py new file mode 100644 index 000000000..4d0b79db7 --- /dev/null +++ b/tests/back/api/test_predict_preview_api.py @@ -0,0 +1,490 @@ +"""End-to-end regression net for ``POST /predict/preview``. + +The synchronous counterpart of ``PredictJob``: same prediction, no persistence, +and errors reported as HTTP responses instead of ``JobError``. Written before +``run_manual_prediction`` is decomposed into the same units the job uses, and +asserted against the pre-refactor implementation, so the refactor has something +to be measured against. + +Every failure mode gets its own test with its exact status code and detail +string, because that is the whole contract this endpoint has with the frontend's +manual-prediction form: eleven distinct responses, four status codes. A +migration that turned any of them into a generic 500 would be a silent UX +regression in the one feature the endpoint exists for. + +This endpoint had no tests at all before this file. + +Lives under ``tests/back/api`` to reuse the ``client`` and ``dataset_1`` +fixtures from this package's ``conftest.py``. +""" + +import json +import shutil +from pathlib import Path + +import pytest +from fastapi.testclient import TestClient + +from DashAI.back.dependencies.database.models import Dataset, ModelSession, Run +from DashAI.back.job.model_job import ModelJob + +INPUT_COLUMNS = [ + "SepalLengthCm", + "SepalWidthCm", + "PetalLengthCm", + "PetalWidthCm", +] +OUTPUT_COLUMN = "Species" + +A_ROW = { + "SepalLengthCm": 5.1, + "SepalWidthCm": 3.5, + "PetalLengthCm": 1.4, + "PetalWidthCm": 0.2, +} + +SPLITS = json.dumps( + { + "train": 0.5, + "test": 0.2, + "validation": 0.3, + "is_random": True, + "has_changed": True, + "seed": 42, + "shuffle": True, + "stratify": False, + } +) + + +@pytest.fixture(scope="module", name="model_session_id") +def create_model_session(client: TestClient, dataset_1: Dataset): + session_factory = client.app.container["session_factory"] + + with session_factory() as db: + model_session = ModelSession( + dataset_id=dataset_1.id, + name="PreviewSession", + task_name="TabularClassificationTask", + input_columns=INPUT_COLUMNS, + output_columns=[OUTPUT_COLUMN], + train_metrics=[], + validation_metrics=[], + test_metrics=[], + splits=SPLITS, + ) + db.add(model_session) + db.commit() + db.refresh(model_session) + return model_session.id + + +@pytest.fixture(scope="module", name="trained_run_id") +def create_trained_run(client: TestClient, model_session_id: int): + """A genuinely trained run: the preview path loads the model from disk.""" + session_factory = client.app.container["session_factory"] + + with session_factory() as db: + run = Run( + model_session_id=model_session_id, + optimizer_name="OptunaOptimizer", + optimizer_parameters={ + "n_trials": 1, + "sampler": "TPESampler", + "pruner": "None", + }, + model_name="KNeighborsClassifier", + parameters={}, + name="PreviewRun", + goal_metric="Accuracy", + ) + db.add(run) + db.commit() + db.refresh(run) + run_id = run.id + + ModelJob(run_id=run_id).run() + + with session_factory() as db: + assert db.get(Run, run_id).run_path, "the fixture did not save a model" + return run_id + + +def _preview(client, run_id, rows=None): + return client.post( + "/api/v1/predict/preview", + data={ + "run_id": str(run_id), + "manual_input_data": json.dumps(rows if rows is not None else [A_ROW]), + }, + ) + + +@pytest.fixture(name="restore_run") +def fixture_restore_run(client: TestClient, trained_run_id: int): + """Let a test corrupt the module-scoped Run row and put it back after. + + The run is module scoped because training is slow; without this the + error-branch tests would poison every test after them. + """ + session_factory = client.app.container["session_factory"] + with session_factory() as db: + run = db.get(Run, trained_run_id) + original = { + "model_name": run.model_name, + "run_path": run.run_path, + "model_session_id": run.model_session_id, + } + + yield + + with session_factory() as db: + run = db.get(Run, trained_run_id) + for key, value in original.items(): + setattr(run, key, value) + db.commit() + + +@pytest.fixture(name="restore_model_session") +def fixture_restore_model_session(client: TestClient, model_session_id: int): + """Same idea for the module-scoped ModelSession row.""" + session_factory = client.app.container["session_factory"] + with session_factory() as db: + row = db.get(ModelSession, model_session_id) + original = { + "dataset_id": row.dataset_id, + "task_name": row.task_name, + "input_columns": list(row.input_columns), + "output_columns": list(row.output_columns), + } + + yield + + with session_factory() as db: + row = db.get(ModelSession, model_session_id) + for key, value in original.items(): + setattr(row, key, value) + db.commit() + + +# --- the happy path ----------------------------------------------------- + + +def test_the_preview_returns_the_inputs_plus_the_prediction(client, trained_run_id): + response = _preview(client, trained_run_id) + + assert response.status_code == 200, response.text + body = response.json() + assert body["columns"] == INPUT_COLUMNS + [OUTPUT_COLUMN] + assert len(body["rows"]) == 1 + assert body["rows"][0][:4] == [5.1, 3.5, 1.4, 0.2] + # The label is decoded against the training dataset, not left as an index. + assert body["rows"][0][4] in {"Iris-setosa", "Iris-versicolor", "Iris-virginica"} + + +def test_the_preview_handles_several_rows_at_once(client, trained_run_id): + rows = [A_ROW, {**A_ROW, "PetalLengthCm": 5.9, "PetalWidthCm": 2.1}] + + response = _preview(client, trained_run_id, rows) + + assert response.status_code == 200, response.text + body = response.json() + assert len(body["rows"]) == 2 + assert body["rows"][1][2] == 5.9 + + +def test_the_preview_and_the_job_agree_on_the_same_input(client, trained_run_id): + """The point of sharing units: the two paths cannot answer differently. + + Same hand-typed row, one predicted synchronously for the preview and one + through ``PredictJob``. They used to run separate copies of the same three + steps, so nothing stopped them from drifting apart; this is what would fail + if the prediction were ever fixed in one place only. + """ + from DashAI.back.dataloaders.classes.dashai_dataset import load_dataset + from DashAI.back.dependencies.database.models import Prediction + from DashAI.back.job.predict_job import PredictJob + + row = {**A_ROW, "PetalLengthCm": 4.7, "PetalWidthCm": 1.4} + + previewed = _preview(client, trained_run_id, [row]) + assert previewed.status_code == 200, previewed.text + preview_label = previewed.json()["rows"][0][4] + + created = client.post( + "/api/v1/predict/", json={"run_id": trained_run_id, "dataset_id": None} + ) + assert created.status_code == 200, created.text + prediction_id = created.json()["id"] + + PredictJob(prediction_id=prediction_id, manual_input_data=[row]).run() + + session_factory = client.app.container["session_factory"] + with session_factory() as db: + results_path = db.get(Prediction, prediction_id).results_path + + saved = load_dataset(str(Path(results_path) / "dataset")) + assert saved[OUTPUT_COLUMN] == [preview_label] + + +def test_the_preview_persists_nothing(client, trained_run_id): + """It is a preview: no Prediction row, no results folder.""" + from DashAI.back.dependencies.database.models import Prediction + + session_factory = client.app.container["session_factory"] + with session_factory() as db: + before = db.query(Prediction).count() + + assert _preview(client, trained_run_id).status_code == 200 + + with session_factory() as db: + assert db.query(Prediction).count() == before + + +# --- the four 404/422 checks the endpoint owns -------------------------- + + +def test_a_missing_run_is_a_404(client): + response = _preview(client, 999999) + + assert response.status_code == 404 + assert response.json()["detail"] == "Run not found for id 999999" + + +def test_a_missing_model_session_is_a_404(client, trained_run_id, restore_run): + session_factory = client.app.container["session_factory"] + with session_factory() as db: + db.get(Run, trained_run_id).model_session_id = 999999 + db.commit() + + response = _preview(client, trained_run_id) + + assert response.status_code == 404 + assert response.json()["detail"] == "Model session not found" + + +def test_a_missing_training_dataset_row_is_a_404( + client, trained_run_id, restore_model_session, model_session_id +): + session_factory = client.app.container["session_factory"] + with session_factory() as db: + db.get(ModelSession, model_session_id).dataset_id = 999999 + db.commit() + + response = _preview(client, trained_run_id) + + assert response.status_code == 404 + assert response.json()["detail"] == "Training dataset not found" + + +def test_no_input_columns_is_a_422( + client, trained_run_id, restore_model_session, model_session_id +): + session_factory = client.app.container["session_factory"] + with session_factory() as db: + db.get(ModelSession, model_session_id).input_columns = [] + db.commit() + + response = _preview(client, trained_run_id) + + assert response.status_code == 422 + assert response.json()["detail"] == "Model session has no input columns configured" + + +def test_no_output_columns_is_a_422( + client, trained_run_id, restore_model_session, model_session_id +): + session_factory = client.app.container["session_factory"] + with session_factory() as db: + db.get(ModelSession, model_session_id).output_columns = [] + db.commit() + + response = _preview(client, trained_run_id) + + assert response.status_code == 422 + assert response.json()["detail"] == "Model session has no output columns configured" + + +# --- the registry and loading failures ---------------------------------- + + +def test_an_unknown_task_is_a_500_naming_the_task( + client, trained_run_id, restore_model_session, model_session_id +): + """The task is resolved before the model, so this wins when both are wrong. + + That ordering is behaviour: it decides which of the two the user is told + about, and a decomposition that resolves the model first would change it. + """ + session_factory = client.app.container["session_factory"] + with session_factory() as db: + db.get(ModelSession, model_session_id).task_name = "NoSuchTask" + db.commit() + + response = _preview(client, trained_run_id) + + assert response.status_code == 500 + assert response.json()["detail"] == "Task NoSuchTask not found in the registry" + + +def test_an_unknown_model_is_a_500_naming_the_model( + client, trained_run_id, restore_run +): + session_factory = client.app.container["session_factory"] + with session_factory() as db: + db.get(Run, trained_run_id).model_name = "NoSuchModel" + db.commit() + + response = _preview(client, trained_run_id) + + assert response.status_code == 500 + assert response.json()["detail"] == "Model NoSuchModel not found in the registry" + + +def test_an_unreadable_model_is_a_500_naming_model_and_path( + client, trained_run_id, restore_run +): + session_factory = client.app.container["session_factory"] + with session_factory() as db: + db.get(Run, trained_run_id).run_path = "nowhere/at/all" + db.commit() + + response = _preview(client, trained_run_id) + + assert response.status_code == 500 + assert response.json()["detail"] == ( + "Failed to load model KNeighborsClassifier from path nowhere/at/all" + ) + + +def test_an_unreadable_training_dataset_is_a_500_without_the_path( + client, trained_run_id, dataset_1, tmp_path +): + """The detail is deliberately bare here — no path, unlike the model error. + + Pinned exactly because the unit that replaces this step reports a message of + its own that *does* carry the path; the endpoint has to keep saying this. + """ + stored = Path(dataset_1.file_path) / "dataset" + backup = tmp_path / "training-dataset-backup" + shutil.copytree(stored, backup) + shutil.rmtree(stored) + try: + response = _preview(client, trained_run_id) + + assert response.status_code == 500 + assert response.json()["detail"] == "Cannot load training dataset" + finally: + shutil.copytree(backup, stored) + + +# --- the input and prediction failures ---------------------------------- + + +def test_an_unknown_input_column_is_a_400(client, trained_run_id): + response = _preview(client, trained_run_id, [{"NotAColumn": 1.0}]) + + assert response.status_code == 400 + detail = response.json()["detail"] + assert detail.startswith("Invalid input data: ") + assert "NotAColumn" in detail + + +def test_a_value_error_while_predicting_is_a_400(client, trained_run_id, monkeypatch): + from DashAI.back.models.scikit_learn.k_neighbors_classifier import ( + KNeighborsClassifier, + ) + + def _bad_value(self, x): + raise ValueError("a value the model cannot use") + + monkeypatch.setattr(KNeighborsClassifier, "predict", _bad_value) + + response = _preview(client, trained_run_id) + + assert response.status_code == 400 + assert response.json()["detail"] == ( + "Invalid input data: a value the model cannot use" + ) + + +def test_a_type_error_while_predicting_is_also_a_400_invalid_input( + client, trained_run_id, monkeypatch +): + """The sync path merges ``TypeError`` into the ``ValueError`` message. + + ``PredictJob`` keeps them apart ("Type validation failed" vs "Invalid input + data"). That divergence is pinned on both sides so neither drifts into the + other while they share units. + """ + from DashAI.back.models.scikit_learn.k_neighbors_classifier import ( + KNeighborsClassifier, + ) + + def _wrong_type(self, x): + raise TypeError("bad type somewhere in the input") + + monkeypatch.setattr(KNeighborsClassifier, "predict", _wrong_type) + + response = _preview(client, trained_run_id) + + assert response.status_code == 400 + assert response.json()["detail"] == ( + "Invalid input data: bad type somewhere in the input" + ) + + +def test_any_other_prediction_failure_is_a_500(client, trained_run_id, monkeypatch): + from DashAI.back.models.scikit_learn.k_neighbors_classifier import ( + KNeighborsClassifier, + ) + + def _explode(self, x): + raise RuntimeError("the model itself blew up") + + monkeypatch.setattr(KNeighborsClassifier, "predict", _explode) + + response = _preview(client, trained_run_id) + + assert response.status_code == 500 + assert response.json()["detail"] == "Model prediction failed" + + +# --- the request-shape checks, owned by the endpoint itself ------------- + + +@pytest.mark.parametrize( + ("payload", "detail"), + [ + ({}, "Missing run_id or manual_input_data"), + ( + {"run_id": "notanint", "manual_input_data": "[]"}, + "Invalid run_id: notanint", + ), + ( + {"run_id": "1", "manual_input_data": "[]"}, + "manual_input_data must be a non-empty JSON array of objects (list[dict]).", + ), + ( + {"run_id": "1", "manual_input_data": "[1, 2]"}, + "Each item in manual_input_data must be a JSON object (dict).", + ), + ], +) +def test_the_request_shape_is_validated_before_anything_is_loaded( + client, payload, detail +): + response = client.post("/api/v1/predict/preview", data=payload) + + assert response.status_code == 422 + assert response.json()["detail"] == detail + + +def test_malformed_manual_input_json_is_a_422(client): + response = client.post( + "/api/v1/predict/preview", + data={"run_id": "1", "manual_input_data": "{not json"}, + ) + + assert response.status_code == 422 + assert response.json()["detail"].startswith("Invalid manual_input_data JSON: ")